Fix issues resolving poly proc fields and improve hover information of poly structs

3 files changed, 68 insertions, 12 deletions

diff --git a/src/server/documentation.odin b/src/server/documentation.odin
index 9e23b52..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)
 		}
@@ -275,18 +275,38 @@ write_struct_hover :: proc(ast_context: ^AstContext, sb: ^strings.Builder, v: Sy
 	using_index := -1
 
 	strings.write_string(sb, "struct")
+	poly_name_index := 0
 	if v.poly != nil {
 		strings.write_string(sb, "(")
-		for field in v.poly.list {
-			for name in field.names {
-				build_string_node(name, sb, false)
+		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, ": ")
-			build_string_node(field.type, sb, false)
 		}
 		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 c34a0ce..43e2787 100644
--- a/src/server/generics.odin
+++ b/src/server/generics.odin
@@ -654,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
 			}
@@ -668,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]
 					}
 				}
@@ -682,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 {
@@ -698,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:
@@ -707,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
 				}
@@ -715,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[:],
 	}
 }