Merge pull request #984 from BradLewis/feat/improve-union-enum-completions

Provide full path for union enum completion labels

2 files changed, 113 insertions, 15 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,