Simplify BTree implementation

3 files changed, 422 insertions, 263 deletions

diff --git a/include/zix/btree.h b/include/zix/btree.h
index 5212ec7..d19da3c 100644
--- a/include/zix/btree.h
+++ b/include/zix/btree.h
@@ -74,7 +74,15 @@ static const ZixBTreeIter zix_btree_end_iter = {
   {0u, 0u, 0u, 0u, 0u, 0u},
   0u};
 
-/// Create a new (empty) B-Tree
+/**
+   Create a new (empty) B-Tree.
+
+   The given comparator must be a total ordering and is used to internally
+   organize the tree and look for values exactly.
+
+   Searching can be done with a custom comparator that supports wildcards, see
+   zix_btree_lower_bound() for details.
+*/
 ZIX_API
 ZixBTree*
 zix_btree_new(ZixComparator cmp, const void* cmp_data);
@@ -126,7 +134,7 @@ ZixStatus
 zix_btree_remove(ZixBTree* t, const void* e, void** out, ZixBTreeIter* next);
 
 /**
-   Set `ti` to an element equal to `e` in `t`.
+   Set `ti` to an element exactly equal to `e` in `t`.
 
    If no such item exists, `ti` is set to the end.
 */
@@ -137,15 +145,26 @@ zix_btree_find(const ZixBTree* t, const void* e, ZixBTreeIter* ti);
 /**
    Set `ti` to the smallest element in `t` that is not less than `e`.
 
-   Wildcards are supported, so if the search key `e` compares equal to many
-   values in the tree, `ti` will be set to the least such element.  The search
-   key `e` is always passed as the second argument to the comparator.
+   The given comparator must be compatible with the tree comparator, that is,
+   any two values must have the same ordering according to both.  Within this
+   constraint, it may implement fuzzier searching by handling special search
+   key values, for example with wildcards.
+
+   If the search key `e` compares equal to many values in the tree, then `ti`
+   will be set to the least such element.
+
+   The comparator is always called with an actual value in the tree as the
+   first argument, and `key` as the second argument.
 */
 ZIX_API
 ZixStatus
-zix_btree_lower_bound(const ZixBTree* t, const void* e, ZixBTreeIter* ti);
+zix_btree_lower_bound(const ZixBTree* t,
+                      ZixComparator   compare_key,
+                      const void*     compare_key_user_data,
+                      const void*     key,
+                      ZixBTreeIter*   ti);
 
-/// Return the data associated with the given tree item
+/// Return the data at the given position in the tree
 ZIX_PURE_API
 void*
 zix_btree_get(ZixBTreeIter ti);
@@ -177,6 +196,11 @@ ZIX_API
 ZixStatus
 zix_btree_iter_increment(ZixBTreeIter* i);
 
+/// Return an iterator one past `iter`
+ZIX_API
+ZixBTreeIter
+zix_btree_iter_next(ZixBTreeIter iter);
+
 /**
    @}
    @}
diff --git a/src/btree.c b/src/btree.c
index 483cfa8..21db8d6 100644
--- a/src/btree.c
+++ b/src/btree.c
@@ -56,7 +56,7 @@ struct ZixBTreeNodeImpl {
 
     struct {
       void*         vals[ZIX_BTREE_INODE_VALS];
-      ZixBTreeNode* children[ZIX_BTREE_INODE_VALS + 1];
+      ZixBTreeNode* children[ZIX_BTREE_INODE_VALS + 1u];
     } inode;
   } data;
 };
@@ -66,17 +66,6 @@ struct ZixBTreeNodeImpl {
 static_assert(sizeof(ZixBTreeNode) == ZIX_BTREE_PAGE_SIZE, "");
 #endif
 
-typedef struct {
-  ZixBTreeNode* node;
-  unsigned      index;
-} ZixBTreeIterFrame;
-
-struct ZixBTreeIterImpl {
-  unsigned          n_levels; ///< Maximum depth of stack
-  unsigned          level;    ///< Current level in stack
-  ZixBTreeIterFrame stack[];  ///< Position stack
-};
-
 static ZixBTreeNode*
 zix_btree_node_new(const bool leaf)
 {
@@ -87,20 +76,13 @@ zix_btree_node_new(const bool leaf)
          ZIX_BTREE_PAGE_SIZE - 2u * sizeof(ZixBTreeNode*));
 #endif
 
-  ZixBTreeNode* node = (ZixBTreeNode*)malloc(sizeof(ZixBTreeNode));
+  ZixBTreeNode* const node = (ZixBTreeNode*)malloc(sizeof(ZixBTreeNode));
   if (node) {
     node->is_leaf = leaf;
-    node->n_vals  = 0;
+    node->n_vals  = 0u;
   }
-  return node;
-}
 
-ZIX_PURE_FUNC
-static void*
-zix_btree_value(const ZixBTreeNode* const node, const unsigned i)
-{
-  assert(i < node->n_vals);
-  return node->is_leaf ? node->data.leaf.vals[i] : node->data.inode.vals[i];
+  return node;
 }
 
 ZIX_PURE_FUNC
@@ -115,17 +97,20 @@ zix_btree_child(const ZixBTreeNode* const node, const unsigned i)
 ZixBTree*
 zix_btree_new(const ZixComparator cmp, const void* const cmp_data)
 {
-  ZixBTree* t = (ZixBTree*)malloc(sizeof(ZixBTree));
-  if (t) {
-    t->root     = zix_btree_node_new(true);
-    t->cmp      = cmp;
-    t->cmp_data = cmp_data;
-    t->size     = 0;
-    if (!t->root) {
-      free(t);
-      return NULL;
-    }
+  ZixBTree* const t = (ZixBTree*)malloc(sizeof(ZixBTree));
+  if (!t) {
+    return NULL;
   }
+
+  if (!(t->root = zix_btree_node_new(true))) {
+    free(t);
+    return NULL;
+  }
+
+  t->cmp      = cmp;
+  t->cmp_data = cmp_data;
+  t->size     = 0;
+
   return t;
 }
 
@@ -180,16 +165,16 @@ zix_btree_size(const ZixBTree* const t)
   return t->size;
 }
 
-static uint16_t
+static ZixShort
 zix_btree_max_vals(const ZixBTreeNode* const node)
 {
   return node->is_leaf ? ZIX_BTREE_LEAF_VALS : ZIX_BTREE_INODE_VALS;
 }
 
-static uint16_t
+static ZixShort
 zix_btree_min_vals(const ZixBTreeNode* const node)
 {
-  return (uint16_t)(((zix_btree_max_vals(node) + 1U) / 2U) - 1U);
+  return (ZixShort)(((zix_btree_max_vals(node) + 1u) / 2u) - 1u);
 }
 
 /// Shift pointers in `array` of length `n` right starting at `i`
@@ -223,7 +208,7 @@ zix_btree_split_child(ZixBTreeNode* const n,
   assert(i < n->n_vals + 1U);
   assert(zix_btree_child(n, i) == lhs);
 
-  const uint16_t max_n_vals = zix_btree_max_vals(lhs);
+  const ZixShort max_n_vals = zix_btree_max_vals(lhs);
   ZixBTreeNode*  rhs        = zix_btree_node_new(lhs->is_leaf);
   if (!rhs) {
     return NULL;
@@ -231,7 +216,7 @@ zix_btree_split_child(ZixBTreeNode* const n,
 
   // LHS and RHS get roughly half, less the middle value which moves up
   lhs->n_vals = max_n_vals / 2U;
-  rhs->n_vals = (uint16_t)(max_n_vals - lhs->n_vals - 1);
+  rhs->n_vals = (ZixShort)(max_n_vals - lhs->n_vals - 1);
 
   if (lhs->is_leaf) {
     // Copy large half from LHS to new RHS node
@@ -265,20 +250,23 @@ zix_btree_split_child(ZixBTreeNode* const n,
 #ifdef ZIX_BTREE_SORTED_CHECK
 /// Check that `n` is sorted with respect to search key `e`
 static bool
-zix_btree_node_is_sorted_with_respect_to(const ZixBTree* const     t,
-                                         const ZixBTreeNode* const n,
-                                         const void* const         e)
+zix_btree_node_is_sorted_with_respect_to(const ZixComparator compare,
+                                         const void* const   compare_user_data,
+                                         void* const* const  values,
+                                         const unsigned      n_values,
+                                         const void* const   key)
 {
-  if (n->n_vals <= 1) {
+  if (n_values <= 1u) {
     return true;
   }
 
-  int cmp = t->cmp(zix_btree_value(n, 0), e, t->cmp_data);
-  for (uint16_t i = 1; i < n->n_vals; ++i) {
-    const int next_cmp = t->cmp(zix_btree_value(n, i), e, t->cmp_data);
+  int cmp = compare(values[0], key, compare_user_data);
+  for (unsigned i = 1u; i < n_values; ++i) {
+    const int next_cmp = compare(values[i], key, compare_user_data);
     if ((cmp >= 0 && next_cmp < 0) || (cmp > 0 && next_cmp <= 0)) {
       return false;
     }
+
     cmp = next_cmp;
   }
 
@@ -286,95 +274,199 @@ zix_btree_node_is_sorted_with_respect_to(const ZixBTree* const     t,
 }
 #endif
 
-/// Find the first value in `n` that is not less than `e` (lower bound)
 static unsigned
-zix_btree_node_find(const ZixBTree* const     t,
-                    const ZixBTreeNode* const n,
-                    const void* const         e,
-                    bool* const               equal)
+zix_btree_find_value(const ZixComparator compare,
+                     const void* const   compare_user_data,
+                     void* const* const  values,
+                     const unsigned      n_values,
+                     const void* const   key,
+                     bool* const         equal)
+{
+  unsigned first = 0u;
+  unsigned count = n_values;
+
+  while (count > 0u) {
+    const unsigned half  = count >> 1u;
+    const unsigned i     = first + half;
+    void* const    value = values[i];
+    const int      cmp   = compare(value, key, compare_user_data);
+
+    if (!cmp) {
+      *equal = true;
+      return i;
+    }
+
+    if (cmp < 0) {
+      first += half + 1u;
+      count -= half + 1u;
+    } else {
+      count = half;
+    }
+  }
+
+  assert(first == n_values || compare(values[first], key, compare_user_data));
+  *equal = false;
+  return first;
+}
+
+static unsigned
+zix_btree_find_pattern(const ZixComparator compare_key,
+                       const void* const   compare_key_user_data,
+                       void* const* const  values,
+                       const unsigned      n_values,
+                       const void* const   key,
+                       bool* const         equal)
 {
 #ifdef ZIX_BTREE_SORTED_CHECK
-  assert(zix_btree_node_is_sorted_with_respect_to(t, n, e));
+  assert(zix_btree_node_is_sorted_with_respect_to(
+    compare_key, compare_key_user_data, values, n_values, key));
 #endif
 
-  unsigned first = 0U;
-  unsigned len   = n->n_vals;
-  while (len > 0) {
-    const unsigned half = len >> 1U;
-    const unsigned i    = first + half;
-    const int      cmp  = t->cmp(zix_btree_value(n, i), e, t->cmp_data);
+  unsigned first = 0u;
+  unsigned count = n_values;
+
+  while (count > 0u) {
+    const unsigned half  = count >> 1u;
+    const unsigned i     = first + half;
+    void* const    value = values[i];
+    const int      cmp   = compare_key(value, key, compare_key_user_data);
+
     if (cmp == 0) {
+      // Found a match, but keep searching for the leftmost one
       *equal = true;
-      len    = half; // Keep searching for wildcard matches
+      count  = half;
+
     } else if (cmp < 0) {
-      const unsigned chop = half + 1U;
-      first += chop;
-      len -= chop;
+      // Search right half
+      first += half + 1u;
+      count -= half + 1u;
+
     } else {
-      len = half;
+      // Search left half
+      count = half;
     }
   }
 
-  assert(!*equal || t->cmp(zix_btree_value(n, first), e, t->cmp_data) == 0);
+  assert(!*equal ||
+         (compare_key(values[first], key, compare_key_user_data) == 0 &&
+          (first == 0u ||
+           (compare_key(values[first - 1u], key, compare_key_user_data) < 0))));
+
   return first;
 }
 
+/// Convenience wrapper to find a value in an internal node
+static unsigned
+zix_btree_inode_find(const ZixBTree* const     t,
+                     const ZixBTreeNode* const n,
+                     const void* const         e,
+                     bool* const               equal)
+{
+  assert(!n->is_leaf);
+
+  return zix_btree_find_value(
+    t->cmp, t->cmp_data, n->data.inode.vals, n->n_vals, e, equal);
+}
+
+/// Convenience wrapper to find a value in a leaf node
+static unsigned
+zix_btree_leaf_find(const ZixBTree* const     t,
+                    const ZixBTreeNode* const n,
+                    const void* const         e,
+                    bool* const               equal)
+{
+  assert(n->is_leaf);
+
+  return zix_btree_find_value(
+    t->cmp, t->cmp_data, n->data.leaf.vals, n->n_vals, e, equal);
+}
+
+static inline bool
+zix_btree_can_remove_from(const ZixBTreeNode* const n)
+{
+  assert(n->n_vals >= zix_btree_min_vals(n));
+  return n->n_vals > zix_btree_min_vals(n);
+}
+
+static inline bool
+zix_btree_is_full(const ZixBTreeNode* const n)
+{
+  assert(n->n_vals <= zix_btree_max_vals(n));
+  return n->n_vals == zix_btree_max_vals(n);
+}
+
+static ZixStatus
+zix_btree_grow_up(ZixBTree* const t)
+{
+  ZixBTreeNode* const new_root = zix_btree_node_new(false);
+  if (!new_root) {
+    return ZIX_STATUS_NO_MEM;
+  }
+
+  // Set old root as the only child of the new root
+  new_root->data.inode.children[0] = t->root;
+
+  // Split the old root to get two balanced siblings
+  zix_btree_split_child(new_root, 0, t->root);
+  t->root = new_root;
+
+  return ZIX_STATUS_SUCCESS;
+}
+
 ZixStatus
 zix_btree_insert(ZixBTree* const t, void* const e)
 {
-  ZixBTreeNode* parent = NULL;    // Parent of n
-  ZixBTreeNode* n      = t->root; // Current node
-  unsigned      i      = 0;       // Index of n in parent
-  while (n) {
-    if (n->n_vals == zix_btree_max_vals(n)) {
-      // Node is full, split to ensure there is space for a leaf split
-      if (!parent) {
-        // Root is full, grow tree upwards
-        if (!(parent = zix_btree_node_new(false))) {
-          return ZIX_STATUS_NO_MEM;
-        }
-        t->root                        = parent;
-        parent->data.inode.children[0] = n;
-      }
+  ZixStatus st = ZIX_STATUS_SUCCESS;
+
+  // Grow up if necessary to ensure the root is not full
+  if (zix_btree_is_full(t->root)) {
+    if ((st = zix_btree_grow_up(t))) {
+      return st;
+    }
+  }
 
-      ZixBTreeNode* const rhs = zix_btree_split_child(parent, i, n);
+  // Walk down from the root until we reach a suitable leaf
+  ZixBTreeNode* node = t->root;
+  while (!node->is_leaf) {
+    // Search for the value in this node
+    bool           equal = false;
+    const unsigned i     = zix_btree_inode_find(t, node, e, &equal);
+    if (equal) {
+      return ZIX_STATUS_EXISTS;
+    }
+
+    // Value not in this node, but may be in the ith child
+    ZixBTreeNode* child = node->data.inode.children[i];
+    if (zix_btree_is_full(child)) {
+      // The child is full, split it before continuing
+      ZixBTreeNode* const rhs = zix_btree_split_child(node, i, child);
       if (!rhs) {
         return ZIX_STATUS_NO_MEM;
       }
 
-      const int cmp = t->cmp(parent->data.inode.vals[i], e, t->cmp_data);
-      if (cmp == 0) {
-        return ZIX_STATUS_EXISTS;
-      }
-
+      // Compare with new split value to determine which side to use
+      const int cmp = t->cmp(node->data.inode.vals[i], e, t->cmp_data);
       if (cmp < 0) {
-        // Move to new RHS
-        n = rhs;
-        ++i;
+        child = rhs; // Split value is less than the new value, move right
+      } else if (cmp == 0) {
+        return ZIX_STATUS_EXISTS; // Split value is exactly the value to insert
       }
     }
 
-    assert(!parent || zix_btree_child(parent, i) == n);
-
-    bool equal = false;
-    i          = zix_btree_node_find(t, n, e, &equal);
-    if (equal) {
-      return ZIX_STATUS_EXISTS;
-    }
+    // Descend to child node and continue
+    node = child;
+  }
 
-    if (!n->is_leaf) {
-      // Descend to child node left of value
-      parent = n;
-      n      = zix_btree_child(n, i);
-    } else {
-      // Insert into internal node
-      zix_btree_ainsert(n->data.leaf.vals, n->n_vals++, i, e);
-      break;
-    }
+  // Search for the value in the leaf
+  bool           equal = false;
+  const unsigned i     = zix_btree_leaf_find(t, node, e, &equal);
+  if (equal) {
+    return ZIX_STATUS_EXISTS;
   }
 
+  // The value is not in the tree, insert into the leaf
+  zix_btree_ainsert(node->data.leaf.vals, node->n_vals++, i, e);
   ++t->size;
-
   return ZIX_STATUS_SUCCESS;
 }
 
@@ -383,8 +475,6 @@ zix_btree_iter_set_frame(ZixBTreeIter* const ti,
                          ZixBTreeNode* const n,
                          const ZixShort      i)
 {
-  assert(i <= ZIX_BTREE_LEAF_VALS);
-
   ti->nodes[ti->level]   = n;
   ti->indexes[ti->level] = (uint16_t)i;
 }
@@ -409,14 +499,6 @@ zix_btree_iter_pop(ZixBTreeIter* const ti)
   --ti->level;
 }
 
-ZIX_PURE_FUNC
-static bool
-zix_btree_node_is_minimal(ZixBTreeNode* const n)
-{
-  assert(n->n_vals >= zix_btree_min_vals(n));
-  return n->n_vals == zix_btree_min_vals(n);
-}
-
 /// Enlarge left child by stealing a value from its right sibling
 static ZixBTreeNode*
 zix_btree_rotate_left(ZixBTreeNode* const parent, const unsigned i)
@@ -451,7 +533,7 @@ zix_btree_rotate_left(ZixBTreeNode* const parent, const unsigned i)
   return lhs;
 }
 
-/// Enlarge right child by stealing a value from its left sibling
+/// Enlarge a child by stealing a value from its left sibling
 static ZixBTreeNode*
 zix_btree_rotate_right(ZixBTreeNode* const parent, const unsigned i)
 {
@@ -493,7 +575,6 @@ zix_btree_merge(ZixBTree* const t, ZixBTreeNode* const n, const unsigned i)
   ZixBTreeNode* const rhs = zix_btree_child(n, i + 1);
 
   assert(lhs->is_leaf == rhs->is_leaf);
-  assert(zix_btree_node_is_minimal(lhs));
   assert(lhs->n_vals + rhs->n_vals < zix_btree_max_vals(lhs));
 
   // Move parent value to end of LHS
@@ -522,7 +603,7 @@ zix_btree_merge(ZixBTree* const t, ZixBTreeNode* const n, const unsigned i)
            (rhs->n_vals + 1U) * sizeof(void*));
   }
 
-  lhs->n_vals = (uint16_t)(lhs->n_vals + rhs->n_vals);
+  lhs->n_vals = (ZixShort)(lhs->n_vals + rhs->n_vals);
 
   if (--n->n_vals == 0) {
     // Root is now empty, replace it with its only child
@@ -539,19 +620,14 @@ zix_btree_merge(ZixBTree* const t, ZixBTreeNode* const n, const unsigned i)
 static void*
 zix_btree_remove_min(ZixBTree* const t, ZixBTreeNode* n)
 {
+  assert(zix_btree_can_remove_from(n));
+
   while (!n->is_leaf) {
-    if (zix_btree_node_is_minimal(zix_btree_child(n, 0))) {
-      // Leftmost child is minimal, must expand
-      if (!zix_btree_node_is_minimal(zix_btree_child(n, 1))) {
-        // Child's right sibling has at least one key to steal
-        n = zix_btree_rotate_left(n, 0);
-      } else {
-        // Both child and right sibling are minimal, merge
-        n = zix_btree_merge(t, n, 0);
-      }
-    } else {
-      n = zix_btree_child(n, 0);
-    }
+    ZixBTreeNode* const* const children = n->data.inode.children;
+
+    n = zix_btree_can_remove_from(children[0])   ? children[0]
+        : zix_btree_can_remove_from(children[1]) ? zix_btree_rotate_left(n, 0)
+                                                 : zix_btree_merge(t, n, 0);
   }
 
   return zix_btree_aerase(n->data.leaf.vals, --n->n_vals, 0);
@@ -561,24 +637,80 @@ zix_btree_remove_min(ZixBTree* const t, ZixBTreeNode* n)
 static void*
 zix_btree_remove_max(ZixBTree* const t, ZixBTreeNode* n)
 {
+  assert(zix_btree_can_remove_from(n));
+
   while (!n->is_leaf) {
-    if (zix_btree_node_is_minimal(zix_btree_child(n, n->n_vals))) {
-      // Leftmost child is minimal, must expand
-      if (!zix_btree_node_is_minimal(zix_btree_child(n, n->n_vals - 1u))) {
-        // Child's left sibling has at least one key to steal
-        n = zix_btree_rotate_right(n, n->n_vals);
-      } else {
-        // Both child and left sibling are minimal, merge
-        n = zix_btree_merge(t, n, n->n_vals - 1U);
-      }
-    } else {
-      n = zix_btree_child(n, n->n_vals);
-    }
+    ZixBTreeNode* const* const children = n->data.inode.children;
+
+    const unsigned y = n->n_vals - 1u;
+    const unsigned z = n->n_vals;
+
+    n = zix_btree_can_remove_from(children[z])   ? children[z]
+        : zix_btree_can_remove_from(children[y]) ? zix_btree_rotate_right(n, z)
+                                                 : zix_btree_merge(t, n, y);
   }
 
   return n->data.leaf.vals[--n->n_vals];
 }
 
+static ZixBTreeNode*
+zix_btree_fatten_child(ZixBTree* const t, ZixBTreeIter* const iter)
+{
+  ZixBTreeNode* const n = iter->nodes[iter->level];
+  const ZixShort      i = iter->indexes[iter->level];
+
+  assert(n);
+  assert(!n->is_leaf);
+  ZixBTreeNode* const* const children = n->data.inode.children;
+
+  if (i > 0 && zix_btree_can_remove_from(children[i - 1u])) {
+    return zix_btree_rotate_right(n, i); // Steal a key from left sibling
+  }
+
+  if (i < n->n_vals && zix_btree_can_remove_from(children[i + 1u])) {
+    return zix_btree_rotate_left(n, i); // Steal a key from right sibling
+  }
+
+  // Both child's siblings are minimal, merge them
+
+  if (i == n->n_vals) {
+    --iter->indexes[iter->level];
+    return zix_btree_merge(t, n, i - 1u); // Merge last two children
+  }
+
+  return zix_btree_merge(t, n, i); // Merge left and right siblings
+}
+
+/// Replace the ith value in `n` with one from a child if possible
+static ZixStatus
+zix_btree_replace_value(ZixBTree* const     t,
+                        ZixBTreeNode* const n,
+                        const unsigned      i,
+                        void** const        out)
+{
+  ZixBTreeNode* const lhs = zix_btree_child(n, i);
+  ZixBTreeNode* const rhs = zix_btree_child(n, i + 1);
+  if (!zix_btree_can_remove_from(lhs) && !zix_btree_can_remove_from(rhs)) {
+    return ZIX_STATUS_NOT_FOUND;
+  }
+
+  // Stash the value for the caller before it is replaced
+  *out = n->data.inode.vals[i];
+
+  n->data.inode.vals[i] =
+    // Left child has more values, steal its largest
+    (lhs->n_vals > rhs->n_vals) ? zix_btree_remove_max(t, lhs)
+
+    // Right child has more values, steal its smallest
+    : (rhs->n_vals > lhs->n_vals) ? zix_btree_remove_min(t, rhs)
+
+    // Children are balanced, use index parity as a low-bias tie breaker
+    : (i & 1u) ? zix_btree_remove_max(t, lhs)
+               : zix_btree_remove_min(t, rhs);
+
+  return ZIX_STATUS_SUCCESS;
+}
+
 ZixStatus
 zix_btree_remove(ZixBTree* const     t,
                  const void* const   e,
@@ -587,106 +719,80 @@ zix_btree_remove(ZixBTree* const     t,
 {
   ZixBTreeNode* n  = t->root;
   ZixBTreeIter* ti = next;
+  ZixStatus     st = ZIX_STATUS_SUCCESS;
 
   *ti = zix_btree_end_iter;
 
-  while (true) {
-    /* To remove in a single walk down, the tree is adjusted along the way
-       so that the current node always has at least one more value than the
-       minimum required in general. Thus, there is always room to remove
-       without adjusting on the way back up. */
-    assert(n == t->root || !zix_btree_node_is_minimal(n));
+  /* To remove in a single walk down, the tree is adjusted along the way so
+     that the current node always has at least one more value than the
+     minimum.  This ensures that there is always room to remove, without
+     having to merge nodes again on a traversal back up. */
+
+  if (!n->is_leaf && n->n_vals == 1u &&
+      !zix_btree_can_remove_from(n->data.inode.children[0u]) &&
+      !zix_btree_can_remove_from(n->data.inode.children[1u])) {
+    // Root has only two children, both minimal, merge them into a new root
+    n = zix_btree_merge(t, n, 0);
+  }
+
+  while (!n->is_leaf) {
+    assert(n == t->root || zix_btree_can_remove_from(n));
 
+    // Search for the value in the current node and update the iterator
     bool           equal = false;
-    const unsigned i     = zix_btree_node_find(t, n, e, &equal);
-    zix_btree_iter_set_frame(ti, n, i);
-    if (n->is_leaf) {
-      if (equal) {
-        // Found in leaf node
-        *out = zix_btree_aerase(n->data.leaf.vals, --n->n_vals, i);
-        if (ti && i == n->n_vals) {
-          if (i == 0) {
-            ti->level    = 0;
-            ti->nodes[0] = NULL;
-          } else {
-            --ti->indexes[ti->level];
-            zix_btree_iter_increment(ti);
-          }
-        }
-        --t->size;
-        return ZIX_STATUS_SUCCESS;
-      }
+    const unsigned i     = zix_btree_inode_find(t, n, e, &equal);
 
-      // Not found in leaf node, or tree
-      *next = zix_btree_end_iter;
-      return ZIX_STATUS_NOT_FOUND;
-    }
+    zix_btree_iter_set_frame(ti, n, i);
 
     if (equal) {
       // Found in internal node
-      ZixBTreeNode* const lhs    = zix_btree_child(n, i);
-      ZixBTreeNode* const rhs    = zix_btree_child(n, i + 1);
-      const size_t        l_size = lhs->n_vals;
-      const size_t        r_size = rhs->n_vals;
-      if (zix_btree_node_is_minimal(lhs) && zix_btree_node_is_minimal(rhs)) {
-        // Both preceding and succeeding child are minimal
-        n = zix_btree_merge(t, n, i);
-      } else if (l_size >= r_size) {
-        // Left child can remove without merge
-        assert(!zix_btree_node_is_minimal(lhs));
-        *out                  = n->data.inode.vals[i];
-        n->data.inode.vals[i] = zix_btree_remove_max(t, lhs);
+      if (!(st = zix_btree_replace_value(t, n, i, out))) {
+        // Replaced hole with a value from a direct child
         --t->size;
-        return ZIX_STATUS_SUCCESS;
-      } else {
-        // Right child can remove without merge
-        assert(!zix_btree_node_is_minimal(rhs));
-        *out                  = n->data.inode.vals[i];
-        n->data.inode.vals[i] = zix_btree_remove_min(t, rhs);
-        --t->size;
-        return ZIX_STATUS_SUCCESS;
+        return st;
       }
+
+      // Both preceding and succeeding child are minimal, merge and continue
+      n = zix_btree_merge(t, n, i);
+
     } else {
-      // Not found in internal node, key is in/under children[i]
-      if (zix_btree_node_is_minimal(zix_btree_child(n, i))) {
-        if (i > 0 && !zix_btree_node_is_minimal(zix_btree_child(n, i - 1))) {
-          // Steal a key from child's left sibling
-          n = zix_btree_rotate_right(n, i);
-        } else if (i < n->n_vals &&
-                   !zix_btree_node_is_minimal(zix_btree_child(n, i + 1))) {
-          // Steal a key from child's right sibling
-          n = zix_btree_rotate_left(n, i);
-        } else if (n == t->root && n->n_vals == 1) {
-          // Root has two children, both minimal, delete it
-          assert(i == 0 || i == 1);
-          const ZixShort counts[2] = {zix_btree_child(n, 0)->n_vals,
-                                      zix_btree_child(n, 1)->n_vals};
-
-          n                      = zix_btree_merge(t, n, 0);
-          ti->nodes[ti->level]   = n;
-          ti->indexes[ti->level] = (uint16_t)counts[i];
-        } else {
-          // Both child's siblings are minimal, merge them
-          if (i < n->n_vals) {
-            n = zix_btree_merge(t, n, i);
-          } else {
-            n = zix_btree_merge(t, n, i - 1U);
-            if (ti) {
-              --ti->indexes[ti->level];
-            }
-          }
-        }
-      } else {
-        n = zix_btree_child(n, i);
-      }
-    }
-    if (ti) {
-      ++ti->level;
+      // Not found in internal node, is in the ith child if anywhere
+      n = zix_btree_can_remove_from(zix_btree_child(n, i))
+            ? zix_btree_child(n, i)
+            : zix_btree_fatten_child(t, ti);
     }
+
+    ++ti->level;
+  }
+
+  // We're at the leaf the value may be in, search for the value in it
+  bool           equal = false;
+  const unsigned i     = zix_btree_leaf_find(t, n, e, &equal);
+
+  if (!equal) { // Not found in tree
+    *ti = zix_btree_end_iter;
+    return ZIX_STATUS_NOT_FOUND;
+  }
+
+  // Erase from leaf node
+  *out = zix_btree_aerase(n->data.leaf.vals, --n->n_vals, i);
+
+  // Update next iterator
+  if (n->n_vals == 0u) {
+    // Removed the last element in the tree
+    assert(n == t->root);
+    assert(t->size == 1u);
+    *ti = zix_btree_end_iter;
+  } else if (i == n->n_vals) {
+    // Removed the largest element in this leaf, increment to the next
+    zix_btree_iter_set_frame(ti, n, i - 1);
+    zix_btree_iter_increment(ti);
+  } else {
+    zix_btree_iter_set_frame(ti, n, i);
   }
 
-  assert(false); // Not reached
-  return ZIX_STATUS_ERROR;
+  --t->size;
+  return ZIX_STATUS_SUCCESS;
 }
 
 ZixStatus
@@ -698,67 +804,86 @@ zix_btree_find(const ZixBTree* const t,
 
   *ti = zix_btree_end_iter;
 
-  while (n) {
+  while (!n->is_leaf) {
     bool           equal = false;
-    const unsigned i     = zix_btree_node_find(t, n, e, &equal);
+    const unsigned i     = zix_btree_inode_find(t, n, e, &equal);
 
-    zix_btree_iter_set_frame(ti, n, (uint16_t)i);
+    zix_btree_iter_set_frame(ti, n, i);
 
     if (equal) {
       return ZIX_STATUS_SUCCESS;
     }
 
-    if (n->is_leaf) {
-      break;
-    }
-
     ++ti->level;
     n = zix_btree_child(n, i);
   }
 
+  bool           equal = false;
+  const unsigned i     = zix_btree_leaf_find(t, n, e, &equal);
+  if (equal) {
+    zix_btree_iter_set_frame(ti, n, i);
+    return ZIX_STATUS_SUCCESS;
+  }
+
   *ti = zix_btree_end_iter;
   return ZIX_STATUS_NOT_FOUND;
 }
 
 ZixStatus
 zix_btree_lower_bound(const ZixBTree* const t,
-                      const void* const     e,
+                      const ZixComparator   compare_key,
+                      const void* const     compare_key_user_data,
+                      const void* const     key,
                       ZixBTreeIter* const   ti)
 {
-  ZixBTreeNode* n           = t->root;
-  bool          found       = false;
-  unsigned      found_level = 0;
-
   *ti = zix_btree_end_iter;
 
-  while (n) {
-    bool           equal = false;
-    const unsigned i     = zix_btree_node_find(t, n, e, &equal);
+  ZixBTreeNode* n           = t->root; // Current node
+  uint16_t      found_level = 0u;      // Lowest level a match was found at
+  bool          found       = false;   // True if a match was ever found
+
+  // Search down until we reach a leaf
+  while (!n->is_leaf) {
+    bool equal = false;
 
-    zix_btree_iter_set_frame(ti, n, (uint16_t)i);
+    const unsigned i = zix_btree_find_pattern(compare_key,
+                                              compare_key_user_data,
+                                              n->data.inode.vals,
+                                              n->n_vals,
+                                              key,
+                                              &equal);
 
+    zix_btree_iter_set_frame(ti, n, i);
     if (equal) {
       found_level = ti->level;
       found       = true;
     }
 
-    if (n->is_leaf) {
-      break;
-    }
-
     ++ti->level;
     n = zix_btree_child(n, i);
-    assert(n);
   }
 
-  if (ti->indexes[ti->level] == ti->nodes[ti->level]->n_vals) {
-    if (found) {
-      // Found on a previous level but went too far
-      ti->level = (uint16_t)found_level;
-    } else {
-      // Reached end (key is greater than everything in tree)
-      *ti = zix_btree_end_iter;
-    }
+  bool equal = false;
+
+  const unsigned i = zix_btree_find_pattern(compare_key,
+                                            compare_key_user_data,
+                                            n->data.leaf.vals,
+                                            n->n_vals,
+                                            key,
+                                            &equal);
+
+  zix_btree_iter_set_frame(ti, n, i);
+  if (equal) {
+    return ZIX_STATUS_SUCCESS;
+  }
+
+  assert(ti->indexes[ti->level] == ti->nodes[ti->level]->n_vals);
+  if (found) {
+    // Found on a previous level but went too far
+    ti->level = found_level;
+  } else {
+    // Reached end (key is greater than everything in tree)
+    *ti = zix_btree_end_iter;
   }
 
   return ZIX_STATUS_SUCCESS;
@@ -848,3 +973,13 @@ zix_btree_iter_increment(ZixBTreeIter* const i)
 
   return ZIX_STATUS_SUCCESS;
 }
+
+ZixBTreeIter
+zix_btree_iter_next(const ZixBTreeIter iter)
+{
+  ZixBTreeIter next = iter;
+
+  zix_btree_iter_increment(&next);
+
+  return next;
+}
diff --git a/test/btree_test.c b/test/btree_test.c
index 9c4eab8..77771c4 100644
--- a/test/btree_test.c
+++ b/test/btree_test.c
@@ -234,7 +234,7 @@ stress(const unsigned test_num, const size_t n_elems)
   // Find the lower bound of all elements and ensure it's exact
   for (size_t i = 0; i < n_elems; ++i) {
     r = ith_elem(test_num, n_elems, i);
-    if (zix_btree_lower_bound(t, (void*)r, &ti)) {
+    if (zix_btree_lower_bound(t, int_cmp, NULL, (void*)r, &ti)) {
       return test_fail(
         t, "Lower bound %" PRIuPTR " @ %" PRIuPTR " failed\n", (uintptr_t)r, i);
     }
@@ -467,7 +467,7 @@ stress(const unsigned test_num, const size_t n_elems)
 
   // Find lower bound of wildcard
   const uintptr_t wildcard = 0;
-  if (zix_btree_lower_bound(t, (void*)wildcard, &ti)) {
+  if (zix_btree_lower_bound(t, wildcard_cmp, &ctx, (void*)wildcard, &ti)) {
     return test_fail(t, "Lower bound failed\n");
   }
 
@@ -490,7 +490,7 @@ stress(const unsigned test_num, const size_t n_elems)
 
   // Find lower bound of value past end
   const uintptr_t max = (uintptr_t)-1;
-  if (zix_btree_lower_bound(t, (void*)max, &ti)) {
+  if (zix_btree_lower_bound(t, wildcard_cmp, &ctx, (void*)max, &ti)) {
     return test_fail(t, "Lower bound failed\n");
   }