Background#

For context, see my previous post about defining nodes and why I’m doing this.

Basic Tree Operations#

In the previous post we defined our basic node structure and now we will use the nodes to build a tree.

For this tree I wanted to define some basic constraints that will influence my design.

  1. The tree must own all of its own allocations. This will make more sense when we talk about how Zig manages heap memory.
  2. Formally we are building an arborescence where m=2 and each node is unique, but informally this is a binary tree that will be unbalanced yet sorted.
  3. Duplicate values cannot be inserted into the tree.

If we remember the node implementation, each node has a left_child and a right_child that are optional pointers to other nodes. This means the linked structure of our tree with a depth of 3 would look like this.

                      ┌──────────────────────┐
                      │   root : ?*Node(T)   │
                      └───────────┬──────────┘
                                  │
             ┌────────────────────┴────────────────────┐
             ▼                                         ▼
     ┌──────────────────────┐                ┌──────────────────────┐
     │ left_child : ?*Node  │                │ right_child : ?*Node │
     └───────────┬──────────┘                └───────────┬──────────┘
                 │                                       │
      ┌──────────┴──────────┐                 ┌──────────┴──────────┐
      ▼                     ▼                 ▼                     ▼
┌──────────────┐     ┌──────────────┐   ┌──────────────┐     ┌──────────────┐
│ left_child   │     │ right_child  │   │ left_child   │     │ right_child  │
│ (ll) : ?*Node│     │ (lr) : null  │   │ (rl) : ?*Node│     │ (rr) : ?*Node│
└──────────────┘     └──────────────┘   └──────────────┘     └──────────────┘

Each node will point to its child, but that pointer is wrapped in an option. You can see from the root.left_child.right_child that when an option has no underlying value it is null. Pointers in Zig cannot be null, so I terminate node pointers with a null option.

Let’s jump into stubbing out the tree in code. I will start with our type definition and type properties.

pub fn Tree(comptime T: type) type {
  return struct {
    root: ?*Node(T),
    allocator:  std.mem.allocator,

    const Self = @this();

  }
}

The code above defines the following things

  • A Tree of type T, where T is not known until compile time
  • A root that follows the same typing as our Nodes by being optional, and pointing to a Node.
  • An Allocator as part of my type, more on these in just a second.
  • Lastly my function scoped Self that will allow me to define instance methods with brevity.

Heap Memory#

Now that we have defined an allocator, what is it and how is it used? At a basic level, allocators will allow us to create heap memory allocations. Zig provides a number of different allocator implementations that utilize different strategies for memory allocation. These different implementations all exist to optimize how efficiently you can manage memory for your use case, we are going to stick to the General Purpose and Debug allocators in this post.

One of the key pillars of Zig is that there are no hidden memory allocations. When writing Zig code all of the memory management is done explicit and there is no guessing what memory is allocated. The abstraction given to use by Zig to manage heap memory is the Allocator. Allocators provide an easy to understand interface for seeing exactly when heap memory is used.

You may wonder how is this different from other languages, in C I can allocate memory with malloc, and in languages like Go I don’t even have to deal with this. The answer comes down to what you value and the constraints of your program. Lets compare Zig to other languages and how memory is managed.

Go#

In Zig you will know exactly what and how much data in in the heap. In Go heap allocations are decisions that the programmer is not always privy to, the Go compiler make heap allocation decisions based on the size, lifetime and usage of program data. While sometimes the heap allocations are easy to identify, you have no control over this in Go and are subject to whatever decision the go compiler makes for you.

Go can also lead to memory problems, for example a go routine’s function calls can end up creating heap allocations do to the limited size of a goroutines call stack. This means we are getting memory allocations when we wouldn’t really expect to have them. Go’s defer statements can also cause memory issues because each defer takes a segment of memory from the stack, meaning you can accidentally create out of memory errors by using defer in a loop.

If Go hides a lot of memory management and heap allocations what about other languages?

Note: I’m picking on Go a little here, but I do actually like Go and think is a very useful language.

C++ and Rust#

I know, why on earth would I lump Rust and C++ together. Clearly they are very very very different languages with very different goals, ecosystems and syntax. Modern C++ and Rust actually share some common traits for memory management in direct contrast to Zig’s approach.

C++ and Rust both have the concept of smart pointers, and while there are some key differences in the two the underlying result of dynamic heap management is common between the two. Allocation and deallocation depend on references, usage and compiler decisions. While smart pointers are great, they do create yet another situation where memory management is not explicit. Rust has gone as far as trying to implement allocators inside the language to prevent some of the standard library’s hidden allocations from failing with this RFC.

C++ has even further hidden allocations, for example, coroutines dynamically create heap allocations when they are created. Another primitive built into the language that is doing hidden memory management.

The Case for Zig#

In contrast to the languages I have discussed, Zig aims to make all of the memory management clear, explicit and easy to identify. There is no new(T) or any built in function to create a heap allocation. Heap memory is not even part of the language spec or definition, it is added by the standard library. If you do not have an allocator then you are not using heap memory, that is Zig’s guarantee.

To learn more on this topic I would highly recommend watching this talk on allocators in Zig by Benjamin Feng.

Back to our Tree#

Our tree needs an init function, so we can create the empty tree itself. We have talked at length about the need to use an allocator, so we are going to define the init function to accept an allocator for the function caller. This means all heap allocations and deallocations are owned by the code leveraging this tree.

pub fn Tree(comptime T: type) type {
    return struct {
        root: ?*Node(T),
        allocator: std.mem.Allocator,

        const Self = @This();

        pub fn init(allocator: std.mem.Allocator) @This() {
            return .{
                .root = null,
                .allocator = allocator,
            };
        }
    }
}

We set the root to null, since the tree is empty, and set our instance variable to the user provided allocator. In the sample code I shared we are actually just declaring that Tree.allocator is of type std.mem.allocator, this expression you can compare to saying that my allocator property must implement the allocator interface from the standard library.

We can use the above code in a simple test just to confirm it works.

 test "tree_allocation" {
    var gp_allocator = std.heap.GeneralPurposeAllocator(.{}){};
    defer _ = gp_allocator.deinit();
    const allocator = gp_allocator.allocator();
    const test_tree = Tree(i32).init(allocator);

    try std.testing.expectEqual(null, test_tree.root);
}

I create a test case, define a general purpose allocator, and create an empty tree. Notice how even though we didn’t have to de-initialize our tree, since we manage the memory allocator in the caller scope we will always be able to clean up anything the tree allocates by cleaning up our allocator. We know that memory can’t escape my scope because no matter what the Tree will do, that memory is manged by the caller.

You may notice that deinit is a non void function with a return type and value. When de-initializing allocators, especially the general purpose and debug allocators, they can return any found memory leaks but we will cover that later.

We can test individual unit tests by calling zig test like so.

zig test src/tree.zig --test-filter "tree_allocation"

We would see the following results:

All 1 tests passed

Success, our tests pass and now we can move on

Insert#

Our first non trivial function is going to be inserting nodes into our tree. The algorithm to insert some n value into our tree will be as follows:

  1. Start at the root node, if root is empty insert value as the root nodes value, else compare n to root.
  2. If n is greater than the node value, chose the right child as our next comparison.
  3. If n is less than the node value, choose the left child as our next comparison.
  4. Repeat steps 2 and 3 until we find a child node that is null, allocate a new node and set n as it’s value.

Starting with Step 1, let’s define our function signature and the base case.

pub fn insert(self: *Self, value: T) !void {
  // Handle empty tree, set root node and return
  if (self.root == null) {
      self.root = try self.allocator.create(Node(T));
      self.root.?.* = Node(T).init(value);
      return;
  }
}

This part of insert defines a few things we have seen, and some we haven’t.

  • *Self provides us a way to create instance methods and manage instance state.
  • value: T is some generic value we will accept
  • !void is the first time we have used a return type with !. This is some sugar syntax for a function that could return some error, or return a type. This could be expanded to return specific errors such as MemoryError!void, but I left it without an explicit return which is also acceptable.
  • We use our allocator to create.(Node(T)). This will create our Node of type T on the heap, which is managed by the allocator instance we provided when initializing our Tree. create() returns a pointer to the memory that was created. Remember our self.root is a ?*Node(T) so now we actually have that implemented with this line.
  • try is used, try will call a function with a ! in the return type, meaning it will attempt a function that could error. In this case, allocator instance methods could fail if our system is out of memory.
  • Our instance of Tree would have root set, and a value of root initialized to the value parameter in our function.
  • self.root.?.* = Node(T).init(value) is just setting root, which we have to access via unwrapping its option with .? and then setting it’s pointer value with .*.

Sweet, now lets establish how to extend past the base case. Regardless of how many nodes are in the tree, the traversal to add a new node is the same.

The first step is to create a variable to track where we are in the tree, to treat as a comparison point for value. Then we can create our new node with the value parameter.

var current = self.root;

const new_node: ?*Node(T) = try self.allocator.create(Node(T));
new_node.?.* = Node(T).init(value);

Nothing special going on in the above code, everything going on here we have covered already. Now we have a way to track where we are in the tree, and the new node defined. We need to implement a loop that use each iteration to compare the current node’s value to our value parameter. Our current node could have none, one or both child nodes. We will use step 2 or 3 from our insertion algorithm depending on how our value parameter compares to our current node’s value. Once we know what child node we are looking at, check if that node exists and if it exists we use that as the next iteration of our current node. If the child node does not exist, then we set that child as the node we allocated as new_node. Putting our function all together would look like this.

pub fn insert(self: *Self, value: T) !void {
    // Handle empty tree, set root node and bail
    if (self.root == null) {
        self.root = try self.allocator.create(Node(T));
        self.root.?.* = Node(T).init(value);
        return;
    }

    var current = self.root;

    const new_node: ?*Node(T) = try self.allocator.create(Node(T));
    new_node.?.* = Node(T).init(value);

    while (true) {
        // check if option has value since every node is optional
        if (current) |current_node| {
            if (value < current_node.*.value) {
                if (current_node.*.left_child == null) {
                    current_node.*.left_child = new_node;
                    break;
                } else {
                    current = current_node.*.left_child;
                }
            } else if (value > current_node.*.value) {
                if (current_node.*.right_child == null) {
                    current_node.*.right_child = new_node;
                    break;
                } else {
                    current = current_node.*.right_child;
                }
            } else {
                // the value == current node value, we don't insert duplicates
                self.allocator.destroy(new_node.?);
                break;
            }
        } else {
            // we have encountered an indeterminate state
            return;
        }
    }
}

Voila, we have a tree we can initialize and insert values into. Let’s write a test to verify we are inserting values as we expect. We do the same setup as before, by creating an allocator, initializing our tree, and we will insert some values into our tree. We will insert 5, 3, 4, 6 and 3 again in that order. Our tree should look like this:

 5
/ \

3 6
4

Where 5 is our root, it’s left child is 4, it’s right child is 6. The root’s left child has one child node with a value of 4 as it’s right child. The second 3 to be inserted should be a noop since we don’t insert duplicates. Let’s write out our test to verify this behavior.

test "tree_basics" {
    var gp_allocator = std.heap.GeneralPurposeAllocator(.{}){};
    defer _ = gp_allocator.deinit();
    const allocator = gp_allocator.allocator();
    var test_tree = Tree(i32).init(allocator);

    try std.testing.expectEqual(null, test_tree.root);

    try test_tree.insert(5); // root Node()
    try std.testing.expectEqual(5, test_tree.root.?.*.value);
    try std.testing.expectEqual(null, test_tree.root.?.*.left_child);

    try test_tree.insert(3);
    try std.testing.expectEqual(3, test_tree.root.?.*.left_child.?.*.value);
    try std.testing.expectEqual(null, test_tree.root.?.*.right_child);

    try test_tree.insert(4);
    try test_tree.insert(6);
    try std.testing.expectEqual(4, test_tree.root.?.*.left_child.?.*.right_child.?.*.value);
    try std.testing.expectEqual(6, test_tree.root.?.*.right_child.?.*.value);

    // Can't add duplicate
    try test_tree.insert(3);
    try std.testing.expectEqual(3, test_tree.root.?.*.left_child.?.*.value);
    try std.testing.expectEqual(null, test_tree.root.?.*.left_child.?.*.left_child);
}

Everything looks solid, to run this test we can use

zig test src/tree.zig --test-filter "tree_basics"

Now we have run into one of Zig’s memory safeguards, we get the following error:

 zig test src/tree.zig --test-filter "tree_basics"
[gpa] (err): memory address 0x7f6703060000 leaked: 
/.../src/tree.zig:77:54: 0x103f2cf in insert (tree.zig)
                self.root = try self.allocator.create(Node(T));
                                                     ^
/.../src/tree.zig:298:25: 0x103fd25 in test.tree_basics (tree.zig)
    try test_tree.insert(5); // root Node()
                        ^
/usr/lib/zig/compiler/test_runner.zig:66:28: 0x11800f1 in main (test_runner.zig)
        return mainTerminal();
                           ^
/usr/lib/zig/std/start.zig:618:22: 0x1179e8d in posixCallMainAndExit (std.zig)
            root.main();
                     ^
/usr/lib/zig/std/start.zig:232:5: 0x1179721 in _start (std.zig)
    asm volatile (switch (native_arch) {
    ^

[gpa] (err): memory address 0x7f6703060020 leaked: 
/.../src/tree.zig:84:66: 0x103f45d in insert (tree.zig)
            const new_node: ?*Node(T) = try self.allocator.create(Node(T));
                                                                 ^
/.../src/tree.zig:302:25: 0x10400bd in test.tree_basics (tree.zig)
    try test_tree.insert(3);
                        ^
/usr/lib/zig/compiler/test_runner.zig:218:25: 0x1186ed0 in mainTerminal (test_runner.zig)
        if (test_fn.func()) |_| {
                        ^
/usr/lib/zig/compiler/test_runner.zig:66:28: 0x11800f1 in main (test_runner.zig)
        return mainTerminal();
                           ^
/usr/lib/zig/std/start.zig:618:22: 0x1179e8d in posixCallMainAndExit (std.zig)
            root.main();
                     ^
/usr/lib/zig/std/start.zig:232:5: 0x1179721 in _start (std.zig)
    asm volatile (switch (native_arch) {
    ^

[gpa] (err): memory address 0x7f6703060040 leaked: 
/.../src/tree.zig:84:66: 0x103f45d in insert (tree.zig)
            const new_node: ?*Node(T) = try self.allocator.create(Node(T));
                                                                 ^
/.../src/tree.zig:306:25: 0x10404cc in test.tree_basics (tree.zig)
    try test_tree.insert(4);
                        ^
/usr/lib/zig/compiler/test_runner.zig:218:25: 0x1186ed0 in mainTerminal (test_runner.zig)
        if (test_fn.func()) |_| {
                        ^
/usr/lib/zig/compiler/test_runner.zig:66:28: 0x11800f1 in main (test_runner.zig)
        return mainTerminal();
                           ^
/usr/lib/zig/std/start.zig:618:22: 0x1179e8d in posixCallMainAndExit (std.zig)
            root.main();
                     ^
/usr/lib/zig/std/start.zig:232:5: 0x1179721 in _start (std.zig)
    asm volatile (switch (native_arch) {
    ^

[gpa] (err): memory address 0x7f6703060060 leaked: 
/.../src/tree.zig:84:66: 0x103f45d in insert (tree.zig)
            const new_node: ?*Node(T) = try self.allocator.create(Node(T));
                                                                 ^
/.../src/tree.zig:307:25: 0x1040573 in test.tree_basics (tree.zig)
    try test_tree.insert(6);
                        ^
/usr/lib/zig/compiler/test_runner.zig:218:25: 0x1186ed0 in mainTerminal (test_runner.zig)
        if (test_fn.func()) |_| {
                        ^
/usr/lib/zig/compiler/test_runner.zig:66:28: 0x11800f1 in main (test_runner.zig)
        return mainTerminal();
                           ^
/usr/lib/zig/std/start.zig:618:22: 0x1179e8d in posixCallMainAndExit (std.zig)
            root.main();
                     ^
/usr/lib/zig/std/start.zig:232:5: 0x1179721 in _start (std.zig)
    asm volatile (switch (native_arch) {
    ^

All 1 tests passed.
4 errors were logged.
error: the following test command failed with exit code 1:
.zig-cache/o/0da37dfd92864bc60afc2a5be8ef06fd/test --seed=0xe84b9fd1

Some important context from this error, our unit test actually passed, our code executed and ran. The errors thrown are returned from the general purpose allocator. Each memory allocation created by inserting a new node, which consequentially allocated a node on the heap, thew an error when our execution ended with that memory not being cleaned up. The errors call out that we have a memory leak, meaning memory is left on the heap while the context using that memory has ended. We did not see this error earlier when we had our tree initialization because we weren’t allocating any heap memory yet.

This warning looks complex but it is actually just letting us know that we need to implement a de-initialization function for our tree. We need a function that can free all our node memory when we are done using the tree.

Destroying Allocated Memory#

For cleaning up a tree, we actually have to use a second data structure and a breadth first traversal. We will visit each node, add those nodes children to a list of things to be destroyed, and destroy the last visited node. We will use an array list to store the values we delete since an array list can be extended in memory as we wont know the number of values in our tree at any given time. The de-initialization of our tree should run in O(nlogn), with a space complexity of O(2n) since we are copying values from the heap to another data structure until they are deleted. We will also check preemptively to see if we are calling deinit on an emtpy tree and we can just return.

pub fn deinit(self: *Self) void {
    if (self.root == null) {
        return;
    }

    var item_stack = std.ArrayList(*Node(T)).empty;
    defer item_stack.deinit(self.allocator);

    item_stack.append(self.allocator, self.root.?) catch unreachable;

    while (item_stack.items.len > 0) {
        const node_opt = item_stack.pop();

        if (node_opt) |node| {
            if (node.*.right_child) |righ_node| {
                item_stack.append(self.allocator, righ_node) catch unreachable;
            }
            if (node.*.left_child) |left_node| {
                item_stack.append(self.allocator, left_node) catch unreachable;
            }

            self.allocator.destroy(node);
        }
    }
}

The idea here is to use the array list, append each node we visit starting with our root and keep adding the node’s children to the array list and deleting the previous node. To verify our logic we can re purpose our exiting insertion test, because we will get errors if any nodes are not cleaned up. We just add defer test_tree.deinit() to our test right after we initialize the tree. Re run our test again and let’s see what happens now.

> zig test src/tree.zig --test-filter "tree_basics"
All 1 tests passed

Now we have a working tree!

All Together#

Our tree now looks like this when we put all our code together.

const std = @import("std");

pub fn Node(comptime T: type) type {
    return struct {
        value: T,
        left_child: ?*Node(T),
        right_child: ?*Node(T),

        const Self = @This();

        pub fn init(value: T) @This() {
            return .{
                .left_child = null,
                .right_child = null,
                .value = value,
            };
        }

        pub fn is_leaf(self: *Self) bool {
            if (self.left_child || self.right_child) {
                return false;
            }
            return true;
        }
    };
}

pub fn Tree(comptime T: type) type {
    return struct {
        root: ?*Node(T),
        allocator: std.mem.Allocator,

        const Self = @This();

        pub fn init(allocator: std.mem.Allocator) @This() {
            return .{
                .root = null,
                .allocator = allocator,
            };
        }

        pub fn deinit(self: *Self) void {
            if (self.root == null) {
                return;
            }

            var item_stack = std.ArrayList(*Node(T)).empty;
            defer item_stack.deinit(self.allocator);

            item_stack.append(self.allocator, self.root.?) catch unreachable;

            while (item_stack.items.len > 0) {
                const node_opt = item_stack.pop();

                if (node_opt) |node| {
                    if (node.*.right_child) |righ_node| {
                        item_stack.append(self.allocator, righ_node) catch unreachable;
                    }
                    if (node.*.left_child) |left_node| {
                        item_stack.append(self.allocator, left_node) catch unreachable;
                    }

                    self.allocator.destroy(node);
                }
            }
        }

        pub fn insert(self: *Self, value: T) !void {
            // Handle empty tree, set root node and bail
            if (self.root == null) {
                self.root = try self.allocator.create(Node(T));
                self.root.?.* = Node(T).init(value);
                return;
            }

            var current = self.root;

            const new_node: ?*Node(T) = try self.allocator.create(Node(T));
            new_node.?.* = Node(T).init(value);

            while (true) {
                // check if option has value since every node is optional
                if (current) |current_node| {
                    if (value < current_node.*.value) {
                        if (current_node.*.left_child == null) {
                            current_node.*.left_child = new_node;
                            break;
                        } else {
                            current = current_node.*.left_child;
                        }
                    } else if (value > current_node.*.value) {
                        if (current_node.*.right_child == null) {
                            current_node.*.right_child = new_node;
                            break;
                        } else {
                            current = current_node.*.right_child;
                        }
                    } else {
                        self.allocator.destroy(new_node.?);
                        break;
                    }
                } else {
                    break;
                }
            }
        }
}

test "tree_allocation" {
    var gp_allocator = std.heap.GeneralPurposeAllocator(.{}){};
    defer _ = gp_allocator.deinit();
    const allocator = gp_allocator.allocator();
    const test_tree = Tree(i32).init(allocator);

    try std.testing.expectEqual(null, test_tree.root);
}

test "tree_basics" {
    var gp_allocator = std.heap.GeneralPurposeAllocator(.{}){};
    defer _ = gp_allocator.deinit();
    const allocator = gp_allocator.allocator();
    var test_tree = Tree(i32).init(allocator);

    try std.testing.expectEqual(null, test_tree.root);
    defer test_tree.deinit();

    try test_tree.insert(5); // root Node()
    try std.testing.expectEqual(5, test_tree.root.?.*.value);
    try std.testing.expectEqual(null, test_tree.root.?.*.left_child);

    try test_tree.insert(3);
    try std.testing.expectEqual(3, test_tree.root.?.*.left_child.?.*.value);
    try std.testing.expectEqual(null, test_tree.root.?.*.right_child);

    try test_tree.insert(4);
    try test_tree.insert(6);
    try std.testing.expectEqual(4, test_tree.root.?.*.left_child.?.*.right_child.?.*.value);
    try std.testing.expectEqual(6, test_tree.root.?.*.right_child.?.*.value);

    // Can't add duplicate
    try test_tree.insert(3);
    try std.testing.expectEqual(3, test_tree.root.?.*.left_child.?.*.value);
    try std.testing.expectEqual(null, test_tree.root.?.*.left_child.?.*.left_child);
}

The next blog will cover the most challenging tree operation, the removal of nodes while preserving the order of the tree. We will learn how to debug zig programs with lldb and dap in neovim.