From 94f493ec2e522f2e780f00ea9e25d5576b6ad30c Mon Sep 17 00:00:00 2001
From: Christopher Haster <geky@geky.net>
Date: Wed, 22 Jan 2025 15:01:18 -0600
Subject: [PATCH] rbyd: Added LFS_ASSERTRBYDBALANCE for asserting rbyds are
 balanced

This is _not_ free, so an opt-in define is needed.

This is also expected to fail right now due to balance issues with
diverging yellow nodes in range-removals, but it's the first step
towards trying to fix said balance issues.

I've been trying to avoid adding debug machinery to lfs.c itself, since
it tends to hurt readability, but I really don't know how you could
assert rbyd balance out-of-tree short of trying to fetch every block in
emubd somehow...

Oh well, finding and asserting rbyd balance in lfsr_rbyd_fetch at least
doesn't require that much code.

No code changes, thanks to const-propagation.
---
 lfs.c      | 59 ++++++++++++++++++++++++++++++++++++++++++++++++++++--
 lfs_util.h |  6 ++++++
 2 files changed, 63 insertions(+), 2 deletions(-)

diff --git a/lfs.c b/lfs.c
index 73617c22..fb0248f8 100644
--- a/lfs.c
+++ b/lfs.c
@@ -2717,6 +2717,12 @@ static int lfsr_rbyd_ckecksum(lfs_t *lfs, const lfsr_rbyd_t *rbyd,
     return (ecksum_ == ecksum->cksum) ? 0 : LFS_ERR_CORRUPT;
 }
 
+// needed in lfsr_rbyd_fetch_ if asserting rbyd balance
+static int lfsr_rbyd_lookupnext_(lfs_t *lfs, const lfsr_rbyd_t *rbyd,
+        lfsr_srid_t rid, lfsr_tag_t tag,
+        lfsr_srid_t *rid_, lfsr_tag_t *tag_, lfsr_rid_t *weight_,
+        lfsr_data_t *data_, lfs_size_t *height_);
+
 // fetch an rbyd
 static int lfsr_rbyd_fetch_(lfs_t *lfs,
         lfsr_rbyd_t *rbyd, uint32_t *gcksumdelta,
@@ -2931,6 +2937,32 @@ static int lfsr_rbyd_fetch_(lfs_t *lfs,
         rbyd->eoff = -1;
     }
 
+    // asserting rbyd balance? check that all branches in the rbyd have
+    // the same height
+    #ifdef LFS_ASSERTRBYDBALANCE
+    lfsr_srid_t rid = -1;
+    lfsr_tag_t tag = 0;
+    lfs_ssize_t height = -1;
+    while (true) {
+        lfs_size_t height_;
+        int err = lfsr_rbyd_lookupnext_(lfs, rbyd,
+                rid, tag+1,
+                &rid, &tag, NULL, NULL, &height_);
+        if (err) {
+            if (err == LFS_ERR_NOENT) {
+                break;
+            }
+            return err;
+        }
+
+        // all branches should have the same height
+        if (height != -1) {
+            LFS_ASSERT(height_ == height);
+        }
+        height = height_;
+    }
+    #endif
+
     return 0;
 }
 
@@ -2973,10 +3005,10 @@ static int lfsr_rbyd_fetchck(lfs_t *lfs, lfsr_rbyd_t *rbyd,
 }
 
 
-static int lfsr_rbyd_lookupnext(lfs_t *lfs, const lfsr_rbyd_t *rbyd,
+static int lfsr_rbyd_lookupnext_(lfs_t *lfs, const lfsr_rbyd_t *rbyd,
         lfsr_srid_t rid, lfsr_tag_t tag,
         lfsr_srid_t *rid_, lfsr_tag_t *tag_, lfsr_rid_t *weight_,
-        lfsr_data_t *data_) {
+        lfsr_data_t *data_, lfs_size_t *height_) {
     // these bits should be clear at this point
     LFS_ASSERT(lfsr_tag_mode(tag) == 0);
 
@@ -2989,6 +3021,11 @@ static int lfsr_rbyd_lookupnext(lfs_t *lfs, const lfsr_rbyd_t *rbyd,
         return LFS_ERR_NOENT;
     }
 
+    // optionally find height for asserting rbyd balance
+    if (height_) {
+        *height_ = 0;
+    }
+
     // keep track of bounds as we descend down the tree
     lfs_size_t branch = lfsr_rbyd_trunk(rbyd);
     lfsr_srid_t lower_rid = 0;
@@ -3011,6 +3048,16 @@ static int lfsr_rbyd_lookupnext(lfs_t *lfs, const lfsr_rbyd_t *rbyd,
         if (lfsr_tag_isalt(alt)) {
             lfs_size_t branch_ = branch + d;
 
+            // only count black alts and followed alts towards height
+            if (height_
+                    && (lfsr_tag_isblack(alt)
+                        || lfsr_tag_follow(
+                            alt, weight,
+                            lower_rid, upper_rid,
+                            rid, tag))) {
+                *height_ += 1;
+            }
+
             // take alt?
             if (lfsr_tag_follow(
                     alt, weight,
@@ -3061,6 +3108,14 @@ static int lfsr_rbyd_lookupnext(lfs_t *lfs, const lfsr_rbyd_t *rbyd,
     }
 }
 
+static int lfsr_rbyd_lookupnext(lfs_t *lfs, const lfsr_rbyd_t *rbyd,
+        lfsr_srid_t rid, lfsr_tag_t tag,
+        lfsr_srid_t *rid_, lfsr_tag_t *tag_, lfsr_rid_t *weight_,
+        lfsr_data_t *data_) {
+    return lfsr_rbyd_lookupnext_(lfs, rbyd, rid, tag,
+            rid_, tag_, weight_, data_, NULL);
+}
+
 static int lfsr_rbyd_lookup(lfs_t *lfs, const lfsr_rbyd_t *rbyd,
         lfsr_srid_t rid, lfsr_tag_t tag,
         lfsr_data_t *data_) {
diff --git a/lfs_util.h b/lfs_util.h
index 2520f522..d2392b3f 100644
--- a/lfs_util.h
+++ b/lfs_util.h
@@ -168,6 +168,12 @@ extern "C"
 #define LFS_IFDEF_GC(a, b) (b)
 #endif
 
+#ifdef LFS_ASSERTRBYDBALANCE
+#define LFS_IFDEF_ASSERTRBYDBALANCE(a, b) (a)
+#else
+#define LFS_IFDEF_ASSERTRBYDBALANCE(a, b) (b)
+#endif
+
 
 // Builtin functions, these may be replaced by more efficient
 // toolchain-specific implementations. LFS_NO_BUILTINS falls back to a more