Autocompletion trie for English words in zig

Question

I am trying to write a program that creates a Trie from a list of all English words, and then when run provides completions for a given prefix. Basically for each node I allocate a list of 255 null chars, and if a node has a child at a certain character, I set the node at index of the character as an int to have a value and it's own children.

It's pretty fast right now as I'm embedding the list of words into the program to avoid reading from a file. It currently has a max memory footprint of 2G when running. Is there some optimization I could make to decrease the amount of memory I use?

Heres the link for the word list: link

const std = @import("std");
const NodePool = std.ArrayList(Trie.Node);
const data = @embedFile("words_alpha.txt");

const Trie = struct {
    pool: NodePool,
    allocator: std.mem.Allocator,

    pub const Node = struct {
        end: bool,
        children: [256]?*Node,
    };

    pub fn init(allocator: std.mem.Allocator) !Trie {
        return Trie{
            .allocator = allocator,
            .pool = try NodePool.initCapacity(allocator, 3494707),
        };
    }

    pub fn deinit(self: *Trie) void {
        self.pool.deinit();
    }

    pub fn alloc_node(self: *Trie) !*Node {
        const node = Node{
            .end = false,
            .children = [_]?*Node{null} ** 256,
        };
        try self.pool.append(node);
        return &(self.pool.items[self.pool.items.len - 1]);
    }

    pub fn add(self: *Trie, root: *Node, slice: []const u8) !void {
        if (slice.len == 0) {
            return;
        }

        const index: usize = @intCast(slice[0]);
        if (root.children[index] == null) {
            var node = try self.alloc_node();
            if (slice.len == 1) {
                node.end = true;
            }
            root.children[index] = node;
        }

        try self.add(root.children[index].?, slice[1..]);
    }

    pub fn get_sub_tree(self: *Trie, root: *Node, slice: []const u8) *Node {
        var node = root;
        var str = slice[0..];
        while (str.len > 0) {
            const index: usize = @intCast(str[0]);
            if (node.children[index] == null) {
                break;
            }
            node = node.children[index].?;
            str = str[1..];
        }

        _ = self;

        return node;
    }

    pub fn get_autocompletion(self: *Trie, root: ?*Node, writer: anytype, ac_buffer: *std.ArrayList(u8)) !void {
        if (root) |r| {
            if (r.end) {
                try writer.print("{s}\n", .{ac_buffer.items});
            }

            var i: usize = 0;
            while (i < r.children.len) : (i += 1) {
                try ac_buffer.append(@intCast(i));
                try self.get_autocompletion(r.children[i], writer, ac_buffer);
                _ = ac_buffer.pop();
            }
        }
    }
};

pub fn main() !void {
    var gpa = std.heap.GeneralPurposeAllocator(.{}){};
    defer _ = gpa.deinit();
    const allocator = gpa.allocator();

    var stdout = std.io.getStdOut();
    var stdin = std.io.getStdIn();
    defer stdin.close();
    defer stdout.close();

    const writer = stdout.writer();
    const reader = stdin.reader();
    var trie = try Trie.init(allocator);
    defer trie.deinit();

    const root = try trie.alloc_node();

    // Read wordlist and create trie
    var itr = std.mem.splitScalar(u8, data, '\r');
    var next = itr.next();
    while (next) |n| {
        try trie.add(root, n[1..]);
        //next = try file.reader().readUntilDelimiterOrEof(&buf, '\r');
        next = itr.next();
    }

    // Event Loop
    while (true) {
        try writer.print("Enter prefix: ", .{});

        const slice = try reader.readUntilDelimiterAlloc(
            allocator,
            '\n',
            20,
        );
        defer allocator.free(slice);

        // If nothing is entered exit
        if (slice.len == 0) {
            break;
        }

        const node = trie.get_sub_tree(root, slice);

        // Buffer that writer !!
        var words = std.ArrayList(u8).init(allocator);
        var temp = std.ArrayList(u8).init(allocator);
        defer temp.deinit();
        defer words.deinit();

        try trie.get_autocompletion(node, words.writer(), &temp);
        try writer.print("{s}", .{words.items});
        try writer.print("\n", .{});
    }
}

Answer 1

One way you can cut down on memory usage is by only using the set of printable characters instead of the whole ASCII charset. In the children array, 32 should map to 0, 65 should map to 33, and so on. This brings down the size for each Node's children from 256 to 95.

The mapping itself is quite simple, simply subtract 32 when storing the character, and add 32 when printing it. And of course, 32 should be some named, not naked, constant in your code. You could also look at using a dynamically growing hash table.

As the requirement is to only print English words, you can further eliminate most special characters (judging from this), and some other printable symbols which may not be part of your data-set, and adjust the mapping accordingly. (If it gets too complicated, use some sort of a map).

---

Side-note: I did something similar in C some time ago that might be of interest (algorithm-wise): Trie Implementation, Autocompletion, and Graph Visualization

In my application, changing the children array from 255 to 95 brought down the struct's size (the struct contained a bool field and the children array) from 1028 bytes to 384 bytes (I used 32-bit integers to store the indices in a dynamically growing pool of nodes. Previously, I was using 64-bit pointers, and the struct size had went down from 2056 bytes to 768 bytes). I should expect a similar reduction in your implementation too.