yabai wrapper

main.zig

const std = @import("std");
// const log = std.log;
const print = std.debug.print;

var gpa: std.heap.GeneralPurposeAllocator(.{}) = .{};
const gpa_alloc = gpa.allocator();

var arena = std.heap.ArenaAllocator.init(gpa_alloc);
const arena_alloc = arena.allocator();

var yabai_path: []const u8 = undefined;

pub fn main() !void {
    defer arena.deinit();
    defer _ = gpa.deinit();

    const user_name = std.os.getenv("USER") orelse return error.GetUsernameFailed;
    const y3_path = try std.fmt.allocPrint(gpa_alloc, "/tmp/y3_{s}.socket", .{user_name});
    yabai_path = try std.fmt.allocPrint(gpa_alloc, "/tmp/yabai_{s}.socket", .{user_name});

    // cheap temporary way of cleaning up at startup instead of handling sigkill
    std.fs.deleteFileAbsolute(y3_path) catch |err|
        if (err != error.FileNotFound) return print("error deleting socket fd: {}", .{err});

    const listener_address = try std.net.Address.initUnix(y3_path);
    var server = std.net.StreamServer.init(.{});
    defer server.deinit();

    try server.listen(listener_address);
    print("started listener...\n", .{});

    var command_buf: [32]u8 = undefined;
    while (true) {
        defer _ = arena.reset(.retain_capacity);

        const conn = try server.accept();
        defer conn.stream.close();

        const n = try conn.stream.readAll(&command_buf);
        run(command_buf[0..n]) catch |err| print("failed to run command: {}\n", .{err});
        // print("arena: {}\n", .{arena.queryCapacity()});
    }
}

fn run(message: []const u8) !void {
    var args = std.mem.splitSequence(u8, message, " ");
    const arg1 = args.next() orelse return error.MissingCommandArg;
    const arg2 = args.next() orelse return error.MissingDomainArg;
    if (args.next()) |_| return error.TooManyArgs;

    const CommandArg = enum { focus, move };
    const command = std.meta.stringToEnum(CommandArg, arg1) orelse return error.InvalidCommand;
    const direction = std.meta.stringToEnum(Direction, arg2[0 .. arg2.len - 1]) orelse return error.InvalidTarget;

    switch (command) {
        .focus => try focus(direction),
        .move => try move(direction),
    }
}

fn focus(input_dir: Direction) !void {
    sendCommand(&.{ "window", "--focus", @tagName(input_dir) }) catch |window_err| switch (window_err) {
        // try focus in other display
        error.LocateNorth, error.LocateSouth => focusInOtherDisplay(input_dir) catch |display_err| {
            // try cycling stack
            if (input_dir == .north) {
                sendCommand(&.{ "window", "--focus", "stack.prev" }) catch |stack_err| switch (stack_err) {
                    // try wrapping around
                    error.LocatePrevStacked => sendCommand(&.{ "window", "--focus", "stack.last" }) catch |err|
                        if (err != error.LocateLastStacked) return err,
                    else => return stack_err,
                };
            } else if (input_dir == .south) {
                sendCommand(&.{ "window", "--focus", "stack.next" }) catch |stack_err| switch (stack_err) {
                    // try wrapping around
                    error.LocateNextStacked => sendCommand(&.{ "window", "--focus", "stack.first" }) catch |err|
                        if (err != error.LocateFirstStacked) return err,
                    else => return stack_err,
                };
            }
            switch (display_err) {
                error.LocateNorthDisplay, error.LocateSouthDisplay => {},
                else => return display_err,
            }
        },
        error.LocateEast, error.LocateWest => focusInOtherDisplay(input_dir) catch |display_err| switch (display_err) {
            error.LocateEastDisplay, error.LocateWestDisplay => {},
            else => return display_err,
        },
        error.LocateSelected => {},
        else => return window_err,
    };
}

fn focusInOtherDisplay(input_dir: Direction) !void {
    try sendCommand(&.{ "display", "--focus", @tagName(input_dir) });

    const target_win = switch (input_dir) {
        .east => "first",
        .west => "last",
        else => "last",
    };
    sendCommand(&.{ "window", "--focus", target_win }) catch |err| switch (err) {
        error.LocateFirst, error.LocateLast => try sendCommand(&.{ "window", "--focus", "prev" }),
        else => return err,
    };
}

fn move(input_dir: Direction) !void {
    const space_wins = try query([]Window, &.{ "--windows", "--space" });
    const active_win = blk: {
        for (space_wins) |win| {
            if (win.@"has-focus" == true) break :blk win;
        } else return error.NoFocusedWindow;
    };
    if (active_win.isStacked()) return try unstack(active_win, space_wins, input_dir);

    const target_win = query(Window, &.{ "--windows", "--window", @tagName(input_dir) }) catch null;
    if (target_win) |tw| {
        if (active_win.isSimilar(tw)) {
            try stack(active_win, tw, input_dir);
        } else try sendCommand(&.{ "window", "--warp", @tagName(input_dir) });
    } else {
        const target_disp = query(Display, &.{ "--displays", "--display", @tagName(input_dir) }) catch null;
        if (target_disp) |_| {
            try moveToDisplay(input_dir);
        } else {
            const active_disp = try query(Display, &.{ "--displays", "--display" });
            if (space_wins.len == 2) try rotate(active_win, active_disp, input_dir);
        }
    }
}

fn rotate(active_win: Window, active_disp: Display, input_dir: Direction) !void {
    switch (input_dir) {
        .north, .south => if (active_win.isFullHeight(active_disp)) {
            if (@abs(active_win.frame.x) == 0.0) {
                try insertAndWarp(.east, input_dir);
            } else try insertAndWarp(.west, input_dir);
        },
        .east, .west => {
            if (active_win.isTopRow(active_disp)) {
                try insertAndWarp(.south, input_dir);
            } else if (!active_win.isFullHeight(active_disp))
                try insertAndWarp(.north, input_dir);
        },
    }
}

/// Inserts active window into the side of the target window
fn insertAndWarp(target_win_dir: Direction, insert_side: Direction) !void {
    const stream = try sendCommandUnchecked(&.{ "window", @tagName(target_win_dir), "--insert", @tagName(insert_side) });
    defer stream.close();

    try sendCommand(&.{ "window", "--warp", @tagName(target_win_dir) });
    try checkResponse(stream);
}

fn stack(active_win: Window, target_win: Window, input_dir: Direction) !void {
    const stream = try sendCommandUnchecked(&.{ "window", @tagName(input_dir), "--stack", intStr(active_win.id).slice() });
    defer stream.close();

    try sendCommand(&.{ "window", intStr(target_win.id).slice(), "--layer", "normal" }); // workaround for "below" bug
    try checkResponse(stream);
}

fn unstack(active_win: Window, space_wins: []Window, input_dir: Direction) !void {
    const under_win_id = blk: {
        if (active_win.@"stack-index" > 0) {
            for (space_wins) |win| {
                if (std.meta.eql(win.frame, active_win.frame) and @abs(win.@"stack-index" - active_win.@"stack-index") == 1)
                    break :blk win.id;
            } else return error.NoWindowUnderFocused;
        } else return error.ActiveWindowNotStacked;
    };
    const stream = try sendCommandUnchecked(&.{ "window", intStr(under_win_id).slice(), "--insert", @tagName(input_dir) });
    defer stream.close();

    const active_win_id_str = intStr(active_win.id).slice();
    const stream2 = try sendCommandUnchecked(&.{ "window", active_win_id_str, "--toggle", "float", "--toggle", "float" });
    defer stream2.close();

    try sendCommand(&.{ "window", active_win_id_str, "--layer", "normal" }); // workaround for "below" bug
    try checkResponse(stream);
    try checkResponse(stream2);
}

fn moveToDisplay(input_dir: Direction) !void {
    print("running move to display", .{});
    const spaces = try query([]Space, &.{ "--spaces", "--display", @tagName(input_dir) });
    var streams: std.BoundedArray(std.net.Stream, 2) = .{};

    for (spaces) |sp| if (sp.@"is-visible") {
        if (sp.@"first-window" == 0) {
            const stream = try sendCommandUnchecked(&.{ "window", "--display", @tagName(input_dir) });
            defer stream.close();
            break streams.appendAssumeCapacity(stream);
        }

        const stream = switch (input_dir) {
            .east => try sendCommandUnchecked(&.{ "window", intStr(sp.@"first-window").slice(), "--insert", "west" }),
            .west => try sendCommandUnchecked(&.{ "window", intStr(sp.@"last-window").slice(), "--insert", "east" }),
            .north => try sendCommandUnchecked(&.{ "window", intStr(sp.@"last-window").slice(), "--insert", "south" }),
            .south => try sendCommandUnchecked(&.{ "window", intStr(sp.@"last-window").slice(), "--insert", "north" }),
        };
        defer stream.close();
        streams.appendAssumeCapacity(stream);

        const stream2 = try sendCommandUnchecked(&.{ "window", "--display", @tagName(input_dir) });
        defer stream2.close();
        streams.appendAssumeCapacity(stream2);
    };

    try focusInOtherDisplay(input_dir);
    for (streams.slice()) |st| try checkResponse(st);
}

/// Formats the command, writes to yabai's socket, and checks the reply for error messages.
fn sendCommand(args: []const []const u8) !void {
    var buf: std.BoundedArray(u8, 64) = .{};
    buf.appendSliceAssumeCapacity(&.{ 0, 0, 0, 0 });

    for (args) |arg| {
        buf.appendSliceAssumeCapacity(arg);
        buf.appendAssumeCapacity(0);
    }
    buf.appendAssumeCapacity(0);
    buf.slice()[0] = @intCast(buf.len - 4); // payload len

    const stream = try std.net.connectUnixSocket(yabai_path);
    defer stream.close();

    try stream.writeAll(buf.slice());
    try checkResponse(stream);
}

/// Similar to `sendCommand()` but doesn't read yabai's reply to allow the caller to defer the delay.
/// Caller is responsible for the returned stream.
fn sendCommandUnchecked(args: []const []const u8) !std.net.Stream {
    var buf: std.BoundedArray(u8, 64) = .{};
    buf.appendSliceAssumeCapacity(&.{ 0, 0, 0, 0 });

    for (args) |arg| {
        buf.appendSliceAssumeCapacity(arg);
        buf.appendAssumeCapacity(0);
    }
    buf.appendAssumeCapacity(0);
    buf.slice()[0] = @intCast(buf.len - 4); // payload len

    const stream = try std.net.connectUnixSocket(yabai_path);
    errdefer stream.close();

    try stream.writeAll(buf.slice());
    return stream;
}

/// Sends a query command to yabai and parses the JSON response into struct `T`
fn query(comptime T: type, query_args: []const []const u8) !T {
    switch (T) {
        Window, []Window, Display, []Display, Space, []Space => {},
        else => @compileError("Invalid query type"),
    }

    var args: std.BoundedArray([]const u8, 48) = .{};
    args.appendAssumeCapacity("query");
    args.appendSliceAssumeCapacity(query_args);

    const stream = try sendCommandUnchecked(args.slice());
    defer stream.close();

    var reader = std.json.reader(arena_alloc, stream.reader());
    return try std.json.parseFromTokenSourceLeaky(T, arena_alloc, &reader, .{ .ignore_unknown_fields = true });
}

fn IntStr(comptime T: type) type {
    const is_signed: comptime_int = @intFromBool(@typeInfo(T).Int.signedness == .signed);
    return std.BoundedArray(u8, std.fmt.count("{d}", .{std.math.maxInt(T) + is_signed}));
}

fn intStr(int: anytype) IntStr(@TypeOf(int)) {
    var result: IntStr(@TypeOf(int)) = .{};
    result.writer().print("{d}", .{int}) catch unreachable;
    return result;
}

/// Check yabai's response for error messages. Adds noticeable delay when chaining commands. Caller must close stream.
fn checkResponse(stream: std.net.Stream) !void {
    const error_map = std.ComptimeStringMap(Error, .{
        .{ "could not locate the selected window.", error.LocateSelected },
        .{ "could not locate a northward managed window.", error.LocateNorth },
        .{ "could not locate a southward managed window.", error.LocateSouth },
        .{ "could not locate a eastward managed window.", error.LocateEast },
        .{ "could not locate a westward managed window.", error.LocateWest },
        .{ "could not locate the first managed window.", error.LocateFirst },
        .{ "could not locate the last managed window.", error.LocateLast },
        .{ "could not locate the next managed window.", error.LocateNext },
        .{ "could not locate the prev managed window.", error.LocatePrev },
        .{ "could not locate the next stacked window.", error.LocateNextStacked },
        .{ "could not locate the prev stacked window.", error.LocatePrevStacked },
        .{ "could not locate the first stacked window.", error.LocateFirstStacked },
        .{ "could not locate the last stacked window.", error.LocateLastStacked },
        .{ "could not locate a northward display.", error.LocateNorthDisplay },
        .{ "could not locate a southward display.", error.LocateSouthDisplay },
        .{ "could not locate a eastward display.", error.LocateEastDisplay },
        .{ "could not locate a westward display.", error.LocateWestDisplay },
    });

    var buf: [64]u8 = undefined;
    const n = stream.readAll(&buf) catch |err| {
        print("checkResponse readAll failed", .{});
        return err;
    };

    if (n > 0) return error_map.get(buf[1 .. n - 1]) orelse switch (buf[0]) {
        '{', '[', '"' => return error.UnexpectedJsonResponse,
        else => {
            print("unknown error: {s}", .{buf[0..n]});
            return error.ErrorNotFound;
        },
    };
}

const Error = error{
    LocateSelected,
    LocateNorth,
    LocateSouth,
    LocateEast,
    LocateWest,
    LocateFirst,
    LocateLast,
    LocateNext,
    LocatePrev,
    LocateNextStacked,
    LocatePrevStacked,
    LocateFirstStacked,
    LocateLastStacked,
    LocateNorthDisplay,
    LocateSouthDisplay,
    LocateEastDisplay,
    LocateWestDisplay,
};

const Window = struct {
    const notch_height: f16 = 40.0;
    frame: Frame,
    id: u32,
    display: u8,
    space: u8,
    @"stack-index": i8,
    @"has-focus": bool,
    @"is-visible": bool,

    fn isStacked(self: Window) bool {
        return self.@"stack-index" > 0;
    }
    fn isFullHeight(self: Window, active_display: Display) bool {
        return ((active_display.frame.h - self.frame.h) <= notch_height);
    }
    fn isTopRow(self: Window, active_display: Display) bool {
        return self.frame.y < notch_height and self.frame.h < (active_display.frame.h - notch_height);
    }
    fn isSimilar(self: Window, target_window: Window) bool {
        const similar_w = (self.frame.w / target_window.frame.w) > 0.95 and (self.frame.w / target_window.frame.w) < 1.05;
        const similar_h = (self.frame.h / target_window.frame.h) > 0.95 and (self.frame.h / target_window.frame.h) < 1.05;
        return similar_w and similar_h;
    }
};

const Display = struct {
    frame: Frame,
};

const Space = struct {
    @"first-window": u32,
    @"last-window": u32,
    @"is-visible": bool,
};

const Frame = struct {
    x: f16,
    y: f16,
    w: f16,
    h: f16,
};

const Direction = enum {
    east,
    west,
    north,
    south,
};