semantic token fix on package selector, more union completion

5 files changed, 79 insertions, 48 deletions

diff --git a/src/index/collector.odin b/src/index/collector.odin
index cd524f4..4bf2aca 100644
--- a/src/index/collector.odin
+++ b/src/index/collector.odin
@@ -20,8 +20,9 @@ SymbolCollection :: struct {
 }
 
 get_index_unique_string :: proc {
-get_index_unique_string_collection, 
-get_index_unique_string_collection_raw};
+	get_index_unique_string_collection,
+	get_index_unique_string_collection_raw,
+};
 
 get_index_unique_string_collection :: proc (collection: ^SymbolCollection, s: string) -> string {
 	return get_index_unique_string_collection_raw(&collection.unique_strings, collection.allocator, s);
@@ -147,16 +148,25 @@ collect_enum_fields :: proc (collection: ^SymbolCollection, fields: []^ast.Expr,
 collect_union_fields :: proc (collection: ^SymbolCollection, union_type: ast.Union_Type, package_map: map[string]string) -> SymbolUnionValue {
 
 	names := make([dynamic]string, 0, collection.allocator);
+	types := make([dynamic]^ast.Expr, 0, collection.allocator);
 
 	for variant in union_type.variants {
 
 		if ident, ok := variant.derived.(ast.Ident); ok {
 			append(&names, get_index_unique_string(collection, ident.name));
+		} else if selector, ok := variant.derived.(ast.Selector_Expr); ok {
+
+			if ident, ok := selector.field.derived.(ast.Ident); ok {
+				append(&names, ident.name);
+			}
 		}
+
+		append(&types, variant);
 	}
 
 	value := SymbolUnionValue {
-		names = names[:]
+		names = names[:],
+		types = types[:],
 	};
 
 	return value;
@@ -216,16 +226,16 @@ collect_symbols :: proc (collection: ^SymbolCollection, file: ast.File, uri: str
 
 		switch v in col_expr.derived {
 		case ast.Proc_Lit:
-			token = v;
+			token      = v;
 			token_type = .Function;
 
 			if v.type.params != nil {
-				symbol.signature = strings.concatenate({"(", string(file.src[v.type.params.pos.offset:v.type.params.end.offset]), ")"}, 
+				symbol.signature = strings.concatenate({"(", string(file.src[v.type.params.pos.offset:v.type.params.end.offset]), ")"},
 					collection.allocator);
 			}
 
 			if v.type.results != nil {
-				symbol.returns = strings.concatenate({"(", string(file.src[v.type.results.pos.offset:v.type.results.end.offset]), ")"}, 
+				symbol.returns = strings.concatenate({"(", string(file.src[v.type.results.pos.offset:v.type.results.end.offset]), ")"},
 					collection.allocator);
 			}
 
@@ -233,57 +243,57 @@ collect_symbols :: proc (collection: ^SymbolCollection, file: ast.File, uri: str
 				symbol.value = collect_procedure_fields(collection, v.type, v.type.params, v.type.results, package_map);
 			}
 		case ast.Proc_Type:
-			token = v;
+			token      = v;
 			token_type = .Function;
 
 			if v.params != nil {
-				symbol.signature = strings.concatenate({"(", string(file.src[v.params.pos.offset:v.params.end.offset]), ")"}, 
+				symbol.signature = strings.concatenate({"(", string(file.src[v.params.pos.offset:v.params.end.offset]), ")"},
 					collection.allocator);
 			}
 
 			if v.results != nil {
-				symbol.returns = strings.concatenate({"(", string(file.src[v.results.pos.offset:v.results.end.offset]), ")"}, 
+				symbol.returns = strings.concatenate({"(", string(file.src[v.results.pos.offset:v.results.end.offset]), ")"},
 					collection.allocator);
 			}
 
 			symbol.value = collect_procedure_fields(collection, cast(^ast.Proc_Type)col_expr, v.params, v.results, package_map);
 		case ast.Proc_Group:
-			token = v;
-			token_type = .Function;
+			token        = v;
+			token_type   = .Function;
 			symbol.value = SymbolProcedureGroupValue {
 					group = clone_type(col_expr, collection.allocator, &collection.unique_strings)
 				};
 		case ast.Struct_Type:
-			token = v;
-			token_type = .Struct;
-			symbol.value = collect_struct_fields(collection, v, package_map);
+			token            = v;
+			token_type       = .Struct;
+			symbol.value     = collect_struct_fields(collection, v, package_map);
 			symbol.signature = "struct";
 		case ast.Enum_Type:
-			token = v;
-			token_type = .Enum;
-			symbol.value = collect_enum_fields(collection, v.fields, package_map);
+			token            = v;
+			token_type       = .Enum;
+			symbol.value     = collect_enum_fields(collection, v.fields, package_map);
 			symbol.signature = "enum";
 		case ast.Union_Type:
-			token = v;
-			token_type = .Enum;
-			symbol.value = collect_union_fields(collection, v, package_map);
+			token            = v;
+			token_type       = .Enum;
+			symbol.value     = collect_union_fields(collection, v, package_map);
 			symbol.signature = "union";
 		case ast.Bit_Set_Type:
-			token = v;
-			token_type = .Enum;
-			symbol.value = collect_bitset_field(collection, v, package_map);
+			token            = v;
+			token_type       = .Enum;
+			symbol.value     = collect_bitset_field(collection, v, package_map);
 			symbol.signature = "bitset";
 		case ast.Basic_Lit:
-			token = v;
+			token        = v;
 			symbol.value = collect_generic(collection, col_expr, package_map);
 		case ast.Ident:
-			token = v;
-			token_type = .Variable;
+			token        = v;
+			token_type   = .Variable;
 			symbol.value = collect_generic(collection, col_expr, package_map);
 		case: // default
 			symbol.value = collect_generic(collection, col_expr, package_map);
-			token_type = .Variable;
-			token = expr.expr;
+			token_type   = .Variable;
+			token        = expr.expr;
 			break;
 		}
 
@@ -294,8 +304,7 @@ collect_symbols :: proc (collection: ^SymbolCollection, file: ast.File, uri: str
 
 		when ODIN_OS == "windows" {
 			symbol.uri = get_index_unique_string(collection, strings.to_lower(uri, context.temp_allocator));
-		} else 
-		{
+		} else {
 			symbol.uri = get_index_unique_string(collection, uri);
 		}
 
@@ -391,10 +400,11 @@ get_package_mapping :: proc (file: ast.File, config: ^common.Config, uri: string
 */
 
 replace_package_alias :: proc {
-replace_package_alias_node, 
-replace_package_alias_expr, 
-replace_package_alias_array, 
-replace_package_alias_dynamic_array};
+	replace_package_alias_node,
+	replace_package_alias_expr,
+	replace_package_alias_array,
+	replace_package_alias_dynamic_array,
+};
 
 replace_package_alias_array :: proc (array: $A/[]^$T, package_map: map[string]string, collection: ^SymbolCollection) {
 
diff --git a/src/index/symbol.odin b/src/index/symbol.odin
index 94ebbc5..59b47fb 100644
--- a/src/index/symbol.odin
+++ b/src/index/symbol.odin
@@ -42,6 +42,7 @@ SymbolEnumValue :: struct {
 
 SymbolUnionValue :: struct {
 	names: []string,
+	types: []^ast.Expr,
 }
 
 SymbolBitSetValue :: struct {
@@ -78,7 +79,7 @@ Symbol :: struct {
 	value:     SymbolValue,
 }
 
-SymbolType :: enum 
+SymbolType :: enum
 
 //set by ast symbol
 
@@ -97,8 +98,8 @@ SymbolType :: enum
 
 free_symbol :: proc (symbol: Symbol, allocator: mem.Allocator) {
 
-	if symbol.signature != "" && symbol.signature != "struct" && 
-	symbol.signature != "union" && symbol.signature != "enum" && 
+	if symbol.signature != "" && symbol.signature != "struct" &&
+	symbol.signature != "union" && symbol.signature != "enum" &&
 	symbol.signature != "bitset" && symbol.signature != "bitfield" {
 		delete(symbol.signature, allocator);
 	}
@@ -126,6 +127,7 @@ free_symbol :: proc (symbol: Symbol, allocator: mem.Allocator) {
 		delete(v.names, allocator);
 	case SymbolUnionValue:
 		delete(v.names, allocator);
+		common.free_ast(v.types, allocator);
 	case SymbolBitSetValue:
 		common.free_ast(v.expr, allocator);
 	}
diff --git a/src/server/analysis.odin b/src/server/analysis.odin
index 2e95c3a..c8c1827 100644
--- a/src/server/analysis.odin
+++ b/src/server/analysis.odin
@@ -1164,11 +1164,19 @@ make_symbol_union_from_ast :: proc (ast_context: ^AstContext, v: ast.Union_Type,
 		if ident, ok := variant.derived.(ast.Ident); ok {
 			append(&names, ident.name);
 		}
+
+		else if selector, ok := variant.derived.(ast.Selector_Expr); ok {
+
+			if ident, ok := selector.field.derived.(ast.Ident); ok {
+				append(&names, ident.name);
+			}
+		}
 	}
 
 	symbol.value = index.SymbolUnionValue {
-			names = names[:]
-		};
+			names = names[:],
+			types = v.variants,
+	};
 
 	return symbol;
 }
diff --git a/src/server/completion.odin b/src/server/completion.odin
index 40f633c..e19bd0a 100644
--- a/src/server/completion.odin
+++ b/src/server/completion.odin
@@ -242,12 +242,22 @@ get_selector_completion :: proc (ast_context: ^AstContext, position_context: ^Do
 	case index.SymbolUnionValue:
 		list.isIncomplete = false;
 
-		for name in v.names {
-			symbol: index.Symbol;
-			symbol.name = fmt.aprintf("(%v)", name);
-			symbol.pkg  = selector.name;
-			symbol.type = .EnumMember;
-			append(&symbols, symbol);
+		for name, i in v.names {
+
+			if symbol, ok := resolve_type_expression(ast_context, v.types[i]); ok {
+
+				if symbol.pkg == ast_context.document_package {
+					symbol.name = fmt.aprintf("(%v)", name);
+				}
+
+				else {
+					symbol.name = fmt.aprintf("(%v.%v)", path.base(symbol.pkg, false, context.temp_allocator), name);
+				}
+
+				symbol.pkg  = selector.name;
+				symbol.type = .EnumMember;
+				append(&symbols, symbol);
+			}
 		}
 
 	case index.SymbolEnumValue:
diff --git a/src/server/semantic_tokens.odin b/src/server/semantic_tokens.odin
index 8e65841..ab0d6a2 100644
--- a/src/server/semantic_tokens.odin
+++ b/src/server/semantic_tokens.odin
@@ -139,10 +139,11 @@ write_semantic_token_pos :: proc (builder: ^SemanticTokenBuilder, pos: tokenizer
 	builder.current_start = pos.offset;
 }
 
-resolve_and_write_ident :: proc (node: ^ast.Node, builder: ^SemanticTokenBuilder, ast_context: ^AstContext) -> (is_member: bool, is_package: bool) {
+resolve_and_write_ident :: proc (node: ^ast.Node, builder: ^SemanticTokenBuilder, ast_context: ^AstContext) -> (is_member: bool, is_package: bool, package_name: string) {
 
 	n := node.derived.(ast.Ident);
 
+	package_name = ast_context.document_package;
 	ast_context.current_package = ast_context.document_package;
 	ast_context.use_globals     = true;
 	ast_context.use_locals      = true;
@@ -151,11 +152,11 @@ resolve_and_write_ident :: proc (node: ^ast.Node, builder: ^SemanticTokenBuilder
 		write_semantic_node(builder, node, ast_context.file.src, .Variable, .None);
 		is_member = true;
 	} else if symbol, ok := resolve_type_identifier(ast_context, n); ok {
-
 		#partial switch v in symbol.value {
 		case index.SymbolPackageValue:
 			write_semantic_node(builder, node, ast_context.file.src, .Namespace, .None);
 			is_package = true;
+			package_name = symbol.pkg;
 		case index.SymbolStructValue:
 			write_semantic_node(builder, node, ast_context.file.src, .Struct, .None);
 		case index.SymbolEnumValue:
@@ -531,7 +532,7 @@ write_semantic_selector :: proc (selector: ^ast.Selector_Expr, builder: ^Semanti
 
 	if ident, ok := selector.expr.derived.(Ident); ok {
 		get_locals_at(builder.current_function, selector.expr, ast_context);
-		builder.selector_member, builder.selector_package = resolve_and_write_ident(selector.expr, builder, ast_context); //base
+		builder.selector_member, builder.selector_package, ast_context.current_package = resolve_and_write_ident(selector.expr, builder, ast_context); //base
 
 		if builder.selector_package && selector.field != nil && resolve_ident_is_variable(ast_context, selector.field^) {
 			builder.selector_member = true;