math_shell/src/main.zig

pub fn main() !void {
    clear();

    var gpa = std.heap.GeneralPurposeAllocator(.{}).init;
    defer _ = gpa.deinit();

    const allocator = gpa.allocator();

    var ln = Linenoise.init(allocator);
    defer ln.deinit();

    var prng = Random.DefaultPrng.init(blk: {
        var seed: u64 = undefined;
        try posix.getrandom(mem.asBytes(&seed));
        break :blk seed;
    });
    const rand = prng.random();

    const uname_cmd = try tokenizeCommand("uname -a", allocator);
    defer allocator.free(uname_cmd);

    var shell_state: ShellState = .{
        .env_map = try process.getEnvMap(allocator),
    };
    defer shell_state.env_map.deinit();

    _ = try runCommand(uname_cmd, &shell_state, allocator);
    var p = try getPrompt(allocator);

    var limit: i32 = 100;

    while (try ln.linenoise(p)) |input| {
        allocator.free(p);
        defer allocator.free(input);

        if (input.len > 0) {
            try ln.history.add(input);
            const command = try tokenizeCommand(input, allocator);
            defer allocator.free(command);
            if (!shell_state.should_test or
                eql(u8, command[0], "nomorenumbers") or
                try mathTest(rand, limit, shell_state.iteration))
            {
                _ = execCommand(command, &shell_state, allocator) catch |err| switch (err) {
                    error.FileNotFound => print("mash: {s}: command not found\n", .{command[0]}),
                    error.AccessDenied => print("mash: {s}: Permission denied\n", .{command[0]}),
                    else => print("Unkown error: {}\n", .{err}),
                };
            } else {
                print(
                    \\
                    \\How can you expect to execute commands without knowing any math???
                    \\
                , .{});
            }

            shell_state.iteration += 1;
            if (shell_state.iteration % 3 == 0) {
                limit *= 10;
            }
        }
        if (shell_state.should_exit) break;
        p = try getPrompt(allocator);
    } else allocator.free(p); // free on crtl+d
}

const ShellState = struct {
    iteration: u32 = 0,
    should_test: bool = true,
    should_exit: bool = false,
    env_map: process.EnvMap,
};

const Operator = enum(u8) {
    add = 0,
    subtract,
    multiply,
    divide,

    const Self = @This();

    fn toString(s: Self) []const u8 {
        return switch (s) {
            .add => "+",
            .subtract => "-",
            .multiply => "*",
            .divide => "/",
        };
    }

    fn apply(s: Self, a: i32, b: i32) i32 {
        return switch (s) {
            .add => a + b,
            .subtract => a - b,
            .multiply => a * b,
            .divide => @divTrunc(a, b),
        };
    }

    /// Gets add or subtract for the first three commands.
    /// Includes multiplication and division after the first three.
    fn rand(r: Random, iteration: u32) Self {
        if (iteration > 3) {
            return r.enumValue(Self);
        } else {
            return @enumFromInt(r.uintAtMost(u8, 1));
        }
    }
};

/// return true if the user passed, false if the user did not.
fn mathTest(rand: Random, limit: i32, iteration: u32) !bool {
    const op = Operator.rand(rand, iteration);
    const num1 = rand.intRangeAtMost(i32, 1, limit);
    const num2 = rand.intRangeAtMost(i32, 1, limit);
    const answer = op.apply(num1, num2);

    if (comptime builtin.mode == .Debug) {
        print("Answer: {d}\n", .{answer});
    }
    print("{d} {s} {d} = ", .{ num1, op.toString(), num2 });

    defer {
        print("Calculating", .{});
        for (0..3) |_| {
            print(".", .{});
            if (comptime builtin.mode != .Debug) {
                sleep(2 * 1000 * 1000 * 1000);
            }
        }
        print("\n", .{});
    }

    return answer == while (true) {
        break readInt(stdin) catch |err| switch (err) {
            error.InvalidCharacter, error.EndOfStream => continue,
            else => return err,
        };
    };
}

const BuiltinFunc = *const fn (*ShellState, [][]const u8, Allocator) anyerror!void;

fn exitFn(shell_state: *ShellState, command: [][]const u8, allocator: Allocator) !void {
    _ = command;
    _ = allocator;
    shell_state.*.should_exit = true;
}

fn nomorenumbersFn(shell_state: *ShellState, command: [][]const u8, allocator: Allocator) !void {
    shell_state.*.should_test = false;
    if (command.len > 1) {
        _ = execCommand(command[1..], shell_state, allocator) catch |err| {
            return err;
        };
    }
}

fn cdFn(shell_state: *ShellState, command: [][]const u8, allocator: Allocator) !void {
    const cwd = fs.cwd();
    var env_map = shell_state.*.env_map;
    const pwd = try cwd.realpathAlloc(allocator, ".");
    const old_pwd = env_map.get("OLDPWD");
    const home = env_map.get("HOME");
    defer allocator.free(pwd);
    var err: ?anyerror = null;

    var target: ?[]const u8 = null;
    if (command.len == 1) {
        if (home) |h| {
            target = h;
        }
    } else if (eql(u8, command[1], "-")) {
        if (old_pwd) |old| {
            target = old;
        } else {
            print("mash: cd: OLDPWD not set\n", .{});
            return;
        }
    } else {
        var buf: [fs.max_path_bytes]u8 = undefined;
        target = cwd.realpath(command[1], &buf) catch |e| blk: {
            err = e;
            break :blk null;
        };
    }

    if (target) |t| {
        if (posix.chdir(t)) {
            try env_map.put("OLDPWD", pwd);
            try env_map.put("PWD", t);
        } else |e| {
            err = e;
        }
    }

    if (err) |e| {
        switch (e) {
            error.FileNotFound => {
                print("mash: cd: {s}: No such file or directory\n", .{target orelse command[1]});
                return;
            },
            error.AccessDenied => {
                print("mash: cd: {s}: Permission denied\n", .{target orelse command[1]});
                return;
            },
            error.NameTooLong => {
                print("mash: cd: {s}: File name too long\n", .{target orelse command[1]});
                return;
            },
            else => return e,
        }
    }
}

fn authorFn(shell_state: *ShellState, command: [][]const u8, allocator: Allocator) !void {
    _ = shell_state;
    _ = command;
    _ = allocator;
    print(
        \\Author: Robby
        \\Description: A shell to promote math!
        \\Ask for the source code!
        \\
        \\Derived from an earlier version written in C by Spencer.
        \\
    , .{});
}

fn execCommand(command: [][]const u8, shell_state: *ShellState, allocator: Allocator) !u8 {
    const builtins = comptime .{
        .{ "exit", exitFn },
        .{ "nomorenumbers", nomorenumbersFn },
        .{ "cd", cdFn },
        .{ "author", authorFn },
    };
    const builtinMap = StaticStringMap(BuiltinFunc).initComptime(builtins);

    if (builtinMap.has(command[0])) {
        try builtinMap.get(command[0]).?(shell_state, command, allocator);
        return 0;
    } else {
        return runCommand(command, shell_state, allocator);
    }
}

fn runCommand(command: [][]const u8, shell_state: *ShellState, allocator: Allocator) !u8 {
    var child = process.Child.init(command, allocator);
    child.env_map = &shell_state.env_map;
    const result = try child.spawnAndWait();
    return result.Exited;
}

fn getPrompt(allocator: Allocator) ![]const u8 {
    const script = @embedFile("./get_prompt.bash");
    const result = try process.Child.run(.{
        .allocator = allocator,
        .argv = &[_][]const u8{ "bash", "-i", "-c", script },
    });
    allocator.free(result.stderr);
    return result.stdout;
}

fn tokenizeCommand(command: []const u8, allocator: Allocator) ![][]const u8 {
    var argv_array_list = ArrayList([]const u8).init(allocator);
    defer argv_array_list.deinit();

    var tokens = tokenizeAny(u8, command, "\t\r\n ");
    while (tokens.next()) |token| {
        try argv_array_list.append(token);
    }
    return argv_array_list.toOwnedSlice();
}

fn clear() void {
    print("\u{001b}[H\u{001b}[J", .{});
}

fn print(comptime fmt: []const u8, args: anytype) void {
    stdout.print(fmt, args) catch unreachable;
}

fn readInt(reader: anytype) !i32 {
    var buf: [1024]u8 = undefined;
    const buf2 = try reader.readUntilDelimiter(&buf, '\n');
    return std.fmt.parseInt(i32, buf2, 0);
}

const std = @import("std");
const posix = std.posix;

const fs = std.fs;
const path = std.fs.path;

const mem = std.mem;
const Allocator = mem.Allocator;
const eql = mem.eql;

const Reader = std.io.Reader;
const stdout = std.io.getStdOut().writer();
const stdin = std.io.getStdIn().reader();
const tokenizeAny = mem.tokenizeAny;

const process = std.process;

const sleep = std.Thread.sleep;

const Random = std.Random;
const ArrayList = std.ArrayList;
const StaticStringMap = std.static_string_map.StaticStringMap;

const Linenoise = @import("linenoise").Linenoise;

const builtin = @import("builtin");