Added label support

src/analysis.zig

@@ -889,6 +889,7 @@ pub const PositionContext = union(enum) {
     var_access: SourceRange,
     global_error_set,
     enum_literal,
+    label,
     other,
     empty,
 
@@ -900,6 +901,7 @@ pub const PositionContext = union(enum) {
             .field_access => |r| r,
             .var_access => |r| r,
             .enum_literal => null,
+            .label => null,
             .other => null,
             .empty => null,
             .global_error_set => null,
@@ -980,6 +982,7 @@ pub fn documentPositionContext(allocator: *std.mem.Allocator, document: types.Te
                     .field_access = tokenRangeAppend(curr_ctx.ctx.range().?, tok),
                 },
             },
+            .Colon => curr_ctx.ctx = .label,
             .QuestionMark => switch (curr_ctx.ctx) {
                 .field_access => {},
                 else => curr_ctx.ctx = .empty,
@@ -1112,6 +1115,7 @@ pub const Declaration = union(enum) {
         node: *ast.Node.PointerPayload,
         items: []const *ast.Node,
     },
+    label_decl: *ast.Node, // .id is While, For or Block (firstToken will be the label)
 };
 
 pub const DeclWithHandle = struct {
@@ -1129,6 +1133,7 @@ pub const DeclWithHandle = struct {
             .pointer_payload => |pp| tree.tokenLocation(0, pp.node.value_symbol.firstToken()),
             .array_payload => |ap| tree.tokenLocation(0, ap.identifier.firstToken()),
             .switch_payload => |sp| tree.tokenLocation(0, sp.node.value_symbol.firstToken()),
+            .label_decl => |ld| tree.tokenLocation(0, ld.firstToken()),
         };
     }
 
@@ -1174,6 +1179,7 @@ pub const DeclWithHandle = struct {
                 .Single,
                 bound_type_params,
             ),
+            .label_decl => return null,
             // TODO Resolve switch payload types
             .switch_payload => |pay| return null,
         };
@@ -1218,6 +1224,7 @@ pub fn iterateSymbolsContainer(
         var decl_it = container_scope.decls.iterator();
         while (decl_it.next()) |entry| {
             if (!include_fields and entry.value == .ast_node and entry.value.ast_node.id == .ContainerField) continue;
+            if (entry.value == .label_decl) continue;
             const decl = DeclWithHandle{ .decl = &entry.value, .handle = handle };
             if (handle != orig_handle and !decl.isPublic()) continue;
             try callback(context, decl);
@@ -1233,6 +1240,27 @@ pub fn iterateSymbolsContainer(
     std.debug.warn("Did not find container scope when iterating container {} (name: {})\n", .{ container, getDeclName(handle.tree, container) });
 }
 
+pub fn iterateLabels(
+    handle: *DocumentStore.Handle,
+    source_index: usize,
+    comptime callback: var,
+    context: var,
+) error{OutOfMemory}!void {
+    for (handle.document_scope.scopes) |scope| {
+        if (source_index >= scope.range.start and source_index < scope.range.end) {
+            var decl_it = scope.decls.iterator();
+            while (decl_it.next()) |entry| {
+                switch (entry.value) {
+                    .label_decl => {},
+                    else => continue,
+                }
+                try callback(context, DeclWithHandle{ .decl = &entry.value, .handle = handle });
+            }
+        }
+        if (scope.range.start >= source_index) return;
+    }
+}
+
 pub fn iterateSymbolsGlobal(
     store: *DocumentStore,
     arena: *std.heap.ArenaAllocator,
@@ -1246,6 +1274,7 @@ pub fn iterateSymbolsGlobal(
             var decl_it = scope.decls.iterator();
             while (decl_it.next()) |entry| {
                 if (entry.value == .ast_node and entry.value.ast_node.id == .ContainerField) continue;
+                if (entry.value == .label_decl) continue;
                 try callback(context, DeclWithHandle{ .decl = &entry.value, .handle = handle });
             }
 
@@ -1294,6 +1323,29 @@ fn resolveUse(
     return null;
 }
 
+pub fn lookupLabel(
+    handle: *DocumentStore.Handle,
+    symbol: []const u8,
+    source_index: usize,
+) error{OutOfMemory}!?DeclWithHandle {
+    for (handle.document_scope.scopes) |scope| {
+        if (source_index >= scope.range.start and source_index < scope.range.end) {
+            if (scope.decls.get(symbol)) |candidate| {
+                switch (candidate.value) {
+                    .label_decl => {},
+                    else => continue,
+                }
+                return DeclWithHandle{
+                    .decl = &candidate.value,
+                    .handle = handle,
+                };
+            }
+        }
+        if (scope.range.start > source_index) return null;
+    }
+    return null;
+}
+
 pub fn lookupSymbolGlobal(
     store: *DocumentStore,
     arena: *std.heap.ArenaAllocator,
@@ -1308,6 +1360,7 @@ pub fn lookupSymbolGlobal(
                     .ast_node => |node| {
                         if (node.id == .ContainerField) continue;
                     },
+                    .label_decl => continue,
                     else => {},
                 }
                 return DeclWithHandle{
@@ -1341,6 +1394,7 @@ pub fn lookupSymbolContainer(
                 .ast_node => |node| {
                     if (node.id == .ContainerField and !accept_fields) return null;
                 },
+                .label_decl => unreachable,
                 else => {},
             }
             return DeclWithHandle{ .decl = &candidate.value, .handle = handle };
@@ -1525,6 +1579,27 @@ fn makeScopeInternal(
             return try makeScopeInternal(allocator, scopes, tree, node.cast(ast.Node.TestDecl).?.body_node);
         },
         .Block => {
+            const block = node.cast(ast.Node.Block).?;
+            if (block.label) |label| {
+                std.debug.assert(tree.token_ids[label] == .Identifier);
+                var scope = try scopes.addOne();
+                scope.* = .{
+                    .range = .{
+                        .start = tree.token_locs[block.lbrace].start,
+                        .end = tree.token_locs[block.rbrace].end,
+                    },
+                    .decls = std.StringHashMap(Declaration).init(allocator),
+                    .uses = &[0]*ast.Node.Use{},
+                    .tests = &[0]*ast.Node{},
+                    .data = .other,
+                };
+                errdefer scope.decls.deinit();
+
+                try scope.decls.putNoClobber(tree.tokenSlice(label), .{
+                    .label_decl = node,
+                });
+            }
+
             (try scopes.addOne()).* = .{
                 .range = nodeSourceRange(tree, node),
                 .decls = std.StringHashMap(Declaration).init(allocator),
@@ -1618,6 +1693,26 @@ fn makeScopeInternal(
         },
         .While => {
             const while_node = node.cast(ast.Node.While).?;
+            if (while_node.label) |label| {
+                std.debug.assert(tree.token_ids[label] == .Identifier);
+                var scope = try scopes.addOne();
+                scope.* = .{
+                    .range = .{
+                        .start = tree.token_locs[while_node.while_token].start,
+                        .end = tree.token_locs[while_node.lastToken()].end,
+                    },
+                    .decls = std.StringHashMap(Declaration).init(allocator),
+                    .uses = &[0]*ast.Node.Use{},
+                    .tests = &[0]*ast.Node{},
+                    .data = .other,
+                };
+                errdefer scope.decls.deinit();
+
+                try scope.decls.putNoClobber(tree.tokenSlice(label), .{
+                    .label_decl = node,
+                });
+            }
+
             if (while_node.payload) |payload| {
                 std.debug.assert(payload.id == .PointerPayload);
                 var scope = try scopes.addOne();
@@ -1671,6 +1766,26 @@ fn makeScopeInternal(
         },
         .For => {
             const for_node = node.cast(ast.Node.For).?;
+            if (for_node.label) |label| {
+                std.debug.assert(tree.token_ids[label] == .Identifier);
+                var scope = try scopes.addOne();
+                scope.* = .{
+                    .range = .{
+                        .start = tree.token_locs[for_node.for_token].start,
+                        .end = tree.token_locs[for_node.lastToken()].end,
+                    },
+                    .decls = std.StringHashMap(Declaration).init(allocator),
+                    .uses = &[0]*ast.Node.Use{},
+                    .tests = &[0]*ast.Node{},
+                    .data = .other,
+                };
+                errdefer scope.decls.deinit();
+
+                try scope.decls.putNoClobber(tree.tokenSlice(label), .{
+                    .label_decl = node,
+                });
+            }
+
             std.debug.assert(for_node.payload.id == .PointerIndexPayload);
             const ptr_idx_payload = for_node.payload.cast(ast.Node.PointerIndexPayload).?;
             std.debug.assert(ptr_idx_payload.value_symbol.id == .Identifier);

src/main.zig

@@ -482,6 +482,13 @@ fn hoverSymbol(id: types.RequestId, arena: *std.heap.ArenaAllocator, decl_handle
             try std.fmt.allocPrint(&arena.allocator, "```zig\n{}\n```", .{handle.tree.tokenSlice(payload.node.value_symbol.firstToken())})
         else
             try std.fmt.allocPrint(&arena.allocator, "{}", .{handle.tree.tokenSlice(payload.node.value_symbol.firstToken())}),
+        .label_decl => |label_decl| block: {
+            const source = handle.tree.source[handle.tree.token_locs[label_decl.firstToken()].start..handle.tree.token_locs[label_decl.lastToken()].end];
+            break :block if (hover_kind == .Markdown)
+                try std.fmt.allocPrint(&arena.allocator, "```zig\n{}\n```", .{source})
+            else
+                try std.fmt.allocPrint(&arena.allocator, "```{}```", .{source});
+        },
     };
 
     try send(types.Response{
@@ -494,6 +501,13 @@ fn hoverSymbol(id: types.RequestId, arena: *std.heap.ArenaAllocator, decl_handle
     });
 }
 
+fn getLabelGlobal(pos_index: usize, handle: *DocumentStore.Handle) !?analysis.DeclWithHandle {
+    const name = identifierFromPosition(pos_index, handle.*);
+    if (name.len == 0) return null;
+
+    return try analysis.lookupLabel(handle, name, pos_index);
+}
+
 fn getSymbolGlobal(arena: *std.heap.ArenaAllocator, pos_index: usize, handle: *DocumentStore.Handle) !?analysis.DeclWithHandle {
     const name = identifierFromPosition(pos_index, handle.*);
     if (name.len == 0) return null;
@@ -501,6 +515,14 @@ fn getSymbolGlobal(arena: *std.heap.ArenaAllocator, pos_index: usize, handle: *D
     return try analysis.lookupSymbolGlobal(&document_store, arena, handle, name, pos_index);
 }
 
+fn gotoDefinitionLabel(id: types.RequestId, pos_index: usize, handle: *DocumentStore.Handle, config: Config) !void {
+    var arena = std.heap.ArenaAllocator.init(allocator);
+    defer arena.deinit();
+
+    const decl = (try getLabelGlobal(pos_index, handle)) orelse return try respondGeneric(id, null_result_response);
+    return try gotoDefinitionSymbol(id, &arena, decl);
+}
+
 fn gotoDefinitionGlobal(id: types.RequestId, pos_index: usize, handle: *DocumentStore.Handle, config: Config) !void {
     var arena = std.heap.ArenaAllocator.init(allocator);
     defer arena.deinit();
@@ -509,6 +531,14 @@ fn gotoDefinitionGlobal(id: types.RequestId, pos_index: usize, handle: *Document
     return try gotoDefinitionSymbol(id, &arena, decl);
 }
 
+fn hoverDefinitionLabel(id: types.RequestId, pos_index: usize, handle: *DocumentStore.Handle, config: Config) !void {
+    var arena = std.heap.ArenaAllocator.init(allocator);
+    defer arena.deinit();
+
+    const decl = (try getLabelGlobal(pos_index, handle)) orelse return try respondGeneric(id, null_result_response);
+    return try hoverSymbol(id, &arena, decl);
+}
+
 fn hoverDefinitionGlobal(id: types.RequestId, pos_index: usize, handle: *DocumentStore.Handle, config: Config) !void {
     var arena = std.heap.ArenaAllocator.init(allocator);
     defer arena.deinit();
@@ -639,9 +669,41 @@ fn declToCompletion(context: DeclToCompletionContext, decl_handle: analysis.Decl
                 .kind = .Variable,
             });
         },
+        .label_decl => |label_decl| {
+            try context.completions.append(.{
+                .label = tree.tokenSlice(label_decl.firstToken()),
+                .kind = .Variable,
+            });
+        },
     }
 }
 
+fn completeLabel(id: types.RequestId, pos_index: usize, handle: *DocumentStore.Handle, config: Config) !void {
+    // We use a local arena allocator to deallocate all temporary data without iterating
+    var arena = std.heap.ArenaAllocator.init(allocator);
+    var completions = std.ArrayList(types.CompletionItem).init(&arena.allocator);
+    // Deallocate all temporary data.
+    defer arena.deinit();
+
+    const context = DeclToCompletionContext{
+        .completions = &completions,
+        .config = &config,
+        .arena = &arena,
+        .orig_handle = handle,
+    };
+    try analysis.iterateLabels(handle, pos_index, declToCompletion, context);
+
+    try send(types.Response{
+        .id = id,
+        .result = .{
+            .CompletionList = .{
+                .isIncomplete = false,
+                .items = completions.items,
+            },
+        },
+    });
+}
+
 fn completeGlobal(id: types.RequestId, pos_index: usize, handle: *DocumentStore.Handle, config: Config) !void {
     // We use a local arena allocator to deallocate all temporary data without iterating
     var arena = std.heap.ArenaAllocator.init(allocator);
@@ -1017,6 +1079,7 @@ fn processJsonRpc(parser: *std.json.Parser, json: []const u8, config: Config) !v
                         },
                     },
                 }),
+                .label => try completeLabel(id, pos_index, handle, this_config),
                 else => try respondGeneric(id, no_completions_response),
             }
         } else {
@@ -1051,20 +1114,10 @@ fn processJsonRpc(parser: *std.json.Parser, json: []const u8, config: Config) !v
             const pos_context = try analysis.documentPositionContext(allocator, handle.document, pos);
 
             switch (pos_context) {
-                .var_access => try gotoDefinitionGlobal(
+                .var_access => try gotoDefinitionGlobal(id, pos_index, handle, configFromUriOr(uri, config)),
-                    id,
+                .field_access => |range| try gotoDefinitionFieldAccess(id, handle, pos, range, configFromUriOr(uri, config)),
-                    pos_index,
-                    handle,
-                    configFromUriOr(uri, config),
-                ),
-                .field_access => |range| try gotoDefinitionFieldAccess(
-                    id,
-                    handle,
-                    pos,
-                    range,
-                    configFromUriOr(uri, config),
-                ),
                 .string_literal => try gotoDefinitionString(id, pos_index, handle, config),
+                .label => try gotoDefinitionLabel(id, pos_index, handle, configFromUriOr(uri, config)),
                 else => try respondGeneric(id, null_result_response),
             }
         } else {
@@ -1090,19 +1143,9 @@ fn processJsonRpc(parser: *std.json.Parser, json: []const u8, config: Config) !v
             const pos_context = try analysis.documentPositionContext(allocator, handle.document, pos);
 
             switch (pos_context) {
-                .var_access => try hoverDefinitionGlobal(
+                .var_access => try hoverDefinitionGlobal(id, pos_index, handle, configFromUriOr(uri, config)),
-                    id,
+                .field_access => |range| try hoverDefinitionFieldAccess(id, handle, pos, range, configFromUriOr(uri, config)),
-                    pos_index,
+                .label => try hoverDefinitionLabel(id, pos_index, handle, configFromUriOr(uri, config)),
-                    handle,
-                    configFromUriOr(uri, config),
-                ),
-                .field_access => |range| try hoverDefinitionFieldAccess(
-                    id,
-                    handle,
-                    pos,
-                    range,
-                    configFromUriOr(uri, config),
-                ),
                 else => try respondGeneric(id, null_result_response),
             }
         } else {