Adding support to the union enums

2 files changed, 55 insertions, 0 deletions

diff --git a/src/server/analysis.odin b/src/server/analysis.odin
index 228eafa..61855ef 100644
--- a/src/server/analysis.odin
+++ b/src/server/analysis.odin
@@ -3992,12 +3992,37 @@ unwrap_enum :: proc(
 	if enum_symbol, ok := resolve_type_expression(ast_context, node); ok {
 		if enum_value, ok := enum_symbol.value.(SymbolEnumValue); ok {
 			return enum_value, true
+		} else if union_value, ok := enum_symbol.value.(SymbolUnionValue); ok {
+			return unwrap_super_enum(ast_context, union_value)
 		}
 	}
 
 	return {}, false
 }
 
+unwrap_super_enum :: proc(
+	ast_context: ^AstContext,
+	symbol_union: SymbolUnionValue,
+) -> (
+	ret_value: SymbolEnumValue,
+	ret_ok: bool,
+) {
+	names := make([dynamic]string, 0, 20, ast_context.allocator)
+	ranges := make([dynamic]common.Range, 0, 20, ast_context.allocator)
+
+	for type in symbol_union.types {
+		symbol := resolve_type_expression(ast_context, type) or_return
+		value := symbol.value.(SymbolEnumValue) or_return
+		append(&names, ..value.names)
+		append(&ranges, ..value.ranges)
+	}
+
+	ret_value.names = names[:]
+	ret_value.ranges = ranges[:]
+
+	return ret_value, true
+}
+
 unwrap_union :: proc(
 	ast_context: ^AstContext,
 	node: ^ast.Expr,
@@ -4030,6 +4055,9 @@ unwrap_bitset :: proc(
 		); ok {
 			if enum_value, ok := enum_symbol.value.(SymbolEnumValue); ok {
 				return enum_value, true
+			} else if union_value, ok := enum_symbol.value.(SymbolUnionValue);
+			   ok {
+				return unwrap_super_enum(ast_context, union_value)
 			}
 		}
 	}
diff --git a/tests/completions_test.odin b/tests/completions_test.odin
index d74de1b..ea7a8ab 100644
--- a/tests/completions_test.odin
+++ b/tests/completions_test.odin
@@ -2665,6 +2665,33 @@ ast_simple_bit_field_completion :: proc(t: ^testing.T) {
 	)
 }
 
+@(test)
+ast_simple_union_of_enums_completion :: proc(t: ^testing.T) {
+	source := test.Source {
+		main = `package test
+			Sub_Enum_1 :: enum {
+				ONE,
+			}
+			Sub_Enum_2 :: enum {
+				TWO,
+			}
+
+			Super_Enum :: union {
+				Sub_Enum_1,
+				Sub_Enum_2,
+			}
+
+			fn :: proc(mode: Super_Enum) {}
+
+			main :: proc() {
+				fn(.{*})
+			}
+		`,
+	}
+
+	test.expect_completion_labels(t, &source, ".", {"ONE", "TWO"})
+}
+
 
 @(test)
 ast_generics_function_with_struct_same_pkg :: proc(t: ^testing.T) {