Merge pull request #702 from BradLewis/fix/poly-struct-autocomplete

Improve resolution of poly fields and parameters

7 files changed, 385 insertions, 14 deletions

diff --git a/src/server/analysis.odin b/src/server/analysis.odin
index 9ac88fe..611bf4b 100644
--- a/src/server/analysis.odin
+++ b/src/server/analysis.odin
@@ -193,6 +193,31 @@ set_ast_package_from_symbol_scoped :: proc(ast_context: ^AstContext, symbol: Sym
 	}
 }
 
+set_ast_package_from_node_deferred :: proc(ast_context: ^AstContext, node: ast.Node) {
+	if ast_context.deferred_count <= 0 {
+		return
+	}
+	ast_context.deferred_count -= 1
+	ast_context.current_package = ast_context.deferred_package[ast_context.deferred_count]
+}
+
+@(deferred_in = set_ast_package_from_node_deferred)
+set_ast_package_from_node_scoped :: proc(ast_context: ^AstContext, node: ast.Node) {
+	if ast_context.deferred_count >= DeferredDepth {
+		return
+	}
+
+	ast_context.deferred_package[ast_context.deferred_count] = ast_context.current_package
+	ast_context.deferred_count += 1
+	pkg := get_package_from_node(node)
+
+	if pkg != "" && pkg != "." {
+		ast_context.current_package = 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)
@@ -875,7 +900,7 @@ internal_resolve_type_expression :: proc(ast_context: ^AstContext, node: ^ast.Ex
 		return {}, false
 	}
 
-	set_ast_package_scoped(ast_context)
+	set_ast_package_from_node_scoped(ast_context, node)
 
 	if check_node_recursion(ast_context, node) {
 		return {}, false
@@ -1127,12 +1152,12 @@ resolve_selector_expression :: proc(ast_context: ^AstContext, node: ^ast.Selecto
 				)
 				selector_expr.expr = s.return_types[0].type
 				selector_expr.field = node.field
-
 				return internal_resolve_type_expression(ast_context, selector_expr)
 			}
 		case SymbolStructValue:
 			for name, i in s.names {
 				if node.field != nil && name == node.field.name {
+					set_ast_package_from_node_scoped(ast_context, s.types[i])
 					ast_context.field_name = node.field^
 					symbol, ok := internal_resolve_type_expression(ast_context, s.types[i])
 					symbol.type = .Variable
@@ -1300,6 +1325,7 @@ internal_resolve_type_identifier :: proc(ast_context: ^AstContext, node: ast.Ide
 
 	set_ast_package_scoped(ast_context)
 
+
 	if v, ok := keyword_map[node.name]; ok {
 		//keywords
 		ident := new_type(Ident, node.pos, node.end, ast_context.allocator)
diff --git a/src/server/completion.odin b/src/server/completion.odin
index 12c4205..de8274a 100644
--- a/src/server/completion.odin
+++ b/src/server/completion.odin
@@ -562,8 +562,6 @@ get_selector_completion :: proc(
 				continue
 			}
 
-			set_ast_package_from_symbol_scoped(ast_context, selector)
-
 			if symbol, ok := resolve_type_expression(ast_context, v.types[i]); ok {
 				if expr, ok := position_context.selector.derived.(^ast.Selector_Expr); ok {
 					if expr.op.text == "->" && symbol.type != .Function {
diff --git a/src/server/documentation.odin b/src/server/documentation.odin
index 3be174f..c0853d7 100644
--- a/src/server/documentation.odin
+++ b/src/server/documentation.odin
@@ -124,7 +124,7 @@ get_short_signature :: proc(ast_context: ^AstContext, symbol: Symbol) -> string
 		)
 	case SymbolProcedureValue:
 		sb := strings.builder_make(context.temp_allocator)
-		if symbol.type_pkg != "" {
+		if symbol.type_pkg != "" && symbol.type_name != "" {
 			pkg_name := get_pkg_name(ast_context, symbol.type_pkg)
 			fmt.sbprintf(&sb, "%s%s.%s :: ", pointer_prefix, pkg_name, symbol.type_name)
 		}
@@ -274,7 +274,39 @@ write_struct_hover :: proc(ast_context: ^AstContext, sb: ^strings.Builder, v: Sy
 
 	using_index := -1
 
-	strings.write_string(sb, "struct {\n")
+	strings.write_string(sb, "struct")
+	poly_name_index := 0
+	if v.poly != nil {
+		strings.write_string(sb, "(")
+		for field, i in v.poly.list {
+			write_type := true
+			for name, j in field.names {
+				if poly_name_index < len(v.poly_names) {
+					poly_name := v.poly_names[poly_name_index]
+					if !strings.starts_with(poly_name, "$") {
+						write_type = false
+					}
+					strings.write_string(sb, poly_name)
+				} else {
+					build_string_node(name, sb, false)
+				}
+				if j != len(field.names) - 1 {
+					strings.write_string(sb, ", ")
+				}
+				poly_name_index += 1
+			}
+			if write_type {
+				strings.write_string(sb, ": ")
+				build_string_node(field.type, sb, false)
+			}
+			if i != len(v.poly.list) - 1 {
+				strings.write_string(sb, ", ")
+			}
+		}
+		strings.write_string(sb, ")")
+	}
+	strings.write_string(sb, " {\n")
+
 	for i in 0 ..< len(v.names) {
 		if i < len(v.from_usings) {
 			if index := v.from_usings[i]; index != using_index {
diff --git a/src/server/generics.odin b/src/server/generics.odin
index fcb1397..43e2787 100644
--- a/src/server/generics.odin
+++ b/src/server/generics.odin
@@ -493,7 +493,11 @@ resolve_generic_function_symbol :: proc(
 			ast_context.current_package = ast_context.document_package
 
 			if symbol, ok := resolve_type_expression(ast_context, call_expr.args[i]); ok {
-				symbol_expr := symbol_to_expr(symbol, call_expr.args[i].pos.file, context.temp_allocator)
+				file := strings.trim_prefix(symbol.uri, "file://")
+				if file == "" {
+					file = call_expr.args[i].pos.file
+				}
+				symbol_expr := symbol_to_expr(symbol, file, context.temp_allocator)
 
 				if symbol_expr == nil {
 					return {}, false
@@ -650,9 +654,11 @@ resolve_poly_struct :: proc(ast_context: ^AstContext, b: ^SymbolStructValueBuild
 	i := 0
 
 	poly_map := make(map[string]^ast.Expr, 0, context.temp_allocator)
+	clear(&b.poly_names)
 
 	for param in poly_params.list {
 		for name in param.names {
+			append(&b.poly_names, node_to_string(name))
 			if len(ast_context.call.args) <= i {
 				break
 			}
@@ -664,8 +670,10 @@ resolve_poly_struct :: proc(ast_context: ^AstContext, b: ^SymbolStructValueBuild
 			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]
+					b.poly_names[i] = node_to_string(ast_context.call.args[i])
 				} else if poly, ok := name.derived.(^ast.Poly_Type); ok {
 					if poly.type != nil {
+						b.poly_names[i] = node_to_string(ast_context.call.args[i])
 						poly_map[poly.type.name] = ast_context.call.args[i]
 					}
 				}
@@ -678,11 +686,12 @@ resolve_poly_struct :: proc(ast_context: ^AstContext, b: ^SymbolStructValueBuild
 	}
 
 	Visit_Data :: struct {
-		poly_map:     map[string]^ast.Expr,
+		poly_map:             map[string]^ast.Expr,
 		symbol_value_builder: ^SymbolStructValueBuilder,
-		parent:       ^ast.Node,
-		i:            int,
-		poly_index:   int,
+		parent:               ^ast.Node,
+		parent_proc:	      ^ast.Proc_Type,
+		i:                    int,
+		poly_index:           int,
 	}
 
 	visit :: proc(visitor: ^ast.Visitor, node: ^ast.Node) -> ^ast.Visitor {
@@ -694,6 +703,29 @@ resolve_poly_struct :: proc(ast_context: ^AstContext, b: ^SymbolStructValueBuild
 
 		if ident, ok := node.derived.(^ast.Ident); ok {
 			if expr, ok := data.poly_map[ident.name]; ok {
+				if data.parent_proc != nil {
+					// If the field is a parapoly procedure, we check to see if any of the params or return types
+					// need to be updated
+					if data.parent_proc.params != nil {
+						for &param in data.parent_proc.params.list {
+							if param_ident, ok := param.type.derived.(^ast.Ident); ok {
+								if param_ident.name == ident.name {
+									param.type = expr
+								}
+							}
+						}
+					}
+					if data.parent_proc.results != nil {
+						for &return_value in data.parent_proc.results.list {
+							if return_ident, ok := return_value.type.derived.(^ast.Ident); ok {
+								if return_ident.name == ident.name {
+									return_value.type = expr
+								}
+							}
+						}
+					}
+				}
+
 				if data.parent != nil {
 					#partial switch &v in data.parent.derived {
 					case ^ast.Array_Type:
@@ -703,7 +735,7 @@ resolve_poly_struct :: proc(ast_context: ^AstContext, b: ^SymbolStructValueBuild
 					case ^ast.Pointer_Type:
 						v.elem = expr
 					}
-				} else {
+				} else if data.parent_proc == nil {
 					data.symbol_value_builder.types[data.i] = expr
 					data.poly_index += 1
 				}
@@ -711,8 +743,10 @@ resolve_poly_struct :: proc(ast_context: ^AstContext, b: ^SymbolStructValueBuild
 		}
 
 		#partial switch v in node.derived {
-		case ^ast.Array_Type, ^ast.Dynamic_Array_Type, ^ast.Selector_Expr, ^ast.Pointer_Type:
+		case ^ast.Array_Type, ^ast.Dynamic_Array_Type, ^ast.Selector_Expr, ^ast.Pointer_Type: 
 			data.parent = node
+		case ^ast.Proc_Type:
+			data.parent_proc = v
 		}
 
 		return visitor
diff --git a/src/server/symbol.odin b/src/server/symbol.odin
index b795a96..2d99d22 100644
--- a/src/server/symbol.odin
+++ b/src/server/symbol.odin
@@ -32,6 +32,7 @@ SymbolStructValue :: struct {
 	from_usings:       []int,
 	unexpanded_usings: []int,
 	poly:              ^ast.Field_List,
+	poly_names:        []string, // The resolved names for the poly fields
 	args:              []^ast.Expr, //The arguments in the call expression for poly
 	docs:              []^ast.Comment_Group,
 	comments:          []^ast.Comment_Group,
@@ -208,6 +209,7 @@ SymbolStructValueBuilder :: struct {
 	from_usings:       [dynamic]int,
 	unexpanded_usings: [dynamic]int,
 	poly:              ^ast.Field_List,
+	poly_names:        [dynamic]string,
 }
 
 symbol_struct_value_builder_make_none :: proc(allocator := context.allocator) -> SymbolStructValueBuilder {
@@ -221,6 +223,7 @@ symbol_struct_value_builder_make_none :: proc(allocator := context.allocator) ->
 		usings = make(map[int]struct{}, allocator),
 		from_usings = make([dynamic]int, allocator),
 		unexpanded_usings = make([dynamic]int, allocator),
+		poly_names = make([dynamic]string, allocator),
 	}
 }
 
@@ -239,6 +242,7 @@ symbol_struct_value_builder_make_symbol :: proc(
 		usings = make(map[int]struct{}, allocator),
 		from_usings = make([dynamic]int, allocator),
 		unexpanded_usings = make([dynamic]int, allocator),
+		poly_names = make([dynamic]string, allocator),
 	}
 }
 
@@ -258,6 +262,7 @@ symbol_struct_value_builder_make_symbol_symbol_struct_value :: proc(
 		usings = v.usings,
 		from_usings = slice.to_dynamic(v.from_usings, allocator),
 		unexpanded_usings = slice.to_dynamic(v.unexpanded_usings, allocator),
+		poly_names = slice.to_dynamic(v.poly_names, allocator),
 	}
 }
 
@@ -285,6 +290,7 @@ to_symbol_struct_value :: proc(b: SymbolStructValueBuilder) -> SymbolStructValue
 		from_usings = b.from_usings[:],
 		unexpanded_usings = b.unexpanded_usings[:],
 		poly = b.poly,
+		poly_names = b.poly_names[:],
 	}
 }
 
diff --git a/tests/completions_test.odin b/tests/completions_test.odin
index 2a3d7c2..ae3b023 100644
--- a/tests/completions_test.odin
+++ b/tests/completions_test.odin
@@ -1586,7 +1586,7 @@ ast_maybe_index_completion :: proc(t: ^testing.T) {
 		packages = packages[:],
 	}
 
-	test.expect_completion_labels(t, &source, ".", {"(my_package.int)"})
+	test.expect_completion_labels(t, &source, ".", {"(int)"})
 }
 
 @(test)
@@ -3362,3 +3362,122 @@ ast_completion_inline_using :: proc(t: ^testing.T) {
 
 	test.expect_completion_details(t, &source, "", {"Foo.a: int", "Foo.b: int"})
 }
+
+@(test)
+ast_completion_poly_struct_another_package :: proc(t: ^testing.T) {
+	packages := make([dynamic]test.Package, context.temp_allocator)
+
+	append(
+		&packages,
+		test.Package {
+			pkg = "my_package",
+			source = `package my_package
+		Runner :: struct($TState: typeid) {
+			state: TState, // state
+		}
+		`,
+		},
+	)
+	source := test.Source {
+		main = `package test
+
+		import "my_package"
+
+		app: my_package.Runner(State) = {
+			state = {score = 55},
+		}
+
+		main :: proc() {
+			app.{*}
+		}
+
+		State :: struct {
+			score: int,
+		}
+		`,
+		packages = packages[:],
+	}
+
+	test.expect_completion_details(t, &source, "", {"Runner.state: test.State // state"})
+}
+
+@(test)
+ast_completion_poly_struct_another_package_field :: proc(t: ^testing.T) {
+	packages := make([dynamic]test.Package, context.temp_allocator)
+
+	append(
+		&packages,
+		test.Package {
+			pkg = "my_package",
+			source = `package my_package
+		Runner :: struct($TState: typeid) {
+			state: TState, // state
+		}
+		`,
+		},
+	)
+	source := test.Source {
+		main = `package test
+
+		import "my_package"
+
+		app: my_package.Runner(State) = {
+			state = {score = 55},
+		}
+
+		main :: proc() {
+			app.state.{*}
+		}
+
+		State :: struct {
+			score: int,
+		}
+		`,
+		packages = packages[:],
+	}
+
+	test.expect_completion_details(t, &source, "", {"State.score: int"})
+}
+
+@(test)
+ast_completion_poly_proc_mixed_packages :: proc(t: ^testing.T) {
+	packages := make([dynamic]test.Package, context.temp_allocator)
+
+	append(
+		&packages,
+		test.Package {
+			pkg = "foo_package",
+			source = `package foo_package
+		foo :: proc(t: $T) -> T {
+			return t
+		}
+		`,
+		},
+		test.Package {
+			pkg = "bar_package",
+			source = `package bar_package
+			Bar :: struct {
+				bar: int,
+			}
+		`,
+		},
+	)
+
+	source := test.Source {
+		main = `package test
+
+		import "foo_package"
+		import "bar_package"
+
+		main :: proc() {
+			b := bar_package.Bar{}
+			f := foo_package.foo(b)
+			f.{*}
+		}
+	}
+		`,
+		packages = packages[:],
+	}
+
+	test.expect_completion_details(t, &source, "", {"Bar.bar: int"})
+}
diff --git a/tests/hover_test.odin b/tests/hover_test.odin
index 1511af6..a4a60e7 100644
--- a/tests/hover_test.odin
+++ b/tests/hover_test.odin
@@ -1595,6 +1595,7 @@ ast_hover_distinct_definition_external_package :: proc(t: ^testing.T) {
 
 	test.expect_hover(t, &source, "my_package.A: distinct u64")
 }
+
 @(test)
 ast_hover_poly_type :: proc(t: ^testing.T) {
 	source := test.Source {
@@ -1723,6 +1724,161 @@ ast_hover_poly_type_external_package_with_external_type :: proc(t: ^testing.T) {
 
 	test.expect_hover(t, &source, "test.foo: small_array.Foo :: struct {}")
 }
+
+@(test)
+ast_hover_struct_poly_type :: proc(t: ^testing.T) {
+	source := test.Source {
+		main = `package test
+		F{*}oo :: struct($T: typeid) {
+			foo: T,
+		}
+		`,
+	}
+
+	test.expect_hover(t, &source, "test.Foo: struct($T: typeid) {\n\tfoo: T,\n}")
+}
+
+@(test)
+ast_hover_poly_proc_mixed_packages :: proc(t: ^testing.T) {
+	packages := make([dynamic]test.Package, context.temp_allocator)
+
+	append(
+		&packages,
+		test.Package {
+			pkg = "foo_package",
+			source = `package foo_package
+		foo :: proc(t: $T) -> T {
+			return t
+		}
+		`,
+		},
+		test.Package {
+			pkg = "bar_package",
+			source = `package bar_package
+			Bar :: struct {
+				bar: int,
+			}
+		`,
+		},
+	)
+
+	source := test.Source {
+		main = `package test
+
+		import "foo_package"
+		import "bar_package"
+
+		main :: proc() {
+			b := bar_package.Bar{}
+			f{*} := foo_package.foo(b)
+		}
+	}
+		`,
+		packages = packages[:],
+	}
+
+	test.expect_hover(t, &source, "test.f: bar_package.Bar :: struct {\n\tbar: int,\n}")
+}
+
+@(test)
+ast_hover_poly_struct_proc_field :: proc(t: ^testing.T) {
+	packages := make([dynamic]test.Package, context.temp_allocator)
+
+	append(
+		&packages,
+		test.Package {
+			pkg = "my_package",
+			source = `package my_package
+		Foo :: struct($T: typeid) {
+			foo: proc(t: ^T) -> T,
+		}
+		`,
+		},
+	)
+
+	source := test.Source {
+		main = `package test
+
+		import "my_package"
+
+		Bar :: struct{}
+
+		my_proc :: proc(b: ^Bar) -> Bar {}
+
+		main :: proc() {
+			foo: my_package.Foo(Bar) = {
+				foo = my_proc,
+			}
+			foo.f{*}oo()
+
+		}
+	}
+		`,
+		packages = packages[:],
+	}
+
+	test.expect_hover(t, &source, "Foo.foo: proc(t: ^Bar) -> Bar")
+}
+
+@(test)
+ast_hover_poly_struct_poly_proc_fields :: proc(t: ^testing.T) {
+	source := test.Source {
+		main = `package test
+
+		F{*}oo :: struct($S: typeid, $T: typeid) {
+			my_proc1: proc(s: S) -> ^S,
+			my_proc2: proc(t: ^T) -> T,
+			my_proc3: proc(s: ^S, t: T) -> T,
+			my_proc4: proc() -> T,
+			my_proc5: proc(t: T),
+			foo1: T,
+			foo2: ^S,
+		}
+	}
+		`,
+	}
+
+	test.expect_hover(t, &source, "test.Foo: struct($S: typeid, $T: typeid) {\n\tmy_proc1: proc(s: S) -> ^S,\n\tmy_proc2: proc(t: ^T) -> T,\n\tmy_proc3: proc(s: ^S,t: T) -> T,\n\tmy_proc4: proc() -> T,\n\tmy_proc5: proc(t: T),\n\tfoo1:     T,\n\tfoo2:     ^S,\n}")
+}
+
+@(test)
+ast_hover_poly_struct_poly_proc_fields_resolved :: proc(t: ^testing.T) {
+	packages := make([dynamic]test.Package, context.temp_allocator)
+
+	append(
+		&packages,
+		test.Package {
+			pkg = "my_package",
+			source = `package my_package
+			Bazz :: struct{}
+		`,
+		},
+	)
+
+	source := test.Source {
+		main = `package test
+		import "my_package"
+
+		Foo :: struct($S: typeid, $T: typeid) {
+			my_proc1: proc(s: S) -> ^S,
+			my_proc2: proc(t: T) -> T,
+			my_proc3: proc(s: ^S, t: T) -> T,
+			foo1: T,
+			foo2: ^S,
+		}
+
+		Bar :: struct{}
+
+		main :: proc() {
+			foo := Fo{*}o(Bar, my_package.Bazz){}
+		}
+	}
+		`,
+		packages = packages[:],
+	}
+
+	test.expect_hover(t, &source, "test.Foo: struct(Bar, my_package.Bazz) {\n\tmy_proc1: proc(s: Bar) -> ^Bar,\n\tmy_proc2: proc(t: my_package.Bazz) -> my_package.Bazz,\n\tmy_proc3: proc(s: ^my_package.Bazz,t: my_package.Bazz) -> my_package.Bazz,\n\tfoo1:     my_package.Bazz,\n\tfoo2:     ^Bar,\n}")
+}
 /*
 
 Waiting for odin fix