Improve reference handling for variables part of selector expr

3 files changed, 60 insertions, 9 deletions

diff --git a/src/server/file_resolve.odin b/src/server/file_resolve.odin
index dd2eabb..4e2ce5d 100644
--- a/src/server/file_resolve.odin
+++ b/src/server/file_resolve.odin
@@ -196,7 +196,8 @@ resolve_node :: proc(node: ^ast.Node, data: ^FileResolveData) {
 				}
 			}
 
-			if _, ok := n.expr.derived.(^ast.Selector_Expr); ok {
+			#partial switch v in n.expr.derived {
+			case ^ast.Selector_Expr, ^ast.Index_Expr:
 				resolve_node(n.expr, data)
 			}
 		} else {
diff --git a/src/server/references.odin b/src/server/references.odin
index 944a337..b62efee 100644
--- a/src/server/references.odin
+++ b/src/server/references.odin
@@ -136,13 +136,7 @@ prepare_references :: proc(
 
 		resolve_flag = .Field
 	} else if position_context.selector_expr != nil {
-		resolve_flag = .Field
-
-		base: ^ast.Ident
-		base, ok = position_context.selector.derived.(^ast.Ident)
-
-		if position_in_node(base, position_context.position) && position_context.identifier != nil && ok {
-
+		if position_in_node(position_context.selector, position_context.position) && position_context.identifier != nil {
 			ident := position_context.identifier.derived.(^ast.Ident)
 
 			symbol, ok = resolve_location_identifier(ast_context, ident^)
@@ -151,7 +145,7 @@ prepare_references :: proc(
 				return
 			}
 
-			resolve_flag = .Base
+			resolve_flag = .Identifier
 		} else {
 			symbol, ok = resolve_location_selector(ast_context, position_context.selector_expr)
 
diff --git a/tests/references_test.odin b/tests/references_test.odin
index 8c7c422..1d0fbde 100644
--- a/tests/references_test.odin
+++ b/tests/references_test.odin
@@ -3,6 +3,8 @@ package tests
 import "core:fmt"
 import "core:testing"
 
+import "src:common"
+
 import test "src:testing"
 
 @(test)
@@ -217,3 +219,57 @@ reference_struct_field :: proc(t: ^testing.T) {
 		},
 	)
 }
+
+@(test)
+ast_reference_variable_declaration_with_selector_expr :: proc(t: ^testing.T) {
+	source := test.Source {
+		main = `package test
+
+		Bar :: struct {
+			foo: int,
+		}
+
+		main :: proc() {
+			bar: [2]Bar
+			bar[0].foo = 5
+			b{*}ar[1].foo = 6
+		}
+		`,
+		packages = {},
+	}
+
+	locations := []common.Location{
+		{range = { start = {line = 7, character = 3}, end = {line = 7, character = 6}}},
+		{range = { start = {line = 8, character = 3}, end = {line = 8, character = 6}}},
+		{range = { start = {line = 9, character = 3}, end = {line = 9, character = 6}}},
+	}
+
+	test.expect_reference_locations(t, &source, locations[:])
+}
+
+@(test)
+ast_reference_variable_declaration_field_with_selector_expr :: proc(t: ^testing.T) {
+	source := test.Source {
+		main = `package test
+
+		Bar :: struct {
+			foo: int,
+		}
+
+		main :: proc() {
+			bar: [2]Bar
+			bar[0].foo = 5
+			bar[1].f{*}oo = 6
+		}
+		`,
+		packages = {},
+	}
+
+	locations := []common.Location{
+		{range = { start = {line = 3, character = 3}, end = {line = 3, character = 6}}},
+		{range = { start = {line = 8, character = 10}, end = {line = 8, character = 13}}},
+		{range = { start = {line = 9, character = 10}, end = {line = 9, character = 13}}},
+	}
+
+	test.expect_reference_locations(t, &source, locations[:])
+}