Fix/improve union(enum) completions

zigtools · Feb 29, 2024 · 87bd0b1 · 87bd0b1
1 parent 7d6a9e2
commit 87bd0b1
Show file tree

Hide file tree

Showing 2 changed files with 110 additions and 5 deletions.
diff --git a/src/features/completions.zig b/src/features/completions.zig
@@ -701,10 +701,7 @@ fn completeDot(builder: *Builder, loc: offsets.Loc) error{OutOfMemory}!void {
             dot_context,
         );
         for (containers) |container| {
-            if (dot_context.likely == .enum_arg and !container.isEnumType()) continue;
-            if (dot_context.likely != .struct_field)
-                if (!container.isEnumType() and !container.isUnionType()) continue;
-            try collectContainerFields(builder, container);
+            try collectContainerFields(builder, dot_context.likely, container);
         }
     }
 
@@ -917,6 +914,7 @@ pub fn completionAtIndex(
                 else
                     .{ .TextEdit = .{ .newText = item.insertText orelse item.label, .range = insert_range } };
             }
+            item.insertText = null;
         }
     }
 
@@ -1227,8 +1225,10 @@ fn getSwitchOrStructInitContext(
 /// Given a Type that is a container, adds it's `.container_field*`s to completions
 pub fn collectContainerFields(
     builder: *Builder,
+    likely: EnumLiteralContext.Likely,
     container: Analyser.Type,
 ) error{OutOfMemory}!void {
+    const use_snippets = builder.server.config.enable_snippets and builder.server.client_capabilities.supports_snippets;
     const node_handle = switch (container.data) {
         .other => |n| n,
         else => return,
@@ -1240,10 +1240,25 @@ pub fn collectContainerFields(
     for (container_decl.ast.members) |member| {
         const field = handle.tree.fullContainerField(member) orelse continue;
         const name = handle.tree.tokenSlice(field.ast.main_token);
-        try builder.completions.append(builder.arena, .{
+        if (likely != .struct_field and !field.ast.tuple_like) {
+            try builder.completions.append(builder.arena, .{
+                .label = name,
+                .kind = if (field.ast.tuple_like) .EnumMember else .Field,
+                .detail = Analyser.getContainerFieldSignature(handle.tree, field),
+                .insertText = if (use_snippets)
+                    try std.fmt.allocPrint(builder.arena, "{{ .{s} = $1 }}$0", .{name})
+                else
+                    try std.fmt.allocPrint(builder.arena, "{{ .{s} = ", .{name}),
+                .insertTextFormat = if (use_snippets) .Snippet else .PlainText,
+            });
+        } else try builder.completions.append(builder.arena, .{
             .label = name,
             .kind = if (field.ast.tuple_like) .EnumMember else .Field,
             .detail = Analyser.getContainerFieldSignature(handle.tree, field),
+            .insertText = if (field.ast.tuple_like)
+                name
+            else
+                try std.fmt.allocPrint(builder.arena, "{s} = ", .{name}),
         });
     }
 }

diff --git a/tests/lsp_features/completion.zig b/tests/lsp_features/completion.zig
@@ -1269,6 +1269,26 @@ test "enum" {
     });
 }
 
+test "tagged union" {
+    try testCompletion(
+        \\const Birdie = enum {
+        \\    canary,
+        \\};
+        \\const Ue = union(enum) {
+        \\    alpha,
+        \\    beta: []const u8,
+        \\};
+        \\const S = struct{ foo: Ue };
+        \\test {
+        \\    const s = S{};
+        \\    s.foo = .<cursor>
+        \\}
+    , &.{
+        .{ .label = "alpha", .kind = .EnumMember },
+        .{ .label = "beta", .kind = .Field },
+    });
+}
+
 test "global enum set" {
     try testCompletion(
         \\const SomeError = error{ e };
@@ -1691,6 +1711,30 @@ test "struct init" {
         .{ .label = "beta", .kind = .Field, .detail = "u32" },
         .{ .label = "alpha", .kind = .Field, .detail = "*const S" },
     });
+    try testCompletion(
+        \\const S = struct {
+        \\    alpha: *const S,
+        \\    beta: u32,
+        \\    gamma: ?S = null,
+        \\};
+        \\test {
+        \\    const foo: S = undefined;
+        \\    foo.gamma = .<cursor>
+        \\}
+    , &.{
+        .{ .label = "alpha", .kind = .Field, .detail = "*const S" },
+        .{ .label = "beta", .kind = .Field, .detail = "u32" },
+        .{ .label = "gamma", .kind = .Field, .detail = "?S = null" },
+    });
+    try testCompletion(
+        \\const S = struct { alpha: u32 };
+        \\fn foo(s: S) void {}
+        \\test {
+        \\    foo(.<cursor>)
+        \\}
+    , &.{
+        .{ .label = "alpha", .kind = .Field, .detail = "u32" },
+    });
     try testCompletion(
         \\const S = struct { alpha: u32 };
         \\fn foo(s: *S) void { s = .{.<cursor>} }
@@ -2786,6 +2830,52 @@ test "insert replace behaviour - function with partial argument placeholders" {
     });
 }
 
+test "insert replace behaviour - struct literal" {
+    try testCompletionTextEdit(.{
+        .source =
+        \\const S = struct { alpha: u32 };
+        \\const foo: S = .<cursor>
+        ,
+        .label = "alpha",
+        .expected_insert_line = "const foo: S = .{ .alpha = ",
+        .expected_replace_line = "const foo: S = .{ .alpha = ",
+    });
+    try testCompletionTextEdit(.{
+        .source =
+        \\const S = struct { alpha: u32 };
+        \\const foo: S = .<cursor>
+        ,
+        .label = "alpha",
+        .expected_insert_line = "const foo: S = .{ .alpha = $1 }$0",
+        .expected_replace_line = "const foo: S = .{ .alpha = $1 }$0",
+        .enable_snippets = true,
+    });
+}
+
+test "insert replace behaviour - tagged union" {
+    try testCompletionTextEdit(.{
+        .source =
+        \\const Birdie = enum { canary };
+        \\const U = union(enum) { alpha: []const u8 };
+        \\const foo: U = .<cursor>
+        ,
+        .label = "alpha",
+        .expected_insert_line = "const foo: U = .{ .alpha = $1 }$0",
+        .expected_replace_line = "const foo: U = .{ .alpha = $1 }$0",
+        .enable_snippets = true,
+    });
+    try testCompletionTextEdit(.{
+        .source =
+        \\const Birdie = enum { canary };
+        \\const U = union(enum) { alpha: []const u8 };
+        \\const foo: U = .<cursor>
+        ,
+        .label = "alpha",
+        .expected_insert_line = "const foo: U = .{ .alpha = ",
+        .expected_replace_line = "const foo: U = .{ .alpha = ",
+    });
+}
+
 test "insert replace behaviour - doc test name" {
     if (true) return error.SkipZigTest; // TODO
     try testCompletionTextEdit(.{