Store the identifier rather than the specialization for poly types when the types match

2 files changed, 35 insertions, 1 deletions

diff --git a/src/server/generics.odin b/src/server/generics.odin
index 33543cf..cec9926 100644
--- a/src/server/generics.odin
+++ b/src/server/generics.odin
@@ -61,7 +61,20 @@ resolve_poly :: proc(
 		return true
 	} else if type != nil {
 		if ident, ok := unwrap_ident(type); ok {
-			save_poly_map(ident, specialization, poly_map)
+			call_node_id := reflect.union_variant_typeid(call_node.derived)
+			specialization_id := reflect.union_variant_typeid(specialization.derived)
+			if call_node_id == specialization_id {
+				// if the specialization type matches the type of the parameter passed to the proc
+				// we store that rather than the specialization so we can follow it correctly
+				// for things like `textDocument/typeDefinition`
+				save_poly_map(
+					ident,
+					make_ident_ast(ast_context, call_node.pos, call_node.end, call_symbol.name),
+					poly_map,
+				)
+			} else {
+				save_poly_map(ident, specialization, poly_map)
+			}
 		}
 	}
 
diff --git a/tests/type_definition_test.odin b/tests/type_definition_test.odin
index 3e818bf..e0cfae9 100644
--- a/tests/type_definition_test.odin
+++ b/tests/type_definition_test.odin
@@ -1010,3 +1010,24 @@ ast_type_definition_variable_in_comp_lit :: proc (t: ^testing.T) {
 
 	test.expect_type_definition_locations(t, &source, locations[:])
 }
+
+@(test)
+ast_type_definition_polymorphic_type_with_specialization :: proc (t: ^testing.T) {
+	source := test.Source {
+		main     = `package test
+		Vec :: [2]f32
+		foo :: proc (a: $T/[$N]$E) -> T {return a}
+
+		main :: proc () {
+			a: Vec
+			f{*} := foo(a)
+		}
+		`,
+	}
+
+	locations := []common.Location {
+		{range = {start = {line = 1, character = 2}, end = {line = 1, character = 5}}},
+	}
+
+	test.expect_type_definition_locations(t, &source, locations[:])
+}