Add btree_remove_v to selectively remove bindings that have a key

2 files changed, 57 insertions, 0 deletions

diff --git a/src/common/include/btree.h b/src/common/include/btree.h
index 34fd2b3..9975885 100644
--- a/src/common/include/btree.h
+++ b/src/common/include/btree.h
@@ -7,6 +7,7 @@
 //	- btree_lower returns any item with matching key, or if none returns last item with smaller key
 //	- btree_upper returns any item with matching key, or if none returns first item with bigger key
 //	- btree_remove removes *all bindings* with matching key
+//	- btree_remove_v removes bindings with matching *key and value*
 //	- btree_iter_on calls iterator function on all bindings with matching key
 
 // Memory management is same as for hashtbl (a kv_iter_fun_t is called when an item is released)
@@ -19,6 +20,7 @@ void delete_btree(btree_t *t);
 
 bool btree_add(btree_t *t, void* key, void* val);
 void btree_remove(btree_t *t, const void* key);
+void btree_remove_v(btree_t *t, const void* key, const void* value);
 
 void* btree_find(btree_t *i, const void* key);
 void* btree_lower(btree_t *i, const void* key);
diff --git a/src/common/libalgo/btree.c b/src/common/libalgo/btree.c
index 20323aa..d3a030a 100644
--- a/src/common/libalgo/btree.c
+++ b/src/common/libalgo/btree.c
@@ -193,6 +193,61 @@ void btree_remove(btree_t *t, const void* key) {
 	t->root = remove_aux(t, t->root, key);
 }
 
+void btree_remove_v(btree_t *t, const void* key, const void* val) {
+	btree_item_t *extract_smallest(btree_item_t *i, btree_item_t **out_smallest) {
+		ASSERT(i != 0);
+		if (i->left == 0) {
+			*out_smallest = i;
+			return i->right;
+		} else {
+			i->left = extract_smallest(i->left, out_smallest);
+			recalc_height(i);
+			return equilibrate(i);
+		}
+	}
+
+	btree_item_t *remove_aux(btree_t *t, btree_item_t *i, const void* key, const void* val) {
+		if (i == 0) return 0;
+
+		int c = t->cf(key, i->key);
+		if (c < 0) {
+			i->left = remove_aux(t, i->left, key, val);
+			return equilibrate(i);
+		} else if (c > 0) {
+			i->right = remove_aux(t, i->right, key, val);
+			return equilibrate(i);
+		} else if (i->val == i) {
+			// remove item i
+
+			btree_item_t *new_i;
+			if (i->right == 0 || i->left == 0) {
+				new_i = (i->right == 0 ? i->left : i->right);
+			} else {
+				btree_item_t *new_i_right = extract_smallest(i->right, &new_i);
+
+				new_i->left = i->left;
+				new_i->right = new_i_right;
+
+				recalc_height(new_i);
+				new_i = equilibrate(new_i);
+			}
+
+			if (t->releasef) t->releasef(i->key, i->val);
+			free(i);
+			t->nitems--;
+
+			return remove_aux(t, new_i, key, val);	// loop because several elements may correspond
+		} else {
+			i->left = remove_aux(t, i->left, key, val);
+			i->right = remove_aux(t, i->right, key, val);
+			recalc_height(i);
+			return equilibrate(i);
+		}
+	}
+
+	t->root = remove_aux(t, t->root, key, val);
+}
+
 // ======================== //
 // LOOKING UP AND ITERATING //
 // ======================== //