Support comp lit inference from return

2 files changed, 67 insertions, 0 deletions

diff --git a/src/server/analysis.odin b/src/server/analysis.odin
index 7b4d125..b6ae251 100644
--- a/src/server/analysis.odin
+++ b/src/server/analysis.odin
@@ -1853,6 +1853,34 @@ resolve_comp_literal :: proc(
 				value.arg_types[arg_index].type,
 			) or_return
 		}
+	} else if position_context.returns != nil {
+		return_index: int
+
+		if position_context.returns.results == nil {
+			return {}, false
+		}
+
+		for result, i in position_context.returns.results {
+			if position_in_node(result, position_context.position) {
+				return_index = i
+				break
+			}
+		}
+
+		if position_context.function.type == nil {
+			return {}, false
+		}
+
+		if position_context.function.type.results == nil {
+			return {}, false
+		}
+
+		if len(position_context.function.type.results.list) > return_index {
+			symbol = resolve_type_expression(
+				ast_context,
+				position_context.function.type.results.list[return_index].type,
+			) or_return
+		}
 	}
 
 	set_ast_package_set_scoped(ast_context, symbol.pkg)
diff --git a/tests/references_test.odin b/tests/references_test.odin
index cadaff4..61c9cd0 100644
--- a/tests/references_test.odin
+++ b/tests/references_test.odin
@@ -179,3 +179,42 @@ reference_field_comp_lit_infer_from_function :: proc(t: ^testing.T) {
 		},
 	)
 }
+
+@(test)
+reference_field_comp_lit_infer_from_return :: proc(t: ^testing.T) {
+	source := test.Source {
+		main     = `package test
+		Foo :: struct {
+			soo_many_cases: int,
+		}
+
+		My_Struct :: struct {
+			foo: Foo,
+		}
+
+		my_function :: proc() -> My_Struct {
+			return {foo = {soo_many_cases{*} = 2}}
+		}
+		`,
+		packages = {},
+	}
+
+	test.expect_reference_locations(
+		t,
+		&source,
+		{
+			{
+				range = {
+					start = {line = 2, character = 3},
+					end = {line = 2, character = 17},
+				},
+			},
+			{
+				range = {
+					start = {line = 10, character = 18},
+					end = {line = 10, character = 32},
+				},
+			},
+		},
+	)
+}