From 0f907769e3c348694be658f706a93e56fdaea198 Mon Sep 17 00:00:00 2001
From: Humza Shahid <humzasaur@gmail.com>
Date: Thu, 5 Feb 2026 20:39:01 +0000
Subject: [PATCH] add a function to persistent-vector.sml to check if all
 leaves are at the same depth (which is how we want to balance)

---
 fcore/persistent-vector.sml | 36 ++++++++++++++++++++++++++++++++++--
 1 file changed, 34 insertions(+), 2 deletions(-)

diff --git a/fcore/persistent-vector.sml b/fcore/persistent-vector.sml
index 5849bad..fda1f83 100644
--- a/fcore/persistent-vector.sml
+++ b/fcore/persistent-vector.sml
@@ -711,9 +711,17 @@ struct
       if leftDepth = rightDepth then
         mergeSameDepth (left, right)
       else if leftDepth < rightDepth then
-        mergeWhenRightDepthIsGreater (left, right, leftDepth, 0)
+        let
+          val targetDepth = rightDepth - leftDepth
+        in
+          mergeWhenRightDepthIsGreater (left, right, leftDepth, 0)
+        end
       else
-        mergeWhenLeftDepthIsGreater (left, right, rightDepth, 0)
+        let
+          val targetDepth = leftDepth - rightDepth
+        in
+          mergeWhenLeftDepthIsGreater (left, right, rightDepth, 0)
+        end
     end
 
   fun delete (start, length, tree) =
@@ -827,6 +835,30 @@ struct
     end
 
   (* functions only for testing *)
+  fun allLeavesAtSameDepth tree =
+    case tree of
+      BRANCH (nodes, _) =>
+        let
+          fun loop (pos, expectedDepth) =
+            if pos = Vector.length nodes then
+              true
+            else
+              let
+                val node = Vector.sub (nodes, pos)
+                val nodeDepth = countDepth node
+              in
+                if nodeDepth = expectedDepth then
+                  loop (pos + 1, expectedDepth)
+                else
+                  false
+              end
+
+          val expectedDepth = countDepth (Vector.sub (nodes, 0))
+        in
+          loop (0, expectedDepth)
+        end
+    | LEAF _ => true
+
   fun fromListLoop (lst, acc) =
     case lst of
       {start, finish} :: tl =>