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.

filament/src/Froxelizer.cpp

@@ -939,7 +939,7 @@ void Froxelizer::computeLightTree(
     auto const* UTILS_RESTRICT spheres = lightData.data<FScene::POSITION_RADIUS>() + 1;
     BinaryTreeArray::traverse(h,
             [&array, &camera, spheres, indices, count]
-            (size_t index, size_t parent, size_t col, size_t next) {
+            (size_t index, size_t col, size_t next) {
                 float min = 1.0;
                 float max = 0.0;
                 if (col < count) {
@@ -954,7 +954,7 @@ void Froxelizer::computeLightTree(
                 }
                 array[index] = Node{ min, max, (uint16_t)next };
             },
-            [&array](size_t index, size_t parent, size_t l, size_t r, size_t next) {
+            [&array](size_t index, size_t l, size_t r, size_t next) {
                 array[index] = Node{
                         std::min(array[l].min, array[r].min),
                         std::max(array[l].max, array[r].max),

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;
         }
     }
 };

libs/utils/test/test_BinaryTreeArray.cpp

@@ -30,10 +30,10 @@ TEST(BinaryTreeArray, basic) {
     Node array[256];
 
     BinaryTreeArray::traverse(1,
-            [&array](size_t index, size_t parent, size_t col, size_t next) {
+            [&array](size_t index, size_t col, size_t next) {
                 array[index] = Node{ index, next };
             },
-            [&array](size_t index, size_t parent, size_t l, size_t r, size_t next) {
+            [&array](size_t index, size_t l, size_t r, size_t next) {
                 array[index] = Node{ index, next };
             });
 
@@ -42,10 +42,10 @@ TEST(BinaryTreeArray, basic) {
 
 
     BinaryTreeArray::traverse(2,
-            [&array](size_t index, size_t parent, size_t col, size_t next) {
+            [&array](size_t index, size_t col, size_t next) {
                 array[index] = Node{ index, next };
             },
-            [&array](size_t index, size_t parent, size_t l, size_t r, size_t next) {
+            [&array](size_t index, size_t l, size_t r, size_t next) {
                 array[index] = Node{ index, next };
             });
 
@@ -60,10 +60,10 @@ TEST(BinaryTreeArray, basic) {
 
 
     BinaryTreeArray::traverse(3,
-            [&array](size_t index, size_t parent, size_t col, size_t next) {
+            [&array](size_t index, size_t col, size_t next) {
                 array[index] = Node{ index, next };
             },
-            [&array](size_t index, size_t parent, size_t l, size_t r, size_t next) {
+            [&array](size_t index, size_t l, size_t r, size_t next) {
                 array[index] = Node{ index, next };
             });
 
@@ -90,10 +90,10 @@ TEST(BinaryTreeArray, basic) {
 
 
     BinaryTreeArray::traverse(4,
-            [&array](size_t index, size_t parent, size_t col, size_t next) {
+            [&array](size_t index, size_t col, size_t next) {
                 array[index] = Node{ index, next };
             },
-            [&array](size_t index, size_t parent, size_t l, size_t r, size_t next) {
+            [&array](size_t index, size_t l, size_t r, size_t next) {
                 array[index] = Node{ index, next };
             });