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

3 files changed, 41 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
 				}