Correctly resolve comp lits in map keys

2 files changed, 49 insertions, 10 deletions

diff --git a/src/server/analysis.odin b/src/server/analysis.odin
index f19df28..f10ed28 100644
--- a/src/server/analysis.odin
+++ b/src/server/analysis.odin
@@ -1999,18 +1999,36 @@ internal_resolve_comp_literal :: proc(
 		symbol = resolve_proc(ast_context, position_context.value_decl.type) or_return
 	} else if position_context.assign != nil {
 		if len(position_context.assign.lhs) > 0 {
-			index := 0
-			for value, i in position_context.assign.rhs {
-				if position_in_node(value, position_context.position) {
-					index = i
-					break
+			if position_in_exprs(position_context.assign.lhs, position_context.position) {
+				for n in position_context.assign.lhs {
+					if position_in_node(n, position_context.position) {
+						// check if we're a comp literal of a map key
+						if index_expr, ok := n.derived.(^ast.Index_Expr); ok {
+							if s, ok := resolve_proc(ast_context, index_expr.expr); ok {
+								if value, ok := s.value.(SymbolMapValue); ok {
+									symbol = resolve_proc(ast_context, value.key) or_return
+								}
+							}
+						} else {
+							symbol = resolve_proc(ast_context, n) or_return
+						}
+						break
+					}
 				}
+			} else {
+				index := 0
+				for value, i in position_context.assign.rhs {
+					if position_in_node(value, position_context.position) {
+						index = i
+						break
+					}
+				}
+				// Just to be safe
+				if index >= len(position_context.assign.lhs) {
+					index = 0
+				}
+				symbol = resolve_proc(ast_context, position_context.assign.lhs[index]) or_return
 			}
-			// Just to be safe
-			if index >= len(position_context.assign.lhs) {
-				index = 0
-			}
-			symbol = resolve_proc(ast_context, position_context.assign.lhs[index]) or_return
 		}
 	}
 
diff --git a/tests/hover_test.odin b/tests/hover_test.odin
index a4769cb..8660eb7 100644
--- a/tests/hover_test.odin
+++ b/tests/hover_test.odin
@@ -4629,6 +4629,27 @@ ast_hover_assign_comp_lit_with_multiple_assigns_second :: proc(t: ^testing.T) {
 	}
 	test.expect_hover(t, &source, "Bar.b: int")
 }
+
+@(test)
+ast_hover_comp_lit_map_key :: proc(t: ^testing.T) {
+	source := test.Source {
+		main     = `package test
+		Foo :: struct {
+			a: int,
+		}
+
+		Bar :: struct {
+			b: int,
+		}
+
+		main :: proc() {
+			m: map[Foo]Bar
+			m[{a{*} = 1}] = {b = 2}
+		}
+		`,
+	}
+	test.expect_hover(t, &source, "Foo.a: int")
+}
 /*
 
 Waiting for odin fix