Sparse Arrays and type composition in Zig

One of the lovely things about Zig is the way that you can just create types at compile time and then use these types elsewhere in your code. The other day, I needed a sparse array, but the standard library didn’t really have what I wanted. I could have used one of the hash-map types, but that wasn’t really what I wanted, so I created my own.

pub fn SparseArray(IndexType: type, DataType: type) type {
}

The first thing I did was create a function that returns a type and takes in the types that I’d be using for the index and data types.

I then defined a few constants that would help with the computation of the actual data structure.

const page_size = 64 * 1024;

const data_size = @sizeOf(DataType);
const data_num = @divFloor(page_size, data_size);

const max_index = std.math.maxInt(IndexType);
const page_table_size = max_index / data_num + 1;

const DataTable = [data_num]DataType;

I defined a page size, that defines the size of each page of data within this sparse array. Essentially, what I want is an array that can hold lots of data, indexed by an normal index value, but that if there were large gaps in the data then no memory was wasted. The smaller the page size, the more efficient this is with regards to the memory waste, but the more overhead we had. For example, if we had a page size that was just the size of a single bit of data then we just have an array that allocates memory for each index when it was used.

I then calculate the number of data items that can be stored in a single page. Obviously, this is the page size divided by the size of a bit of data. I then determine the maximum index value based off the IndexType that was passed in. For example, if this was just a 16-bit integer, then the maximum index would be 65535. From this, we can determine the page table size.

We also define a new type, only used within the context of this function called DataTable, that holds the data for a single page.

return struct {
    page_table: [page_table_size]?*DataTable = .{null} ** page_table_size,

    const This = @This();

After this, our function returns a structure. Remember, we are in a function that is returning a type and therefore we return a struct type back to the caller.

We also define that this structure has a single page_table, of the appropriate size that optionally points to a DataTable, that is all initialized to null.

Finally, we define that within this struct, the constant This is defined to be this structure using the @This() builtin.

We have a page table, all initialized to null, we need a few member functions that can operate on this new type.

pub fn deinit(self: This, alloc: std.mem.Allocator) void {
    for (self.page_table) |pt| {
        if (pt) |p| {
            alloc.free(p);
        }
    }
}

The deinit function takes a This, and an allocator, and frees up any memory that was allocated. Obviously, if you just create an instance of this type and then immediately .deiniting it, this would just be a no-operation.

The for syntax just iterates over all the elements in the page_table this this instance, and for each it captures the value into pt. It then checks that this optional value is a value or null. If the value is not-null, then this is captured into p that can then be used to free that page_table entry.

pub fn set(
    self: *This,
    alloc: std.mem.Allocator,
    index: IndexType,
    data: DataType,
) !void {
    const page_num = @divFloor(index, data_num);
    if (self.page_table[page_num]) |page| {
        const offset = @mod(index, data_num);
        page[offset] = data;
    } else {
        const page = try alloc.create(DataTable);
        @memset(page, 0);
        const offset = @mod(index, data_num);
        page[offset] = data;
        self.page_table[page_num] = page;
    }
}

To set a value into the sparse array, we use the set function. This is a member function that takes an allocator, an index value, and a data value. Note that the index is of type IndexType, and data is of type DataType, that were defined in the SparseArray function we are defining the return type for.

The function is fairly simple. We first compute the page number that our data should be in. If that page table entry is not-null, then we just index into that page entry and set the data.

Otherwise, we allocate a new page table entry using the allocator.create and the type DataTable. We then use the builtin @memset function to clear this memory to all zeros. This is not necessarily required, and may be a performance problem if we are constantly creating new page table entries, but frankly, I’d prefer zeroed memory to weird bugs in the future. Note also, that the try alloc.create will pass any allocator error, like an out of memory error, up to the caller to work out what the do. Then we just use the same offset calculation, set the data, and make sure that the page_table entry is correct.

pub fn get(self: This, index: IndexType) ?DataType {
    const page_num = @divFloor(index, data_num);
    if (self.page_table[page_num]) |page| {
        const offset = @mod(index, data_num);
        return page[offset];
    }
    return null;
}

The get function is even easier. Check if the page table entry for the page number we require is not-null, then just capture that page, determine the offset, and return the data. Otherwise, we just return null.

Yes, this does not know if that data was ever set in the past, but it is an array not a hash map. Arrays are simple, hash maps are more complex. Arrays are faster, although sometimes you’ll get data back that was never set, which is why I zero’ed the memory in the first place.

Can we optimize this a little? Perhaps. We could define a put function that sets a value but knows that the page table should be set and therefore would never allocator memory or fail due to an out of memory error. Note that this doesn’t error in the unhappy case. This is a design decision that may not be valid for everybody, but it is fine for my usecase.

pub fn put(self: *This, index: IndexType, data: DataType) void {
    const page_num = @divFloor(index, data_num);
    if (self.page_table[page_num]) |page| {
        const offset = @mod(index, data_num);
        page[offset] = data;
    }
}

And we should never forget testing.

test "sparse creation" {
    var data: SparseArray(u16, u32) = .{};
    defer data.deinit(std.testing.allocator);

    try std.testing.expectEqual(null, data.get(0));
}

My first standard test is the ‘can you create and destroy it without leaking memory’ test. That’s what the sparse creation test is doing.

test "sparse set indexes" {
    var data: SparseArray(u24, u8) = .{};
    defer data.deinit(std.testing.allocator);

    try std.testing.expectEqual(null, data.get(10));

    try data.set(std.testing.allocator, 10, 37);
    try data.set(std.testing.allocator, 11, 39);
    try data.set(std.testing.allocator, 300000, 42);

    try std.testing.expectEqual(42, data.get(300000));
    try std.testing.expectEqual(37, data.get(10));
    try std.testing.expectEqual(39, data.get(11));

    data.put(10, 79);
    try std.testing.expectEqual(79, data.get(10));
}

The next test is attempting to test the functions we’ve defined above. Create a SparseArray. Note, the types used are different from the first test. Is this deliberate? No. But it could be a better test because we did that, so it stays. We first check that we can’t read an empty page for index 10. Then we write some data into the array at indexes 10, 11, and 300000. This will create a couple of pages of data. We then check that we can get that data back out. And then we test to see if we can put an index we’ve already used without passing in an allocator, and check that it worked. A very simple test.

The testing is much more through later on, but only because I use this data structure as part of a larger system and the testing of that system does much more than just exercise a couple of sets and puts and gets.

How much memory did the above test allocate? Well, each page has 64k of memory, and given that we use it to hold u8 data types, that means that we have 65536 entries per table. This means that the data_num is 65536 and that indexes 10 and 11 are in page 0, and index 300000 is in page 4. This means we have to have 2 page table entries used, about 128 kbytes.

Could we make this more efficient? Perhaps. We could pass the page size into the function. We could just define a smaller page size to start with. However, this is sufficient for my needs, so I’m sticking with what I have.

Could I have done this in C++ with Templates? Yes, but I’m not sure that I’d really want to see all those < and > brackets everywhere. Passing in types into a function and returning a type feels elegant, the whole concept of having a different mechanism to do this now feels alien. Yes, I understand that languages like C++ are stuck in the legacy of the past design decisions, but that doesn’t mean that new languages like Zig have to follow along the same way.