Merge pull request #1263 from pnarimani/overload_resolution

Procedure Overload Resolution

8 files changed, 190 insertions, 7 deletions

diff --git a/README.md b/README.md
index 75b51ff..92ee383 100644
--- a/README.md
+++ b/README.md
@@ -77,6 +77,8 @@ Options:
 
 - `enable_fake_methods`: Turn on fake methods completion. This is currently highly experimental.
 
+- `enable_overload_resolution`: Enable go-to-definition to resolve overloaded procedures from procedure groups based on call arguments.
+
 - `enable_references`: Turns on finding references for a symbol. _(Enabled by default)_
 
 - `enable_document_highlights`: Turns on highlighting of symbol references in file. _(Enabled by default)_
diff --git a/misc/ols.schema.json b/misc/ols.schema.json
index c95d8e6..81bc0cd 100644
--- a/misc/ols.schema.json
+++ b/misc/ols.schema.json
@@ -87,6 +87,11 @@
 			"description": "Turn on fake methods completion.",
 			"default": false
 		},
+		"enable_overload_resolution": {
+			"type": "boolean",
+			"description": "Enable go-to-definition to resolve overloaded procedures from procedure groups based on call arguments.",
+			"default": false
+		},
 		"enable_document_links": {
 			"type": "boolean",
 			"description": "Follow links when opening documentation.",
diff --git a/src/common/config.odin b/src/common/config.odin
index 441e8a9..f2ae68f 100644
--- a/src/common/config.odin
+++ b/src/common/config.odin
@@ -33,6 +33,7 @@ Config :: struct {
 	enable_std_references:                   bool,
 	enable_import_fixer:                     bool,
 	enable_fake_method:                      bool,
+	enable_overload_resolution:              bool,
 	enable_procedure_snippet:                bool,
 	enable_checker_only_saved:               bool,
 	enable_auto_import:                      bool,
diff --git a/src/server/definition.odin b/src/server/definition.odin
index c793f32..869816c 100644
--- a/src/server/definition.odin
+++ b/src/server/definition.odin
@@ -42,7 +42,7 @@ get_all_package_file_locations :: proc(
 	return true
 }
 
-get_definition_location :: proc(document: ^Document, position: common.Position) -> ([]common.Location, bool) {
+get_definition_location :: proc(document: ^Document, position: common.Position, config: ^common.Config) -> ([]common.Location, bool) {
 	locations := make([dynamic]common.Location, context.temp_allocator)
 
 	location: common.Location
@@ -100,6 +100,14 @@ get_definition_location :: proc(document: ^Document, position: common.Position)
 		}
 
 		if resolved, ok := resolve_location_selector(&ast_context, position_context.selector_expr); ok {
+			if config.enable_overload_resolution {
+				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 +147,12 @@ get_definition_location :: proc(document: ^Document, position: common.Position)
 			&ast_context,
 			position_context.identifier.derived.(^ast.Ident)^,
 		); ok {
+			if config.enable_overload_resolution {
+				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 +175,114 @@ 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 result, ok := get_full_symbol_from_selector(ast_context, selector_expr, symbol); ok {
+		full_symbol = result
+	} else if result, ok := get_full_symbol_from_identifier(ast_context, position_context, symbol); ok {
+		full_symbol = result
+	}
+
+	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
+}
+
+get_full_symbol_from_selector :: proc(
+	ast_context: ^AstContext,
+	selector_expr: ^ast.Node,
+	symbol: Symbol,
+) -> (
+	full_symbol: Symbol,
+	ok: bool,
+) {
+	if selector_expr == nil do return
+
+	selector := selector_expr.derived.(^ast.Selector_Expr) or_return
+
+	_, is_pkg := symbol.value.(SymbolPackageValue)
+	if !is_pkg && symbol.value != nil do return
+
+	if selector.field == nil do return
+
+	ident := selector.field.derived.(^ast.Ident) or_return
+
+	return lookup(ident.name, symbol.pkg, ast_context.fullpath);
+}
+
+get_full_symbol_from_identifier :: proc(
+	ast_context: ^AstContext,
+	position_context: ^DocumentPositionContext,
+	symbol: Symbol,
+) -> (
+	full_symbol: Symbol,
+	ok: bool,
+) {
+	if position_context.identifier == nil || symbol.value != nil do return
+
+	// 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.
+	ident := position_context.identifier.derived.(^ast.Ident) or_return
+
+	pkg := symbol.pkg if symbol.pkg != "" else ast_context.document_package
+
+	if pkg_symbol, ok := lookup(ident.name, pkg, ast_context.fullpath); ok {
+		return pkg_symbol, true
+	}
+
+	// If lookup fails (e.g., in tests without full indexing), try checking if it's a proc group
+
+	global := ast_context.globals[ident.name] or_return
+	if proc_group, is_proc_group := global.expr.derived.(^ast.Proc_Group); is_proc_group {
+		full_symbol = symbol
+		full_symbol.value = SymbolProcedureGroupValue {
+			group = global.expr,
+		}
+		return full_symbol, true
+	}
+
+	return Symbol{}, false
+}
diff --git a/src/server/requests.odin b/src/server/requests.odin
index 0d6e025..64142ed 100644
--- a/src/server/requests.odin
+++ b/src/server/requests.odin
@@ -454,6 +454,8 @@ read_ols_initialize_options :: proc(config: ^common.Config, ols_config: OlsConfi
 		ols_config.enable_inlay_hints_implicit_return.(bool) or_else config.enable_inlay_hints_implicit_return
 
 	config.enable_fake_method = ols_config.enable_fake_methods.(bool) or_else config.enable_fake_method
+	config.enable_overload_resolution =
+		ols_config.enable_overload_resolution.(bool) or_else config.enable_overload_resolution
 
 	// Delete overriding collections.
 	for it in ols_config.collections {
@@ -861,7 +863,7 @@ request_definition :: proc(
 		return .InternalError
 	}
 
-	locations, ok2 := get_definition_location(document, definition_params.position)
+	locations, ok2 := get_definition_location(document, definition_params.position, config)
 
 	if !ok2 {
 		log.warn("Failed to get definition location")
diff --git a/src/server/types.odin b/src/server/types.odin
index c0bcf9e..c65e181 100644
--- a/src/server/types.odin
+++ b/src/server/types.odin
@@ -419,6 +419,7 @@ OlsConfig :: struct {
 	enable_hover:                            Maybe(bool),
 	enable_document_symbols:                 Maybe(bool),
 	enable_fake_methods:                     Maybe(bool),
+	enable_overload_resolution:              Maybe(bool),
 	enable_references:                       Maybe(bool),
 	enable_document_highlights:              Maybe(bool),
 	enable_document_links:                   Maybe(bool),
diff --git a/src/testing/testing.odin b/src/testing/testing.odin
index 373ef62..2013f79 100644
--- a/src/testing/testing.odin
+++ b/src/testing/testing.odin
@@ -346,7 +346,7 @@ expect_definition_locations :: proc(t: ^testing.T, src: ^Source, expect_location
 	setup(src)
 	defer teardown(src)
 
-	locations, ok := server.get_definition_location(src.document, src.position)
+	locations, ok := server.get_definition_location(src.document, src.position, &src.config)
 
 	if !ok {
 		log.error("Failed get_definition_location")
diff --git a/tests/definition_test.odin b/tests/definition_test.odin
index 4810361..fe4cd3b 100644
--- a/tests/definition_test.odin
+++ b/tests/definition_test.odin
@@ -723,3 +723,56 @@ 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[:],
+		config = {enable_overload_resolution = true},
+	}
+	// 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)
+		}
+	`,
+		config = {enable_overload_resolution = true},
+	}
+	// 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