Improve Binary Tree Array by getting rid of the recursion

We use a stack instead, it’s not terribly different
from the recursion but it saves function calls
which use a lot more stack than needed.

libs/utils/include/utils/BinaryTreeArray.h

@@ -17,40 +17,93 @@
 #ifndef TNT_UTILS_BINARYTREEARRAY_H
 #define TNT_UTILS_BINARYTREEARRAY_H
 
+#include <utils/compiler.h>
+
+#include <type_traits>
+
 #include <stddef.h>
 #include <stdint.h>
 
 namespace utils {
 
 class BinaryTreeArray {
+
+    // Simple fixed capacity stack
+    template<typename TYPE, size_t CAPACITY,
+            typename = typename std::enable_if<std::is_pod<TYPE>::value>::type>
+    class stack {
+        TYPE mElements[CAPACITY];
+        size_t mSize = 0;
+    public:
+        bool empty() const noexcept { return mSize == 0; }
+        void push(TYPE const& v) noexcept {
+            assert(mSize < CAPACITY);
+            mElements[mSize++] = v;
+        }
+        void pop() noexcept {
+            assert(mSize > 0);
+            --mSize;
+        }
+        const TYPE& back() const noexcept {
+            return mElements[mSize - 1];
+        }
+    };
+
 public:
     static size_t count(size_t height) noexcept { return  (1u << height) - 1; }
     static size_t left(size_t i, size_t height) noexcept { return i + 1; }
     static size_t right(size_t i, size_t height) noexcept { return i + (1u << (height - 1)); }
 
+    // this builds the depth-first binary tree array top down (post-order)
     template<typename Leaf, typename Node>
     static void traverse(size_t height, Leaf leaf, Node node) noexcept {
-        traverse(0, 0, 0, height, count(height), leaf, node);
-    }
 
-private:
-    // this builds the depth-first binary tree array top down, so we need to evaluate the
-    // light ranges only once.
-    template<typename Leaf, typename Node>
-    static void traverse(size_t index, size_t parent,
-                         size_t col, size_t height,
-                         size_t next,
-                         Leaf& leaf, Node& node) noexcept {
-        if (height > 1) {
-            size_t l = left(index, height);
-            size_t r = right(index, height);
-            // the 'next' node of our left node's right descendants is our right child
-            traverse(l, index, 2 * col,     height - 1, r, leaf, node);
-            // the 'next' node of our right child is our own 'next' sibling
-            traverse(r, index, 2 * col + 1, height - 1, next, leaf, node);
-            node(index, parent, l, r, next);
-        } else {
-            leaf(index, parent, col, next);
+        struct TNode {
+            uint32_t index;
+            uint32_t col;
+            uint32_t height;
+            uint32_t next;
+
+            bool isLeaf() const noexcept { return height == 1; }
+            size_t left() const noexcept { return BinaryTreeArray::left(index, height); }
+            size_t right() const noexcept { return BinaryTreeArray::right(index, height); }
+        };
+
+        stack<TNode, 16> stack;
+        stack.push(TNode{ 0, 0, (uint32_t)height, (uint32_t)count(height) });
+
+        uint32_t prevLeft = 0;
+        uint32_t prevRight = 0;
+        uint32_t prevIndex = 0;
+        while (!stack.empty()) {
+            TNode const* const UTILS_RESTRICT curr = &stack.back();
+            const bool isLeaf = curr->isLeaf();
+            const uint32_t index = curr->index;
+            const uint32_t l = (uint32_t)curr->left();
+            const uint32_t r = (uint32_t)curr->right();
+
+            if (prevLeft == index || prevRight == index) {
+                if (!isLeaf) {
+                    // the 'next' node of our left node's right descendants is our right child
+                    stack.push({ l, 2 * curr->col, curr->height - 1, r });
+                }
+            } else if (l == prevIndex) {
+                if (!isLeaf) {
+                    // the 'next' node of our right child is our own 'next' sibling
+                    stack.push({ r, 2 * curr->col + 1, curr->height - 1, curr->next });
+                }
+            } else {
+                if (!isLeaf) {
+                    node(index, l, r, curr->next);
+                } else {
+                    leaf(index, curr->col, curr->next);
+                }
+                stack.pop();
+            }
+
+            prevLeft  = l;
+            prevRight = r;
+            prevIndex = index;
         }
     }
 };