Correctly resolve references when returning a struct

2 files changed, 169 insertions, 92 deletions

diff --git a/src/server/analysis.odin b/src/server/analysis.odin
index b91840e..c1d3038 100644
--- a/src/server/analysis.odin
+++ b/src/server/analysis.odin
@@ -1787,6 +1787,96 @@ resolve_comp_literal :: proc(
 	return symbol, true
 }
 
+// Used to get the name of the field for resolving the implicit selectors
+get_field_value_name :: proc(field_value: ^ast.Field_Value) ->(string, bool) {
+	if field, ok := field_value.field.derived.(^ast.Ident); ok {
+		return field.name, true
+	} else if field, ok := field_value.field.derived.(^ast.Implicit_Selector_Expr); ok {
+		return field.field.name, true
+	}
+	return "", false
+}
+
+resolve_implicit_selector_comp_literal :: proc(
+	ast_context: ^AstContext,
+	position_context: ^DocumentPositionContext,
+	symbol: Symbol,
+	field_name: string,
+) -> (Symbol, bool) {
+	if comp_symbol, comp_lit, ok := resolve_type_comp_literal(
+		ast_context,
+		position_context,
+		symbol,
+		position_context.parent_comp_lit,
+	); ok {
+		if s, ok := comp_symbol.value.(SymbolStructValue); ok {
+			set_ast_package_set_scoped(ast_context, comp_symbol.pkg)
+
+			//We can either have the final
+			elem_index := -1
+
+			for elem, i in comp_lit.elems {
+				if position_in_node(elem, position_context.position) {
+					elem_index = i
+				}
+			}
+
+			type: ^ast.Expr
+
+			for name, i in s.names {
+				if name != field_name {
+					continue
+				}
+
+				type = s.types[i]
+				break
+			}
+
+			if type == nil && len(s.types) > elem_index {
+				type = s.types[elem_index]
+			}
+
+			return resolve_type_expression(ast_context, type)
+		} else if s, ok := comp_symbol.value.(SymbolBitFieldValue); ok {
+			set_ast_package_set_scoped(ast_context, comp_symbol.pkg)
+
+			//We can either have the final
+			elem_index := -1
+
+			for elem, i in comp_lit.elems {
+				if position_in_node(elem, position_context.position) {
+					elem_index = i
+				}
+			}
+
+			type: ^ast.Expr
+
+			for name, i in s.names {
+				if name != field_name {
+					continue
+				}
+
+				type = s.types[i]
+				break
+			}
+
+			if type == nil && len(s.types) > elem_index {
+				type = s.types[elem_index]
+			}
+
+			return resolve_type_expression(ast_context, type)
+		} else if s, ok := comp_symbol.value.(SymbolFixedArrayValue); ok {
+			//This will be a comp_lit for an enumerated array
+			//EnumIndexedArray :: [TestEnum]u32 {
+			//	.valueOne = 1,
+			//	.valueTwo = 2,
+			//}
+			return resolve_type_expression(ast_context, s.len)
+		}
+	}
+	return {}, false
+}
+
 resolve_implicit_selector :: proc(
 	ast_context: ^AstContext,
 	position_context: ^DocumentPositionContext,
@@ -1856,102 +1946,19 @@ resolve_implicit_selector :: proc(
 		}
 	}
 
-	if position_context.value_decl != nil && position_context.value_decl.type != nil {
-		return resolve_type_expression(ast_context, position_context.value_decl.type)
-	}
-
-	if position_context.comp_lit != nil {
-		if position_context.parent_comp_lit.type == nil {
+	if position_context.comp_lit != nil && position_context.parent_comp_lit.type != nil {
+		if position_context.field_value == nil {
 			return {}, false
 		}
-
-		field_name: string
-
-		if position_context.field_value != nil {
-			if field, ok := position_context.field_value.field.derived.(^ast.Ident); ok {
-				field_name = field.name
-			} else if field, ok := position_context.field_value.field.derived.(^ast.Implicit_Selector_Expr); ok {
-				field_name = field.field.name
-			} else {
-				return {}, false
-			}
+		field_name, ok := get_field_value_name(position_context.field_value)
+		if !ok {
+			return {}, false
 		}
 
-
 		if symbol, ok := resolve_type_expression(ast_context, position_context.parent_comp_lit.type); ok {
-			if comp_symbol, comp_lit, ok := resolve_type_comp_literal(
-				ast_context,
-				position_context,
-				symbol,
-				position_context.parent_comp_lit,
-			); ok {
-				if s, ok := comp_symbol.value.(SymbolStructValue); ok {
-					set_ast_package_set_scoped(ast_context, comp_symbol.pkg)
-
-					//We can either have the final
-					elem_index := -1
-
-					for elem, i in comp_lit.elems {
-						if position_in_node(elem, position_context.position) {
-							elem_index = i
-						}
-					}
-
-					type: ^ast.Expr
-
-					for name, i in s.names {
-						if name != field_name {
-							continue
-						}
-
-						type = s.types[i]
-						break
-					}
-
-					if type == nil && len(s.types) > elem_index {
-						type = s.types[elem_index]
-					}
-
-					return resolve_type_expression(ast_context, type)
-				} else if s, ok := comp_symbol.value.(SymbolBitFieldValue); ok {
-					set_ast_package_set_scoped(ast_context, comp_symbol.pkg)
-
-					//We can either have the final
-					elem_index := -1
-
-					for elem, i in comp_lit.elems {
-						if position_in_node(elem, position_context.position) {
-							elem_index = i
-						}
-					}
-
-					type: ^ast.Expr
-
-					for name, i in s.names {
-						if name != field_name {
-							continue
-						}
-
-						type = s.types[i]
-						break
-					}
-
-					if type == nil && len(s.types) > elem_index {
-						type = s.types[elem_index]
-					}
-
-					return resolve_type_expression(ast_context, type)
-				} else if s, ok := comp_symbol.value.(SymbolFixedArrayValue); ok {
-					/*
-					This will be a comp_lit for an enumerated array
-					EnumIndexedArray :: [TestEnum]u32 {
-						.valueOne = 1,
-						.valueTwo = 2,
-					}
-					*/
-					return resolve_type_expression(ast_context, s.len)
-				}
-			}
+			return resolve_implicit_selector_comp_literal(
+				ast_context, position_context, symbol, field_name,
+			)
 		}
 	}
 
@@ -1978,7 +1985,18 @@ resolve_implicit_selector :: proc(
 		}
 
 		if len(position_context.function.type.results.list) > return_index {
-			return resolve_type_expression(ast_context, position_context.function.type.results.list[return_index].type)
+			current_symbol, ok := resolve_type_expression(ast_context, position_context.function.type.results.list[return_index].type)
+			if !ok {
+				return {}, false
+			}
+			if position_context.parent_comp_lit != nil && position_context.field_value != nil {
+				if field_name, ok := get_field_value_name(position_context.field_value); ok {
+					return resolve_implicit_selector_comp_literal(
+						ast_context, position_context, current_symbol, field_name,
+					)
+				}
+			}
+			return current_symbol, ok
 		}
 	}
 
diff --git a/tests/references_test.odin b/tests/references_test.odin
index c4665d8..5298a12 100644
--- a/tests/references_test.odin
+++ b/tests/references_test.odin
@@ -823,3 +823,62 @@ ast_reference_enum_variants_comp_lit_return :: proc(t: ^testing.T) {
 
 	test.expect_reference_locations(t, &source, locations[:])
 }
+
+@(test)
+ast_reference_enum_variants_comp_lit_return_implicit :: proc(t: ^testing.T) {
+	source := test.Source {
+		main = `package test
+
+		Foo :: enum {
+			A,
+			B,
+		}
+
+		Bar :: struct {
+			foo: Foo,
+		}
+
+		foo :: proc() -> Bar {
+			return {
+				foo = .A{*},
+			}
+		}
+
+		`,
+	}
+
+	locations := []common.Location {
+		{range = {start = {line = 3, character = 3}, end = {line = 3, character = 4}}},
+		{range = {start = {line = 13, character = 11}, end = {line = 13, character = 12}}},
+	}
+
+	test.expect_reference_locations(t, &source, locations[:])
+}
+
+@(test)
+ast_reference_enum_indexed_array_return_value :: proc(t: ^testing.T) {
+	source := test.Source {
+		main = `package test
+
+		Foo :: enum {
+			A,
+			B,
+		}
+
+		foo :: proc() -> [Foo]int {
+			return {
+				.A{*} = 2,
+				.B = 1,
+			}
+		}
+
+		`,
+	}
+
+	locations := []common.Location {
+		{range = {start = {line = 3, character = 3}, end = {line = 3, character = 4}}},
+		{range = {start = {line = 9, character = 5}, end = {line = 9, character = 6}}},
+	}
+
+	test.expect_reference_locations(t, &source, locations[:])
+}