Format Zig 0.16.0 source code with 2 spaces

zig_fmt_2_spaces.zig

// Format Zig 0.16.0 source code with 2 spaces
// zig build-exe --zig-lib-dir /ABSOLUTE/PATH/TO/zig/lib -O ReleaseSmall ./fmt_16.zig

const std = @import("std");
const Io = std.Io;
const mem = std.mem;
const fs = std.fs;
const process = std.process;
const Allocator = std.mem.Allocator;
const Color = std.zig.Color;
const fatal = std.process.fatal;

const usage_fmt =
    \\Usage: zig fmt [file]...
    \\
    \\   Formats the input files and modifies them in-place to use 2-space indentation.
    \\   Arguments can be files or directories, which are searched
    \\   recursively.
    \\
    \\Options:
    \\  -h, --help             Print this help and exit
    \\  --color [auto|off|on]  Enable or disable colored error messages
    \\  --stdin                Format code from stdin; output to stdout
    \\  --check                List non-conforming files and exit with an error
    \\                         if the list is non-empty
    \\  --ast-check            Run zig ast-check on every file
    \\  --exclude [file]       Exclude file or directory from formatting
    \\  --zon                  Treat all input files as ZON, regardless of file extension
    \\
    \\
;

const Fmt = struct {
  seen: SeenMap,
  any_error: bool,
  check_ast: bool,
  force_zon: bool,
  color: Color,
  gpa: Allocator,
  arena: Allocator,
  io: Io,
  stdout_writer: *Io.File.Writer,

  const SeenMap = std.AutoHashMap(Io.File.INode, void);
};

pub fn main(init: std.process.Init) !void {
  const gpa = init.gpa;
  const arena = init.arena.allocator();
  const io = init.io;

  const args = try init.minimal.args.toSlice(arena);
  const positional_args = if (args.len > 1) args[1..] else &[_][]const u8{};

  var color: Color = .auto;
  var stdin_flag = false;
  var check_flag = false;
  var check_ast_flag = false;
  var force_zon = false;
  var input_files = std.array_list.Managed([]const u8).init(gpa);
  defer input_files.deinit();
  var excluded_files = std.array_list.Managed([]const u8).init(gpa);
  defer excluded_files.deinit();

  {
    var i: usize = 0;
    while (i < positional_args.len) : (i += 1) {
      const arg = positional_args[i];
      if (mem.startsWith(u8, arg, "-")) {
        if (mem.eql(u8, arg, "-h") or mem.eql(u8, arg, "--help")) {
          try Io.File.stdout().writeStreamingAll(io, usage_fmt);
          return process.cleanExit(io);
        } else if (mem.eql(u8, arg, "--color")) {
          if (i + 1 >= positional_args.len) {
            fatal("expected [auto|on|off] after --color", .{});
          }
          i += 1;
          const next_arg = positional_args[i];
          color = std.meta.stringToEnum(Color, next_arg) orelse {
            fatal("expected [auto|on|off] after --color, found '{s}'", .{next_arg});
          };
        } else if (mem.eql(u8, arg, "--stdin")) {
          stdin_flag = true;
        } else if (mem.eql(u8, arg, "--check")) {
          check_flag = true;
        } else if (mem.eql(u8, arg, "--ast-check")) {
          check_ast_flag = true;
        } else if (mem.eql(u8, arg, "--exclude")) {
          if (i + 1 >= positional_args.len) {
            fatal("expected parameter after --exclude", .{});
          }
          i += 1;
          const next_arg = positional_args[i];
          try excluded_files.append(next_arg);
        } else if (mem.eql(u8, arg, "--zon")) {
          force_zon = true;
        } else {
          fatal("unrecognized parameter: '{s}'", .{arg});
        }
      } else {
        try input_files.append(arg);
      }
    }
  }

  if (stdin_flag) {
    if (input_files.items.len != 0) {
      fatal("cannot use --stdin with positional arguments", .{});
    }

    const stdin: Io.File = .stdin();
    var stdio_buffer = [_]u8{0} ** 1024;
    var file_reader: Io.File.Reader = stdin.reader(io, &stdio_buffer);
    const source_code = std.zig.readSourceFileToEndAlloc(gpa, &file_reader) catch |err| {
      fatal("unable to read stdin: {}", .{err});
    };
    defer gpa.free(source_code);

    var tree = std.zig.Ast.parse(gpa, source_code, if (force_zon) .zon else .zig) catch |err| {
      fatal("error parsing stdin: {}", .{err});
    };
    defer tree.deinit(gpa);

    if (check_ast_flag) {
      if (!force_zon) {
        var zir = try std.zig.AstGen.generate(gpa, tree);
        defer zir.deinit(gpa);

        if (zir.hasCompileErrors()) {
          var wip_errors: std.zig.ErrorBundle.Wip = undefined;
          try wip_errors.init(gpa);
          defer wip_errors.deinit();
          try wip_errors.addZirErrorMessages(zir, tree, source_code, "<stdin>");
          var error_bundle = try wip_errors.toOwnedBundle("");
          defer error_bundle.deinit(gpa);
          error_bundle.renderToStderr(io, .{}, color) catch {};
          process.exit(2);
        }
      } else {
        const zoir = try std.zig.ZonGen.generate(gpa, tree, .{});
        defer zoir.deinit(gpa);

        if (zoir.hasCompileErrors()) {
          var wip_errors: std.zig.ErrorBundle.Wip = undefined;
          try wip_errors.init(gpa);
          defer wip_errors.deinit();
          try wip_errors.addZoirErrorMessages(zoir, tree, source_code, "<stdin>");
          var error_bundle = try wip_errors.toOwnedBundle("");
          defer error_bundle.deinit(gpa);
          error_bundle.renderToStderr(io, .{}, color) catch {};
          process.exit(2);
        }
      }
    } else if (tree.errors.len != 0) {
      std.zig.printAstErrorsToStderr(gpa, io, tree, "<stdin>", color) catch {};
      process.exit(2);
    }

    var out_stream = std.Io.Writer.Allocating.init(gpa);
    defer out_stream.deinit();

    try tree.render(gpa, &out_stream.writer, .{});
    const four_space_raw = try out_stream.toOwnedSlice();
    defer gpa.free(four_space_raw);

    const formatted = try convertFourToTwoSpaces(gpa, four_space_raw);
    defer gpa.free(formatted);

    if (check_flag) {
      const code: u8 = @intFromBool(!mem.eql(u8, formatted, source_code));
      process.exit(code);
    }

    return Io.File.stdout().writeStreamingAll(io, formatted);
  }

  if (input_files.items.len == 0) {
    fatal("expected at least one file or directory argument", .{});
  }

  var stdout_buffer = [_]u8{0} ** 4096;
  var stdout_writer = Io.File.stdout().writer(io, &stdout_buffer);

  var fmt: Fmt = .{
    .gpa = gpa,
    .arena = arena,
    .io = io,
    .seen = .init(gpa),
    .any_error = false,
    .check_ast = check_ast_flag,
    .force_zon = force_zon,
    .color = color,
    .stdout_writer = &stdout_writer,
  };
  defer fmt.seen.deinit();

  for (excluded_files.items) |file_path| {
    const stat = Io.Dir.cwd().statFile(io, file_path, .{}) catch |err| switch (err) {
      error.FileNotFound => continue,
      error.IsDir => dir: {
        var dir = try Io.Dir.cwd().openDir(io, file_path, .{});
        defer dir.close(io);
        break :dir try dir.stat(io);
      },
      else => |e| return e,
    };
    try fmt.seen.put(stat.inode, {});
  }

  for (input_files.items) |file_path| {
    try fmtPath(&fmt, file_path, check_flag, Io.Dir.cwd(), file_path);
  }
  try fmt.stdout_writer.interface.flush();
  if (fmt.any_error) {
    process.exit(1);
  }
}

fn fmtPath(fmt: *Fmt, file_path: []const u8, check_mode: bool, dir: Io.Dir, sub_path: []const u8) anyerror!void {
  fmtPathFile(fmt, file_path, check_mode, dir, sub_path) catch |err| switch (err) {
    error.IsDir, error.AccessDenied => return fmtPathDir(fmt, file_path, check_mode, dir, sub_path),
    else => |e| return e,
  };
}

fn fmtPathFile(fmt: *Fmt, file_path: []const u8, check_mode: bool, dir: Io.Dir, sub_path: []const u8) anyerror!void {
  const is_zig = mem.endsWith(u8, sub_path, ".zig");
  const is_zon = mem.endsWith(u8, sub_path, ".zon");
  if (!is_zig and !is_zon and !fmt.force_zon) return;

  const source_code = code: {
    var io_file = try dir.openFile(fmt.io, sub_path, .{ .mode = .read_only });
    defer io_file.close(fmt.io);

    const stat = try io_file.stat(fmt.io);
    const gget = try fmt.seen.getOrPut(stat.inode);
    if (gget.found_existing) return;

    var read_buf = [_]u8{0} ** 1024;
    var reader = io_file.reader(fmt.io, &read_buf);
    break :code try std.zig.readSourceFileToEndAlloc(fmt.gpa, &reader);
  };
  defer fmt.gpa.free(source_code);

  var tree = try std.zig.Ast.parse(fmt.gpa, source_code, if (is_zon or fmt.force_zon) .zon else .zig);
  defer tree.deinit(fmt.gpa);

  if (tree.errors.len != 0) {
    try std.zig.printAstErrorsToStderr(fmt.gpa, fmt.io, tree, file_path, fmt.color);
    fmt.any_error = true;
    return;
  }

  var out_stream = std.Io.Writer.Allocating.init(fmt.gpa);
  defer out_stream.deinit();

  try tree.render(fmt.gpa, &out_stream.writer, .{});
  const four_space_raw = try out_stream.toOwnedSlice();
  defer fmt.gpa.free(four_space_raw);

  const formatted = try convertFourToTwoSpaces(fmt.gpa, four_space_raw);
  defer fmt.gpa.free(formatted);

  if (mem.eql(u8, source_code, formatted)) return;

  if (check_mode) {
    try Io.File.stdout().writeStreamingAll(fmt.io, file_path);
    try Io.File.stdout().writeStreamingAll(fmt.io, "\n");
    fmt.any_error = true;
    return;
  }

  var write_file = try dir.createFile(fmt.io, sub_path, .{});
  defer write_file.close(fmt.io);

  try write_file.writeStreamingAll(fmt.io, formatted);

  try Io.File.stdout().writeStreamingAll(fmt.io, file_path);
  try Io.File.stdout().writeStreamingAll(fmt.io, "\n");
}

fn fmtPathDir(fmt: *Fmt, file_path: []const u8, check_mode: bool, dir: Io.Dir, sub_path: []const u8) anyerror!void {
  var child_dir = try dir.openDir(fmt.io, sub_path, .{});
  defer child_dir.close(fmt.io);

  const stat = try child_dir.stat(fmt.io);
  const gget = try fmt.seen.getOrPut(stat.inode);
  if (gget.found_existing) return;

  var iterator = child_dir.iterate();
  while (try iterator.next(fmt.io)) |entry| {
    const child_path = try fs.path.join(fmt.arena, &[_][]const u8{ file_path, entry.name });
    try fmtPath(fmt, child_path, check_mode, child_dir, entry.name);
  }
}

fn convertFourToTwoSpaces(gpa: Allocator, input: []const u8) anyerror![]u8 {
  var result = std.array_list.Managed(u8).init(gpa);
  errdefer result.deinit();
  var line_it = mem.splitScalar(u8, input, '\n');
  while (line_it.next()) |line| {
    var spaces: usize = 0;
    while (spaces < line.len and line[spaces] == ' ') : (spaces += 1) {}
    const remaining = line[spaces..];
    if (mem.startsWith(u8, remaining, "\\")) {
      try result.appendSlice(line);
    } else {
      const indent_level = spaces / 4;
      const remainder_spaces = spaces % 4;
      const target_spaces = (indent_level * 2) + remainder_spaces;
      try result.appendNTimes(' ', target_spaces);
      try result.appendSlice(remaining);
    }
    if (line_it.rest().len > 0 or (input.len > 0 and input[input.len - 1] == '\n')) {
      try result.append('\n');
    }
  }
  return result.toOwnedSlice();
}