Rockin' around the BGP Trie

pub const TrieNode = struct {
    // 0 or 1 to lead to the next bit (node)
    zero: ?*TrieNode,
    one: ?*TrieNode,

    // an IP can have more than 1 ASN
    asns: ?[]i32,

    fn blank() @This() {
        return .{ .zero = null, .one = null, .asns = null };
    }
};

const std = @import("std");

const ArenaAllocator = std.heap.ArenaAllocator;
const Allocator = std.mem.Allocator;

pub const TrieParent = struct {
    arena: ArenaAllocator,
    root: TrieNode,

    // a beginner mistake in Zig is passing another arena's 
    // `arena.allocator()` into a new ArenaAllocator. 
    // don't do that! it'll segfault immediately.
    pub fn init(allocator: Allocator) @This() {
        return .{
            .arena = ArenaAllocator.init(allocator),
            .root = TrieNode.blank(),
        };
    }

    pub fn deinit(self: @This()) void {
        self.arena.deinit();
    }
};

// insert an item into the trie given an ip 
// and a length the asn is associated data
pub fn insert(self: *@This(), ip: u32, length: u5, asn: i32) !void {
    const allocator = self.arena.allocator();

    var current_node = &self.root;
    // the shift is how many bits from the left we are in
    // 0 meaning we are starting at the leftmost bit and going
    // rightward until we reach the desired length
    var shift: u5 = 0;

    while (shift < length) {
        // checks if the shift(th) bit from the left is 1 or 0, 
        // get the corresponding node from the current node
        // as a pointer (so that if it is null we can set the 
        // pointer's value to a pointer to a valid new TrieNode)
        const node = if (ip >> (31 - shift) & 1 == 0)
            &current_node.zero 
        else 
            &current_node.one;

        // if the node isn't null, set the current node and 
        // shift another bit rightwards and continue
        if (node.*) |n| {
            current_node = n;
            shift += 1;
            continue;
        }

        // otherwise we have to create a new node to place there
        const new_node = try allocator.create(TrieNode);
        new_node.* = TrieNode.blank();
        node.* = new_node;

        current_node = new_node;
        shift += 1;
    }

    // now that we've reached the desired node we can add our data

    // if there's already an ASN there, we need to expand the array
    if (current_node.asns) |asns| {
        // allocate the same array but one longer
        const new_asns = try allocator.alloc(i32, asns.len + 1);
        // copy the old content over
        @memcpy(new_asns[0..asns.len], asns);
        // set the last item to the new ASN
        new_asns[asns.len] = asn;
        // free the old array
        allocator.free(asns);
        // sort the array for future optimizations...
        std.sort.insertion(i32, new_asns, {}, std.sort.asc(i32));
        // set the asns to the new array
        current_node.asns = new_asns;

        return;
    }

    // otherwise since there's no array here we can just allocate
    // one and set the array's singlular value to the asn and
    // set the node's asn array to the pointer.
    const new_asns = try allocator.alloc(i32, 1);
    new_asns[0] = asn;
    current_node.asns = new_asns;
}

pub fn query(self: @This(), ip: u32) ?[]i32 {
    var current_node = &self.root;

    // same concept as the last but since we can count up 
    // instead of counting down we can just count backwards
    var shift: u5 = 31;

    var least_specific = current_node.asns;

    while (shift > 0) {
        const node = if (ip >> shift & 1 == 0)
            current_node.zero
        else
            current_node.one;

        // if the node exists go deeper
        // otherwise give up
        if (node) |n| {
            if (n.asns) |asns| least_specific = asns;
            current_node = n;
        } else {
            return least_specific;
        }

        shift -= 1;
    }

    return least_specific;
}