3 files changed, 59 insertions, 1 deletions

diff --git a/rbtree.c b/rbtree.c
index ce3ed36..94f53e3 100644
--- a/rbtree.c
+++ b/rbtree.c
@@ -254,6 +254,40 @@ rbtree_insert_rebalance(struct rbtree *tree, struct rbtree_node *parent,
     assert(rbtree_node_is_black(tree->root));
 }
 
+void *
+rbtree_replace_slot(struct rbtree *tree, rbtree_slot_t slot,
+                    struct rbtree_node *node)
+{
+    struct rbtree_node *parent, *prev;
+    int index;
+
+    parent = rbtree_slot_parent(slot);
+
+    if (!parent) {
+        prev = tree->root;
+        tree->root = node;
+    } else {
+        index = rbtree_slot_index(slot);
+        assert(rbtree_check_index(index));
+        prev = parent->children[index];
+        parent->children[index] = node;
+    }
+
+    assert(prev);
+    *node = *prev;
+
+    for (size_t i = 0; i < ARRAY_SIZE(node->children); i++) {
+        if (!node->children[i]) {
+            continue;
+        }
+
+
+        rbtree_node_set_parent(node->children[i], node);
+    }
+
+    return prev;
+}
+
 void
 rbtree_remove(struct rbtree *tree, struct rbtree_node *node)
 {
diff --git a/rbtree.h b/rbtree.h
index 3327a70..e8ae915 100644
--- a/rbtree.h
+++ b/rbtree.h
@@ -278,6 +278,15 @@ rbtree_insert_slot(struct rbtree *tree, rbtree_slot_t slot,
 }
 
 /*
+ * Replace a node at an insertion point in a tree.
+ *
+ * The given node must compare strictly equal to the previous node,
+ * which is returned on completion.
+ */
+void * rbtree_replace_slot(struct rbtree *tree, rbtree_slot_t slot,
+                           struct rbtree_node *node);
+
+/*
  * Remove a node from a tree.
  *
  * After completion, the node is stale.
diff --git a/test/test_rbtree.c b/test/test_rbtree.c
index 328e1f7..6c7d903 100644
--- a/test/test_rbtree.c
+++ b/test/test_rbtree.c
@@ -27,6 +27,7 @@
 #include <stdlib.h>
 #include <string.h>
 
+#include "../check.h"
 #include "../hash.h"
 #include "../macros.h"
 #include "../rbtree.h"
@@ -88,7 +89,7 @@ main(int argc, char *argv[])
 {
     struct rbtree tree;
     struct rbtree_node *node, *tmp;
-    struct obj *obj;
+    struct obj *obj, *prev;
     rbtree_slot_t slot;
     int i, id;
 
@@ -111,6 +112,20 @@ main(int argc, char *argv[])
         rbtree_insert_slot(&tree, slot, &obj->node);
     }
 
+    id = get_id(0);
+    node = rbtree_lookup_slot(&tree, id, obj_cmp_lookup, slot);
+    check(node);
+    obj = malloc(sizeof(*obj));
+    check(obj);
+    obj->id = id;
+    printf("replacing: %d ", obj->id);
+    node = rbtree_replace_slot(&tree, slot, &obj->node);
+    check(node != &obj->node);
+    prev = rbtree_entry(node, struct obj, node);
+    free(prev);
+    node = rbtree_lookup(&tree, id, obj_cmp_lookup);
+    check(node == &obj->node);
+
     printf("\n");
     print_tree(&tree);