From 975b823e3bb068efd430630f0d343d476a2b1bdd Mon Sep 17 00:00:00 2001
From: LordMZTE <lord@mzte.de>
Date: Tue, 8 Aug 2023 19:09:43 +0200
Subject: [PATCH] feat: add alecor script

---
 .config/fish/conf.d/50-alecor.fish |   1 +
 scripts/alecor/.gitignore          |   1 +
 scripts/alecor/build.zig           |  25 ++++
 scripts/alecor/src/cache.zig       |  60 ++++++++
 scripts/alecor/src/correct.zig     | 217 +++++++++++++++++++++++++++++
 scripts/alecor/src/main.zig        |  68 +++++++++
 scripts/alecor/src/util.zig        | 100 +++++++++++++
 setup/commands/install-scripts.rkt |   1 +
 8 files changed, 473 insertions(+)
 create mode 100644 .config/fish/conf.d/50-alecor.fish
 create mode 100644 scripts/alecor/.gitignore
 create mode 100644 scripts/alecor/build.zig
 create mode 100644 scripts/alecor/src/cache.zig
 create mode 100644 scripts/alecor/src/correct.zig
 create mode 100644 scripts/alecor/src/main.zig
 create mode 100644 scripts/alecor/src/util.zig

diff --git a/.config/fish/conf.d/50-alecor.fish b/.config/fish/conf.d/50-alecor.fish
new file mode 100644
index 0000000..5399027
--- /dev/null
+++ b/.config/fish/conf.d/50-alecor.fish
@@ -0,0 +1 @@
+alecor printfish | source
diff --git a/scripts/alecor/.gitignore b/scripts/alecor/.gitignore
new file mode 100644
index 0000000..fe95f8d
--- /dev/null
+++ b/scripts/alecor/.gitignore
@@ -0,0 +1 @@
+/zig-*
diff --git a/scripts/alecor/build.zig b/scripts/alecor/build.zig
new file mode 100644
index 0000000..6361516
--- /dev/null
+++ b/scripts/alecor/build.zig
@@ -0,0 +1,25 @@
+const std = @import("std");
+
+pub fn build(b: *std.Build) void {
+    const target = b.standardTargetOptions(.{});
+    const optimize = b.standardOptimizeOption(.{});
+
+    const exe = b.addExecutable(.{
+        .name = "alecor",
+        .root_source_file = .{ .path = "src/main.zig" },
+        .target = target,
+        .optimize = optimize,
+    });
+
+    b.installArtifact(exe);
+
+    const run_cmd = b.addRunArtifact(exe);
+    run_cmd.step.dependOn(b.getInstallStep());
+
+    if (b.args) |args| {
+        run_cmd.addArgs(args);
+    }
+
+    const run_step = b.step("run", "Run the app");
+    run_step.dependOn(&run_cmd.step);
+}
diff --git a/scripts/alecor/src/cache.zig b/scripts/alecor/src/cache.zig
new file mode 100644
index 0000000..5aeeec7
--- /dev/null
+++ b/scripts/alecor/src/cache.zig
@@ -0,0 +1,60 @@
+const std = @import("std");
+
+pub fn commandsCachePath(alloc: std.mem.Allocator) ![]const u8 {
+    return try std.fs.path.join(alloc, &.{
+        std.os.getenv("HOME") orelse return error.HomeNotSet,
+        ".cache",
+        "alecor",
+        "commands",
+    });
+}
+
+pub fn generate(alloc: std.mem.Allocator) !void {
+    const cache_path = try commandsCachePath(alloc);
+    defer alloc.free(cache_path);
+
+    if (std.fs.path.dirname(cache_path)) |cache_dir| {
+        try std.fs.cwd().makePath(cache_dir);
+    }
+
+    var cache_file = try std.fs.cwd().createFile(cache_path, .{});
+    defer cache_file.close();
+
+    const pipefds = try std.os.pipe();
+    defer std.os.close(pipefds[0]);
+
+    var stdout_buf_reader = std.io.bufferedReader((std.fs.File{ .handle = pipefds[0] }).reader());
+
+    // ChildProcess being useless again...
+    const pid = try std.os.fork();
+    if (pid == 0) {
+        errdefer std.os.exit(1);
+        try std.os.dup2(pipefds[1], 1);
+        std.os.close(pipefds[0]);
+        std.os.close(pipefds[1]);
+        return std.os.execvpeZ(
+            "fish",
+            &[_:null]?[*:0]const u8{ "fish", "-c", "complete -C ''" },
+            @ptrCast(std.os.environ.ptr),
+        );
+    }
+
+    std.os.close(pipefds[1]);
+
+    var cmd_buf: [1024]u8 = undefined;
+    var fbs = std.io.fixedBufferStream(&cmd_buf);
+    while (true) {
+        fbs.reset();
+        stdout_buf_reader.reader().streamUntilDelimiter(fbs.writer(), '\n', null) catch |e| switch (e) {
+            error.EndOfStream => break,
+            else => return e,
+        };
+
+        // FBS will have <cmd>\tgarbage here
+        var spliter = std.mem.tokenize(u8, fbs.getWritten(), "\t");
+        try cache_file.writeAll(spliter.next() orelse continue);
+        try cache_file.writer().writeByte('\n');
+    }
+
+    _ = std.os.waitpid(pid, 0);
+}
diff --git a/scripts/alecor/src/correct.zig b/scripts/alecor/src/correct.zig
new file mode 100644
index 0000000..5933554
--- /dev/null
+++ b/scripts/alecor/src/correct.zig
@@ -0,0 +1,217 @@
+const std = @import("std");
+
+const util = @import("util.zig");
+
+/// Commands which are prioritized for correction
+const priority_commands = [_][]const u8{
+    "git",
+};
+
+/// Commands which wrap another
+const wrapper_commands = std.ComptimeStringMap(void, .{
+    .{ "doas", {} },
+    .{ "sudo", {} },
+    .{ "rbg", {} },
+    .{ "rbgd", {} },
+    .{ "pkexec", {} },
+});
+
+fn populateArgMap(map: *ArgMap) !void {
+    try map.put(&.{"git"}, .{ .subcommand = &.{ "push", "pull", "reset", "checkout" } });
+    try map.put(&.{ "git", "checkout" }, .file_or_directory);
+}
+
+const ArgRequirement = union(enum) {
+    subcommand: []const []const u8,
+    file,
+    directory,
+    file_or_directory,
+};
+
+const ArgMap = std.HashMap(
+    []const []const u8,
+    ArgRequirement,
+    struct {
+        pub fn hash(self: @This(), v: []const []const u8) u64 {
+            _ = self;
+            var hasher = std.hash.Wyhash.init(0);
+            hasher.update(std.mem.asBytes(&v.len));
+            for (v) |s| {
+                hasher.update(std.mem.asBytes(&v.len));
+                hasher.update(s);
+            }
+
+            return hasher.final();
+        }
+
+        pub fn eql(self: @This(), a: []const []const u8, b: []const []const u8) bool {
+            _ = self;
+            if (a.len != b.len)
+                return false;
+
+            for (a, b) |va, vb|
+                if (!std.mem.eql(u8, va, vb))
+                    return false;
+
+            return true;
+        }
+    },
+    std.hash_map.default_max_load_percentage,
+);
+
+pub fn correctCommand(
+    arena: *std.heap.ArenaAllocator,
+    /// Command to correct in-place
+    cmd: [][]const u8,
+    /// Set of all valid commands
+    commands: *std.StringHashMap(void),
+) !void {
+    const alloc = arena.child_allocator;
+
+    var subslice = cmd;
+
+    // skip wrapper commands
+    while (subslice.len > 0 and wrapper_commands.has(subslice[0]))
+        subslice = subslice[1..];
+
+    // empty command
+    if (subslice.len == 0)
+        return;
+
+    if (!commands.contains(subslice[0])) {
+        // correct command
+        var best: ?struct { []const u8, usize } = null;
+
+        // do priority commands first and sub 1 from distance
+        for (priority_commands) |possible_cmd| {
+            const dist = try util.dist(alloc, subslice[0], possible_cmd) -| 1; // prioritize by subtracting 1
+            if (best == null or best.?.@"1" > dist)
+                best = .{ possible_cmd, dist };
+        }
+
+        if (best != null and best.?.@"1" != 0) {
+            var iter = commands.keyIterator();
+            while (iter.next()) |possible_cmd| {
+                const dist = try util.dist(alloc, subslice[0], possible_cmd.*);
+                if (best == null or best.?.@"1" > dist)
+                    best = .{ possible_cmd.*, dist };
+            }
+        }
+
+        if (best) |b| {
+            if (!std.mem.eql(u8, subslice[0], b.@"0")) {
+                std.log.info("[C] {s} => {s}", .{ subslice[0], b.@"0" });
+                subslice[0] = b.@"0";
+            }
+        }
+    }
+
+    if (subslice.len < 2)
+        return;
+
+    var arg_map = ArgMap.init(alloc);
+    defer arg_map.deinit();
+    try populateArgMap(&arg_map);
+
+    // correct args. loop as long as corrections are made
+    while (true) {
+        var req: ArgRequirement = .file_or_directory;
+
+        var cmd_slice_end = subslice.len - 1;
+
+        while (cmd_slice_end > 1) : (cmd_slice_end -= 1) {
+            if (arg_map.get(subslice[0..cmd_slice_end])) |r| {
+                req = r;
+                break;
+            }
+        }
+
+        var new_arg = subslice[cmd_slice_end];
+        try correctArgForReq(arena, req, &new_arg);
+        if (!std.mem.eql(u8, subslice[cmd_slice_end], new_arg)) {
+            std.log.info("[A] {s} => {s}", .{ subslice[cmd_slice_end], new_arg });
+            subslice[cmd_slice_end] = new_arg;
+        } else break;
+    }
+}
+
+fn correctArgForReq(arena: *std.heap.ArenaAllocator, req: ArgRequirement, arg: *[]const u8) !void {
+    const alloc = arena.child_allocator;
+    switch (req) {
+        .subcommand => |subcmds| {
+            var best: ?struct { []const u8, usize } = null;
+            for (subcmds) |possible_cmd| {
+                const dist = try util.dist(alloc, arg.*, possible_cmd);
+                if (best == null or best.?.@"1" > dist)
+                    best = .{ possible_cmd, dist };
+            }
+
+            if (best) |b|
+                arg.* = b.@"0";
+        },
+        .file, .directory, .file_or_directory => {
+            if (arg.len == 0)
+                return;
+
+            var path_spliter = std.mem.tokenize(u8, arg.*, "/");
+            var path_splits = std.ArrayList([]const u8).init(alloc);
+            defer path_splits.deinit();
+
+            // path is absolute
+            if (arg.*[0] == '/')
+                try path_splits.append("/");
+
+            while (path_spliter.next()) |split| {
+                if (std.mem.eql(u8, split, "~")) {
+                    try path_splits.append(std.os.getenv("HOME") orelse return error.HomeNotSet);
+                } else {
+                    try path_splits.append(split);
+                }
+            }
+
+            for (path_splits.items, 0..) |*split, cur_idx| {
+                const dirs = path_splits.items[0..cur_idx];
+                const dir_subpath = try std.fs.path.join(arena.allocator(), if (dirs.len == 0) &.{"."} else dirs);
+
+                var iterable_dir = try std.fs.cwd().openIterableDir(dir_subpath, .{});
+                defer iterable_dir.close();
+
+                // if the given file already exists, there's no point in iterating the dir
+                if (iterable_dir.dir.statFile(split.*)) |_| continue else |e| switch (e) {
+                    error.FileNotFound => {},
+                    else => return e,
+                }
+
+                var best: ?struct { []const u8, usize } = null;
+
+                var dir_iter = iterable_dir.iterate();
+
+                var best_buf: [1024]u8 = undefined;
+
+                while (try dir_iter.next()) |entry| {
+                    switch (req) {
+                        .file => if (entry.kind == .directory) continue,
+                        .directory => if (entry.kind != .directory and
+                            entry.kind != .sym_link) continue,
+                        else => {},
+                    }
+
+                    const dist = try util.dist(alloc, split.*, entry.name);
+                    if (best == null or best.?.@"1" > dist) {
+                        if (entry.name.len > best_buf.len)
+                            return error.OutOfMemory;
+                        const buf_slice = best_buf[0..entry.name.len];
+                        @memcpy(buf_slice, entry.name);
+                        best = .{ buf_slice, dist };
+                    }
+                }
+
+                if (best) |b| {
+                    split.* = try arena.allocator().dupe(u8, b.@"0");
+                } else break;
+            }
+
+            arg.* = try std.fs.path.join(arena.allocator(), path_splits.items);
+        },
+    }
+}
diff --git a/scripts/alecor/src/main.zig b/scripts/alecor/src/main.zig
new file mode 100644
index 0000000..87c1a48
--- /dev/null
+++ b/scripts/alecor/src/main.zig
@@ -0,0 +1,68 @@
+const std = @import("std");
+
+const cache = @import("cache.zig");
+const util = @import("util.zig");
+
+pub fn main() !void {
+    if (std.os.argv.len < 2)
+        return error.NotEnoughArguments;
+
+    var gpa = std.heap.GeneralPurposeAllocator(.{}){};
+    defer _ = gpa.deinit();
+    const alloc = gpa.allocator();
+
+    const subcmd = std.mem.span(std.os.argv[1]);
+
+    if (std.mem.eql(u8, subcmd, "doalec")) {
+        if (std.os.argv.len < 3)
+            return error.NotEnoughArguments;
+
+        var args = std.ArrayList([]const u8).init(alloc);
+        defer args.deinit();
+
+        var spliter = std.mem.tokenize(u8, std.mem.span(std.os.argv[2]), "\n");
+        while (spliter.next()) |arg|
+            try args.append(arg);
+
+        // open and map cache
+        const cache_path = try cache.commandsCachePath(alloc);
+        defer alloc.free(cache_path);
+
+        var cache_file = try std.fs.cwd().openFile(cache_path, .{});
+        defer cache_file.close();
+
+        const cache_content = try std.os.mmap(
+            null,
+            (try cache_file.stat()).size,
+            std.os.PROT.READ,
+            std.os.MAP.PRIVATE,
+            cache_file.handle,
+            0,
+        );
+        defer std.os.munmap(cache_content);
+
+        var command_set = std.StringHashMap(void).init(alloc);
+        defer command_set.deinit();
+
+        var cache_tok = std.mem.tokenize(u8, cache_content, "\n");
+        while (cache_tok.next()) |tok|
+            if (tok.len != 0)
+                try command_set.put(tok, {});
+
+        var arena = std.heap.ArenaAllocator.init(alloc);
+        defer arena.deinit();
+
+        try @import("correct.zig").correctCommand(&arena, args.items, &command_set);
+        try std.io.getStdOut().writer().print("{}\n", .{util.fmtCommand(args.items)});
+    } else if (std.mem.eql(u8, subcmd, "printfish")) {
+        try std.io.getStdOut().writer().print(
+            \\function alec --description 'ALEC'
+            \\    commandline (builtin history search -n 1)
+            \\    commandline ({s} doalec (commandline -o | string split0))
+            \\end
+            \\
+        , .{std.os.argv[0]});
+    } else if (std.mem.eql(u8, subcmd, "mkcache")) {
+        try cache.generate(alloc);
+    } else return error.UnknownCommand;
+}
diff --git a/scripts/alecor/src/util.zig b/scripts/alecor/src/util.zig
new file mode 100644
index 0000000..d11fc50
--- /dev/null
+++ b/scripts/alecor/src/util.zig
@@ -0,0 +1,100 @@
+const std = @import("std");
+
+/// A 2-dimensional heap-allocated matrix.
+pub fn Matrix2D(comptime T: type) type {
+    return struct {
+        data: []T,
+        width: usize,
+
+        const Self = @This();
+
+        pub inline fn init(alloc: std.mem.Allocator, height: usize, width: usize) !Self {
+            return .{ .data = try alloc.alloc(T, height * width), .width = width };
+        }
+
+        pub inline fn deinit(self: Self, alloc: std.mem.Allocator) void {
+            alloc.free(self.data);
+        }
+
+        pub inline fn el(self: *Self, row: usize, col: usize) *T {
+            return &self.data[row * self.width + col];
+        }
+    };
+}
+
+/// Calculates the Damerau-Levenshtein distance between 2 strings
+pub fn dist(alloc: std.mem.Allocator, a: []const u8, b: []const u8) !usize {
+    var d = try Matrix2D(usize).init(alloc, a.len + 1, b.len + 1);
+    defer d.deinit(alloc);
+
+    @memset(d.data, 0);
+
+    var i: usize = 0;
+    var j: usize = 0;
+
+    while (i <= a.len) : (i += 1) {
+        d.el(i, 0).* = i;
+    }
+
+    while (j <= b.len) : (j += 1) {
+        d.el(0, j).* = j;
+    }
+
+    i = 1;
+    while (i <= a.len) : (i += 1) {
+        j = 1;
+        while (j <= b.len) : (j += 1) {
+            const cost = @intFromBool(a[i - 1] != b[j - 1]);
+            d.el(i, j).* = @min(
+                d.el(i - 1, j).* + 1, // deletion
+                @min(
+                    d.el(i, j - 1).* + 1, // insertion
+                    d.el(i - 1, j - 1).* + cost, // substitution
+                ),
+            );
+
+            // transposition
+            if (i > 1 and j > 1 and a[i - 1] == b[j - 2] and a[i - 2] == b[j - 1])
+                d.el(i, j).* = @min(d.el(i, j).*, d.el(i - 2, j - 2).* + cost);
+        }
+    }
+
+    return d.el(a.len, b.len).*;
+}
+
+fn formatCommand(
+    cmd: []const []const u8,
+    comptime fmt: []const u8,
+    options: std.fmt.FormatOptions,
+    writer: anytype,
+) !void {
+    _ = options;
+    _ = fmt;
+
+    var first = true;
+    for (cmd) |arg| {
+        defer first = false;
+        var needs_quote = false;
+        for (arg) |ch| {
+            if (!std.ascii.isPrint(ch) or ch == '\'' or ch == ' ' or ch == '*' or ch == '$') {
+                needs_quote = true;
+                break;
+            }
+        }
+
+        if (!first)
+            try writer.writeByte(' ');
+
+        if (needs_quote) {
+            try writer.writeByte('"');
+            try writer.print("{}", .{std.fmt.fmtSliceEscapeUpper(arg)});
+            try writer.writeByte('"');
+        } else {
+            try writer.writeAll(arg);
+        }
+    }
+}
+
+pub fn fmtCommand(cmd: []const []const u8) std.fmt.Formatter(formatCommand) {
+    return .{ .data = cmd };
+}
diff --git a/setup/commands/install-scripts.rkt b/setup/commands/install-scripts.rkt
index 7a25759..63e1203 100644
--- a/setup/commands/install-scripts.rkt
+++ b/setup/commands/install-scripts.rkt
@@ -20,6 +20,7 @@
   (mklink "scripts/use-country-mirrors.sh" (bin-path "use-country-mirrors"))
 
   ;; Compile scripts
+  (install-zig "scripts/alecor")
   (install-rust "scripts/i3status")
   (install-zig "scripts/mzteinit")
   (install-zig "scripts/openbrowser")