zig fmt: support aligned ptr with bit fields

4 files changed, 103 insertions, 24 deletions

diff --git a/std/zig/ast.zig b/std/zig/ast.zig
index 1f15046a79..6c848b4a54 100644
--- a/std/zig/ast.zig
+++ b/std/zig/ast.zig
@@ -98,6 +98,7 @@ pub const Error = union(enum) {
     UnattachedDocComment: UnattachedDocComment,
     ExpectedEqOrSemi: ExpectedEqOrSemi,
     ExpectedSemiOrLBrace: ExpectedSemiOrLBrace,
+    ExpectedColonOrRParen: ExpectedColonOrRParen,
     ExpectedLabelable: ExpectedLabelable,
     ExpectedInlinable: ExpectedInlinable,
     ExpectedAsmOutputReturnOrType: ExpectedAsmOutputReturnOrType,
@@ -120,6 +121,7 @@ pub const Error = union(enum) {
             @TagType(Error).UnattachedDocComment => |*x| return x.render(tokens, stream),
             @TagType(Error).ExpectedEqOrSemi => |*x| return x.render(tokens, stream),
             @TagType(Error).ExpectedSemiOrLBrace => |*x| return x.render(tokens, stream),
+            @TagType(Error).ExpectedColonOrRParen => |*x| return x.render(tokens, stream),
             @TagType(Error).ExpectedLabelable => |*x| return x.render(tokens, stream),
             @TagType(Error).ExpectedInlinable => |*x| return x.render(tokens, stream),
             @TagType(Error).ExpectedAsmOutputReturnOrType => |*x| return x.render(tokens, stream),
@@ -144,6 +146,7 @@ pub const Error = union(enum) {
             @TagType(Error).UnattachedDocComment => |x| return x.token,
             @TagType(Error).ExpectedEqOrSemi => |x| return x.token,
             @TagType(Error).ExpectedSemiOrLBrace => |x| return x.token,
+            @TagType(Error).ExpectedColonOrRParen => |x| return x.token,
             @TagType(Error).ExpectedLabelable => |x| return x.token,
             @TagType(Error).ExpectedInlinable => |x| return x.token,
             @TagType(Error).ExpectedAsmOutputReturnOrType => |x| return x.token,
@@ -164,6 +167,7 @@ pub const Error = union(enum) {
     pub const ExpectedAggregateKw = SingleTokenError("Expected " ++ @tagName(Token.Id.Keyword_struct) ++ ", " ++ @tagName(Token.Id.Keyword_union) ++ ", or " ++ @tagName(Token.Id.Keyword_enum) ++ ", found {}");
     pub const ExpectedEqOrSemi = SingleTokenError("Expected '=' or ';', found {}");
     pub const ExpectedSemiOrLBrace = SingleTokenError("Expected ';' or '{{', found {}");
+    pub const ExpectedColonOrRParen = SingleTokenError("Expected ':' or ')', found {}");
     pub const ExpectedLabelable = SingleTokenError("Expected 'while', 'for', 'inline', 'suspend', or '{{', found {}");
     pub const ExpectedInlinable = SingleTokenError("Expected 'while' or 'for', found {}");
     pub const ExpectedAsmOutputReturnOrType = SingleTokenError("Expected '->' or " ++ @tagName(Token.Id.Identifier) ++ ", found {}");
@@ -1487,7 +1491,7 @@ pub const Node = struct {
         op: Op,
         rhs: &Node,
 
-        const Op = union(enum) {
+        pub const Op = union(enum) {
             AddrOf: AddrOfInfo,
             ArrayType: &Node,
             Await,
@@ -1504,12 +1508,20 @@ pub const Node = struct {
             UnwrapMaybe,
         };
 
-        const AddrOfInfo = struct {
-            align_expr: ?&Node,
-            bit_offset_start_token: ?TokenIndex,
-            bit_offset_end_token: ?TokenIndex,
+        pub const AddrOfInfo = struct {
+            align_info: ?Align,
             const_token: ?TokenIndex,
             volatile_token: ?TokenIndex,
+
+            pub const Align = struct {
+                node: &Node,
+                bit_range: ?BitRange,
+
+                pub const BitRange = struct {
+                    start: &Node,
+                    end: &Node,
+                };
+            };
         };
 
         pub fn iterate(self: &PrefixOp, index: usize) ?&Node {
diff --git a/std/zig/parse.zig b/std/zig/parse.zig
index 826ea2c3e1..d60f03c55e 100644
--- a/std/zig/parse.zig
+++ b/std/zig/parse.zig
@@ -1450,9 +1450,7 @@ pub fn parse(allocator: &mem.Allocator, source: []const u8) !ast.Tree {
             State.SliceOrArrayType => |node| {
                 if (eatToken(&tok_it, &tree, Token.Id.RBracket)) |_| {
                     node.op = ast.Node.PrefixOp.Op{ .SliceType = ast.Node.PrefixOp.AddrOfInfo{
-                        .align_expr = null,
-                        .bit_offset_start_token = null,
-                        .bit_offset_end_token = null,
+                        .align_info = null,
                         .const_token = null,
                         .volatile_token = null,
                     } };
@@ -1467,6 +1465,7 @@ pub fn parse(allocator: &mem.Allocator, source: []const u8) !ast.Tree {
                 try stack.append(State{ .Expression = OptionalCtx{ .Required = &node.op.ArrayType } });
                 continue;
             },
+
             State.AddrOfModifiers => |addr_of_info| {
                 const token = nextToken(&tok_it, &tree);
                 const token_index = token.index;
@@ -1474,12 +1473,19 @@ pub fn parse(allocator: &mem.Allocator, source: []const u8) !ast.Tree {
                 switch (token_ptr.id) {
                     Token.Id.Keyword_align => {
                         stack.append(state) catch unreachable;
-                        if (addr_of_info.align_expr != null) {
+                        if (addr_of_info.align_info != null) {
                             ((try tree.errors.addOne())).* = Error{ .ExtraAlignQualifier = Error.ExtraAlignQualifier{ .token = token_index } };
                             return tree;
                         }
-                        try stack.append(State{ .ExpectToken = Token.Id.RParen });
-                        try stack.append(State{ .Expression = OptionalCtx{ .RequiredNull = &addr_of_info.align_expr } });
+                        addr_of_info.align_info = ast.Node.PrefixOp.AddrOfInfo.Align {
+                            .node = undefined,
+                            .bit_range = null,
+                        };
+                        // TODO https://github.com/ziglang/zig/issues/1022
+                        const align_info = &??addr_of_info.align_info;
+
+                        try stack.append(State{ .AlignBitRange = align_info });
+                        try stack.append(State{ .Expression = OptionalCtx{ .Required = &align_info.node } });
                         try stack.append(State{ .ExpectToken = Token.Id.LParen });
                         continue;
                     },
@@ -1508,6 +1514,31 @@ pub fn parse(allocator: &mem.Allocator, source: []const u8) !ast.Tree {
                 }
             },
 
+            State.AlignBitRange => |align_info| {
+                const token = nextToken(&tok_it, &tree);
+                switch (token.ptr.id) {
+                    Token.Id.Colon => {
+                        align_info.bit_range = ast.Node.PrefixOp.AddrOfInfo.Align.BitRange(undefined);
+                        const bit_range = &??align_info.bit_range;
+
+                        try stack.append(State{ .ExpectToken = Token.Id.RParen });
+                        try stack.append(State{ .Expression = OptionalCtx{ .Required = &bit_range.end } });
+                        try stack.append(State{ .ExpectToken = Token.Id.Colon });
+                        try stack.append(State{ .Expression = OptionalCtx{ .Required = &bit_range.start } });
+                        continue;
+                    },
+                    Token.Id.RParen => continue,
+                    else => {
+                        (try tree.errors.addOne()).* = Error{
+                            .ExpectedColonOrRParen = Error.ExpectedColonOrRParen{
+                                .token = token.index,
+                            }
+                        };
+                        return tree;
+                    },
+                }
+            },
+
             State.Payload => |opt_ctx| {
                 const token = nextToken(&tok_it, &tree);
                 const token_index = token.index;
@@ -2801,6 +2832,7 @@ const State = union(enum) {
     SliceOrArrayAccess: &ast.Node.SuffixOp,
     SliceOrArrayType: &ast.Node.PrefixOp,
     AddrOfModifiers: &ast.Node.PrefixOp.AddrOfInfo,
+    AlignBitRange: &ast.Node.PrefixOp.AddrOfInfo.Align,
 
     Payload: OptionalCtx,
     PointerPayload: OptionalCtx,
@@ -3120,9 +3152,7 @@ fn tokenIdToPrefixOp(id: @TagType(Token.Id)) ?ast.Node.PrefixOp.Op {
         Token.Id.Asterisk,
         Token.Id.AsteriskAsterisk => ast.Node.PrefixOp.Op{ .PointerType = void{} },
         Token.Id.Ampersand => ast.Node.PrefixOp.Op{ .AddrOf = ast.Node.PrefixOp.AddrOfInfo{
-            .align_expr = null,
-            .bit_offset_start_token = null,
-            .bit_offset_end_token = null,
+            .align_info = null,
             .const_token = null,
             .volatile_token = null,
         } },
diff --git a/std/zig/parser_test.zig b/std/zig/parser_test.zig
index d114179bf7..a29b8c2547 100644
--- a/std/zig/parser_test.zig
+++ b/std/zig/parser_test.zig
@@ -1,5 +1,14 @@
 test "zig fmt: float literal with exponent" {
     try testCanonical(
+        \\test "bit field alignment" {
+        \\    assert(@typeOf(&blah.b) == &align(1:3:6) const u3);
+        \\}
+        \\
+    );
+}
+
+test "zig fmt: float literal with exponent" {
+    try testCanonical(
         \\test "aoeu" {
         \\    switch (state) {
         \\        TermState.Start => switch (c) {
diff --git a/std/zig/render.zig b/std/zig/render.zig
index f48b13c987..3940a61fc1 100644
--- a/std/zig/render.zig
+++ b/std/zig/render.zig
@@ -253,17 +253,30 @@ fn renderExpression(allocator: &mem.Allocator, stream: var, tree: &ast.Tree, ind
             switch (prefix_op_node.op) {
                 ast.Node.PrefixOp.Op.AddrOf => |addr_of_info| {
                     try renderToken(tree, stream, prefix_op_node.op_token, indent, Space.None); // &
-                    if (addr_of_info.align_expr) |align_expr| {
+                    if (addr_of_info.align_info) |align_info| {
                         const align_token = tree.nextToken(prefix_op_node.op_token);
                         try renderToken(tree, stream, align_token, indent, Space.None); // align
 
-                        const lparen_token = tree.prevToken(align_expr.firstToken());
+                        const lparen_token = tree.prevToken(align_info.node.firstToken());
                         try renderToken(tree, stream, lparen_token, indent, Space.None); // (
 
-                        try renderExpression(allocator, stream, tree, indent, align_expr, Space.None);
+                        try renderExpression(allocator, stream, tree, indent, align_info.node, Space.None);
 
-                        const rparen_token = tree.nextToken(align_expr.lastToken());
-                        try renderToken(tree, stream, rparen_token, indent, Space.Space); // )
+                        if (align_info.bit_range) |bit_range| {
+                            const colon1 = tree.prevToken(bit_range.start.firstToken());
+                            const colon2 = tree.prevToken(bit_range.end.firstToken());
+
+                            try renderToken(tree, stream, colon1, indent, Space.None); // :
+                            try renderExpression(allocator, stream, tree, indent, bit_range.start, Space.None);
+                            try renderToken(tree, stream, colon2, indent, Space.None); // :
+                            try renderExpression(allocator, stream, tree, indent, bit_range.end, Space.None);
+
+                            const rparen_token = tree.nextToken(bit_range.end.lastToken());
+                            try renderToken(tree, stream, rparen_token, indent, Space.Space); // )
+                        } else {
+                            const rparen_token = tree.nextToken(align_info.node.lastToken());
+                            try renderToken(tree, stream, rparen_token, indent, Space.Space); // )
+                        }
                     }
                     if (addr_of_info.const_token) |const_token| {
                         try renderToken(tree, stream, const_token, indent, Space.Space); // const
@@ -272,21 +285,35 @@ fn renderExpression(allocator: &mem.Allocator, stream: var, tree: &ast.Tree, ind
                         try renderToken(tree, stream, volatile_token, indent, Space.Space); // volatile
                     }
                 },
+
                 ast.Node.PrefixOp.Op.SliceType => |addr_of_info| {
                     try renderToken(tree, stream, prefix_op_node.op_token, indent, Space.None); // [
                     try renderToken(tree, stream, tree.nextToken(prefix_op_node.op_token), indent, Space.None); // ]
 
-                    if (addr_of_info.align_expr) |align_expr| {
+                    if (addr_of_info.align_info) |align_info| {
                         const align_token = tree.nextToken(prefix_op_node.op_token);
                         try renderToken(tree, stream, align_token, indent, Space.None); // align
 
-                        const lparen_token = tree.prevToken(align_expr.firstToken());
+                        const lparen_token = tree.prevToken(align_info.node.firstToken());
                         try renderToken(tree, stream, lparen_token, indent, Space.None); // (
 
-                        try renderExpression(allocator, stream, tree, indent, align_expr, Space.None);
+                        try renderExpression(allocator, stream, tree, indent, align_info.node, Space.None);
 
-                        const rparen_token = tree.nextToken(align_expr.lastToken());
-                        try renderToken(tree, stream, rparen_token, indent, Space.Space); // )
+                        if (align_info.bit_range) |bit_range| {
+                            const colon1 = tree.prevToken(bit_range.start.firstToken());
+                            const colon2 = tree.prevToken(bit_range.end.firstToken());
+
+                            try renderToken(tree, stream, colon1, indent, Space.None); // :
+                            try renderExpression(allocator, stream, tree, indent, bit_range.start, Space.None);
+                            try renderToken(tree, stream, colon2, indent, Space.None); // :
+                            try renderExpression(allocator, stream, tree, indent, bit_range.end, Space.None);
+
+                            const rparen_token = tree.nextToken(bit_range.end.lastToken());
+                            try renderToken(tree, stream, rparen_token, indent, Space.Space); // )
+                        } else {
+                            const rparen_token = tree.nextToken(align_info.node.lastToken());
+                            try renderToken(tree, stream, rparen_token, indent, Space.Space); // )
+                        }
                     }
                     if (addr_of_info.const_token) |const_token| {
                         try renderToken(tree, stream, const_token, indent, Space.Space);
@@ -295,6 +322,7 @@ fn renderExpression(allocator: &mem.Allocator, stream: var, tree: &ast.Tree, ind
                         try renderToken(tree, stream, volatile_token, indent, Space.Space);
                     }
                 },
+
                 ast.Node.PrefixOp.Op.ArrayType => |array_index| {
                     try renderToken(tree, stream, prefix_op_node.op_token, indent, Space.None); // [
                     try renderExpression(allocator, stream, tree, indent, array_index, Space.None);