Correctly resolve struct fields that are immediately accessed from a proc return

2 files changed, 44 insertions, 34 deletions

diff --git a/src/server/analysis.odin b/src/server/analysis.odin
index f6f6cb0..4b66717 100644
--- a/src/server/analysis.odin
+++ b/src/server/analysis.odin
@@ -2253,50 +2253,60 @@ resolve_location_selector :: proc(ast_context: ^AstContext, selector_expr: ^ast.
 	set_ast_package_set_scoped(ast_context, ast_context.document_package)
 
 	if selector, ok := selector_expr.derived.(^ast.Selector_Expr); ok {
-
 		symbol = resolve_type_expression(ast_context, selector.expr) or_return
+		return resolve_symbol_selector(ast_context, selector, symbol)
+	}
 
-		field: string
+	return {}, false
+}
 
-		if selector.field != nil {
-			#partial switch v in selector.field.derived {
-			case ^ast.Ident:
-				field = v.name
-			}
+resolve_symbol_selector :: proc(ast_context: ^AstContext, selector: ^ast.Selector_Expr, symbol: Symbol) ->(Symbol, bool) {
+	field: string
+	symbol := symbol
+
+	if selector.field != nil {
+		#partial switch v in selector.field.derived {
+		case ^ast.Ident:
+			field = v.name
 		}
+	}
 
-		#partial switch v in symbol.value {
-		case SymbolEnumValue:
-			for name, i in v.names {
-				if strings.compare(name, field) == 0 {
-					symbol.range = v.ranges[i]
-				}
-			}
-		case SymbolStructValue:
-			for name, i in v.names {
-				if strings.compare(name, field) == 0 {
-					symbol.range = v.ranges[i]
-				}
+	#partial switch v in symbol.value {
+	case SymbolEnumValue:
+		for name, i in v.names {
+			if strings.compare(name, field) == 0 {
+				symbol.range = v.ranges[i]
 			}
-		case SymbolBitFieldValue:
-			for name, i in v.names {
-				if strings.compare(name, field) == 0 {
-					symbol.range = v.ranges[i]
-				}
+		}
+	case SymbolStructValue:
+		for name, i in v.names {
+			if strings.compare(name, field) == 0 {
+				symbol.range = v.ranges[i]
 			}
-		case SymbolPackageValue:
-			if pkg, ok := lookup(field, symbol.pkg); ok {
-				symbol.range = pkg.range
-				symbol.uri = pkg.uri
-			} else {
-				return {}, false
+		}
+	case SymbolBitFieldValue:
+		for name, i in v.names {
+			if strings.compare(name, field) == 0 {
+				symbol.range = v.ranges[i]
 			}
 		}
-
-		return symbol, true
+	case SymbolPackageValue:
+		if pkg, ok := lookup(field, symbol.pkg); ok {
+			symbol.range = pkg.range
+			symbol.uri = pkg.uri
+		} else {
+			return {}, false
+		}
+	case SymbolProcedureValue:
+		if len(v.return_types) != 1 {
+			return {}, false
+		}
+		if s, ok := resolve_type_expression(ast_context, v.return_types[0].type); ok {
+			return resolve_symbol_selector(ast_context, selector, s)
+		}
 	}
 
-	return {}, false
+	return symbol, true
 }
 
 
diff --git a/src/server/file_resolve.odin b/src/server/file_resolve.odin
index 86ef1f5..f6096c5 100644
--- a/src/server/file_resolve.odin
+++ b/src/server/file_resolve.odin
@@ -202,7 +202,7 @@ resolve_node :: proc(node: ^ast.Node, data: ^FileResolveData) {
 
 			#partial switch v in n.expr.derived {
 			// TODO: Should there be more here?
-			case ^ast.Selector_Expr, ^ast.Index_Expr, ^ast.Ident, ^ast.Paren_Expr:
+			case ^ast.Selector_Expr, ^ast.Index_Expr, ^ast.Ident, ^ast.Paren_Expr, ^ast.Call_Expr:
 				resolve_node(n.expr, data)
 			}
 		} else {