add zig fmt fix for async/extern fn

4 files changed, 63 insertions, 20 deletions

diff --git a/lib/std/zig/ast.zig b/lib/std/zig/ast.zig
index 220b0390f1..8dec2de15c 100644
--- a/lib/std/zig/ast.zig
+++ b/lib/std/zig/ast.zig
@@ -880,6 +880,8 @@ pub const Node = struct {
         align_expr: ?*Node, // populated if align(A) is present
         section_expr: ?*Node, // populated if linksection(A) is present
         callconv_expr: ?*Node, // populated if callconv(A) is present
+        is_extern_prototype: bool = false, // TODO: Remove once extern fn rewriting is
+        is_async: bool = false, // TODO: remove once async fn rewriting is
 
         pub const ParamList = SegmentedList(*Node, 2);
 
diff --git a/lib/std/zig/parse.zig b/lib/std/zig/parse.zig
index ac1c1f8b15..031fd9c160 100644
--- a/lib/std/zig/parse.zig
+++ b/lib/std/zig/parse.zig
@@ -337,7 +337,25 @@ fn parseTopLevelDecl(arena: *Allocator, it: *TokenIterator, tree: *Tree) !?*Node
 
 /// FnProto <- KEYWORD_fn IDENTIFIER? LPAREN ParamDeclList RPAREN ByteAlign? LinkSection? EXCLAMATIONMARK? (KEYWORD_var / TypeExpr)
 fn parseFnProto(arena: *Allocator, it: *TokenIterator, tree: *Tree) !?*Node {
-    const fn_token = eatToken(it, .Keyword_fn) orelse return null;
+    // TODO: Remove once extern/async fn rewriting is
+    var is_async = false;
+    var is_extern = false;
+    const cc_token: ?usize = blk: {
+        if (eatToken(it, .Keyword_extern)) |token| {
+            is_extern = true;
+            break :blk token;
+        }
+        if (eatToken(it, .Keyword_async)) |token| {
+            is_async = true;
+            break :blk token;
+        }
+        break :blk null;
+    };
+    const fn_token = eatToken(it, .Keyword_fn) orelse {
+        if (cc_token) |token|
+            putBackToken(it, token);
+        return null;
+    };
     const name_token = eatToken(it, .Identifier);
     const lparen = try expectToken(it, tree, .LParen);
     const params = try parseParamDeclList(arena, it, tree);
@@ -381,6 +399,8 @@ fn parseFnProto(arena: *Allocator, it: *TokenIterator, tree: *Tree) !?*Node {
         .align_expr = align_expr,
         .section_expr = section_expr,
         .callconv_expr = callconv_expr,
+        .is_extern_prototype = is_extern,
+        .is_async = is_async,
     };
 
     return &fn_proto_node.base;
@@ -1175,6 +1195,16 @@ fn parseSuffixExpr(arena: *Allocator, it: *TokenIterator, tree: *Tree) !?*Node {
     const maybe_async = eatToken(it, .Keyword_async);
     if (maybe_async) |async_token| {
         const token_fn = eatToken(it, .Keyword_fn);
+        if (token_fn != null) {
+            // TODO: remove this hack when async fn rewriting is
+            // HACK: If we see the keyword `fn`, then we assume that
+            //       we are parsing an async fn proto, and not a call.
+            //       We therefore put back all tokens consumed by the async
+            //       prefix...
+            putBackToken(it, token_fn.?);
+            putBackToken(it, async_token);
+            return parsePrimaryTypeExpr(arena, it, tree);
+        }
         var res = try expectNode(arena, it, tree, parsePrimaryTypeExpr, .{
             .ExpectedPrimaryTypeExpr = .{ .token = it.index },
         });
diff --git a/lib/std/zig/parser_test.zig b/lib/std/zig/parser_test.zig
index 198ce1ad94..7ab3e86b1a 100644
--- a/lib/std/zig/parser_test.zig
+++ b/lib/std/zig/parser_test.zig
@@ -236,10 +236,10 @@ test "zig fmt: anon list literal syntax" {
 test "zig fmt: async function" {
     try testCanonical(
         \\pub const Server = struct {
-        \\    handleRequestFn: async fn (*Server, *const std.net.Address, File) void,
+        \\    handleRequestFn: fn (*Server, *const std.net.Address, File) callconv(.Async) void,
         \\};
         \\test "hi" {
-        \\    var ptr = @ptrCast(async fn (i32) void, other);
+        \\    var ptr = @ptrCast(fn (i32) callconv(.Async) void, other);
         \\}
         \\
     );
@@ -435,15 +435,6 @@ test "zig fmt: aligned struct field" {
     );
 }
 
-test "zig fmt: preserve space between async fn definitions" {
-    try testCanonical(
-        \\async fn a() void {}
-        \\
-        \\async fn b() void {}
-        \\
-    );
-}
-
 test "zig fmt: comment to disable/enable zig fmt first" {
     try testCanonical(
         \\// Test trailing comma syntax
@@ -1499,7 +1490,7 @@ test "zig fmt: line comments in struct initializer" {
 
 test "zig fmt: first line comment in struct initializer" {
     try testCanonical(
-        \\pub async fn acquire(self: *Self) HeldLock {
+        \\pub fn acquire(self: *Self) HeldLock {
         \\    return HeldLock{
         \\        // guaranteed allocation elision
         \\        .held = self.lock.acquire(),
@@ -2461,8 +2452,7 @@ test "zig fmt: fn type" {
         \\}
         \\
         \\const a: fn (u8) u8 = undefined;
-        \\const b: extern fn (u8) u8 = undefined;
-        \\const c: fn (u8) callconv(.Naked) u8 = undefined;
+        \\const b: fn (u8) callconv(.Naked) u8 = undefined;
         \\const ap: fn (u8) u8 = a;
         \\
     );
@@ -2484,7 +2474,7 @@ test "zig fmt: inline asm" {
 
 test "zig fmt: async functions" {
     try testCanonical(
-        \\async fn simpleAsyncFn() void {
+        \\fn simpleAsyncFn() void {
         \\    const a = async a.b();
         \\    x += 1;
         \\    suspend;
@@ -2920,6 +2910,25 @@ test "zig fmt: noasync to nosuspend" {
         \\pub fn main() void {
         \\    nosuspend call();
         \\}
+    );
+}
+
+test "zig fmt: convert async fn into callconv(.Async)" {
+    try testTransform(
+        \\async fn foo() void {}
+    ,
+        \\fn foo() callconv(.Async) void {}
+        \\
+    );
+}
+
+test "zig fmt: convert extern fn proto into callconv(.C)" {
+    try testTransform(
+        \\extern fn foo0() void {}
+        \\const foo1 = extern fn () void;
+    ,
+        \\extern fn foo0() void {}
+        \\const foo1 = fn () callconv(.C) void;
         \\
     );
 }
@@ -2970,7 +2979,6 @@ fn testParse(source: []const u8, allocator: *mem.Allocator, anything_changed: *b
     anything_changed.* = try std.zig.render(allocator, buffer.outStream(), tree);
     return buffer.toOwnedSlice();
 }
-
 fn testTransform(source: []const u8, expected_source: []const u8) !void {
     const needed_alloc_count = x: {
         // Try it once with unlimited memory, make sure it works
@@ -3018,11 +3026,9 @@ fn testTransform(source: []const u8, expected_source: []const u8) !void {
         }
     }
 }
-
 fn testCanonical(source: []const u8) !void {
     return testTransform(source, source);
 }
-
 fn testError(source: []const u8) !void {
     const tree = try std.zig.parse(std.testing.allocator, source);
     defer tree.deinit();
diff --git a/lib/std/zig/render.zig b/lib/std/zig/render.zig
index 3e2cebcd8a..bcdc08d43a 100644
--- a/lib/std/zig/render.zig
+++ b/lib/std/zig/render.zig
@@ -1414,7 +1414,8 @@ fn renderExpression(
             }
 
             if (fn_proto.extern_export_inline_token) |extern_export_inline_token| {
-                try renderToken(tree, stream, extern_export_inline_token, indent, start_col, Space.Space); // extern/export
+                if (!fn_proto.is_extern_prototype)
+                    try renderToken(tree, stream, extern_export_inline_token, indent, start_col, Space.Space); // extern/export/inline
             }
 
             if (fn_proto.lib_name) |lib_name| {
@@ -1510,6 +1511,10 @@ fn renderExpression(
                 try renderToken(tree, stream, callconv_lparen, indent, start_col, Space.None); // (
                 try renderExpression(allocator, stream, tree, indent, start_col, callconv_expr, Space.None);
                 try renderToken(tree, stream, callconv_rparen, indent, start_col, Space.Space); // )
+            } else if (fn_proto.is_extern_prototype) {
+                try stream.writeAll("callconv(.C) ");
+            } else if (fn_proto.is_async) {
+                try stream.writeAll("callconv(.Async) ");
             }
 
             switch (fn_proto.return_type) {