Provide full path for union enum completion labels

3 files changed, 191 insertions, 37 deletions

diff --git a/src/server/analysis.odin b/src/server/analysis.odin
index ac3680c..5985fea 100644
--- a/src/server/analysis.odin
+++ b/src/server/analysis.odin
@@ -3484,17 +3484,19 @@ unwrap_ident :: proc(node: ^ast.Expr) -> (^ast.Ident, bool) {
 	return {}, false
 }
 
-unwrap_enum :: proc(ast_context: ^AstContext, node: ^ast.Expr) -> (SymbolEnumValue, bool) {
+// Returns the unwrapped enum, whether it unwrapped a super enum, whether it was successful
+unwrap_enum :: proc(ast_context: ^AstContext, node: ^ast.Expr) -> (SymbolEnumValue, bool, bool) {
 	if node == nil {
-		return {}, false
+		return {}, false, false
 	}
 
 	if enum_symbol, ok := resolve_type_expression(ast_context, node); ok {
 		#partial switch value in enum_symbol.value {
 		case SymbolEnumValue:
-			return value, true
+			return value, false, true
 		case SymbolUnionValue:
-			return unwrap_super_enum(ast_context, value)
+			result, ok := unwrap_super_enum(ast_context, value)
+			return result, true, ok
 		case SymbolSliceValue:
 			return unwrap_enum(ast_context, value.expr)
 		case SymbolFixedArrayValue:
@@ -3506,7 +3508,7 @@ unwrap_enum :: proc(ast_context: ^AstContext, node: ^ast.Expr) -> (SymbolEnumVal
 		}
 	}
 
-	return {}, false
+	return {}, false, false
 }
 
 unwrap_super_enum :: proc(
@@ -3522,7 +3524,17 @@ unwrap_super_enum :: proc(
 	for type in symbol_union.types {
 		symbol := resolve_type_expression(ast_context, type) or_return
 		if value, ok := symbol.value.(SymbolEnumValue); ok {
-			append(&names, ..value.names)
+			for name in value.names {
+				if ast_context.current_package != symbol.pkg {
+					pkg_name := get_pkg_name(ast_context, symbol.pkg)
+					append(
+						&names,
+						fmt.aprintf("%s.%s.%s", pkg_name, symbol.name, name, allocator = ast_context.allocator),
+					)
+				} else {
+					append(&names, fmt.aprintf("%s.%s", symbol.name, name, allocator = ast_context.allocator))
+				}
+			}
 			append(&ranges, ..value.ranges)
 		}
 	}
diff --git a/src/server/completion.odin b/src/server/completion.odin
index 802445f..e9b4de3 100644
--- a/src/server/completion.odin
+++ b/src/server/completion.odin
@@ -1115,7 +1115,7 @@ get_implicit_completion :: proc(
 
 	//value decl infer a : My_Enum = .*
 	if position_context.value_decl != nil && position_context.value_decl.type != nil {
-		if enum_value, ok := unwrap_enum(ast_context, position_context.value_decl.type); ok {
+		if enum_value, unwrapped_super_enum, ok := unwrap_enum(ast_context, position_context.value_decl.type); ok {
 			for name in enum_value.names {
 				if position_context.comp_lit != nil && field_exists_in_comp_lit(position_context.comp_lit, name) {
 					continue
@@ -1125,6 +1125,9 @@ get_implicit_completion :: proc(
 					kind   = .EnumMember,
 					detail = name,
 				}
+				if unwrapped_super_enum {
+					add_implicit_selector_remove_edit(position_context, &item, name, enum_value.names)
+				}
 				append(results, CompletionResult{completion_item = item})
 			}
 
@@ -1197,7 +1200,7 @@ get_implicit_completion :: proc(
 			}
 		}
 
-		if enum_value, ok := unwrap_enum(ast_context, position_context.switch_stmt.cond); ok {
+		if enum_value, unwrapped_super_enum, ok := unwrap_enum(ast_context, position_context.switch_stmt.cond); ok {
 			for name in enum_value.names {
 				if name in used_enums {
 					continue
@@ -1208,6 +1211,9 @@ get_implicit_completion :: proc(
 					kind   = .EnumMember,
 					detail = name,
 				}
+				if unwrapped_super_enum {
+					add_implicit_selector_remove_edit(position_context, &item, name, enum_value.names)
+				}
 
 				append(results, CompletionResult{completion_item = item})
 			}
@@ -1280,13 +1286,16 @@ get_implicit_completion :: proc(
 	if position_context.comp_lit != nil {
 		if symbol, ok := resolve_type_expression(ast_context, position_context.comp_lit); ok {
 			if symbol_value, ok := symbol.value.(SymbolFixedArrayValue); ok {
-				if enum_value, ok := unwrap_enum(ast_context, symbol_value.len); ok {
+				if enum_value, unwrapped_super_enum, ok := unwrap_enum(ast_context, symbol_value.len); ok {
 					for enum_name in enum_value.names {
 						item := CompletionItem {
 							label  = enum_name,
 							kind   = .EnumMember,
 							detail = enum_name,
 						}
+						if unwrapped_super_enum {
+							add_implicit_selector_remove_edit(position_context, &item, enum_name, enum_value.names)
+						}
 
 						append(results, CompletionResult{completion_item = item})
 					}
@@ -1344,13 +1353,16 @@ get_implicit_completion :: proc(
 			}
 
 			if context_node != nil && enum_node != nil {
-				if enum_value, ok := unwrap_enum(ast_context, enum_node); ok {
+				if enum_value, unwrapped_super_enum, ok := unwrap_enum(ast_context, enum_node); ok {
 					for name in enum_value.names {
 						item := CompletionItem {
 							label  = name,
 							kind   = .EnumMember,
 							detail = name,
 						}
+						if unwrapped_super_enum {
+							add_implicit_selector_remove_edit(position_context, &item, name, enum_value.names)
+						}
 
 						append(results, CompletionResult{completion_item = item})
 					}
@@ -1386,13 +1398,16 @@ get_implicit_completion :: proc(
 		}
 
 		if len(position_context.assign.lhs) > rhs_index {
-			if enum_value, ok := unwrap_enum(ast_context, position_context.assign.lhs[rhs_index]); ok {
+			if enum_value, unwrapped_super_enum, ok := unwrap_enum(ast_context, position_context.assign.lhs[rhs_index]); ok {
 				for name in enum_value.names {
 					item := CompletionItem {
 						label  = name,
 						kind   = .EnumMember,
 						detail = name,
 					}
+					if unwrapped_super_enum {
+						add_implicit_selector_remove_edit(position_context, &item, name, enum_value.names)
+					}
 
 					append(results, CompletionResult{completion_item = item})
 				}
@@ -1427,7 +1442,7 @@ get_implicit_completion :: proc(
 		}
 
 		if len(position_context.function.type.results.list) > return_index {
-			if enum_value, ok := unwrap_enum(
+			if enum_value, unwrapped_super_enum, ok := unwrap_enum(
 				ast_context,
 				position_context.function.type.results.list[return_index].type,
 			); ok {
@@ -1438,6 +1453,9 @@ get_implicit_completion :: proc(
 						detail = name,
 					}
 
+					if unwrapped_super_enum {
+						add_implicit_selector_remove_edit(position_context, &item, name, enum_value.names)
+					}
 					append(results, CompletionResult{completion_item = item})
 				}
 
@@ -1484,7 +1502,7 @@ get_implicit_completion :: proc(
 						}
 					}
 
-					if enum_value, ok := unwrap_enum(ast_context, type); ok {
+					if enum_value, unwrapped_super_enum, ok := unwrap_enum(ast_context, type); ok {
 						for name in enum_value.names {
 							if position_context.comp_lit != nil &&
 							   field_exists_in_comp_lit(position_context.comp_lit, name) {
@@ -1495,6 +1513,9 @@ get_implicit_completion :: proc(
 								kind   = .EnumMember,
 								detail = name,
 							}
+							if unwrapped_super_enum {
+								add_implicit_selector_remove_edit(position_context, &item, name, enum_value.names)
+							}
 
 							append(results, CompletionResult{completion_item = item})
 						}
@@ -1533,13 +1554,16 @@ get_implicit_completion :: proc(
 
 		#partial switch v in symbol.value {
 		case SymbolFixedArrayValue:
-			if enum_value, ok := unwrap_enum(ast_context, v.len); ok {
+			if enum_value, unwrapped_super_enum, ok := unwrap_enum(ast_context, v.len); ok {
 				for name in enum_value.names {
 					item := CompletionItem {
 						label  = name,
 						kind   = .EnumMember,
 						detail = name,
 					}
+					if unwrapped_super_enum {
+						add_implicit_selector_remove_edit(position_context, &item, name, enum_value.names)
+					}
 
 					append(results, CompletionResult{completion_item = item})
 				}
@@ -1547,13 +1571,16 @@ get_implicit_completion :: proc(
 				return is_incomplete
 			}
 		case SymbolMapValue:
-			if enum_value, ok := unwrap_enum(ast_context, v.key); ok {
+			if enum_value, unwrapped_super_enum, ok := unwrap_enum(ast_context, v.key); ok {
 				for name in enum_value.names {
 					item := CompletionItem {
 						label  = name,
 						kind   = .EnumMember,
 						detail = name,
 					}
+					if unwrapped_super_enum {
+						add_implicit_selector_remove_edit(position_context, &item, name, enum_value.names)
+					}
 
 					append(results, CompletionResult{completion_item = item})
 				}
@@ -1565,6 +1592,65 @@ get_implicit_completion :: proc(
 	return is_incomplete
 }
 
+add_implicit_selector_remove_edit :: proc(
+	position_context: ^DocumentPositionContext,
+	item: ^CompletionItem,
+	name: string,
+	valid_names: []string,
+) {
+	get_name :: proc(full_name: string) -> string {
+		split_name := strings.split(full_name, ".")
+		return split_name[len(split_name) - 1]
+	}
+
+	enum_variant_name := get_name(name)
+	found_match := false
+	for valid_name in valid_names {
+		if name == valid_name {
+			continue
+		}
+		if enum_variant_name == get_name(valid_name) {
+			found_match = true
+			break
+		}
+	}
+
+	if found_match {
+		remove_edit, ok := create_implicit_selector_remove_edit(position_context)
+		if !ok {
+			return
+		}
+		item.additionalTextEdits = remove_edit
+	} else {
+		item.insertText = enum_variant_name
+	}
+}
+
+create_implicit_selector_remove_edit :: proc(position_context: ^DocumentPositionContext) -> ([]TextEdit, bool) {
+	if position_context.implicit_selector_expr != nil {
+		range := common.get_token_range(position_context.implicit_selector_expr, position_context.file.src)
+
+		end := range.start
+		end.character += len(position_context.implicit_selector_expr.field.name)
+		remove_range := common.Range {
+			start = range.start,
+			end   = end,
+		}
+
+		remove_edit := TextEdit {
+			range   = remove_range,
+			newText = "",
+		}
+
+		additionalTextEdits := make([]TextEdit, 1, context.temp_allocator)
+		additionalTextEdits[0] = remove_edit
+
+		return additionalTextEdits, true
+	}
+
+	return nil, false
+}
+
 get_identifier_completion :: proc(
 	ast_context: ^AstContext,
 	position_context: ^DocumentPositionContext,
diff --git a/tests/completions_test.odin b/tests/completions_test.odin
index d6af935..a2c6883 100644
--- a/tests/completions_test.odin
+++ b/tests/completions_test.odin
@@ -2731,7 +2731,7 @@ ast_simple_union_of_enums_completion :: proc(t: ^testing.T) {
 		`,
 	}
 
-	test.expect_completion_labels(t, &source, ".", {"ONE", "TWO"})
+	test.expect_completion_labels(t, &source, ".", {"Sub_Enum_1.ONE", "Sub_Enum_2.TWO"})
 }
 
 
@@ -4471,9 +4471,7 @@ ast_completion_handle_matching_basic_types :: proc(t: ^testing.T) {
 			bar(f{*})
 		}
 		`,
-		config = {
-			enable_completion_matching = true,
-		}
+		config = {enable_completion_matching = true},
 	}
 	test.expect_completion_insert_text(t, &source, "", {"&foo"})
 }
@@ -4492,9 +4490,7 @@ ast_completion_handle_matching_struct :: proc(t: ^testing.T) {
 			bar(f{*})
 		}
 		`,
-		config = {
-			enable_completion_matching = true,
-		}
+		config = {enable_completion_matching = true},
 	}
 	test.expect_completion_insert_text(t, &source, "", {"&foo"})
 }
@@ -4509,9 +4505,7 @@ ast_completion_handle_matching_append :: proc(t: ^testing.T) {
 			append(fo{*})
 		}
 		`,
-		config = {
-			enable_completion_matching = true,
-		}
+		config = {enable_completion_matching = true},
 	}
 	test.expect_completion_insert_text(t, &source, "", {"&foos"})
 }
@@ -4528,9 +4522,7 @@ ast_completion_handle_matching_dynamic_array_to_slice :: proc(t: ^testing.T) {
 			bar(fo{*})
 		}
 		`,
-		config = {
-			enable_completion_matching = true,
-		}
+		config = {enable_completion_matching = true},
 	}
 	test.expect_completion_insert_text(t, &source, "", {"foos[:]"})
 }
@@ -4554,12 +4546,7 @@ ast_completion_proc_bit_set_comp_lit_default_param_with_no_type :: proc(t: ^test
 		`,
 	}
 
-	test.expect_completion_docs(
-		t,
-		&source,
-		"",
-		{"A", "B"},
-	)
+	test.expect_completion_docs(t, &source, "", {"A", "B"})
 }
 
 @(test)
@@ -4593,9 +4580,7 @@ ast_completion_handle_matching_from_overloaded_proc :: proc(t: ^testing.T) {
 			do_foo(f{*})
 		}
 		`,
-		config = {
-			enable_completion_matching = true,
-		},
+		config = {enable_completion_matching = true},
 	}
 	test.expect_completion_insert_text(t, &source, "", {"&foo"})
 }
@@ -4612,3 +4597,74 @@ ast_completion_poly_proc_narrow_type :: proc(t: ^testing.T) {
 	}
 	test.expect_completion_docs(t, &source, "", {"test.Foo: struct {..}"})
 }
+
+@(test)
+ast_completion_union_with_enums :: proc(t: ^testing.T) {
+	source := test.Source {
+		main = `package test
+
+		Foo :: enum {
+			A, B,
+		}
+
+		Bar :: enum {
+			A, C,
+		}
+
+		Bazz :: union #shared_nil {
+			Foo,
+			Bar,
+		}
+
+		main :: proc() {
+			bazz: Bazz
+			if bazz == .F{*} {}
+		}
+		`,
+	}
+	test.expect_completion_docs(t, &source, "", {"Foo.A", "Foo.B", "Bar.A", "Bar.C"})
+}
+
+@(test)
+ast_completion_union_with_enums_from_package :: proc(t: ^testing.T) {
+	packages := make([dynamic]test.Package, context.temp_allocator)
+
+	append(
+		&packages,
+		test.Package {
+			pkg = "my_package",
+			source = `package my_package
+		Foo :: enum {
+			A, B,
+		}
+
+		Bar :: enum {
+			A, C,
+		}
+
+		Bazz :: union #shared_nil {
+			Foo,
+			Bar,
+		}
+		`,
+		},
+	)
+
+	source := test.Source {
+		main     = `package test
+		import "my_package"
+
+		main :: proc() {
+			bazz: my_package.Bazz
+			if bazz == .{*} {}
+		}
+		`,
+		packages = packages[:],
+	}
+	test.expect_completion_docs(
+		t,
+		&source,
+		"",
+		{"my_package.Foo.A", "my_package.Foo.B", "my_package.Bar.A", "my_package.Bar.C"},
+	)
+}