Refractor ast_context.current_package now being called through defered functions

5 files changed, 168 insertions, 76 deletions

diff --git a/src/server/analysis.odin b/src/server/analysis.odin
index ce84579..8562d66 100644
--- a/src/server/analysis.odin
+++ b/src/server/analysis.odin
@@ -86,6 +86,7 @@ AstContext :: struct {
 	imports:          []Package, //imports for the current document
 	current_package:  string,
 	document_package: string,
+	saved_package:    string,
 	use_locals:       bool,
 	local_id:         int,
 	call:             ^ast.Call_Expr, //used to determine the types for generics and the correct function for overloaded functions
@@ -128,6 +129,46 @@ make_ast_context :: proc(
 	return ast_context
 }
 
+set_ast_package_deferred :: proc(ast_context: ^AstContext, pkg: string) {
+	ast_context.current_package = ast_context.saved_package
+}
+
+@(deferred_in = set_ast_package_deferred)
+set_ast_package_set_scoped :: proc(ast_context: ^AstContext, pkg: string) {
+	ast_context.saved_package = ast_context.current_package
+	ast_context.current_package = pkg
+}
+
+set_ast_package_none_deferred :: proc(ast_context: ^AstContext) {
+	ast_context.current_package = ast_context.saved_package
+}
+
+@(deferred_in = set_ast_package_none_deferred)
+set_ast_package_scoped :: proc(ast_context: ^AstContext) {
+	ast_context.saved_package = ast_context.current_package
+}
+
+set_ast_package_from_symbol_deferred :: proc(
+	ast_context: ^AstContext,
+	symbol: Symbol,
+) {
+	ast_context.current_package = ast_context.saved_package
+}
+
+@(deferred_in = set_ast_package_from_symbol_deferred)
+set_ast_package_from_symbol_scoped :: proc(
+	ast_context: ^AstContext,
+	symbol: Symbol,
+) {
+	ast_context.saved_package = ast_context.current_package
+
+	if symbol.pkg != "" {
+		ast_context.current_package = symbol.pkg
+	} else {
+		ast_context.current_package = ast_context.document_package
+	}
+}
+
 reset_ast_context :: proc(ast_context: ^AstContext) {
 	ast_context.use_locals = true
 	clear(&ast_context.recursion_map)
@@ -154,10 +195,7 @@ resolve_type_comp_literal :: proc(
 	}
 
 
-	prev_package := ast_context.current_package
-	ast_context.current_package = current_symbol.pkg
-
-	defer ast_context.current_package = prev_package
+	set_ast_package_set_scoped(ast_context, current_symbol.pkg)
 
 	for elem, element_index in current_comp_lit.elems {
 		if !position_in_node(elem, position_context.position) {
@@ -737,7 +775,7 @@ resolve_basic_directive :: proc(
 			ast_context.allocator,
 		)
 		ident.name = "Source_Code_Location"
-		ast_context.current_package = ast_context.document_package
+		set_ast_package_set_scoped(ast_context, ast_context.document_package)
 		return internal_resolve_type_identifier(ast_context, ident^)
 	}
 
@@ -781,11 +819,7 @@ internal_resolve_type_expression :: proc(
 		return {}, false
 	}
 
-	saved_package := ast_context.current_package
-
-	defer {
-		ast_context.current_package = saved_package
-	}
+	set_ast_package_scoped(ast_context)
 
 	if check_node_recursion(ast_context, node) {
 		return {}, false
@@ -979,7 +1013,7 @@ internal_resolve_type_expression :: proc(
 			return {}, false
 		}
 
-		ast_context.current_package = indexed.pkg
+		set_ast_package_set_scoped(ast_context, indexed.pkg)
 
 		symbol: Symbol
 
@@ -1027,11 +1061,7 @@ internal_resolve_type_expression :: proc(
 		); ok {
 			ast_context.use_locals = false
 
-			if selector.pkg != "" {
-				ast_context.current_package = selector.pkg
-			} else {
-				ast_context.current_package = ast_context.document_package
-			}
+			set_ast_package_from_symbol_scoped(ast_context, selector)
 
 			#partial switch s in selector.value {
 			case SymbolProcedureValue:
@@ -1052,11 +1082,7 @@ internal_resolve_type_expression :: proc(
 		); ok {
 			ast_context.use_locals = false
 
-			if selector.pkg != "" {
-				ast_context.current_package = selector.pkg
-			} else {
-				ast_context.current_package = ast_context.document_package
-			}
+			set_ast_package_from_symbol_scoped(ast_context, selector)
 
 			#partial switch s in selector.value {
 			case SymbolFixedArrayValue:
@@ -1079,12 +1105,8 @@ internal_resolve_type_expression :: proc(
 				}
 
 				if components_count == 1 {
-					if selector.pkg != "" {
-						ast_context.current_package = selector.pkg
-					} else {
-						ast_context.current_package =
-							ast_context.document_package
-					}
+					set_ast_package_from_symbol_scoped(ast_context, selector)
+
 					symbol, ok := internal_resolve_type_expression(
 						ast_context,
 						s.expr,
@@ -1288,11 +1310,7 @@ internal_resolve_type_identifier :: proc(
 		return {}, false
 	}
 
-	saved_package := ast_context.current_package
-
-	defer {
-		ast_context.current_package = saved_package
-	}
+	set_ast_package_scoped(ast_context)
 
 	if v, ok := common.keyword_map[node.name]; ok {
 		//keywords
@@ -1679,7 +1697,7 @@ expand_struct_usings :: proc(
 	ranges := slice.to_dynamic(value.ranges, ast_context.allocator)
 
 	for k, v in value.usings {
-		ast_context.current_package = symbol.pkg
+		set_ast_package_set_scoped(ast_context, symbol.pkg)
 
 		field_expr: ^ast.Expr
 
@@ -1809,8 +1827,7 @@ resolve_comp_literal :: proc(
 		}
 	}
 
-	old_package := ast_context.current_package
-	ast_context.current_package = symbol.pkg
+	set_ast_package_set_scoped(ast_context, symbol.pkg)
 
 	symbol, _ = resolve_type_comp_literal(
 		ast_context,
@@ -1819,8 +1836,6 @@ resolve_comp_literal :: proc(
 		position_context.parent_comp_lit,
 	) or_return
 
-	ast_context.current_package = old_package
-
 	return symbol, true
 }
 
@@ -1964,7 +1979,7 @@ resolve_implicit_selector :: proc(
 				position_context.parent_comp_lit,
 			); ok {
 				if s, ok := comp_symbol.value.(SymbolStructValue); ok {
-					ast_context.current_package = comp_symbol.pkg
+					set_ast_package_set_scoped(ast_context, comp_symbol.pkg)
 
 					//We can either have the final 
 					elem_index := -1
@@ -1993,7 +2008,7 @@ resolve_implicit_selector :: proc(
 					return resolve_type_expression(ast_context, type)
 				} else if s, ok := comp_symbol.value.(SymbolBitFieldValue);
 				   ok {
-					ast_context.current_package = comp_symbol.pkg
+					set_ast_package_set_scoped(ast_context, comp_symbol.pkg)
 
 					//We can either have the final 
 					elem_index := -1
@@ -2140,7 +2155,8 @@ resolve_unresolved_symbol :: proc(
 	case SymbolBitSetValue:
 		symbol.type = .Enum
 	case SymbolGenericValue:
-		ast_context.current_package = symbol.pkg
+		set_ast_package_set_scoped(ast_context, symbol.pkg)
+
 		if ret, ok := resolve_type_expression(ast_context, v.expr); ok {
 			symbol.type = ret.type
 			symbol.signature = ret.signature
@@ -2203,7 +2219,7 @@ resolve_location_comp_lit_field :: proc(
 ) {
 	reset_ast_context(ast_context)
 
-	ast_context.current_package = ast_context.document_package
+	set_ast_package_set_scoped(ast_context, ast_context.document_package)
 
 	symbol = resolve_comp_literal(ast_context, position_context) or_return
 
@@ -2236,7 +2252,7 @@ resolve_location_implicit_selector :: proc(
 ) {
 	reset_ast_context(ast_context)
 
-	ast_context.current_package = ast_context.document_package
+	set_ast_package_set_scoped(ast_context, ast_context.document_package)
 
 	symbol = resolve_implicit_selector(
 		ast_context,
@@ -2264,7 +2280,8 @@ resolve_location_selector :: proc(
 	ok: bool,
 ) {
 	reset_ast_context(ast_context)
-	ast_context.current_package = ast_context.document_package
+
+	set_ast_package_set_scoped(ast_context, ast_context.document_package)
 
 	symbol = resolve_type_expression(ast_context, selector.expr) or_return
 
@@ -3166,7 +3183,8 @@ get_locals_stmt :: proc(
 	save_assign := false,
 ) {
 	reset_ast_context(ast_context)
-	ast_context.current_package = ast_context.document_package
+
+	set_ast_package_set_scoped(ast_context, ast_context.document_package)
 
 	using ast
 
@@ -3817,7 +3835,7 @@ resolve_entire_file :: proc(
 
 	get_globals(document.ast, &ast_context)
 
-	ast_context.current_package = ast_context.document_package
+	set_ast_package_set_scoped(&ast_context, ast_context.document_package)
 
 	symbols := make(map[uintptr]SymbolAndNode, 10000, allocator)
 
@@ -4214,6 +4232,7 @@ unwrap_union :: proc(
 	bool,
 ) {
 	if union_symbol, ok := resolve_type_expression(ast_context, node); ok {
+		//TODO: This current package is sus, it probably shouldn't be there.
 		ast_context.current_package = union_symbol.pkg
 		if union_value, ok := union_symbol.value.(SymbolUnionValue); ok {
 			return union_value, true
diff --git a/src/server/completion.odin b/src/server/completion.odin
index b8fc6ba..948fee5 100644
--- a/src/server/completion.odin
+++ b/src/server/completion.odin
@@ -237,7 +237,6 @@ get_comp_lit_completion :: proc(
 
 
 	if symbol, ok := resolve_comp_literal(ast_context, position_context); ok {
-		//ast_context.current_package = comp_symbol.pkg
 		#partial switch v in symbol.value {
 		case SymbolStructValue:
 			for name, i in v.names {
@@ -245,7 +244,7 @@ get_comp_lit_completion :: proc(
 					continue
 				}
 
-				ast_context.current_package = symbol.pkg
+				set_ast_package_set_scoped(ast_context, symbol.pkg)
 
 				if resolved, ok := resolve_type_expression(
 					ast_context,
@@ -279,7 +278,7 @@ get_comp_lit_completion :: proc(
 					continue
 				}
 
-				ast_context.current_package = symbol.pkg
+				set_ast_package_set_scoped(ast_context, symbol.pkg)
 
 				if resolved, ok := resolve_type_expression(
 					ast_context,
@@ -343,11 +342,7 @@ get_selector_completion :: proc(
 		return
 	}
 
-	if selector.pkg != "" {
-		ast_context.current_package = selector.pkg
-	} else {
-		ast_context.current_package = ast_context.document_package
-	}
+	set_ast_package_from_symbol_scoped(ast_context, selector)
 
 	field: string
 
@@ -601,11 +596,7 @@ get_selector_completion :: proc(
 				continue
 			}
 
-			if selector.pkg != "" {
-				ast_context.current_package = selector.pkg
-			} else {
-				ast_context.current_package = ast_context.document_package
-			}
+			set_ast_package_from_symbol_scoped(ast_context, selector)
 
 			if symbol, ok := resolve_type_expression(ast_context, v.types[i]);
 			   ok {
@@ -668,11 +659,7 @@ get_selector_completion :: proc(
 				continue
 			}
 
-			if selector.pkg != "" {
-				ast_context.current_package = selector.pkg
-			} else {
-				ast_context.current_package = ast_context.document_package
-			}
+			set_ast_package_from_symbol_scoped(ast_context, selector)
 
 			if symbol, ok := resolve_type_expression(ast_context, v.types[i]);
 			   ok {
@@ -792,11 +779,7 @@ get_implicit_completion :: proc(
 
 	reset_ast_context(ast_context)
 
-	if selector.pkg != "" {
-		ast_context.current_package = selector.pkg
-	} else {
-		ast_context.current_package = ast_context.document_package
-	}
+	set_ast_package_from_symbol_scoped(ast_context, selector)
 
 	//value decl infer a : My_Enum = .*
 	if position_context.value_decl != nil &&
@@ -891,7 +874,7 @@ get_implicit_completion :: proc(
 			ast_context,
 			position_context.assign.lhs[0],
 		); ok {
-			ast_context.current_package = symbol.pkg
+			set_ast_package_set_scoped(ast_context, symbol.pkg)
 			if value, ok := unwrap_bitset(ast_context, symbol); ok {
 				for name in value.names {
 
@@ -920,7 +903,7 @@ get_implicit_completion :: proc(
 			ast_context,
 			position_context.parent_binary,
 		); ok {
-			ast_context.current_package = symbol.pkg
+			set_ast_package_set_scoped(ast_context, symbol.pkg)
 			if value, ok := unwrap_bitset(ast_context, symbol); ok {
 				for name in value.names {
 					item := CompletionItem {
@@ -965,7 +948,10 @@ get_implicit_completion :: proc(
 					position_context.parent_comp_lit,
 				); ok {
 					if s, ok := comp_symbol.value.(SymbolStructValue); ok {
-						ast_context.current_package = comp_symbol.pkg
+						set_ast_package_set_scoped(
+							ast_context,
+							comp_symbol.pkg,
+						)
 
 						//We can either have the final 
 						elem_index := -1
@@ -1014,7 +1000,10 @@ get_implicit_completion :: proc(
 							ast_context,
 							type,
 						); ok {
-							ast_context.current_package = bitset_symbol.pkg
+							set_ast_package_set_scoped(
+								ast_context,
+								bitset_symbol.pkg,
+							)
 
 							if value, ok := unwrap_bitset(
 								ast_context,
@@ -1198,7 +1187,7 @@ get_implicit_completion :: proc(
 			)
 			if symbol, ok := resolve_type_expression(ast_context, call.expr);
 			   ok && parameter_ok {
-				ast_context.current_package = symbol.pkg
+				set_ast_package_set_scoped(ast_context, symbol.pkg)
 
 				//Selector call expression always set the first argument to be the type of struct called, so increment it.
 				if position_context.selector_expr != nil {
@@ -1239,7 +1228,10 @@ get_implicit_completion :: proc(
 							ast_context,
 							proc_value.arg_types[parameter_index].type,
 						); ok {
-							ast_context.current_package = bitset_symbol.pkg
+							set_ast_package_set_scoped(
+								ast_context,
+								bitset_symbol.pkg,
+							)
 							if enum_value, ok := unwrap_bitset(
 								ast_context,
 								bitset_symbol,
@@ -1360,6 +1352,7 @@ get_identifier_completion :: proc(
 		}
 
 		reset_ast_context(ast_context)
+
 		ast_context.current_package = ast_context.document_package
 
 		ident := new_type(
diff --git a/src/server/generics.odin b/src/server/generics.odin
index a2393d0..9c44cf5 100644
--- a/src/server/generics.odin
+++ b/src/server/generics.odin
@@ -254,6 +254,24 @@ resolve_poly :: proc(
 				return true
 			}
 		}
+	case ^ast.Comp_Lit:
+		if comp_lit, ok := call_node.derived.(^ast.Comp_Lit); ok {
+			if poly_type, ok := p.type.derived.(^ast.Poly_Type); ok {
+				if ident, ok := unwrap_ident(poly_type.type); ok {
+					poly_map[ident.name] = comp_lit.type
+				}
+
+				if poly_type.specialization != nil {
+					return resolve_poly(
+						ast_context,
+						comp_lit.type,
+						call_symbol,
+						p.type,
+						poly_map,
+					)
+				}
+			}
+		}
 	case ^ast.Ident:
 		if n, ok := call_node.derived.(^ast.Ident); ok {
 			return true
@@ -348,6 +366,12 @@ find_and_replace_poly_type :: proc(
 		poly_map := cast(^map[string]^ast.Expr)visitor.data
 
 		#partial switch v in node.derived {
+		case ^ast.Comp_Lit:
+			if expr, ok := is_in_poly_map(v.type, poly_map); ok {
+				v.type = expr
+				v.pos.file = expr.pos.file
+				v.end.file = expr.end.file
+			}
 		case ^ast.Matrix_Type:
 			if expr, ok := is_in_poly_map(v.elem, poly_map); ok {
 				v.elem = expr
@@ -484,6 +508,8 @@ resolve_generic_function_symbol :: proc(
 
 			reset_ast_context(ast_context)
 
+			ast_context.current_package = ast_context.document_package
+
 			if symbol, ok := resolve_type_expression(
 				ast_context,
 				call_expr.args[i],
diff --git a/src/server/hover.odin b/src/server/hover.odin
index 566b41c..1dbd8a6 100644
--- a/src/server/hover.odin
+++ b/src/server/hover.odin
@@ -165,7 +165,8 @@ get_hover_information :: proc(
 							}
 						}
 					}
-				} else if v, ok := comp_symbol.value.(SymbolBitFieldValue); ok {
+				} else if v, ok := comp_symbol.value.(SymbolBitFieldValue);
+				   ok {
 					for name, i in v.names {
 						if name == field.name {
 							if symbol, ok := resolve_type_expression(
@@ -257,7 +258,7 @@ get_hover_information :: proc(
 				return {}, false, false
 			}
 
-			ast_context.current_package = selector.pkg
+			set_ast_package_set_scoped(&ast_context, selector.pkg)
 
 			if selector, ok = resolve_type_expression(
 				&ast_context,
diff --git a/tests/completions_test.odin b/tests/completions_test.odin
index fc0cc5c..4084c16 100644
--- a/tests/completions_test.odin
+++ b/tests/completions_test.odin
@@ -2766,3 +2766,56 @@ ast_generics_function_with_struct_diff_pkg :: proc(t: ^testing.T) {
 		},
 	)
 }
+
+
+@(test)
+ast_generics_function_with_comp_lit_struct :: proc(t: ^testing.T) {
+	packages := make([dynamic]test.Package)
+
+	append(
+		&packages,
+		test.Package {
+			pkg = "my_package",
+			source = `package my_package
+			DummyFunction :: proc($T: typeid, value: T) -> T
+			{
+				return value;
+			}	
+
+			OtherStruct :: struct 
+			{
+				val1, val2, val3: int,
+			}
+		`,
+		},
+	)
+
+	source := test.Source {
+		main     = `package main
+		import "my_package"
+		
+		CoolStruct :: struct
+		{
+			val1, val2, val3: int,
+		}
+
+		main :: proc()
+		{
+			newValue := my_package.DummyFunction(CoolStruct, CoolStruct{})
+    		newValue.{*}
+		}
+		`,
+		packages = packages[:],
+	}
+
+	test.expect_completion_details(
+		t,
+		&source,
+		".",
+		{
+			"CoolStruct.val1: int",
+			"CoolStruct.val2: int",
+			"CoolStruct.val3: int",
+		},
+	)
+}