From dfac3014f4b03f14b513107bd8733669eac7416e Mon Sep 17 00:00:00 2001
From: Brett <brettmaster1@gmail.com>
Date: Tue, 19 Mar 2024 13:40:14 -0400
Subject: [PATCH] tree depth calculation

---
 CMakeLists.txt         |  2 +-
 include/lilfbtf/tree.h | 34 +++++++++++++++++++++++++++++-----
 src/tree.cpp           | 34 +++++++++++++++++++++++++++++++++-
 3 files changed, 63 insertions(+), 7 deletions(-)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 3ea3f79..4563ba2 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -1,5 +1,5 @@
 cmake_minimum_required(VERSION 3.25)
-project(lilfbtf5 VERSION 0.1.32)
+project(lilfbtf5 VERSION 0.1.33)
 
 option(ENABLE_ADDRSAN "Enable the address sanitizer" OFF)
 option(ENABLE_UBSAN "Enable the ub sanitizer" OFF)
diff --git a/include/lilfbtf/tree.h b/include/lilfbtf/tree.h
index 7a0d080..1c4facc 100644
--- a/include/lilfbtf/tree.h
+++ b/include/lilfbtf/tree.h
@@ -139,18 +139,22 @@ namespace fb
                     tree(tree), engine(info.engine), types(info.types), terminal_chance(info.terminal_chance)
             {}
         };
+        
+        struct tree_eval_t
+        {
+            blt::unsafe::any_t value;
+            type_id contained_type;
+        };
     }
     
     class tree_t
     {
         private:
-            blt::bump_allocator<blt::BLT_2MB_SIZE, false> alloc;
-            type_engine_t& types;
-            detail::node_t* root = nullptr;
-            
             inline blt::bump_allocator<blt::BLT_2MB_SIZE, false>& get_allocator()
             { return alloc; }
             
+            void recalculate_cache();
+            
             static detail::node_t* allocate_non_terminal(detail::node_construction_info_t info, type_id type);
             
             static detail::node_t* allocate_non_terminal_restricted(detail::node_construction_info_t info, type_id type);
@@ -158,6 +162,7 @@ namespace fb
             static detail::node_t* allocate_terminal(detail::node_construction_info_t info, type_id type);
             
             static void grow(detail::node_construction_info_t info, blt::size_t min_depth, blt::size_t max_depth);
+            
             static void brett_grow(detail::node_construction_info_t info, blt::size_t min_depth, blt::size_t max_depth);
             
             static void full(detail::node_construction_info_t info, blt::size_t depth);
@@ -168,7 +173,26 @@ namespace fb
             static tree_t make_tree(detail::tree_construction_info_t tree_info, blt::size_t min_depth, blt::size_t max_depth,
                                     std::optional<type_id> starting_type = {});
             
-            std::pair<blt::unsafe::any_t, type_id> evaluate();
+            detail::tree_eval_t evaluate();
+            
+            blt::size_t depth();
+            
+            // invalidates the internal cache, as the result of tree modification
+            inline void invalidate()
+            {
+                cache.dirty = true;
+            }
+        
+        private:
+            blt::bump_allocator<blt::BLT_2MB_SIZE, false> alloc;
+            type_engine_t& types;
+            detail::node_t* root = nullptr;
+            struct cache_t
+            {
+                blt::size_t depth = 0;
+                blt::size_t node_count = 0;
+                bool dirty = true;
+            } cache;
     };
 }
 
diff --git a/src/tree.cpp b/src/tree.cpp
index 6b2ed7c..06d7158 100644
--- a/src/tree.cpp
+++ b/src/tree.cpp
@@ -63,7 +63,7 @@ namespace fb
         return tree;
     }
     
-    std::pair<blt::unsafe::any_t, type_id> tree_t::evaluate()
+    detail::tree_eval_t tree_t::evaluate()
     {
         using detail::node_t;
         std::stack<node_t*> nodes;
@@ -258,5 +258,37 @@ namespace fb
         }
     }
     
+    blt::size_t tree_t::depth()
+    {
+        if (cache.dirty)
+            recalculate_cache();
+        return cache.depth;
+    }
+    
+    void tree_t::recalculate_cache()
+    {
+        using detail::node_t;
+        blt::size_t depth = 0;
+        blt::size_t node_count = 0;
+        std::stack<std::pair<node_t*, std::size_t>> nodes;
+        
+        nodes.emplace(root, 1);
+        
+        while (!nodes.empty())
+        {
+            auto top = nodes.top();
+            auto* node = top.first;
+            auto d = top.second;
+            node_count++;
+            depth = std::max(d, depth);
+            nodes.pop();
+            for (blt::size_t i = 0; i < node->type.argc(); i++)
+                nodes.emplace(node->children[i], d + 1);
+        }
+        cache.dirty = false;
+        cache.depth = depth;
+        cache.node_count = node_count;
+    }
+    
     
 }
\ No newline at end of file