zig/test/standalone/brace_expansion.zig

const std = @import("std");
const io = std.io;
const mem = std.mem;
const debug = std.debug;
const assert = debug.assert;
const testing = std.testing;
const ArrayList = std.ArrayList;
const maxInt = std.math.maxInt;

const Token = union(enum) {
    Word: []const u8,
    OpenBrace,
    CloseBrace,
    Comma,
    Eof,
};

var gpa = std.heap.GeneralPurposeAllocator(.{}){};
var global_allocator = gpa.allocator();

fn tokenize(input: []const u8) !ArrayList(Token) {
    const State = enum {
        Start,
        Word,
    };

    var token_list = ArrayList(Token).init(global_allocator);
    errdefer token_list.deinit();
    var tok_begin: usize = undefined;
    var state = State.Start;

    for (input, 0..) |b, i| {
        switch (state) {
            .Start => switch (b) {
                'a'...'z', 'A'...'Z' => {
                    state = State.Word;
                    tok_begin = i;
                },
                '{' => try token_list.append(Token.OpenBrace),
                '}' => try token_list.append(Token.CloseBrace),
                ',' => try token_list.append(Token.Comma),
                else => return error.InvalidInput,
            },
            .Word => switch (b) {
                'a'...'z', 'A'...'Z' => {},
                '{', '}', ',' => {
                    try token_list.append(Token{ .Word = input[tok_begin..i] });
                    switch (b) {
                        '{' => try token_list.append(Token.OpenBrace),
                        '}' => try token_list.append(Token.CloseBrace),
                        ',' => try token_list.append(Token.Comma),
                        else => unreachable,
                    }
                    state = State.Start;
                },
                else => return error.InvalidInput,
            },
        }
    }
    switch (state) {
        State.Start => {},
        State.Word => try token_list.append(Token{ .Word = input[tok_begin..] }),
    }
    try token_list.append(Token.Eof);
    return token_list;
}

const Node = union(enum) {
    Scalar: []const u8,
    List: ArrayList(Node),
    Combine: []Node,

    fn deinit(self: Node) void {
        switch (self) {
            .Scalar => {},
            .Combine => |pair| {
                pair[0].deinit();
                pair[1].deinit();
                global_allocator.free(pair);
            },
            .List => |list| {
                for (list.items) |item| {
                    item.deinit();
                }
                list.deinit();
            },
        }
    }
};

const ParseError = error{
    InvalidInput,
    OutOfMemory,
};

fn parse(tokens: *const ArrayList(Token), token_index: *usize) ParseError!Node {
    const first_token = tokens.items[token_index.*];
    token_index.* += 1;

    const result_node = switch (first_token) {
        .Word => |word| Node{ .Scalar = word },
        .OpenBrace => blk: {
            var list = ArrayList(Node).init(global_allocator);
            errdefer {
                for (list.items) |node| node.deinit();
                list.deinit();
            }
            while (true) {
                try list.append(try parse(tokens, token_index));

                const token = tokens.items[token_index.*];
                token_index.* += 1;

                switch (token) {
                    .CloseBrace => break,
                    .Comma => continue,
                    else => return error.InvalidInput,
                }
            }
            break :blk Node{ .List = list };
        },
        else => return error.InvalidInput,
    };

    switch (tokens.items[token_index.*]) {
        .Word, .OpenBrace => {
            const pair = try global_allocator.alloc(Node, 2);
            errdefer global_allocator.free(pair);
            pair[0] = result_node;
            pair[1] = try parse(tokens, token_index);
            return Node{ .Combine = pair };
        },
        else => return result_node,
    }
}

fn expandString(input: []const u8, output: *ArrayList(u8)) !void {
    const tokens = try tokenize(input);
    defer tokens.deinit();
    if (tokens.items.len == 1) {
        return output.resize(0);
    }

    var token_index: usize = 0;
    const root = try parse(&tokens, &token_index);
    defer root.deinit();
    const last_token = tokens.items[token_index];
    switch (last_token) {
        Token.Eof => {},
        else => return error.InvalidInput,
    }

    var result_list = ArrayList(ArrayList(u8)).init(global_allocator);
    defer {
        for (result_list.items) |*buf| buf.deinit();
        result_list.deinit();
    }

    try expandNode(root, &result_list);

    try output.resize(0);
    for (result_list.items, 0..) |buf, i| {
        if (i != 0) {
            try output.append(' ');
        }
        try output.appendSlice(buf.items);
    }
}

const ExpandNodeError = error{OutOfMemory};

fn expandNode(node: Node, output: *ArrayList(ArrayList(u8))) ExpandNodeError!void {
    assert(output.items.len == 0);
    switch (node) {
        .Scalar => |scalar| {
            var list = ArrayList(u8).init(global_allocator);
            errdefer list.deinit();
            try list.appendSlice(scalar);
            try output.append(list);
        },
        .Combine => |pair| {
            const a_node = pair[0];
            const b_node = pair[1];

            var child_list_a = ArrayList(ArrayList(u8)).init(global_allocator);
            defer {
                for (child_list_a.items) |*buf| buf.deinit();
                child_list_a.deinit();
            }
            try expandNode(a_node, &child_list_a);

            var child_list_b = ArrayList(ArrayList(u8)).init(global_allocator);
            defer {
                for (child_list_b.items) |*buf| buf.deinit();
                child_list_b.deinit();
            }
            try expandNode(b_node, &child_list_b);

            for (child_list_a.items) |buf_a| {
                for (child_list_b.items) |buf_b| {
                    var combined_buf = ArrayList(u8).init(global_allocator);
                    errdefer combined_buf.deinit();

                    try combined_buf.appendSlice(buf_a.items);
                    try combined_buf.appendSlice(buf_b.items);
                    try output.append(combined_buf);
                }
            }
        },
        .List => |list| {
            for (list.items) |child_node| {
                var child_list = ArrayList(ArrayList(u8)).init(global_allocator);
                errdefer for (child_list.items) |*buf| buf.deinit();
                defer child_list.deinit();

                try expandNode(child_node, &child_list);

                for (child_list.items) |buf| {
                    try output.append(buf);
                }
            }
        },
    }
}

pub fn main() !void {
    defer _ = gpa.deinit();
    const stdin_file = io.getStdIn();
    const stdout_file = io.getStdOut();

    const stdin = try stdin_file.reader().readAllAlloc(global_allocator, std.math.maxInt(usize));
    defer global_allocator.free(stdin);

    var result_buf = ArrayList(u8).init(global_allocator);
    defer result_buf.deinit();

    try expandString(stdin, &result_buf);
    try stdout_file.writeAll(result_buf.items);
}

test "invalid inputs" {
    global_allocator = std.testing.allocator;

    try expectError("}ABC", error.InvalidInput);
    try expectError("{ABC", error.InvalidInput);
    try expectError("}{", error.InvalidInput);
    try expectError("{}", error.InvalidInput);
    try expectError("A,B,C", error.InvalidInput);
    try expectError("{A{B,C}", error.InvalidInput);
    try expectError("{A,}", error.InvalidInput);

    try expectError("\n", error.InvalidInput);
}

fn expectError(test_input: []const u8, expected_err: anyerror) !void {
    var output_buf = ArrayList(u8).init(global_allocator);
    defer output_buf.deinit();

    try testing.expectError(expected_err, expandString(test_input, &output_buf));
}

test "valid inputs" {
    global_allocator = std.testing.allocator;

    try expectExpansion("{x,y,z}", "x y z");
    try expectExpansion("{A,B}{x,y}", "Ax Ay Bx By");
    try expectExpansion("{A,B{x,y}}", "A Bx By");

    try expectExpansion("{ABC}", "ABC");
    try expectExpansion("{A,B,C}", "A B C");
    try expectExpansion("ABC", "ABC");

    try expectExpansion("", "");
    try expectExpansion("{A,B}{C,{x,y}}{g,h}", "ACg ACh Axg Axh Ayg Ayh BCg BCh Bxg Bxh Byg Byh");
    try expectExpansion("{A,B}{C,C{x,y}}{g,h}", "ACg ACh ACxg ACxh ACyg ACyh BCg BCh BCxg BCxh BCyg BCyh");
    try expectExpansion("{A,B}a", "Aa Ba");
    try expectExpansion("{C,{x,y}}", "C x y");
    try expectExpansion("z{C,{x,y}}", "zC zx zy");
    try expectExpansion("a{b,c{d,e{f,g}}}", "ab acd acef aceg");
    try expectExpansion("a{x,y}b", "axb ayb");
    try expectExpansion("z{{a,b}}", "za zb");
    try expectExpansion("a{b}", "ab");
}

fn expectExpansion(test_input: []const u8, expected_result: []const u8) !void {
    var result = ArrayList(u8).init(global_allocator);
    defer result.deinit();

    expandString(test_input, &result) catch unreachable;

    try testing.expectEqualSlices(u8, expected_result, result.items);
}