Improve finding which element to rename/find references to in struct field types

5 files changed, 173 insertions, 2 deletions

diff --git a/src/server/analysis.odin b/src/server/analysis.odin
index c4de6dd..bc0b215 100644
--- a/src/server/analysis.odin
+++ b/src/server/analysis.odin
@@ -4580,6 +4580,43 @@ fallback_position_context_signature :: proc(
 	//log.error(string(position_context.file.src[begin_offset:end_offset]));
 }
 
+// Used to find which sub-expr is desired by the position.
+// Eg. for map[Key]Value, do we want 'map', 'Key' or 'Value'
+get_desired_expr :: proc(node: ^ast.Expr, position: common.AbsolutePosition) -> ^ast.Expr{
+	#partial switch n in node.derived {
+	case ^ast.Array_Type:
+		if position_in_node(n.tag, position) {
+			return n.tag
+		}
+		if position_in_node(n.elem, position) {
+			return n.elem
+		}
+		if position_in_node(n.len, position) {
+			return n.len
+		}
+	case ^ast.Map_Type:
+		if position_in_node(n.key, position) {
+			return n.key
+		}
+		if position_in_node(n.value, position) {
+			return n.key
+		}
+	case ^ast.Dynamic_Array_Type:
+		if position_in_node(n.tag, position) {
+			return n.tag
+		}
+		if position_in_node(n.elem, position) {
+			return n.elem
+		}
+	case ^ast.Bit_Set_Type:
+		if position_in_node(n.elem, position) {
+			return n.elem
+		}
+	}
+
+	return node
+}
+
 /*
 	All these fallback functions are not perfect and should be fixed. A lot of weird use of the odin tokenizer and parser.
 */
diff --git a/src/server/references.odin b/src/server/references.odin
index f107759..7e791b8 100644
--- a/src/server/references.odin
+++ b/src/server/references.odin
@@ -65,7 +65,8 @@ prepare_references :: proc(
 				}
 			}
 			if position_in_node(field.type, position_context.position) {
-				symbol, ok = resolve_location_type_expression(ast_context, field.type)
+				node := get_desired_expr(field.type, position_context.position)
+				symbol, ok = resolve_location_type_expression(ast_context, node)
 				if !ok {
 					return
 				}
diff --git a/src/server/rename.odin b/src/server/rename.odin
index d97df2b..0b6bdb1 100644
--- a/src/server/rename.odin
+++ b/src/server/rename.odin
@@ -139,7 +139,8 @@ prepare_rename :: proc(
 				}
 			}
 			if position_in_node(field.type, position_context.position) {
-				symbol, ok = get_struct_field_type_position(ast_context, position_context, field.type)
+				node := get_desired_expr(field.type, position_context.position)
+				symbol, ok = get_struct_field_type_position(ast_context, position_context, node)
 				if !ok {
 					return
 				}
diff --git a/tests/references_test.odin b/tests/references_test.odin
index 5e0795f..5a4836a 100644
--- a/tests/references_test.odin
+++ b/tests/references_test.odin
@@ -949,3 +949,55 @@ ast_reference_enum_nested_with_switch :: proc(t: ^testing.T) {
 
 	test.expect_reference_locations(t, &source, locations[:])
 }
+
+@(test)
+ast_reference_struct_field_enumerated_array :: proc(t: ^testing.T) {
+	source := test.Source {
+		main = `package test
+
+		Foo :: enum {
+			A,
+			B,
+		}
+
+		Bar :: struct {
+			foos: [F{*}oo]Bazz
+		}
+
+		Bazz :: struct {}
+		`,
+	}
+
+	locations := []common.Location {
+		{range = {start = {line = 2, character = 2}, end = {line = 2, character = 5}}},
+		{range = {start = {line = 8, character = 10}, end = {line = 8, character = 13}}},
+	}
+
+	test.expect_reference_locations(t, &source, locations[:])
+}
+
+@(test)
+ast_reference_struct_field_map :: proc(t: ^testing.T) {
+	source := test.Source {
+		main = `package test
+
+		Foo :: enum {
+			A,
+			B,
+		}
+
+		Bar :: struct {
+			foos: map[F{*}oo]Bazz
+		}
+
+		Bazz :: struct {}
+		`,
+	}
+
+	locations := []common.Location {
+		{range = {start = {line = 2, character = 2}, end = {line = 2, character = 5}}},
+		{range = {start = {line = 8, character = 13}, end = {line = 8, character = 16}}},
+	}
+
+	test.expect_reference_locations(t, &source, locations[:])
+}
diff --git a/tests/rename_test.odin b/tests/rename_test.odin
index 600bd39..eb9e199 100644
--- a/tests/rename_test.odin
+++ b/tests/rename_test.odin
@@ -242,3 +242,83 @@ ast_prepare_rename_struct_field_ptr :: proc(t: ^testing.T) {
 	range := common.Range{start = {line = 3, character = 9}, end = {line = 3, character = 12}}
 	test.expect_prepare_rename_range(t, &source, range)
 }
+
+@(test)
+ast_prepare_rename_struct_field_enumerated_array :: proc(t: ^testing.T) {
+	source := test.Source {
+		main = `package test
+
+		Foo :: enum {
+			A,
+			B,
+		}
+
+		Bar :: struct {
+			foos: [F{*}oo]int
+		}
+		`,
+	}
+
+	range := common.Range{start = {line = 8, character = 10}, end = {line = 8, character = 13}}
+	test.expect_prepare_rename_range(t, &source, range)
+}
+
+@(test)
+ast_prepare_rename_struct_field_map :: proc(t: ^testing.T) {
+	source := test.Source {
+		main = `package test
+
+		Foo :: enum {
+			A,
+			B,
+		}
+
+		Bar :: struct {
+			foos: map[F{*}oo]int
+		}
+		`,
+	}
+
+	range := common.Range{start = {line = 8, character = 13}, end = {line = 8, character = 16}}
+	test.expect_prepare_rename_range(t, &source, range)
+}
+
+@(test)
+ast_prepare_rename_struct_field_dynamic_array :: proc(t: ^testing.T) {
+	source := test.Source {
+		main = `package test
+
+		Foo :: enum {
+			A,
+			B,
+		}
+
+		Bar :: struct {
+			foos: [dynamic]Fo{*}o
+		}
+		`,
+	}
+
+	range := common.Range{start = {line = 8, character = 18}, end = {line = 8, character = 21}}
+	test.expect_prepare_rename_range(t, &source, range)
+}
+
+@(test)
+ast_prepare_rename_struct_field_bit_set :: proc(t: ^testing.T) {
+	source := test.Source {
+		main = `package test
+
+		Foo :: enum {
+			A,
+			B,
+		}
+
+		Bar :: struct {
+			foos: bit_set[Fo{*}o]
+		}
+		`,
+	}
+
+	range := common.Range{start = {line = 8, character = 17}, end = {line = 8, character = 20}}
+	test.expect_prepare_rename_range(t, &source, range)
+}