Expand roc format to support directories

Also supports formatting multiple passed in args.
Also prints the overall execution time.

Author: bhansconnect

src/collections/utils.zig

@@ -11,10 +11,15 @@ pub fn exitOnOom(err: std.mem.Allocator.Error) noreturn {
             const oom_message =
                 \\I ran out of memory! I can't do anything to recover, so I'm exiting.
                 \\Try reducing memory usage on your machine and then running again.
+                \\
             ;
-
-            std.debug.print(oom_message, .{});
-            std.process.exit(1);
+            fatal(oom_message, .{});
         },
     }
 }
+
+/// Log a fatal error and exit the process with a non-zero code.
+pub fn fatal(comptime format: []const u8, args: anytype) noreturn {
+    std.io.getStdErr().writer().print(format, args) catch unreachable;
+    std.process.exit(1);
+}

src/coordinate/Filesystem.zig

@@ -31,6 +31,10 @@ pub fn default() Self {
     };
 }
 
+/// The max valid file size.
+/// Anything larger will fail due to us using u32 offsets.
+pub const max_file_size = std.math.maxInt(u32);
+
 /// All errors that can occur when reading a file.
 pub const ReadError = std.fs.File.OpenError || std.posix.ReadError || Allocator.Error || error{StreamTooLong};
 
@@ -137,8 +141,7 @@ fn readFileDefault(relative_path: []const u8, allocator: std.mem.Allocator) Read
     const file = try std.fs.cwd().openFile(relative_path, .{});
     defer file.close();
 
-    const max_allowed_file_length = std.math.maxInt(usize);
-    const contents = try file.reader().readAllAlloc(allocator, max_allowed_file_length);
+    const contents = try file.reader().readAllAlloc(allocator, max_file_size);
 
     return contents;
 }

src/fmt.zig

@@ -1,12 +1,15 @@
 //! Formatting logic for Roc modules.
 
 const std = @import("std");
+const parse = @import("check/parse.zig").parse;
 const IR = @import("check/parse/IR.zig");
 const Node = IR.Node;
+const Filesystem = @import("coordinate/Filesystem.zig");
 const tokenizer = @import("check/parse/tokenize.zig");
 const TokenizedBuffer = tokenizer.TokenizedBuffer;
 const TokenIdx = tokenizer.Token.Idx;
 const exitOnOom = @import("./collections/utils.zig").exitOnOom;
+const fatal = @import("./collections/utils.zig").fatal;
 const base = @import("base.zig");
 
 const NodeStore = IR.NodeStore;
@@ -42,6 +45,75 @@ pub fn resetWith(fmt: *Formatter, ast: IR) void {
     fmt.ast = ast;
 }
 
+/// Formats all roc files in the specified path.
+/// Handles both single files and directories
+/// Returns the number of files formatted.
+pub fn formatPath(gpa: std.mem.Allocator, base_dir: std.fs.Dir, path: []const u8) !usize {
+    // TODO: update this to use the filesystem abstraction
+    // When doing so, add a mock filesystem and some tests.
+    var count: usize = 0;
+    // First try as a directory.
+    if (base_dir.openDir(path, .{ .iterate = true })) |const_dir| {
+        var dir = const_dir;
+        defer dir.close();
+        // Walk is recursive.
+        var walker = try dir.walk(gpa);
+        defer walker.deinit();
+        while (try walker.next()) |entry| {
+            if (entry.kind == .file) {
+                if (try formatFilePath(gpa, entry.dir, entry.basename)) {
+                    count += 1;
+                }
+            }
+        }
+    } else |_| {
+        if (try formatFilePath(gpa, base_dir, path)) {
+            count += 1;
+        }
+    }
+
+    return count;
+}
+
+fn formatFilePath(gpa: std.mem.Allocator, base_dir: std.fs.Dir, path: []const u8) !bool {
+    // Skip non ".roc" files.
+    if (!std.mem.eql(u8, std.fs.path.extension(path), ".roc")) {
+        return false;
+    }
+
+    const input_file = try base_dir.openFile(path, .{ .mode = .read_only });
+    defer input_file.close();
+
+    const contents = try input_file.reader().readAllAlloc(gpa, Filesystem.max_file_size);
+    defer gpa.free(contents);
+
+    var module_env = base.ModuleEnv.init(gpa);
+    defer module_env.deinit();
+
+    var parse_ast = parse(&module_env, contents);
+    defer parse_ast.deinit();
+    if (parse_ast.errors.len > 0) {
+        // TODO: pretty print the parse failures.
+        const stderr = std.io.getStdErr().writer();
+        try stderr.print("Failed to parse '{s}' for formatting.\nErrors:\n", .{path});
+        for (parse_ast.errors) |err| {
+            try stderr.print("\t{s}\n", .{@tagName(err.tag)});
+        }
+        fatal("\n", .{});
+    }
+
+    var formatter = init(parse_ast);
+    defer formatter.deinit();
+
+    const formatted_output = formatter.formatFile();
+    defer gpa.free(formatted_output);
+
+    const output_file = try base_dir.createFile(path, .{});
+    defer output_file.close();
+    try output_file.writeAll(formatted_output);
+    return true;
+}
+
 /// Emits a string containing the well-formed source of a Roc parse IR (AST).
 /// The resulting string is owned by the caller.
 pub fn formatFile(fmt: *Formatter) []const u8 {
@@ -728,8 +800,6 @@ fn pushTokenText(fmt: *Formatter, ti: TokenIdx) void {
 }
 
 fn moduleFmtsSame(source: []const u8) !void {
-    const parse = @import("check/parse.zig").parse;
-
     const gpa = std.testing.allocator;
 
     var env = base.ModuleEnv.init(gpa);
@@ -836,8 +906,6 @@ pub fn moduleFmtsStable(gpa: std.mem.Allocator, input: []const u8, debug: bool)
 }
 
 fn parseAndFmt(gpa: std.mem.Allocator, input: []const u8, debug: bool) ![]const u8 {
-    const parse = @import("check/parse.zig").parse;
-
     var module_env = base.ModuleEnv.init(gpa);
     defer module_env.deinit();
 
@@ -1064,3 +1132,32 @@ test "Dot access super test" {
     const expr = "some_fn(arg1)?.static_dispatch_method()?.next_static_dispatch_method()?.record_field?";
     try exprFmtsSame(expr, .no_debug);
 }
+
+// TODO: replace this test with one that doesn't interact with the real filesystem.
+test "format single file" {
+    const gpa = std.testing.allocator;
+    const roc_filename = "test.roc";
+
+    const roc_file = try std.fs.cwd().createFile(roc_filename, .{ .read = true });
+    defer roc_file.close();
+    try roc_file.writeAll(
+        \\module []
+        \\
+        \\foo =      "bar"
+    );
+    defer std.fs.cwd().deleteFile(roc_filename) catch std.debug.panic("Failed to clean up test.roc", .{});
+
+    const count = try formatPath(gpa, std.fs.cwd(), roc_filename);
+    try std.testing.expectEqual(1, count);
+
+    // Reset file position to read formatted roc code
+    try roc_file.seekTo(0);
+    const formatted_code = try roc_file.reader().readAllAlloc(gpa, Filesystem.max_file_size);
+    defer gpa.free(formatted_code);
+
+    try std.testing.expectEqualStrings(
+        \\module []
+        \\
+        \\foo = "bar"
+    , formatted_code);
+}

src/main.zig

@@ -1,6 +1,5 @@
 const std = @import("std");
 const fmt = @import("fmt.zig");
-const parse = @import("check/parse.zig");
 const base = @import("base.zig");
 const cli = @import("cli.zig");
 const collections = @import("collections.zig");
@@ -13,6 +12,7 @@ const RocOpt = cli.RocOpt;
 const Problem = problem_mod.Problem;
 const Allocator = std.mem.Allocator;
 const exitOnOom = collections.utils.exitOnOom;
+const fatal = collections.utils.fatal;
 
 const usage =
     \\Usage:
@@ -36,12 +36,6 @@ const usage =
     \\ -h, --help       Print usage
 ;
 
-/// Log a fatal error and exit the process with a non-zero code.
-pub fn fatal(comptime format: []const u8, args: anytype) noreturn {
-    std.io.getStdErr().writer().print(format, args) catch unreachable;
-    std.process.exit(1);
-}
-
 /// The CLI entrypoint for the Roc compiler.
 pub fn main() !void {
     const gpa = std.heap.c_allocator;
@@ -125,62 +119,21 @@ fn rocRepl(gpa: Allocator, opt: RocOpt, args: []const []const u8) !void {
     fatal("not implemented", .{});
 }
 
-/// Reads, parses, formats, and overwrites a Roc file.
+/// Reads, parses, formats, and overwrites all Roc files at the given paths.
+/// Recurses into directories to search for Roc files.
 fn rocFormat(gpa: Allocator, args: []const []const u8) !void {
-    const roc_file_path = if (args.len > 0) args[0] else "main.roc";
-
-    const input_file = try std.fs.cwd().openFile(roc_file_path, .{ .mode = .read_only });
-    defer input_file.close();
-
-    const contents = try input_file.reader().readAllAlloc(gpa, std.math.maxInt(usize));
-    defer gpa.free(contents);
-
-    var module_env = base.ModuleEnv.init(gpa);
-    defer module_env.deinit();
-
-    var parse_ast = parse.parse(&module_env, contents);
-    defer parse_ast.deinit();
-    if (parse_ast.errors.len > 0) {
-        // TODO: pretty print the parse failures.
-        fatal("Failed to parse '{s}' for formatting.\nErrors:\n{any}\n", .{ roc_file_path, parse_ast.errors });
+    // var timer = try std.time.Timer.start();
+    var count: usize = 0;
+    if (args.len > 0) {
+        for (args) |arg| {
+            count += try fmt.formatPath(gpa, std.fs.cwd(), arg);
+        }
+    } else {
+        count = try fmt.formatPath(gpa, std.fs.cwd(), "main.roc");
     }
 
-    var formatter = fmt.init(parse_ast);
-    defer formatter.deinit();
-
-    const formatted_output = formatter.formatFile();
-    defer gpa.free(formatted_output);
-
-    const output_file = try std.fs.cwd().createFile(roc_file_path, .{});
-    defer output_file.close();
-    try output_file.writeAll(formatted_output);
-}
-
-test "format single file" {
-    const gpa = std.testing.allocator;
-    const roc_filename = "test.roc";
-
-    const roc_file = try std.fs.cwd().createFile(roc_filename, .{ .read = true });
-    defer roc_file.close();
-    try roc_file.writeAll(
-        \\module []
-        \\
-        \\foo =      "bar"
-    );
-    defer std.fs.cwd().deleteFile(roc_filename) catch std.debug.panic("Failed to clean up test.roc", .{});
-
-    try rocFormat(gpa, &.{roc_filename});
-
-    // Reset file position to read formatted roc code
-    try roc_file.seekTo(0);
-    const formatted_code = try roc_file.reader().readAllAlloc(gpa, std.math.maxInt(usize));
-    defer gpa.free(formatted_code);
-
-    try std.testing.expectEqualStrings(
-        \\module []
-        \\
-        \\foo = "bar"
-    , formatted_code);
+    // const elapsed = timer.read() / std.time.ns_per_ms;
+    // try std.io.getStdOut().writer().print("Successfully formatted {} files in {} ms.\n", .{ count, elapsed });
 }
 
 fn rocVersion(gpa: Allocator, args: []const []const u8) !void {