3 files changed, 60 insertions, 2 deletions

diff --git a/avltree.c b/avltree.c
index 0b6ff9e..85d618f 100644
--- a/avltree.c
+++ b/avltree.c
@@ -310,6 +310,40 @@ avltree_insert_rebalance(struct avltree *tree, struct avltree_node *parent,
     avltree_rotate(tree, node, new_balance);
 }
 
+void *
+avltree_replace_slot(struct avltree *tree, avltree_slot_t slot,
+                     struct avltree_node *node)
+{
+    struct avltree_node *parent, *prev;
+    int index;
+
+    parent = avltree_slot_parent(slot);
+
+    if (!parent) {
+        prev = tree->root;
+        tree->root = node;
+    } else {
+        index = avltree_slot_index(slot);
+        assert(avltree_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;
+        }
+
+
+        avltree_node_set_parent(node->children[i], node);
+    }
+
+    return prev;
+}
+
 void
 avltree_remove(struct avltree *tree, struct avltree_node *node)
 {
diff --git a/avltree.h b/avltree.h
index f01718e..019f158 100644
--- a/avltree.h
+++ b/avltree.h
@@ -279,6 +279,15 @@ avltree_insert_slot(struct avltree *tree, avltree_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 * avltree_replace_slot(struct avltree *tree, avltree_slot_t slot,
+                            struct avltree_node *node);
+
+/*
  * Remove a node from a tree.
  *
  * After completion, the node is stale.
diff --git a/test/test_avltree.c b/test/test_avltree.c
index 39afe1f..c6baefb 100644
--- a/test/test_avltree.c
+++ b/test/test_avltree.c
@@ -27,9 +27,10 @@
 #include <stdlib.h>
 #include <string.h>
 
+#include "../avltree.h"
+#include "../check.h"
 #include "../hash.h"
 #include "../macros.h"
-#include "../avltree.h"
 
 #define SIZE 28
 
@@ -103,7 +104,7 @@ main(int argc, char *argv[])
 {
     struct avltree tree;
     struct avltree_node *node, *tmp;
-    struct obj *obj;
+    struct obj *obj, *prev;
     avltree_slot_t slot;
     int i, id;
 
@@ -126,6 +127,20 @@ main(int argc, char *argv[])
         avltree_insert_slot(&tree, slot, &obj->node);
     }
 
+    id = get_id(0);
+    node = avltree_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 = avltree_replace_slot(&tree, slot, &obj->node);
+    check(node != &obj->node);
+    prev = avltree_entry(node, struct obj, node);
+    free(prev);
+    node = avltree_lookup(&tree, id, obj_cmp_lookup);
+    check(node == &obj->node);
+
     printf("\n");
     print_tree(&tree);