From 86d2d2fb8d15fcd3ba36fa10fe592ac47348ed08 Mon Sep 17 00:00:00 2001
From: Leo Collins <collins.leo@proton.me>
Date: Mon, 23 Mar 2026 17:15:40 +0000
Subject: [PATCH] add recursive tree depth function

---
 rtree-capi/include/rtree-capi.h |  2 ++
 rtree-capi/src/rtree.rs         | 54 +++++++++++++++++++++++++++++++--
 rtree-capi/tests/test.c         | 33 ++++++++++++++++++++
 3 files changed, 86 insertions(+), 3 deletions(-)

diff --git a/rtree-capi/include/rtree-capi.h b/rtree-capi/include/rtree-capi.h
index a16d96e..7071f7d 100644
--- a/rtree-capi/include/rtree-capi.h
+++ b/rtree-capi/include/rtree-capi.h
@@ -26,6 +26,8 @@ RTreeError rtree_bulk_load(struct RTreeH **tree,
 
 RTreeError rtree_create(struct RTreeH **tree, uint32_t dim);
 
+RTreeError rtree_depth(const struct RTreeH *tree, size_t *depth_out);
+
 RTreeError rtree_free(struct RTreeH *tree);
 
 RTreeError rtree_free_ids(size_t *ids, size_t n);
diff --git a/rtree-capi/src/rtree.rs b/rtree-capi/src/rtree.rs
index d09265d..17d8651 100644
--- a/rtree-capi/src/rtree.rs
+++ b/rtree-capi/src/rtree.rs
@@ -1,6 +1,6 @@
 use rstar::primitives::{GeomWithData, Rectangle};
-use rstar::RTree;
-use interval_tree::IntervalTree;
+use rstar::{ParentNode, RTree, RTreeNode, RTreeObject};
+use interval_tree::{IntervalTree, IntervalTreeNode};
 
 use crate::error::RTreeError;
 
@@ -144,7 +144,7 @@ pub extern "C" fn rtree_locate_all_at_point(
     let rtree = unsafe { &*(tree as *const RTreeDim) };
     let mut ids: Vec<usize> = match rtree {
         RTreeDim::D1(tree) => {
-            let p: f64 = unsafe { *(point as *const f64) };
+            let p: f64 = unsafe { *point };
             tree.locate_all_at_point(p)
         }
         RTreeDim::D2(tree) => {
@@ -183,6 +183,54 @@ pub extern "C" fn rtree_size(tree: *const RTreeH, size_out: *mut usize) -> RTree
     RTreeError::Success
 }
 
+fn _interval_tree_depth(node: &IntervalTreeNode) -> usize {
+    let left_depth = node
+        .left
+        .as_deref()
+        .map_or(0, |left| 1 + _interval_tree_depth(left));
+    let right_depth = node
+        .right
+        .as_deref()
+        .map_or(0, |right| 1 + _interval_tree_depth(right));
+    left_depth.max(right_depth)
+}
+
+fn _rtree_depth<T: RTreeObject>(node: &ParentNode<T>) -> usize {
+    node
+        .children()
+        .iter()
+        .map(|child| match child {
+            RTreeNode::Leaf(_) => 1,
+            RTreeNode::Parent(parent) => 1 + _rtree_depth(parent),
+        })
+        .max()
+        .unwrap_or(0)
+}
+
+#[no_mangle]
+pub extern "C" fn rtree_depth(tree: *const RTreeH, depth_out: *mut usize) -> RTreeError {
+    if tree.is_null() || depth_out.is_null() {
+        return RTreeError::NullPointer;
+    }
+    let rtree = unsafe { &*(tree as *const RTreeDim) };
+    let size = match rtree {
+        RTreeDim::D1(tree) => tree.size(),
+        RTreeDim::D2(tree) => tree.size(),
+        RTreeDim::D3(tree) => tree.size(),
+    };
+    if size == 0 {
+        unsafe { *depth_out = 0 };
+        return RTreeError::Success;
+    }
+    let depth = match rtree {
+        RTreeDim::D1(tree) => _interval_tree_depth(tree.root().unwrap()),
+        RTreeDim::D2(tree) => _rtree_depth(tree.root()),
+        RTreeDim::D3(tree) => _rtree_depth(tree.root()),
+    };
+    unsafe { *depth_out = depth };
+    RTreeError::Success
+}
+
 #[no_mangle]
 pub extern "C" fn rtree_free_ids(ids: *mut usize, n: usize) -> RTreeError {
     if ids.is_null() {
diff --git a/rtree-capi/tests/test.c b/rtree-capi/tests/test.c
index bd54d87..798d54a 100644
--- a/rtree-capi/tests/test.c
+++ b/rtree-capi/tests/test.c
@@ -103,6 +103,14 @@ bool test_bulk_load(void) {
         return false;
     }
 
+    size_t depth = 0;
+    rtree_depth(tree, &depth);
+    if (depth != 1) {
+        fprintf(stderr, "Expected tree depth 1, got %zu\n", depth);
+        rtree_free(tree);
+        return false;
+    }
+
     double point1[2] = {1.5, 1.5};
     double point2[2] = {0.0, 0.0};
     double point3[2] = {-1.0, 0.0};
@@ -253,6 +261,15 @@ bool test_rtree_1d(void) {
         return false;
     }
 
+    // test rtree_depth
+    size_t depth = 0;
+    rtree_depth(tree, &depth);
+    if (depth != 1) {
+        fprintf(stderr, "Expected tree depth 1, got %zu\n", depth);
+        rtree_free(tree);
+        return false;
+    }
+
     double point1[1] = {0.5};
     double point2[1] = {1.5};
     double point3[1] = {3.5};
@@ -421,6 +438,14 @@ bool test_rtree_empty(void) {
         return false;
     }
 
+    size_t depth = 0;
+    rtree_depth(tree, &depth);
+    if (depth != 0) {
+        fprintf(stderr, "Expected empty tree depth 0, got %zu\n", depth);
+        rtree_free(tree);
+        return false;
+    }
+
     // Query empty tree
     double point[2] = {0.0, 0.0};
     size_t *ids_out = NULL;
@@ -503,6 +528,14 @@ bool test_rtree_empty_1d(void) {
         return false;
     }
 
+    size_t depth = 0;
+    rtree_depth(tree, &depth);
+    if (depth != 0) {
+        fprintf(stderr, "Expected empty tree depth 0, got %zu\n", depth);
+        rtree_free(tree);
+        return false;
+    }
+
     // Query empty tree
     double point[1] = {0.0};
     size_t *ids_out = NULL;