fix completion text edit on escaped identifiers

src/Server.zig

@@ -1459,7 +1459,7 @@ fn prepareRenameHandler(server: *Server, arena: std.mem.Allocator, request: type
     };
     const handle = server.document_store.getHandle(document_uri) orelse return null;
     const source_index = offsets.positionToIndex(handle.tree.source, request.position, server.offset_encoding);
-    const name_loc = Analyser.identifierLocFromIndex(&handle.tree, source_index) orelse return null;
+    const name_loc = offsets.identifierLocFromIndex(&handle.tree, source_index) orelse return null;
     const name = offsets.locToSlice(handle.tree.source, name_loc);
     return .{
         .prepare_rename_placeholder = .{

src/features/completions.zig

@@ -473,26 +473,58 @@ fn populateSnippedCompletions(builder: *Builder, kind: enum { generic, top_level
     }
 }
 
+fn prepareCompletionLoc(tree: *const Ast, source_index: usize) offsets.Loc {
+    const fallback_loc: offsets.Loc = .{ .start = source_index, .end = source_index };
+    const token = switch (offsets.sourceIndexToTokenIndex(tree, source_index)) {
+        .none => return fallback_loc,
+        .one => |token| token,
+        .between => |data| data.left,
+    };
+    switch (tree.tokenTag(token)) {
+        .identifier, .builtin => |tag| {
+            if (tag == .builtin and tree.tokenStart(token) == source_index) return fallback_loc;
+            const token_loc = offsets.tokenToLoc(tree, token);
+            std.debug.assert(token_loc.start <= source_index and source_index <= token_loc.end);
+            return offsets.identifierIndexToLoc(tree.source, token_loc.start, if (tag == .builtin) .name else .full);
+        },
+        else => {
+            const token_start = tree.tokenStart(token);
+
+            var start: usize, var end: usize = start: {
+                if (std.mem.startsWith(u8, tree.source[token_start..], "@\"")) {
+                    break :start .{ token_start, token_start + 2 };
+                } else if (std.mem.startsWith(u8, tree.source[token_start..], "@") or std.mem.startsWith(u8, tree.source[token_start..], ".")) {
+                    if (token_start + 1 < source_index) return fallback_loc;
+                    break :start .{ token_start + 1, token_start + 1 };
+                } else {
+                    break :start .{ token_start, token_start };
+                }
+            };
+            start = @min(start, source_index);
+            end = @max(end, source_index);
+
+            while (end < tree.source.len and offsets.isSymbolChar(tree.source[end])) {
+                end += 1;
+            }
+
+            return .{ .start = start, .end = end };
+        },
+    }
+}
+
 const FunctionCompletionFormat = enum { snippet, only_name };
 const PrepareFunctionCompletionResult = struct { types.Range, types.Range, FunctionCompletionFormat };
 
 fn prepareFunctionCompletion(builder: *Builder) PrepareFunctionCompletionResult {
     if (builder.cached_prepare_function_completion_result) |result| return result;
 
-    const source = builder.orig_handle.tree.source;
-
-    var start_index = builder.source_index;
-    while (start_index > 0 and Analyser.isSymbolChar(source[start_index - 1])) {
-        start_index -= 1;
-    }
-
-    var end_index = builder.source_index;
-    while (end_index < source.len and Analyser.isSymbolChar(source[end_index])) {
-        end_index += 1;
-    }
+    const tree = &builder.orig_handle.tree;
+    const source = tree.source;
+    const source_index = builder.source_index;
 
-    var insert_loc: offsets.Loc = .{ .start = start_index, .end = builder.source_index };
-    var replace_loc: offsets.Loc = .{ .start = start_index, .end = end_index };
+    const identifier_loc = prepareCompletionLoc(tree, source_index);
+    var insert_loc: offsets.Loc = .{ .start = identifier_loc.start, .end = source_index };
+    var replace_loc: offsets.Loc = .{ .start = identifier_loc.start, .end = identifier_loc.end };
 
     var format: FunctionCompletionFormat = .only_name;
 
@@ -505,7 +537,7 @@ fn prepareFunctionCompletion(builder: *Builder) PrepareFunctionCompletionResult
         format = .snippet;
     } else if (insert_can_be_snippet or replace_can_be_snippet) {
         // snippet completions would be possible but insert and replace would need different `newText`
-    } else if (builder.use_snippets and !std.mem.startsWith(u8, source[end_index..], "(")) {
+    } else if (builder.use_snippets and !std.mem.startsWith(u8, source[identifier_loc.end..], "(")) {
         format = .snippet;
     }
 
@@ -949,18 +981,9 @@ pub fn completionAtIndex(
     const completions = builder.completions.items;
     if (completions.len == 0) return null;
 
-    var start_index = source_index;
-    while (start_index > 0 and Analyser.isSymbolChar(source[start_index - 1])) {
-        start_index -= 1;
-    }
-
-    var end_index = source_index;
-    while (end_index < source.len and Analyser.isSymbolChar(source[end_index])) {
-        end_index += 1;
-    }
-
-    const insert_range = offsets.locToRange(source, .{ .start = start_index, .end = source_index }, server.offset_encoding);
-    const replace_range = offsets.locToRange(source, .{ .start = start_index, .end = end_index }, server.offset_encoding);
+    const identifier_loc = prepareCompletionLoc(&handle.tree, source_index);
+    const insert_range = offsets.locToRange(source, .{ .start = identifier_loc.start, .end = source_index }, server.offset_encoding);
+    const replace_range = offsets.locToRange(source, .{ .start = identifier_loc.start, .end = identifier_loc.end }, server.offset_encoding);
 
     for (completions) |*item| {
         if (item.textEdit == null) {
@@ -1741,7 +1764,7 @@ fn collectFieldAccessContainerNodes(
     // inconsistent at returning name_loc for methods, ie
     // `abc.method() == .` => fails, `abc.method(.{}){.}` => ok
     // it also fails for `abc.xyz.*` ... currently we take advantage of this quirk
-    const name_loc = Analyser.identifierLocFromIndex(&handle.tree, loc.end) orelse {
+    const name_loc = offsets.identifierLocFromIndex(&handle.tree, loc.end) orelse {
         const result = try analyser.getFieldAccessType(handle, loc.end, loc) orelse return;
         const container = try analyser.resolveDerefType(result) orelse result;
         if (try analyser.resolveUnwrapErrorUnionType(container, .payload)) |unwrapped| {

src/features/goto.zig

@@ -79,7 +79,7 @@ fn gotoDefinitionLabel(
     const tracy_zone = tracy.trace(@src());
     defer tracy_zone.end();
 
-    const name_loc = Analyser.identifierLocFromIndex(&handle.tree, pos_index) orelse return null;
+    const name_loc = offsets.identifierLocFromIndex(&handle.tree, pos_index) orelse return null;
     const name = offsets.locToSlice(handle.tree.source, name_loc);
     const decl = (try Analyser.lookupLabel(handle, name, pos_index)) orelse return null;
     return try gotoDefinitionSymbol(analyser, offsets.locToRange(handle.tree.source, loc, offset_encoding), decl, kind, offset_encoding);
@@ -95,7 +95,7 @@ fn gotoDefinitionGlobal(
     const tracy_zone = tracy.trace(@src());
     defer tracy_zone.end();
 
-    const name_token, const name_loc = Analyser.identifierTokenAndLocFromIndex(&handle.tree, pos_index) orelse return null;
+    const name_token, const name_loc = offsets.identifierTokenAndLocFromIndex(&handle.tree, pos_index) orelse return null;
     const name = offsets.locToSlice(handle.tree.source, name_loc);
     const decl = (try analyser.lookupSymbolGlobal(handle, name, pos_index)) orelse return null;
     return try gotoDefinitionSymbol(analyser, offsets.tokenToRange(&handle.tree, name_token, offset_encoding), decl, kind, offset_encoding);
@@ -138,7 +138,7 @@ fn gotoDefinitionEnumLiteral(
     const tracy_zone = tracy.trace(@src());
     defer tracy_zone.end();
 
-    const name_token, const name_loc = Analyser.identifierTokenAndLocFromIndex(&handle.tree, source_index) orelse {
+    const name_token, const name_loc = offsets.identifierTokenAndLocFromIndex(&handle.tree, source_index) orelse {
         return gotoDefinitionStructInit(analyser, handle, source_index, kind, offset_encoding);
     };
     const name = offsets.locToSlice(handle.tree.source, name_loc);
@@ -209,7 +209,7 @@ fn gotoDefinitionFieldAccess(
     const tracy_zone = tracy.trace(@src());
     defer tracy_zone.end();
 
-    const name_token, const name_loc = Analyser.identifierTokenAndLocFromIndex(&handle.tree, source_index) orelse return null;
+    const name_token, const name_loc = offsets.identifierTokenAndLocFromIndex(&handle.tree, source_index) orelse return null;
     const name = offsets.locToSlice(handle.tree.source, name_loc);
     const held_loc = offsets.locMerge(loc, name_loc);
     const accesses = (try analyser.getSymbolFieldAccesses(arena, handle, source_index, held_loc, name)) orelse return null;

src/features/hover.zig

@@ -296,7 +296,7 @@ fn hoverDefinitionGlobal(
     const tracy_zone = tracy.trace(@src());
     defer tracy_zone.end();
 
-    const name_token, const name_loc = Analyser.identifierTokenAndLocFromIndex(&handle.tree, pos_index) orelse return null;
+    const name_token, const name_loc = offsets.identifierTokenAndLocFromIndex(&handle.tree, pos_index) orelse return null;
     const name = offsets.locToSlice(handle.tree.source, name_loc);
     const hover_text = blk: {
         const is_escaped_identifier = handle.tree.source[handle.tree.tokenStart(name_token)] == '@';
@@ -372,7 +372,7 @@ fn hoverDefinitionEnumLiteral(
     const tracy_zone = tracy.trace(@src());
     defer tracy_zone.end();
 
-    const name_token, const name_loc = Analyser.identifierTokenAndLocFromIndex(&handle.tree, source_index) orelse {
+    const name_token, const name_loc = offsets.identifierTokenAndLocFromIndex(&handle.tree, source_index) orelse {
         return try hoverDefinitionStructInit(analyser, arena, handle, source_index, markup_kind, offset_encoding);
     };
     const name = offsets.locToSlice(handle.tree.source, name_loc);

src/features/references.zig

@@ -679,7 +679,7 @@ pub fn referencesHandler(server: *Server, arena: std.mem.Allocator, request: Gen
             );
         }
 
-        const name_loc = Analyser.identifierLocFromIndex(&handle.tree, source_index) orelse return null;
+        const name_loc = offsets.identifierLocFromIndex(&handle.tree, source_index) orelse return null;
         const name = offsets.locToSlice(handle.tree.source, name_loc);
 
         const decl = switch (pos_context) {

src/features/signature_help.zig

@@ -287,7 +287,7 @@ pub fn getSignatureInfo(
                     );
                 }
 
-                const name_loc = Analyser.identifierLocFromIndex(&handle.tree, loc.end - 1) orelse {
+                const name_loc = offsets.identifierLocFromIndex(&handle.tree, loc.end - 1) orelse {
                     try symbol_stack.append(arena, .l_paren);
                     continue;
                 };

src/offsets.zig

@@ -227,7 +227,9 @@ pub const IdentifierIndexRange = enum {
     full,
 };
 
-/// Support formats:
+/// The source index must be at the start of the valid identifier.
+///
+/// Supported formats:
 /// - `foo`
 /// - `@"foo"`
 /// - `@foo`
@@ -251,12 +253,7 @@ pub fn identifierIndexToLoc(text: [:0]const u8, source_index: usize, range: Iden
     } else {
         const start: usize = source_index + @intFromBool(text[source_index] == '@');
         var index = start;
-        while (true) : (index += 1) {
-            switch (text[index]) {
-                'a'...'z', 'A'...'Z', '_', '0'...'9' => {},
-                else => break,
-            }
-        }
+        while (isSymbolChar(text[index])) : (index += 1) {}
         return .{ .start = if (range == .full) source_index else start, .end = index };
     }
 }
@@ -304,6 +301,93 @@ pub fn identifierTokenToNameSlice(tree: *const Ast, identifier_token: Ast.TokenI
     return locToSlice(tree.source, identifierTokenToNameLoc(tree, identifier_token));
 }
 
+/// See `identifierTokenAndLocFromIndex`.
+pub fn identifierLocFromIndex(tree: *const Ast, source_index: usize) ?Loc {
+    _, const loc = identifierTokenAndLocFromIndex(tree, source_index) orelse return null;
+    return loc;
+}
+
+/// Returns the source location of `foo` if the source index is on a valid identifier.
+///
+/// Supported formats:
+/// - `foo`    (identifier)
+/// - `@"foo"` (escaped identifier)
+/// - `@foo`   (builtin)
+pub fn identifierTokenAndLocFromIndex(tree: *const Ast, source_index: usize) ?struct { Ast.TokenIndex, Loc } {
+    const token = sourceIndexToTokenIndex(tree, source_index).pickPreferred(&.{ .identifier, .builtin }, tree) orelse return null;
+    switch (tree.tokenTag(token)) {
+        .identifier, .builtin => {},
+        else => return null,
+    }
+    const token_loc = tokenToLoc(tree, token);
+    std.debug.assert(token_loc.start <= source_index and source_index <= token_loc.end);
+    return .{ token, identifierIndexToLoc(tree.source, token_loc.start, .name) };
+}
+
+test identifierLocFromIndex {
+    var tree = try Ast.parse(std.testing.allocator,
+        \\ name  @builtin  @"escaped"  @"s p a c e"  end
+    , .zig);
+    defer tree.deinit(std.testing.allocator);
+
+    try std.testing.expectEqualSlices(
+        std.zig.Token.Tag,
+        &.{ .identifier, .builtin, .identifier, .identifier, .identifier, .eof },
+        tree.tokens.items(.tag),
+    );
+
+    {
+        const expected_loc: Loc = .{ .start = 1, .end = 5 };
+        std.debug.assert(std.mem.eql(u8, "name", locToSlice(tree.source, expected_loc)));
+
+        try std.testing.expectEqual(expected_loc, identifierLocFromIndex(&tree, 1));
+        try std.testing.expectEqual(expected_loc, identifierLocFromIndex(&tree, 2));
+        try std.testing.expectEqual(expected_loc, identifierLocFromIndex(&tree, 5));
+    }
+
+    {
+        const expected_loc: Loc = .{ .start = 8, .end = 15 };
+        std.debug.assert(std.mem.eql(u8, "builtin", locToSlice(tree.source, expected_loc)));
+
+        try std.testing.expectEqual(@as(?Loc, null), identifierLocFromIndex(&tree, 6));
+        try std.testing.expectEqual(expected_loc, identifierLocFromIndex(&tree, 7));
+        try std.testing.expectEqual(expected_loc, identifierLocFromIndex(&tree, 8));
+        try std.testing.expectEqual(expected_loc, identifierLocFromIndex(&tree, 11));
+        try std.testing.expectEqual(expected_loc, identifierLocFromIndex(&tree, 15));
+        try std.testing.expectEqual(@as(?Loc, null), identifierLocFromIndex(&tree, 16));
+    }
+
+    {
+        const expected_loc: Loc = .{ .start = 19, .end = 26 };
+        std.debug.assert(std.mem.eql(u8, "escaped", locToSlice(tree.source, expected_loc)));
+
+        try std.testing.expectEqual(@as(?Loc, null), identifierLocFromIndex(&tree, 16));
+        try std.testing.expectEqual(expected_loc, identifierLocFromIndex(&tree, 17));
+        try std.testing.expectEqual(expected_loc, identifierLocFromIndex(&tree, 18));
+        try std.testing.expectEqual(expected_loc, identifierLocFromIndex(&tree, 19));
+        try std.testing.expectEqual(expected_loc, identifierLocFromIndex(&tree, 23));
+        try std.testing.expectEqual(expected_loc, identifierLocFromIndex(&tree, 27));
+        try std.testing.expectEqual(@as(?Loc, null), identifierLocFromIndex(&tree, 28));
+    }
+
+    {
+        const expected_loc: Loc = .{ .start = 43, .end = 46 };
+        std.debug.assert(std.mem.eql(u8, "end", locToSlice(tree.source, expected_loc)));
+
+        try std.testing.expectEqual(@as(?Loc, null), identifierLocFromIndex(&tree, 42));
+        try std.testing.expectEqual(@as(?Loc, expected_loc), identifierLocFromIndex(&tree, 43));
+        try std.testing.expectEqual(@as(?Loc, expected_loc), identifierLocFromIndex(&tree, 45));
+        try std.testing.expectEqual(@as(?Loc, expected_loc), identifierLocFromIndex(&tree, 46));
+    }
+}
+
+pub fn isSymbolChar(char: u8) bool {
+    return switch (char) {
+        'a'...'z', 'A'...'Z', '_', '0'...'9' => true,
+        else => false,
+    };
+}
+
 pub fn tokensToLoc(tree: *const Ast, first_token: Ast.TokenIndex, last_token: Ast.TokenIndex) Loc {
     return .{ .start = tree.tokenStart(first_token), .end = tokenToLoc(tree, last_token).end };
 }

tests/lsp_features/completion.zig

@@ -4097,6 +4097,36 @@ test "insert replace behaviour - function alias" {
     });
 }
 
+test "insert replace behaviour - escaped identifier" {
+    try testCompletionTextEdit(.{
+        .source =
+        \\const @"foo bar" = 5;
+        \\const foo = @"foo<cursor>
+        ,
+        .label = "@\"foo bar\"",
+        .expected_insert_line = "const foo = @\"foo bar\"",
+        .expected_replace_line = "const foo = @\"foo bar\"",
+    });
+    try testCompletionTextEdit(.{
+        .source =
+        \\fn @"foo bar"() void {}
+        \\const foo = <cursor>@"foo
+        ,
+        .label = "@\"foo bar\"",
+        .expected_insert_line = "const foo = @\"foo bar\"@\"foo",
+        .expected_replace_line = "const foo = @\"foo bar\"",
+    });
+    try testCompletionTextEdit(.{
+        .source =
+        \\fn @"foo bar"() void {}
+        \\const foo = @"foo <cursor>
+        ,
+        .label = "@\"foo bar\"",
+        .expected_insert_line = "const foo = @\"foo bar\"",
+        .expected_replace_line = "const foo = @\"foo bar\"",
+    });
+}
+
 test "insert replace behaviour - decl literal function" {
     try testCompletionTextEdit(.{
         .source =
@@ -4681,7 +4711,6 @@ fn testCompletionTextEdit(
     ctx.server.config_manager.config.enable_snippets = options.enable_snippets;
 
     const test_uri = try ctx.addDocument(.{ .source = text });
-    const handle = ctx.server.document_store.getHandle(test_uri).?;
 
     const cursor_position = offsets.indexToPosition(options.source, cursor_idx, ctx.server.offset_encoding);
     const params: types.completion.Params = .{
@@ -4705,21 +4734,7 @@ fn testCompletionTextEdit(
 
         const TextEditOrInsertReplace = std.meta.Child(@TypeOf(completion_item.textEdit));
 
-        const text_edit_or_insert_replace: TextEditOrInsertReplace = completion_item.textEdit orelse blk: {
-            var start_index: usize = cursor_idx;
-            while (start_index > 0 and zls.Analyser.isSymbolChar(handle.tree.source[start_index - 1])) {
-                start_index -= 1;
-            }
-
-            const start_position = offsets.indexToPosition(text, start_index, ctx.server.offset_encoding);
-
-            break :blk .{
-                .text_edit = .{
-                    .newText = completion_item.insertText orelse completion_item.label,
-                    .range = .{ .start = start_position, .end = cursor_position },
-                },
-            };
-        };
+        const text_edit_or_insert_replace: TextEditOrInsertReplace = completion_item.textEdit.?;
 
         switch (text_edit_or_insert_replace) {
             .text_edit => |text_edit| {