Merge pull request #996 from thetarnav/inlay-hints-fixes

Inlay hints fixes

2 files changed, 280 insertions, 138 deletions

diff --git a/src/server/inlay_hints.odin b/src/server/inlay_hints.odin
index 6dc155d..c87ce3a 100644
--- a/src/server/inlay_hints.odin
+++ b/src/server/inlay_hints.odin
@@ -1,5 +1,6 @@
 package server
 
+import "core:strings"
 import "core:fmt"
 import "core:log"
 import "core:odin/ast"
@@ -14,169 +15,232 @@ get_inlay_hints :: proc(
 	[]InlayHint,
 	bool,
 ) {
-	hints := make([dynamic]InlayHint, context.temp_allocator)
-
-	ast_context := make_ast_context(
-		document.ast,
-		document.imports,
-		document.package_name,
-		document.uri.uri,
-		document.fullpath,
-	)
-
-	Visit_Data :: struct {
-		calls: [dynamic]^ast.Node,
+	Visitor_Data :: struct {
+		hints:    [dynamic]InlayHint,
+		document: ^Document,
+		symbols:  map[uintptr]SymbolAndNode,
+		config:   ^common.Config,
 	}
 
-	data := Visit_Data {
-		calls = make([dynamic]^ast.Node, context.temp_allocator),
+	data := Visitor_Data{
+		hints    = make([dynamic]InlayHint, context.temp_allocator),
+		document = document,
+		symbols  = symbols,
+		config   = config,
 	}
 
-	visit :: proc(visitor: ^ast.Visitor, node: ^ast.Node) -> ^ast.Visitor {
-		if node == nil || visitor == nil {
-			return nil
-		}
-
-		data := cast(^Visit_Data)visitor.data
+	visitor := ast.Visitor{
+		data  = &data,
+		visit = proc(visitor: ^ast.Visitor, node: ^ast.Node) -> ^ast.Visitor {
+			if node == nil || visitor == nil {
+				return nil
+			}
 
-		if call, ok := node.derived.(^ast.Call_Expr); ok {
-			append(&data.calls, node)
-		}
+			if call, is_call := node.derived.(^ast.Call_Expr); is_call {
+				data := (^Visitor_Data)(visitor.data)
+				visit_call(call, data)
+			}
 
-		return visitor
-	}
-
-	visitor := ast.Visitor {
-		data  = &data,
-		visit = visit,
+			return visitor
+		},
 	}
 
 	for decl in document.ast.decls {
 		ast.walk(&visitor, decl)
 	}
 
+	expr_name :: proc (node: ^ast.Node) -> (name: string, ok: bool) {
+		#partial switch v in node.derived {
+		case ^ast.Ident: return v.name, true
+		case ^ast.Poly_Type: return expr_name(v.type)
+		case: return
+		}
+	}
 
-	loop: for node_call in &data.calls {
-
-		is_ellipsis := false
-		has_added_default := false
+	visit_call :: proc (
+		call: ^ast.Call_Expr,
+		data: ^Visitor_Data,
+	) -> (ok: bool) {
 
-		call := node_call.derived.(^ast.Call_Expr)
+		src := string(data.document.text)
+		end_pos := common.token_pos_to_position(call.close, src)
 
 		selector, is_selector_call := call.expr.derived.(^ast.Selector_Expr)
 		is_selector_call &&= selector.op.kind == .Arrow_Right
 
-		end_pos := common.token_pos_to_position(call.close, string(document.text))
-
-		symbol_and_node := symbols[cast(uintptr)call.expr] or_continue
-		symbol_call := symbol_and_node.symbol.value.(SymbolProcedureValue) or_continue
-		
-		positional_arg_idx := 0
-
-		expr_name :: proc (node: ^ast.Node) -> (name: string, ok: bool) {
-			#partial switch v in node.derived {
-			case ^ast.Ident:
-				return v.name, true
-			case ^ast.Poly_Type:
-				ident := v.type.derived.(^ast.Ident) or_return
-				return ident.name, true
-			case:
-				return
-			}
-		}
-		
-		for arg, arg_type_idx in symbol_call.arg_types {
-			if arg_type_idx == 0 && is_selector_call {
-				continue
-			}
+		symbol_and_node := data.symbols[uintptr(call.expr)] or_return // could not resolve symbol
+		proc_symbol := symbol_and_node.symbol.value.(SymbolProcedureValue) or_return // not a procedure call, e.g. type cast
 
-			for name, name_idx in arg.names {
+		param_idx := 1 if is_selector_call else 0
+		label_idx := 0
+		arg_idx   := 0
 
-				arg_call_idx := arg_type_idx + name_idx
-				if is_selector_call do arg_call_idx -= 1
+		if param_idx >= len(proc_symbol.arg_types) do return true // no parameters
 
-				label := expr_name(name) or_continue loop
+		// Positional arguments
+		positional: for ; arg_idx < len(call.args); arg_idx += 1 {
+			arg := call.args[arg_idx]
 
-				is_provided_named, is_provided_positional: bool
-				call_arg: ^ast.Expr
+			// for multi-return function call arguments
+			multi_return: {
+				arg_call := arg.derived.(^ast.Call_Expr) or_break multi_return
+				arg_symbol_and_node := data.symbols[uintptr(arg_call.expr)] or_break multi_return
+				arg_proc_symbol := arg_symbol_and_node.symbol.value.(SymbolProcedureValue) or_break multi_return
 
-				for a, a_i in call.args[positional_arg_idx:] {
-					call_arg_idx := a_i + positional_arg_idx
-					// provided as named
-					if field_value, ok := a.derived.(^ast.Field_Value); ok {
-						ident := field_value.field.derived.(^ast.Ident) or_break
-						if ident.name == label {
-							is_provided_named = true
-							call_arg = a
+				if len(arg_proc_symbol.return_types) <= 1 do break multi_return
+
+				hint_text_sb := strings.builder_make(context.temp_allocator)
+
+				// Collect parameter names for this multi-return call
+				for i in 0..<len(arg_proc_symbol.return_types) {
+					param := proc_symbol.arg_types[param_idx]
+					label := param.names[label_idx]
+
+					label_name, label_has_name := expr_name(label)
+					if data.config.enable_inlay_hints_params && label_has_name {
+						if i > 0 do strings.write_string(&hint_text_sb, ", ")
+						strings.write_string(&hint_text_sb, label_name)
+					}
+
+					// advance to next param
+					label_idx += 1
+					if label_idx >= len(param.names) {
+						param_idx += 1
+						label_idx = 0
+						if param_idx >= len(proc_symbol.arg_types) {
+							return true // end of parameters
 						}
-						break
-					} // provided as positional
-					else if arg_call_idx == call_arg_idx {
-						is_provided_positional = true
-						positional_arg_idx += 1
-						call_arg = a
-						break
 					}
 				}
 
-				if is_ellipsis || (!is_provided_named && !is_provided_positional) {
-					// This parameter is not provided, so it should use default value
-					if arg.default_value == nil {
-						continue loop
-					}
+				// Add combined param name hint
+				if data.config.enable_inlay_hints_params && strings.builder_len(hint_text_sb) > 0 {
+					range := common.get_token_range(arg, src)
+					strings.write_string(&hint_text_sb, " = ")
+					hint_text := strings.to_string(hint_text_sb)
+					append(&data.hints, InlayHint{range.start, .Parameter, hint_text})
+				}
 
-					if !config.enable_inlay_hints_default_params {
-						continue loop
-					}
+				continue positional
+			}
+
+			// provided as named
+			if _, is_field := arg.derived.(^ast.Field_Value); is_field do break
+
+			param := proc_symbol.arg_types[param_idx]
+
+			// param is variadic
+			if param.type != nil {
+				if _, is_variadic := param.type.derived.(^ast.Ellipsis); is_variadic do break
+			}
+
+			// Add param name hint for single-value arg
+			if data.config.enable_inlay_hints_params {
+				label := param.names[label_idx]
+				label_name, label_has_name := expr_name(label)
+				arg_name, arg_has_name := expr_name(arg)
+
+				// Add param name hint (skip idents with same name as param)
+				if label_has_name && (!arg_has_name || arg_name != label_name) {
+					range := common.get_token_range(arg, src)
+					hint_text := fmt.tprintf("%v = ", label_name)
+					append(&data.hints, InlayHint{range.start, .Parameter, hint_text})
+				}
+			}
+
+			// advance to next param
+			label_idx += 1
+			if label_idx >= len(param.names) {
+				param_idx += 1
+				label_idx = 0
+				if param_idx >= len(proc_symbol.arg_types) {
+					return true // end of parameters
+				}
+			}
+		}
+
+		// Variadic arguments
+		variadic: {
+			param := proc_symbol.arg_types[param_idx]
+
+			// param is variadic
+			if param.type == nil do break variadic
+			_ = param.type.derived.(^ast.Ellipsis) or_break variadic
 
-					value := node_to_string(arg.default_value)
+			// skip all provided args
+			init_arg_idx := arg_idx
+			for ; arg_idx < len(call.args); arg_idx += 1 {
+				// provided as named
+				if _, is_field := call.args[arg_idx].derived.(^ast.Field_Value); is_field do break
+			}
+
+			// Add param name hint
+			if arg_idx > init_arg_idx && data.config.enable_inlay_hints_params {
+				if label_name, label_has_name := expr_name(param.names[0]); label_has_name {
+					range := common.get_token_range(call.args[init_arg_idx], src)
+					hint_text := fmt.tprintf("%v = ", label_name)
+					append(&data.hints, InlayHint{range.start, .Parameter, hint_text})
+				}
+			}
+
+			// advance to next param
+			param_idx += 1
+			label_idx = 0
+			if param_idx >= len(proc_symbol.arg_types) {
+				return true // end of parameters
+			}
+		}
 
-					needs_leading_comma := arg_call_idx > 0
+		// Named arguments
+		named: if data.config.enable_inlay_hints_default_params {
 
-					if !has_added_default && needs_leading_comma {
-						till_end := string(document.text[:call.close.offset])
-						#reverse for ch in till_end {
+			init_arg_idx := arg_idx
+			added_default_hint := false
+
+			for ; param_idx < len(proc_symbol.arg_types); param_idx, label_idx = param_idx+1, 0 {
+				param := proc_symbol.arg_types[param_idx]
+
+				label_loop: for ; label_idx < len(param.names); label_idx += 1 {
+					label := param.names[label_idx]
+					label_name := expr_name(label) or_continue
+
+					if param.default_value == nil do continue
+
+					// check if was already provided
+					for arg in call.args[init_arg_idx:] {
+
+						field_value := arg.derived.(^ast.Field_Value) or_break named
+						ident := field_value.field.derived.(^ast.Ident) or_break named
+
+						if ident.name == label_name {
+							continue label_loop
+						}
+					}
+
+					needs_leading_comma := added_default_hint || param_idx > 0 || label_idx > 0
+					if needs_leading_comma && !added_default_hint {
+						// check for existing trailing comma
+						#reverse for ch in string(data.document.text[:call.close.offset]) {
 							switch ch {
-							case ' ', '\t', '\n':
-								continue
-							case ',':
-								needs_leading_comma = false
+							case ' ', '\t', '\n': continue
+							case ',': needs_leading_comma = false
 							}
 							break
 						}
 					}
 
-					hint := InlayHint {
-						kind     = .Parameter,
-						label    = fmt.tprintf("%s%v = %v", needs_leading_comma ? ", " : "", label, value),
-						position = end_pos,
-					}
-					append(&hints, hint)
-
-					has_added_default = true
-				} else if config.enable_inlay_hints_params && is_provided_positional && !is_provided_named {
-					// This parameter is provided via positional argument, show parameter hint
-
-					// if the arg name and param name are the same, don't add it.
-					call_arg_name, _ := expr_name(call_arg)
-					if call_arg_name != label {
-						range := common.get_token_range(call_arg, string(document.text))
-						hint := InlayHint {
-							kind     = .Parameter,
-							label    = fmt.tprintf("%v = ", label),
-							position = range.start,
-						}
-						append(&hints, hint)
-					}
-				}
+					// Add default param hint
+					value := node_to_string(param.default_value)
+					hint_text := fmt.tprintf("%s%v = %v", needs_leading_comma ? ", " : "", label_name, value)
+					append(&data.hints, InlayHint{end_pos, .Parameter, hint_text})
 
-				if arg.type != nil {
-					_, is_current_ellipsis := arg.type.derived.(^ast.Ellipsis)
-					is_ellipsis ||= is_current_ellipsis
+					added_default_hint = true
 				}
 			}
 		}
+
+		return true
 	}
 
-	return hints[:], true
+	return data.hints[:], true
 }
diff --git a/src/testing/testing.odin b/src/testing/testing.odin
index 57354e1..66ed992 100644
--- a/src/testing/testing.odin
+++ b/src/testing/testing.odin
@@ -532,36 +532,114 @@ expect_semantic_tokens :: proc(t: ^testing.T, src: ^Source, expected: []server.S
 	}
 }
 
-expect_inlay_hints :: proc(t: ^testing.T, src: ^Source, expected_hints: []server.InlayHint) {
+expect_inlay_hints :: proc(t: ^testing.T, src: ^Source) {
+
+	src_builder := strings.builder_make(context.temp_allocator)
+	expected_hints := make([dynamic]server.InlayHint, context.temp_allocator)
+
+	HINT_OPEN  :: "[["
+	HINT_CLOSE :: "]]"
+
+	{
+		last, line, col: int
+		saw_brackets: bool
+		for i:= 0; i < len(src.main); i += 1 {
+			if saw_brackets {
+				if i+1 < len(src.main) && src.main[i:][:len(HINT_CLOSE)] == HINT_CLOSE {
+					saw_brackets = false
+					hint_str := src.main[last:i]
+					last = i+len(HINT_CLOSE)
+					i = last-1
+					append(&expected_hints, server.InlayHint{
+						position = {line, col},
+						label    = hint_str,
+						kind     = .Parameter,
+					})
+				}
+			} else {
+				if i+1 < len(src.main) && src.main[i:][:len(HINT_OPEN)] == HINT_OPEN {
+					strings.write_string(&src_builder, src.main[last:i])
+					saw_brackets = true
+					last = i+len(HINT_OPEN)
+					i = last-1
+				} else if src.main[i] == '\n' {
+					line += 1
+					col = 0
+				} else {
+					col += 1
+				}
+			}
+		}
+
+		if saw_brackets {
+			log.error("Unclosed inlay hint marker")
+			return
+		}
+
+		strings.write_string(&src_builder, src.main[last:len(src.main)])
+	}
+
+	src.main = strings.to_string(src_builder)
+
 	setup(src)
 	defer teardown(src)
 
 	resolve_flag: server.ResolveReferenceFlag
-	symbols_and_nodes := server.resolve_entire_file(src.document, resolve_flag, context.temp_allocator)
+	symbols_and_nodes := server.resolve_entire_file(src.document, resolve_flag, )
 
-	hints, ok := server.get_inlay_hints(src.document, symbols_and_nodes, &src.config)
-	if !ok {
+	hints, hints_ok := server.get_inlay_hints(src.document, symbols_and_nodes, &src.config)
+	if !hints_ok {
 		log.error("Failed get_inlay_hints")
 		return
 	}
 
-	if len(expected_hints) == 0 && len(hints) > 0 {
-		log.errorf("Expected empty inlay hints, but received %v", hints)
-		return
-	}
-
-	testing.expectf(
-		t,
+	testing.expectf(t,
 		len(expected_hints) == len(hints),
-		"\nExpected %d inlay hints, but received %d",
+		"Expected %d inlay hints, but received %d",
 		len(expected_hints),
 		len(hints),
 	)
 
-	for i in 0 ..< min(len(expected_hints), len(hints)) {
-		e, a := expected_hints[i], hints[i]
-		if e != a {
-			log.errorf("[%d]: Expected inlay hint\n%v, but received\n%v", i, e, a)
+	lines := strings.split_lines(src.main, context.temp_allocator)
+
+	get_source_line_with_hint :: proc(lines: []string, hint: server.InlayHint) -> string {
+		line := lines[hint.position.line] if hint.position.line >= 0 && hint.position.line < len(lines) else ""
+		if hint.position.character >= 0 && hint.position.character <= len(line) {
+			builder := strings.builder_make(context.temp_allocator)
+			strings.write_string(&builder, line[:hint.position.character])
+			strings.write_string(&builder, HINT_OPEN)
+			strings.write_string(&builder, hint.label)
+			strings.write_string(&builder, HINT_CLOSE)
+			strings.write_string(&builder, line[hint.position.character:])
+			return strings.to_string(builder)
+		}
+		return ""
+	}
+
+	for i in 0 ..< max(len(expected_hints), len(hints)) {
+		expected_text := "---"
+		actual_text   := "---"
+
+		if i < len(expected_hints) {
+			expected := expected_hints[i]
+			expected_line := get_source_line_with_hint(lines, expected)
+			expected_text = fmt.tprintf("\"%s\" at (%d, %d): \"%s\"",
+				expected.label, expected.position.line, expected.position.character, expected_line)
+		}
+
+		if i < len(hints) {
+			actual := hints[i]
+			actual_line := get_source_line_with_hint(lines, actual)
+			actual_text = fmt.tprintf("\"%s\" at (%d, %d): \"%s\"",
+				actual.label, actual.position.line, actual.position.character, actual_line)
+		}
+
+		if i >= len(expected_hints) {
+			log.errorf("[%d]: Unexpected inlay hint\nExpected: %s\nActual:   %s", i, expected_text, actual_text)
+		} else if i >= len(hints) {
+			log.errorf("[%d]: Missing inlay hint\nExpected: %s\nActual:   %s", i, expected_text, actual_text)
+		} else if expected_hints[i] != hints[i] {
+			log.errorf("[%d]: Inlay hint mismatch\nExpected: %s\nActual:   %s", i, expected_text, actual_text)
 		}
 	}
 }