feat: go to resolved proc from a proc overload group when performing go-to-definition

2 files changed, 140 insertions, 4 deletions

diff --git a/src/server/definition.odin b/src/server/definition.odin
index c793f32..34d1315 100644
--- a/src/server/definition.odin
+++ b/src/server/definition.odin
@@ -100,6 +100,12 @@ get_definition_location :: proc(document: ^Document, position: common.Position)
 		}
 
 		if resolved, ok := resolve_location_selector(&ast_context, position_context.selector_expr); ok {
+			resolved = try_resolve_proc_group_overload(
+				&ast_context,
+				&position_context,
+				resolved,
+				position_context.selector_expr,
+			)
 			location.range = resolved.range
 			uri = resolved.uri
 		} else {
@@ -139,12 +145,10 @@ get_definition_location :: proc(document: ^Document, position: common.Position)
 			&ast_context,
 			position_context.identifier.derived.(^ast.Ident)^,
 		); ok {
+			resolved = try_resolve_proc_group_overload(&ast_context, &position_context, resolved)
 			if v, ok := resolved.value.(SymbolAggregateValue); ok {
 				for symbol in v.symbols {
-					append(&locations, common.Location {
-						range = symbol.range,
-						uri = symbol.uri,
-					})
+					append(&locations, common.Location{range = symbol.range, uri = symbol.uri})
 				}
 			}
 			location.range = resolved.range
@@ -167,3 +171,84 @@ get_definition_location :: proc(document: ^Document, position: common.Position)
 
 	return locations[:], true
 }
+
+
+try_resolve_proc_group_overload :: proc(
+	ast_context: ^AstContext,
+	position_context: ^DocumentPositionContext,
+	symbol: Symbol,
+	selector_expr: ^ast.Node = nil,
+) -> Symbol {
+	if position_context.call == nil {
+		return symbol
+	}
+
+	call, is_call := position_context.call.derived.(^ast.Call_Expr)
+	if !is_call {
+		return symbol
+	}
+
+	if position_in_exprs(call.args, position_context.position) {
+		return symbol
+	}
+
+	// For selector expressions, we need to look up the full symbol to check if it's a proc group
+	full_symbol := symbol
+	if selector_expr != nil {
+		if selector, ok := selector_expr.derived.(^ast.Selector_Expr); ok {
+			if _, is_pkg := symbol.value.(SymbolPackageValue); is_pkg || symbol.value == nil {
+				if selector.field != nil {
+					if ident, ok := selector.field.derived.(^ast.Ident); ok {
+						if pkg_symbol, ok := lookup(ident.name, symbol.pkg, ast_context.fullpath); ok {
+							full_symbol = pkg_symbol
+						}
+					}
+				}
+			}
+		}
+	} else if position_context.identifier != nil && symbol.value == nil {
+		// For identifiers (non-selector), the symbol from resolve_location_identifier may not have
+		// value set (e.g., for globals). We need to do a lookup to get the full symbol.
+		if ident, ok := position_context.identifier.derived.(^ast.Ident); ok {
+			pkg := symbol.pkg
+			if pkg == "" do pkg = ast_context.document_package
+
+			if pkg_symbol, ok := lookup(ident.name, pkg, ast_context.fullpath); ok {
+				full_symbol = pkg_symbol
+			} else if global, ok := ast_context.globals[ident.name]; ok {
+				// If lookup fails (e.g., in tests without full indexing), try checking if it's a proc group
+				if proc_group, is_proc_group := global.expr.derived.(^ast.Proc_Group); is_proc_group {
+					full_symbol.value = SymbolProcedureGroupValue {
+						group = global.expr,
+					}
+				}
+			}
+		}
+	}
+
+	proc_group_value, is_proc_group := full_symbol.value.(SymbolProcedureGroupValue)
+	if !is_proc_group {
+		return symbol
+	}
+
+	old_call := ast_context.call
+	ast_context.call = call
+	defer {
+		ast_context.call = old_call
+	}
+
+	if resolved, ok := resolve_function_overload(ast_context, proc_group_value.group.derived.(^ast.Proc_Group)); ok {
+		if resolved.name != "" {
+			if global, ok := ast_context.globals[resolved.name]; ok {
+				resolved.range = common.get_token_range(global.name_expr, ast_context.file.src)
+				resolved.uri = common.create_uri(global.name_expr.pos.file, ast_context.allocator).uri
+			} else if indexed_symbol, ok := lookup(resolved.name, resolved.pkg, ast_context.fullpath); ok {
+				resolved.range = indexed_symbol.range
+				resolved.uri = indexed_symbol.uri
+			}
+		}
+		return resolved
+	}
+
+	return symbol
+}
diff --git a/tests/definition_test.odin b/tests/definition_test.odin
index 4810361..62165e2 100644
--- a/tests/definition_test.odin
+++ b/tests/definition_test.odin
@@ -723,3 +723,54 @@ ast_goto_package_declaration_with_alias :: proc(t: ^testing.T) {
 
 	test.expect_definition_locations(t, &source, locations[:])
 }
+@(test)
+ast_goto_proc_group_overload_with_selector :: proc(t: ^testing.T) {
+	packages := make([dynamic]test.Package, context.temp_allocator)
+
+	append(&packages, test.Package{pkg = "my_package", source = `package my_package
+			push_back :: proc(arr: ^[dynamic]int, val: int) {}
+			push_back_elems :: proc(arr: ^[dynamic]int, vals: ..int) {}
+			append :: proc{push_back, push_back_elems}
+		`})
+	source := test.Source {
+		main = `package test
+		import mp "my_package"
+
+		main :: proc() {
+			arr: [dynamic]int
+			mp.app{*}end(&arr, 1)
+		}
+	`,
+		packages = packages[:],
+	}
+	// Should go to push_back (line 1, character 3) instead of append (line 3)
+	// because push_back is the overload being used with a single value argument
+	locations := []common.Location {
+		{range = {start = {line = 1, character = 3}, end = {line = 1, character = 12}}},
+	}
+
+	test.expect_definition_locations(t, &source, locations[:])
+}
+
+@(test)
+ast_goto_proc_group_overload_identifier :: proc(t: ^testing.T) {
+	source := test.Source {
+		main = `package test
+		push_back :: proc(arr: ^[dynamic]int, val: int) {}
+		push_back_elems :: proc(arr: ^[dynamic]int, vals: ..int) {}
+		append :: proc{push_back, push_back_elems}
+
+		main :: proc() {
+			arr: [dynamic]int
+			app{*}end(&arr, 1)
+		}
+	`,
+	}
+	// Should go to push_back (line 1, character 2) instead of append (line 3)
+	// because push_back is the overload being used with a single value argument
+	locations := []common.Location {
+		{range = {start = {line = 1, character = 2}, end = {line = 1, character = 11}}},
+	}
+
+	test.expect_definition_locations(t, &source, locations[:])
+}
\ No newline at end of file