Fix union poly types

5 files changed, 173 insertions, 61 deletions

diff --git a/src/analysis/analysis.odin b/src/analysis/analysis.odin
index 09bfb3a..29d7760 100644
--- a/src/analysis/analysis.odin
+++ b/src/analysis/analysis.odin
@@ -434,7 +434,6 @@ resolve_generic_function_symbol :: proc(ast_context: ^AstContext, params: []^ast
 	argument_types := make([dynamic]^ast.Field, context.temp_allocator);
 
 	for result in results {
-
 		if result.type == nil {
 			continue;
 		}
@@ -453,7 +452,6 @@ resolve_generic_function_symbol :: proc(ast_context: ^AstContext, params: []^ast
 	}
 
 	for param in params {
-		
 		if len(param.names) == 0 {
 			continue;
 		}
@@ -854,7 +852,7 @@ resolve_type_expression :: proc(ast_context: ^AstContext, node: ^ast.Expr) -> (i
 	case Basic_Lit:
 		return resolve_basic_lit(ast_context, v);
 	case Type_Cast:
-		return resolve_type_expression(ast_context, v.type);
+			return resolve_type_expression(ast_context, v.type);
 	case Auto_Cast:
 		return resolve_type_expression(ast_context, v.expr);
 	case Comp_Lit:
@@ -886,7 +884,20 @@ resolve_type_expression :: proc(ast_context: ^AstContext, node: ^ast.Expr) -> (i
 		ident.name = v.tok.text;
 		return resolve_type_identifier(ast_context, ident^);
 	case Type_Assertion:
-		return resolve_type_expression(ast_context, v.type);
+		if unary, ok := v.type.derived.(ast.Unary_Expr); ok {
+			if unary.op.kind == .Question {
+				if symbol, ok := resolve_type_expression(ast_context, v.expr); ok {
+					if union_value, ok := symbol.value.(index.SymbolUnionValue); ok {
+						if len(union_value.types) != 1 {
+							return {}, false;
+						}
+						return resolve_type_expression(ast_context, union_value.types[0]);
+					}
+				}
+			}
+		} else {
+			return resolve_type_expression(ast_context, v.type);
+		}	
 	case Proc_Lit:
 		if v.type.results != nil {
 			if len(v.type.results.list) == 1 {
@@ -1367,7 +1378,7 @@ resolve_symbol_return :: proc(ast_context: ^AstContext, symbol: index.Symbol, ok
 		resolve_unresolved_symbol(ast_context, &symbol);
 	}
 	
-	#partial switch v in symbol.value {
+	#partial switch v in &symbol.value {
 	case index.SymbolProcedureGroupValue:
 		if symbol, ok := resolve_function_overload(ast_context, v.group.derived.(ast.Proc_Group)); ok {
 			return symbol, true;
@@ -1384,14 +1395,33 @@ resolve_symbol_return :: proc(ast_context: ^AstContext, symbol: index.Symbol, ok
 		} else {
 			return symbol, true;
 		}
+	case index.SymbolUnionValue:
+		if v.poly != nil {
+			//Todo(daniel): Maybe change the function to return a new symbol instead of referencing it.
+			//resolving the poly union means changing the type, so we do a copy of it.
+			types := make([dynamic]^ast.Expr, context.temp_allocator);
+			append_elems(&types, ..v.types);
+			v.types = types[:];
+			resolve_poly_union(ast_context, v.poly, &symbol);
+		}
+		return symbol, ok;
 	case index.SymbolStructValue:
+		if v.poly != nil {
+			//Todo(daniel): Maybe change the function to return a new symbol instead of referencing it.
+			//resolving the struct union means changing the type, so we do a copy of it.
+			types := make([dynamic]^ast.Expr, context.temp_allocator);
+			append_elems(&types, ..v.types);
+			v.types = types[:];
+			resolve_poly_struct(ast_context, v.poly, &symbol);
+		}
+
 		//expand the types and names from the using - can't be done while indexing without complicating everything(this also saves memory)
 		if len(v.usings) > 0 {
 			expanded := symbol;
 			expanded.value = expand_struct_usings(ast_context, symbol, v);
 			return expanded, true;
 		} else {
-			return symbol, true;
+			return symbol, true; 
 		}
 	case index.SymbolGenericValue:
 		ret, ok := resolve_type_expression(ast_context, v.expr);
@@ -1672,7 +1702,7 @@ make_symbol_union_from_ast :: proc(ast_context: ^AstContext, v: ast.Union_Type,
 
 	symbol := index.Symbol {
 		range = common.get_token_range(v, ast_context.file.src),
-		type = .Enum,
+		type = .Union,
 		pkg = get_package_from_node(v.node),
 	};
 
@@ -1696,11 +1726,14 @@ make_symbol_union_from_ast :: proc(ast_context: ^AstContext, v: ast.Union_Type,
 	}
 
 	symbol.value = index.SymbolUnionValue {
-		names = names[:],
 		types = v.variants,
 		union_name = ident,
 	};
 
+	if v.poly_params != nil {
+		resolve_poly_union(ast_context, v.poly_params, &symbol);
+	}
+
 	return symbol;
 }
 
@@ -1778,7 +1811,6 @@ make_symbol_struct_from_ast :: proc(ast_context: ^AstContext, v: ast.Struct_Type
 	usings := make(map[string]bool, 0, context.temp_allocator);
 
 	for field in v.fields.list {
-
 		for n in field.names {
 			if identifier, ok := n.derived.(ast.Ident); ok {
 				append(&names, identifier.name);
@@ -1799,7 +1831,7 @@ make_symbol_struct_from_ast :: proc(ast_context: ^AstContext, v: ast.Struct_Type
 	};
 
 	if v.poly_params != nil {
-		resolve_poly_struct(ast_context, v, &symbol);
+		resolve_poly_struct(ast_context, v.poly_params, &symbol);
 	}
 
 	//TODO change the expand to not double copy the array, but just pass the dynamic arrays
@@ -1810,13 +1842,13 @@ make_symbol_struct_from_ast :: proc(ast_context: ^AstContext, v: ast.Struct_Type
 	return symbol;
 }
 
-resolve_poly_struct :: proc(ast_context: ^AstContext, v: ast.Struct_Type, symbol: ^index.Symbol) {
+resolve_poly_union :: proc(ast_context: ^AstContext, poly_params: ^ast.Field_List, symbol: ^index.Symbol) {
 
 	if ast_context.call == nil {
 		return;
 	}
 
-	symbol_value := &symbol.value.(index.SymbolStructValue);
+	symbol_value := &symbol.value.(index.SymbolUnionValue);
 
 	if symbol_value == nil {
 		return;
@@ -1826,10 +1858,8 @@ resolve_poly_struct :: proc(ast_context: ^AstContext, v: ast.Struct_Type, symbol
 
 	poly_map := make(map[string]^ast.Expr, 0, context.temp_allocator);
 
-	for param in v.poly_params.list {
-
+	for param in poly_params.list {
 		for name in param.names {
-
 			if len(ast_context.call.args) <= i {
 				break;
 			}
@@ -1837,37 +1867,96 @@ resolve_poly_struct :: proc(ast_context: ^AstContext, v: ast.Struct_Type, symbol
 			if param.type == nil {
 				continue;
 			}
-
+			
 			if poly, ok := param.type.derived.(ast.Typeid_Type); ok {
-
 				if ident, ok := name.derived.(ast.Ident); ok {
 					poly_map[ident.name] = ast_context.call.args[i];
+				} else if poly, ok := name.derived.(ast.Poly_Type); ok { 
+					if poly.type != nil {
+						poly_map[poly.type.name] = ast_context.call.args[i];
+					}
 				}
 			}
 
 			i += 1;
 		}
 	}
-
+	
 	for type, i in symbol_value.types {
-
 		if ident, ok := type.derived.(ast.Ident); ok {
-
 			if expr, ok := poly_map[ident.name]; ok {
 				symbol_value.types[i] = expr;
 			}
 		} else if call_expr, ok := type.derived.(ast.Call_Expr); ok {
-
 			if call_expr.args == nil {
 				continue;
 			}
 
 			for arg, i in call_expr.args {
-
 				if ident, ok := arg.derived.(ast.Ident); ok {
+					if expr, ok := poly_map[ident.name]; ok {
+						symbol_value.types[i] = expr;
+					}
+				}
+			}
+		}
+	}
+}
+
+resolve_poly_struct :: proc(ast_context: ^AstContext, poly_params: ^ast.Field_List, symbol: ^index.Symbol) {
+
+	if ast_context.call == nil {
+		return;
+	}
+
+	symbol_value := &symbol.value.(index.SymbolStructValue);
+
+	if symbol_value == nil {
+		return;
+	}
+
+	i := 0;
+
+	poly_map := make(map[string]^ast.Expr, 0, context.temp_allocator);
+
+	for param in poly_params.list {
+		for name in param.names {
+			if len(ast_context.call.args) <= i {
+				break;
+			}
+
+			if param.type == nil {
+				continue;
+			}
+			
+			if poly, ok := param.type.derived.(ast.Typeid_Type); ok {
+				if ident, ok := name.derived.(ast.Ident); ok {
+					poly_map[ident.name] = ast_context.call.args[i];
+				} else if poly, ok := name.derived.(ast.Poly_Type); ok { 
+					if poly.type != nil {
+						poly_map[poly.type.name] = ast_context.call.args[i];
+					}
+				}
+			}
 
+			i += 1;
+		}
+	}
+	
+	for type, i in symbol_value.types {
+		if ident, ok := type.derived.(ast.Ident); ok {
+			if expr, ok := poly_map[ident.name]; ok {
+				symbol_value.types[i] = expr;
+			}
+		} else if call_expr, ok := type.derived.(ast.Call_Expr); ok {
+			if call_expr.args == nil {
+				continue;
+			}
+
+			for arg, i in call_expr.args {
+				if ident, ok := arg.derived.(ast.Ident); ok {
 					if expr, ok := poly_map[ident.name]; ok {
-						call_expr.args[i] = expr;
+						symbol_value.types[i] = expr;
 					}
 				}
 			}
@@ -1894,7 +1983,7 @@ get_generic_assignment :: proc(file: ast.File, value: ^ast.Expr, ast_context: ^A
 	ast_context.use_locals = true;
 	ast_context.use_globals = true;
 
-	switch v in value.derived {
+	switch v in &value.derived {
 	case Call_Expr:
 
 		ast_context.call = cast(^ast.Call_Expr)value;
@@ -1927,8 +2016,16 @@ get_generic_assignment :: proc(file: ast.File, value: ^ast.Expr, ast_context: ^A
 		}
 	case Type_Assertion:
 		if v.type != nil {
-			append(results, v.type);
-			append(results, make_bool_ast());
+			//This is the unique .? that can only be used with maybe
+			if unary, ok := v.type.derived.(ast.Unary_Expr); ok && unary.op.kind == .Question {
+				append(results, cast(^ast.Expr)&v.node);
+			} else {
+				append(results, v.type);
+			}
+
+			b := make_bool_ast();
+			b.pos.file = v.type.pos.file;
+			append(results, b);
 		}
 	case:
 		//log.debugf("default node get_generic_assignment %v", v);
@@ -2304,11 +2401,9 @@ clear_locals :: proc(ast_context: ^AstContext) {
 }
 
 concatenate_symbols_information :: proc(ast_context: ^AstContext, symbol: index.Symbol, is_completion: bool) -> string {
-
 	pkg := path.base(symbol.pkg, false, context.temp_allocator);
 
 	if symbol.type == .Function {
-
 		if symbol.returns != "" {
 			return fmt.tprintf("%v.%v: proc%v -> %v", pkg, symbol.name, symbol.signature, symbol.returns);
 		} else {
diff --git a/src/index/collector.odin b/src/index/collector.odin
index 11e949a..077f6ee 100644
--- a/src/index/collector.odin
+++ b/src/index/collector.odin
@@ -119,6 +119,7 @@ collect_struct_fields :: proc(collection: ^SymbolCollection, struct_type: ast.St
 		types = types[:],
 		usings = usings,
 		struct_name = get_index_unique_string(collection, ident),
+		poly = cast(^ast.Field_List)clone_type(struct_type.poly_params, collection.allocator, &collection.unique_strings),
 	};
 
 	return value;
@@ -151,30 +152,18 @@ 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, ident: 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, get_index_unique_string(collection, ident.name));
-			}
-		}
-
 		cloned := clone_type(variant, collection.allocator, &collection.unique_strings);
 		replace_package_alias(cloned, package_map, collection);
-
 		append(&types, cloned);
 	}
 
 	value := SymbolUnionValue {
-		names = names[:],
 		types = types[:],
 		union_name = get_index_unique_string(collection, ident),
+		poly = cast(^ast.Field_List)clone_type(union_type.poly_params, collection.allocator, &collection.unique_strings),
 	};
 
 	return value;
@@ -320,7 +309,7 @@ collect_symbols :: proc(collection: ^SymbolCollection, file: ast.File, uri: stri
 			symbol.signature = "enum";
 		case ast.Union_Type:
 			token = v;
-			token_type = .Enum;
+			token_type = .Union;
 			symbol.value = collect_union_fields(collection, v, package_map, name);
 			symbol.signature = "union";
 		case ast.Bit_Set_Type:
diff --git a/src/index/symbol.odin b/src/index/symbol.odin
index e2ebad5..5001054 100644
--- a/src/index/symbol.odin
+++ b/src/index/symbol.odin
@@ -15,7 +15,7 @@ SymbolStructValue :: struct {
 	names:   []string,
 	types:   []^ast.Expr,
 	usings:  map[string]bool,
-	generic: bool,
+	poly: ^ast.Field_List,
 }
 
 SymbolPackageValue :: struct {}
@@ -42,8 +42,8 @@ SymbolEnumValue :: struct {
 
 SymbolUnionValue :: struct {
 	union_name: string,
-	names: []string,
 	types: []^ast.Expr,
+	poly: ^ast.Field_List,
 }
 
 SymbolDynamicArrayValue :: struct {
@@ -137,6 +137,7 @@ SymbolType :: enum {
 	EnumMember = 20,
 	Constant   = 21,
 	Struct     = 22,
+	Union      = 9000,
 	Unresolved = 9999,
 }
 
@@ -170,7 +171,6 @@ free_symbol :: proc(symbol: Symbol, allocator: mem.Allocator) {
 	case SymbolEnumValue:
 		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/completion.odin b/src/server/completion.odin
index ce0e872..789d60e 100644
--- a/src/server/completion.odin
+++ b/src/server/completion.odin
@@ -362,13 +362,13 @@ get_selector_completion :: proc(ast_context: ^analysis.AstContext, position_cont
 	case index.SymbolUnionValue:
 		list.isIncomplete = false;
 
-		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);
+		for type in v.types {
+			if symbol, ok := resolve_type_expression(ast_context, type); ok {
+				base := path.base(symbol.pkg, false, context.temp_allocator);
+				if symbol.pkg == ast_context.document_package || base == "runtime" {
+					symbol.name = fmt.aprintf("(%v)", common.node_to_string(type));
 				} else {
-					symbol.name = fmt.aprintf("(%v.%v)", path.base(symbol.pkg, false, context.temp_allocator), name);
+					symbol.name = fmt.aprintf("(%v.%v)", path.base(symbol.pkg, false, context.temp_allocator), common.node_to_string(type));
 				}
 
 				item := CompletionItem {
@@ -1127,7 +1127,6 @@ search_for_packages :: proc(fullpath: string) -> [] string {
 	}
 
 	if files, err := os.read_dir(fh, 0, context.temp_allocator); err == 0 {
-
 		for file in files {
 			if file.is_dir {
 				append(&packages, file.fullpath);
@@ -1149,13 +1148,9 @@ get_type_switch_completion :: proc(ast_context: ^analysis.AstContext, position_c
 	used_unions := make(map[string]bool, 5, context.temp_allocator);
 
 	if block, ok := position_context.switch_type_stmt.body.derived.(ast.Block_Stmt); ok {
-
 		for stmt in block.stmts {
-
 			if case_clause, ok := stmt.derived.(ast.Case_Clause); ok {
-
 				for name in case_clause.list {
-
 					if ident, ok := name.derived.(ast.Ident); ok {
 						used_unions[ident.name] = true;
 					}
@@ -1168,23 +1163,21 @@ get_type_switch_completion :: proc(ast_context: ^analysis.AstContext, position_c
 	ast_context.use_globals = true;
 
 	if assign, ok := position_context.switch_type_stmt.tag.derived.(ast.Assign_Stmt); ok && assign.rhs != nil && len(assign.rhs) == 1 {
-
 		if union_value, ok := unwrap_union(ast_context, assign.rhs[0]); ok {
-
-			for name, i in union_value.names {
+			for type, i in union_value.types {
+				name := common.node_to_string(type);
 
 				if name in used_unions {
 					continue;
 				}
 
 				if symbol, ok := resolve_type_expression(ast_context, union_value.types[i]); ok {
-
 					item := CompletionItem {
 						kind = .EnumMember,
 					};
 
 					if symbol.pkg == ast_context.document_package {
-						item.label  = fmt.aprintf("%v", name);
+						item.label  = fmt.aprintf("%v", common.node_to_string(union_value.types[i]));
 						item.detail = item.label;
 					} else {
 						item.label  = fmt.aprintf("%v.%v", path.base(symbol.pkg, false, context.temp_allocator), name);
diff --git a/src/testing/testing.odin b/src/testing/testing.odin
index 0703708..bd4e39b 100644
--- a/src/testing/testing.odin
+++ b/src/testing/testing.odin
@@ -160,6 +160,41 @@ expect_signature_parameter_position :: proc(t: ^testing.T, src: ^Source, positio
 	}
 }
 
+expect_completion_labels :: proc(t: ^testing.T, src: ^Source, trigger_character: string, expect_labels: []string) {
+	setup(src);
+
+	completion_context := server.CompletionContext {
+		triggerCharacter = trigger_character,
+	};
+
+	completion_list, ok := server.get_completion_list(src.document, src.position, completion_context);
+
+	if !ok {
+		testing.error(t, "Failed get_completion_list");
+	}
+
+	if len(expect_labels) == 0 && len(completion_list.items) > 0 {
+		testing.errorf(t, "Expected empty completion label, but received %v", completion_list.items);
+	}
+
+	flags := make([]int, len(expect_labels));
+
+	for expect_label, i in expect_labels {
+		for completion, j in completion_list.items {
+			if expect_label == completion.label {
+				flags[i] += 1;
+			}
+		}
+	}
+
+	for flag, i in flags {
+		if flag != 1 {
+			testing.errorf(t, "Expected completion detail %v, but received %v", expect_labels[i], completion_list.items);
+		}
+	}
+
+}
+
 expect_completion_details :: proc(t: ^testing.T, src: ^Source, trigger_character: string, expect_details: []string) {
 	setup(src);