From fc89dbc035691859fe72aa62bd9a4b321e50309f Mon Sep 17 00:00:00 2001
From: "Liam R. Howlett" <Liam.Howlett@Oracle.com>
Date: Tue, 7 Jul 2020 11:26:19 -0400
Subject: [PATCH] maple_tree: Add freeing of nodes.

Node freeing is tracked by what was consumed in the previous step and
freed along with what was overwritten in the spanning store that must be
destroyed.

Signed-off-by: Liam R. Howlett <Liam.Howlett@Oracle.com>
---
 include/linux/maple_tree.h |   8 +-
 lib/maple_tree.c           | 510 +++++++++++++++++++++++--------------
 lib/test_maple_tree.c      |  23 +-
 3 files changed, 338 insertions(+), 203 deletions(-)

diff --git a/include/linux/maple_tree.h b/include/linux/maple_tree.h
index 374889c82194..fd2397cdeac5 100644
--- a/include/linux/maple_tree.h
+++ b/include/linux/maple_tree.h
@@ -235,11 +235,11 @@ static inline bool mtree_empty(const struct maple_tree *mt)
 
 struct ma_state {
 	struct maple_tree *tree;	/* The tree we're operating in */
-	unsigned long index;		/* The index we're operating on */
-	unsigned long last;		/* The last index we're operating on */
+	unsigned long index;		/* The index we're operating on - range start */
+	unsigned long last;		/* The last index we're operating on - range end */
 	struct maple_enode *node;	/* The node containing this entry */
-	unsigned long min;		/* The minimum index of this node */
-	unsigned long max;		/* The maximum index of this node */
+	unsigned long min;		/* The minimum index of this node - implied pivot min */
+	unsigned long max;		/* The maximum index of this node - implied pivot max */
 	struct maple_node *alloc;	/* Allocated nodes for this operation */
 	struct maple_enode *span_enode;	/* Pointer to maple parent/slot that set the max */
 	int full_cnt;			/* (+ full nodes) / (- number of almost empty) nodes above */
diff --git a/lib/maple_tree.c b/lib/maple_tree.c
index a8ae9c4374df..321407fb75f4 100644
--- a/lib/maple_tree.c
+++ b/lib/maple_tree.c
@@ -828,14 +828,10 @@ static inline void mas_ascend(struct ma_state *mas)
 	if (_ma_is_root(a_node))
 		goto no_parent;
 
-	printk("ascend: Starting at %p\n", a_node);
-
 	p_enode = mt_mk_node(mte_parent(mas->node),
 			mas_parent_enum(mas, mas->node));
-
 	a_type = mas_parent_enum(mas, mas->node);
 	a_enode = p_enode;
-
 	if (mte_is_root(a_enode)) {
 		a_node = mte_to_node(a_enode);
 		goto no_parent;
@@ -847,7 +843,6 @@ ascend:
 	a_node = mte_parent(mas->node);
 	a_slot = mte_parent_slot(mas->node);
 	a_enode = mt_mk_node(a_node, a_type);
-
 	if (!set_min && a_slot) {
 		set_min = true;
 		min = mte_get_pivot(a_enode, a_slot - 1) + 1;
@@ -873,7 +868,6 @@ no_parent:
 			MT_BUG_ON(mas->tree, mas->node == a_enode);
 		}
 		mas->node = a_enode;
-		printk("ascend again\n");
 		goto ascend;
 	}
 
@@ -1146,7 +1140,7 @@ static inline unsigned char _mas_data_end(const struct ma_state *mas,
 		prev_piv = piv;
 	}
 
-	printk("%s: %u\n", __func__, slot);
+//	printk("%s: %u\n", __func__, slot);
 	*last_piv = piv;
 	return slot;
 }
@@ -1237,6 +1231,7 @@ static inline unsigned long mas_find_gap(struct ma_state *mas)
 		return mas_leaf_max_gap(mas);
 	return mas_max_gap(mas, &slot);
 }
+
 static inline void mas_parent_gap(struct ma_state *mas, unsigned char slot,
 		unsigned long new, bool force)
 {
@@ -1381,20 +1376,24 @@ void mte_destroy_walk(struct maple_enode *mn, struct maple_tree *mtree)
 	unsigned char slot_cnt = mt_slot_count(mn);
 	int i;
 
+	printk("%s: %p type %u\n", __func__, mn, type);
 	switch (type) {
 	case maple_range_16:
 	case maple_range_32:
 	case maple_range_64:
 	case maple_arange_64:
 		for (i = 0; i < slot_cnt; i++) {
+			printk("%s: get slot %u\n", __func__, i);
 			node = mte_get_rcu_slot(mn, i, mtree);
-			if (!mt_is_empty(node) && !xa_is_retry(node))
+			printk("%s: destroy %p\n", __func__, mte_to_node(mn));
+			if (node)
 				mte_destroy_walk(node, mtree);
 		}
 		break;
 	default:
 		break;
 	}
+	printk("%s: Free %p (%p)\n", __func__, mte_to_node(mn), mn);
 	mte_free(mn);
 
 }
@@ -1464,6 +1463,7 @@ static inline void _mas_replace(struct ma_state *mas, bool free, bool push,
 	}
 
 	if (free) {
+		printk("Free %p\n", mte_to_node(prev));
 		mte_free(prev);
 		return;
 	}
@@ -1510,7 +1510,7 @@ static inline struct maple_enode *mas_check_split_parent(struct ma_state *mas,
 	if (!entry)
 		return NULL;
 
-	printk("parent of %p[%u] is %p\n", mas_mn(mas), slot, mte_parent(entry));
+//	printk("parent of %p[%u] is %p\n", mas_mn(mas), slot, mte_parent(entry));
 	if (mte_parent(entry) == entry) {
 		printk("%s: Dead node %p", __func__, entry);
 		MT_BUG_ON(mas->tree, 1);
@@ -1525,11 +1525,11 @@ static inline struct maple_enode *mas_check_split_parent(struct ma_state *mas,
 static inline struct maple_enode *mas_find_l_split(struct ma_state *mas)
 {
 	unsigned char i, end = mas_data_end(mas);
-	printk("data end of %p was %u\n", mas_mn(mas), end);
+//	printk("data end of %p was %u\n", mas_mn(mas), end);
 	struct maple_enode *en = NULL;
 
 	for (i = 0; i <= end; i++) {
-		printk("%s: Checking %p[%i]\n", __func__, mas_mn(mas), i);
+//		printk("%s: Checking %p[%i]\n", __func__, mas_mn(mas), i);
 		if ((en = mas_check_split_parent(mas, i)))
 			break;
 	}
@@ -1552,10 +1552,10 @@ static inline void mab_shift_right(struct maple_big_node *b_node,
 				 unsigned char shift, bool alloc)
 {
 	do {
-		printk("move %u to %u\n", b_end, b_end + shift);
+//		printk("move %u to %u\n", b_end, b_end + shift);
 		b_node->pivot[b_end + shift] = b_node->pivot[b_end];
 		b_node->slot[b_end + shift] = b_node->slot[b_end];
-		printk("piv %u is %lu\n", b_end + shift, b_node->pivot[b_end + shift]);
+//		printk("piv %u is %lu\n", b_end + shift, b_node->pivot[b_end + shift]);
 		if (alloc)
 			b_node->gap[b_end + shift] = b_node->gap[b_end];
 	} while (b_end--);
@@ -1793,6 +1793,7 @@ static inline bool mas_prev_sibling(struct ma_state *mas)
 
 	mas_ascend(mas);
 	mas_set_slot(mas, p_slot - 1);
+	printk("%s: going to go to %p[%u]\n", __func__, mas_mn(mas), p_slot - 1);
 	mas_descend(mas);
 	return true;
 }
@@ -1818,6 +1819,65 @@ static inline bool mas_next_sibling(struct ma_state *mas)
 	mas_descend(mas);
 	return true;
 }
+/* mas_consume() - Add the portions of the tree which will be replaced by the
+ * operation to the removal list; either to be @free or @discard (destroy walk).
+ */
+static inline void
+mas_consume(struct ma_state *l_mas, struct ma_state *r_mas,
+		struct maple_enode **free, unsigned char *f,
+		struct maple_enode **discard, unsigned char *d)
+{
+	unsigned char l_slot, r_slot, slot, end;
+	unsigned long l_min,  range_min, range_max;
+
+	printk("%s: l %p (%lu) r %p (%lu)\n", __func__, mas_mn(l_mas),
+	       l_mas->index, mas_mn(r_mas), r_mas->last);
+	// The left node is consumed, so add to the free list.
+	free[(*f)++] = l_mas->node;
+	printk("free %u -> %p (%p)\n", (*f)-1, mte_to_node(free[(*f)-1]),
+		free[(*f)-1]);
+
+	l_min = l_mas->index;
+	l_mas->index = l_mas->last;
+	printk("OVERWRITE index with %lu for search?\n", l_mas->last);
+	mas_node_walk(l_mas, mte_node_type(l_mas->node), &range_min, &range_max);
+	l_mas->index = l_min;
+	l_slot = mas_get_slot(l_mas);
+	r_slot = mas_get_slot(r_mas);
+	printk("l_slot %u r_slot %u\n", l_slot, r_slot);
+	if (l_mas->node == r_mas->node) {
+		/* If the contents up to l_slot and from r_slot to end are
+		 * still valid and it's the same node, then the rest may need to
+		 * be destroyed.  The exception is if a there is data taken
+		 * from the parent in the previous iteration which is still in
+		 * use.
+		 */
+		printk("Add %p[%u-%u] to discard\n", mas_mn(l_mas), l_slot + 1, r_slot - 1);
+		for (slot = l_slot + 1; slot < r_slot; slot++) {
+			discard[(*d)++] = mas_get_rcu_slot(l_mas, slot);
+	printk("discard %u -> %p\n", (*d)-1, discard[(*d)-1]);
+		}
+		return;
+	}
+	/* r_mas is different and consumed. */
+	free[(*f)++] = r_mas->node;
+	printk("free %u -> %p\n", (*f)-1, free[(*f)-1]);
+	if (mte_is_leaf(r_mas->node))
+		return;
+
+	/* Now free l_slot + 1 -> end and 0 -> r_slot - 1 */
+	end = mas_data_end(l_mas);
+	for (slot = l_slot + 1; slot <= end; slot++) {
+		discard[(*d)++] = mas_get_rcu_slot(l_mas, slot);
+	printk("l discard %u -> %p\n", (*d)-1, discard[(*d)-1]);
+	}
+
+	for (slot = 0; slot < r_slot; slot++) {
+		discard[(*d)++] = mas_get_rcu_slot(r_mas, slot);
+	printk("r discard %u -> %p (%u of %u)\n", (*d)-1, discard[(*d)-1], slot, r_slot);
+	}
+}
+
 /* Private
  *
  * mas_combine_separate() - Follow the tree upwards from @l_mas and @r_mas for
@@ -1825,6 +1885,12 @@ static inline bool mas_next_sibling(struct ma_state *mas)
  * which are inserted into the next iteration of the loop.  @b_node is returned
  * populated with the final iteration. @mas is used to obtain allocations.
  *
+ *
+ * orig_l_mas keeps track of the nodes that will remain active by using
+ * orig_l_mas->index and orig_l_mas->last to account of what has been copied
+ * into the new sub-tree.  The update of orig_l_mas->last is used in mas_consume
+ * to find the slots that will need to be either freed or destroyed.
+ *
  * Returns the @count left.
  */
 static inline int mas_combine_separate(struct ma_state *mas,
@@ -1834,79 +1900,78 @@ static inline int mas_combine_separate(struct ma_state *mas,
 				       unsigned char b_end,
 				       unsigned char count)
 {
-	unsigned char slot_cnt, split;
-	unsigned char child = 0, l_slot, r_slot = 0;
-	struct maple_enode *l, *m, *r; // left, middle, right
+	unsigned char slot_cnt, split, r_end;
+	unsigned char child = 0, l_slot, r_slot = 0, i = 0, d = 0;
+	struct maple_enode *left, *middle, *right;
 	unsigned long range_min, range_max;
+	struct maple_enode *free[100];
+	struct maple_enode *destroy[100];
 
 	MA_STATE(l_mas, mas->tree, mas->index, mas->index);
 	MA_STATE(r_mas, mas->tree, mas->index, mas->index);
 	MA_STATE(m_mas, mas->tree, mas->index, mas->index);
 
 	l_mas.node = r_mas.node = m_mas.node = NULL;
-	l = r = m = NULL;
+	left = right = middle = NULL;
 
+	mt_dump(mas->tree);
 	printk("Starting with orig of %p and %p\n", mas_mn(orig_l_mas),
 	       mas_mn(orig_r_mas));
 	printk("\t\t%d %u\n", __LINE__, b_end);
+	mas_consume(orig_l_mas, orig_r_mas, free, &i, destroy, &d);
 	while(count--) {
 		printk("Iteration %u\n", count);
 		b_end--;
 		slot_cnt = mt_slot_count(orig_l_mas->node);
-		r = NULL;
-		m = NULL;
-		l = mt_mk_node(ma_mnode_ptr(mas_next_alloc(mas)),
+		right = NULL;
+		middle = NULL;
+		left = mt_mk_node(ma_mnode_ptr(mas_next_alloc(mas)),
 			       mte_node_type(orig_l_mas->node));
 
+		///
 		if (b_end < slot_cnt) {
 			split = b_end;
 		} else {
 			split = mab_calc_split(b_node, b_end, slot_cnt,
 					       orig_l_mas->min,
 					       mte_node_type(orig_l_mas->node));
-			r = mt_mk_node(ma_mnode_ptr(mas_next_alloc(mas)),
-				       mte_node_type(orig_l_mas->node));
+			right = mt_mk_node(ma_mnode_ptr(mas_next_alloc(mas)),
+					   mte_node_type(orig_l_mas->node));
 		}
 		if (mab_middle_node(b_node, b_end, slot_cnt))
-			m = mt_mk_node(ma_mnode_ptr(mas_next_alloc(mas)),
-				       mte_node_type(orig_l_mas->node));
+			middle = mt_mk_node(ma_mnode_ptr(mas_next_alloc(mas)),
+					    mte_node_type(orig_l_mas->node));
 		printk("Splitting %u at %u\n", b_end, split);
 
 		/* Set parents from previous run */
 		if (l_mas.node) {
 			if (child <= split)
-				mte_set_parent(l_mas.node, l, child);
+				mte_set_parent(l_mas.node, left, child);
 			else
-				mte_set_parent(l_mas.node, r, child - split);
-			printk("Set parent of %p to %p\n", mas_mn(&l_mas),
-			       mte_parent(l_mas.node));
-			printk("child %u split %u\n", child, split);
+				mte_set_parent(l_mas.node, right,
+					       child - split);
 		}
 		if (m_mas.node) {
 			child++;
 			if (child <= split)
-				mte_set_parent(m_mas.node, l, child);
+				mte_set_parent(m_mas.node, left, child);
 			else
-				mte_set_parent(m_mas.node, r, child - split);
-			printk("Set parent of m %p to %p\n", mas_mn(&m_mas),
-			       mte_parent(m_mas.node));
-			printk("child %u split %u\n", child, split);
+				mte_set_parent(m_mas.node, right,
+					       child - split);
 		}
 		if (r_mas.node) {
 			child++;
 			if (child <= split)
-				mte_set_parent(r_mas.node, l, child);
+				mte_set_parent(r_mas.node, left, child);
 			else
-				mte_set_parent(r_mas.node, r, child - split);
-			printk("Set parent of r %p to %p\n", mas_mn(&r_mas),
-			       mte_parent(r_mas.node));
-			printk("child %u split %u\n", child, split);
+				mte_set_parent(r_mas.node, right,
+					       child - split);
 		}
 
 		/* Copy data from b_node to new nodes */
-		l_mas.node = l;
-		m_mas.node = m;
-		r_mas.node = r;
+		l_mas.node = left;
+		m_mas.node = middle;
+		r_mas.node = right;
 
 		l_mas.min = orig_l_mas->min;
 		mab_mas_cp(b_node, 0, split, &l_mas, 0);
@@ -1915,20 +1980,20 @@ static inline int mas_combine_separate(struct ma_state *mas,
 		r_mas.max = l_mas.max;
 
 
-		if (m) {
+		if (middle) {
 			mab_mas_cp(b_node, 1 + split, split * 2, &m_mas, 0);
 			m_mas.min = b_node->pivot[split] + 1;
 			split *= 2;
 			m_mas.max = b_node->pivot[split];
 		}
-		if (r) {
+		if (right) {
 			printk("there is a r\n");
 			mab_mas_cp(b_node, 1 + split, b_end, &r_mas, 0);
 			r_mas.min = b_node->pivot[split] + 1;
 			r_mas.max = b_node->pivot[b_end];
-			printk("r_max.max from %u\n", b_end);
+			printk("r_mas.max from %u\n", b_end);
 		}
-
+///
 
 		/* Copy data from next level in the tree to b_node from next iteration */
 		memset(b_node, 0, sizeof(struct maple_big_node));
@@ -1940,27 +2005,58 @@ static inline int mas_combine_separate(struct ma_state *mas,
 			mas_mn(&l_mas)->parent = ma_parent_ptr(
 			      ((unsigned long)mas->tree | MA_ROOT_PARENT));
 			mas->depth = orig_l_mas->depth;
+			b_end = 0;
+			printk("orig_l_mas is %p\n", orig_l_mas->node);
+			if (mte_is_root(orig_l_mas->node)) {
+				printk("Already at root, skipping consumption\n");
+				if ((orig_l_mas->node != mas->node) &&
+				    (l_mas.depth > mas->tree->ma_height)) {
+					printk("Adding %p to free list %u\n", mas->node, i);
+					free[i++] = mas->node;
+				}
+			} else {
+			//	mas_set_slot(orig_l_mas, 0);
+				do {
+					mas_ascend(orig_l_mas);
+					mas_ascend(orig_r_mas);
+				mas_set_slot(orig_l_mas, 0);
+				mas_set_slot(orig_r_mas, 0);
+		mas_node_walk(orig_l_mas, mte_node_type(orig_l_mas->node),
+			      &range_min, &range_max);
+		mas_node_walk(orig_r_mas, mte_node_type(orig_r_mas->node),
+				   &range_min, &range_max);
+			//		mas_set_slot(orig_r_mas,
+			//			     mas_data_end(orig_r_mas));
+					mas_consume(orig_l_mas, orig_r_mas,
+						    free, &i, destroy, &d);
+				} while (!mte_is_root(orig_l_mas->node));
+			}
 			mas_dup_state(orig_l_mas, &l_mas);
-			return 0;
+			goto done;
 		}
 
+
 		mas_ascend(orig_l_mas);
 		mas_ascend(orig_r_mas);
+		r_end = mas_data_end(orig_r_mas);
 
-		printk("l_mas min is %lu\n", orig_l_mas->min);
+		printk("orig_l_mas min is %lu\n", orig_l_mas->min);
 		mas_set_slot(orig_r_mas, 0);
-		orig_r_mas->index = r_mas.max + 1;
-		printk("Update r_mas.max to %lu\n", r_mas.max);
-		if (orig_r_mas->index) {
-			printk("R walk %p for %lu\n", mas_mn(orig_r_mas), orig_r_mas->index);
-			orig_r_mas->index++;
+		orig_r_mas->index = r_mas.max;
+		if (orig_r_mas->last < orig_r_mas->index)
 			orig_r_mas->last = orig_r_mas->index;
-			mas_node_walk(orig_r_mas,
-				      mte_node_type(orig_r_mas->node),
-				      &range_min, &range_max);
-			r_slot = mas_get_slot(orig_r_mas);
-			orig_r_mas->index = orig_r_mas->last = r_mas.max;
+		printk("Update r_mas.max to %lu\n", r_mas.max);
+		printk("R walk %p for %lu\n", mas_mn(orig_r_mas), orig_r_mas->index);
+		if (!mas_node_walk(orig_r_mas,
+				   mte_node_type(orig_r_mas->node),
+				   &range_min, &range_max)) {
+			// The node doesn't contain the value so set
+			// slot to ensure all of the nodes contents are
+			// freed or destroyed.
+			mas_set_slot(orig_r_mas, r_end + 1);
 		}
+		r_slot = mas_get_slot(orig_r_mas);
+		printk("r_slot is %u\n", r_slot);
 
 		mas_set_slot(orig_l_mas, 0);
 		orig_l_mas->index = l_mas.min;
@@ -1974,9 +2070,9 @@ static inline int mas_combine_separate(struct ma_state *mas,
 			printk("Copy in orig_l_mas %p\n", mas_mn(orig_l_mas));
 			b_end = mas_mab_cp(orig_l_mas, 0, l_slot - 1, b_node, 0);
 		}
-		// Not a leaf means new l_mas is placed at b_end.
 		printk("Put L %p at %u %lu\n", l_mas.node, b_end, l_mas.max);
 		child = b_end;
+		// Track dead nodes here.
 		b_node->slot[b_end] = l_mas.node;
 		if(mt_is_alloc(mas->tree)) {
 			b_node->gap[b_end] = mas_find_gap(&l_mas);
@@ -1985,10 +2081,9 @@ static inline int mas_combine_separate(struct ma_state *mas,
 		b_node->pivot[b_end++] = l_mas.max;
 		printk("l_mas max is %lu\n", b_node->pivot[b_end - 1]);
 
-		if (m) {
-			// FIXME: What about descend_adopt?
+		if (middle) {
 			printk("Put M! %p at %u %lu\n", m_mas.node, b_end, m_mas.max);
-			b_node->slot[b_end] = m;
+			b_node->slot[b_end] = middle;
 			if(mt_is_alloc(mas->tree)) {
 				b_node->gap[b_end] =
 					mas_find_gap(&m_mas);
@@ -1997,46 +2092,59 @@ static inline int mas_combine_separate(struct ma_state *mas,
 			b_node->pivot[b_end++] = m_mas.max;
 		}
 
-		if (r) {
+		if (right) {
 			printk("Put R %p at %u %lu\n", r_mas.node, b_end, r_mas.max);
-			b_node->slot[b_end] = r;
+			b_node->slot[b_end] = right;
 			if(mt_is_alloc(mas->tree)) {
 				b_node->gap[b_end] =
 					mas_find_gap(&r_mas);
 				printk("gap %u is %lu\n", b_end, b_node->gap[b_end]);
 			}
 			b_node->pivot[b_end++] = r_mas.max;
-			printk("r_mas max is %lu\n", b_node->pivot[b_end]);
+			printk("r_mas max is %lu\n", b_node->pivot[b_end - 1]);
 		}
 
-		// Add orig_r_mas->node to clean_list.
-		if ((orig_r_mas->last) &&
-		    (b_node->pivot[b_end - 1] < ULONG_MAX)) {
-			printk("Copy in orig_r_mas %p\n", mas_mn(orig_r_mas));
-			b_end = mas_mab_cp(orig_r_mas, r_slot,
-					   mas_data_end(orig_r_mas),
-					   b_node, b_end);
+		// Copy anything necessary out of the right node.
+		if (b_node->pivot[b_end - 1] < orig_r_mas->max) {
+			printk("Copy in orig_r_mas %p %u-%u\n", mas_mn(orig_r_mas),
+			       r_slot + 1, r_end);
+			b_end = mas_mab_cp(orig_r_mas, r_slot + 1, r_end, b_node,
+					   b_end);
+			orig_r_mas->last = orig_r_mas->max;
 		}
+		mas_consume(orig_l_mas, orig_r_mas, free, &i, destroy, &d);
+		orig_l_mas->last = orig_l_mas->max;
+
 
 		// Attempt to balance from this parent
 		if (b_end - 1 < mt_min_slot_cnt(orig_l_mas->node)) {
 			unsigned char end;
-			printk("\t merge with other node\n");
+			printk("\t merge with other node l is %p r is %p\n",
+			       mas_mn(orig_l_mas), mas_mn(orig_r_mas));
 			if (mas_next_sibling(orig_r_mas)) {
 				printk("r sibling merger\n");
 				end = mas_data_end(orig_r_mas);
 				b_end = mas_mab_cp(orig_r_mas, 0, end, b_node,
 						   b_end);
+				printk("Add free r %p (%p)\n", mas_mn(orig_r_mas),
+				       orig_r_mas->node);
+				free[i++] = orig_r_mas->node;
+				orig_r_mas->last = orig_r_mas->max;
 				if (!count)
 					count++;
 			} else if (mas_prev_sibling(orig_l_mas)) {
-				printk("l sibling merger\n");
+				printk("l sibling merger %p\n", mas_mn(orig_l_mas));
 				end = mas_data_end(orig_l_mas);
 				// shift b_node by prev size
 				mab_shift_right(b_node, b_end - 1, end + 1,
 						(mt_is_alloc(mas->tree) ? true : false));
 				// copy in prev.
 				mas_mab_cp(orig_l_mas, 0, end, b_node, 0);
+				printk("Add free l %p (%p)\n", mas_mn(orig_l_mas),
+				       orig_l_mas->node);
+				free[i++] = orig_l_mas->node;
+				l_mas.min = orig_l_mas->min;
+				orig_l_mas->index = orig_l_mas->min;
 				b_end += end + 1;
 				child += end + 1;
 				if (!count)
@@ -2059,6 +2167,8 @@ static inline int mas_combine_separate(struct ma_state *mas,
 				end = mas_data_end(orig_r_mas);
 				b_end = mas_mab_cp(orig_r_mas, 0, end,
 						   b_node, b_end);
+				orig_r_mas->last = orig_r_mas->max;
+				free[i++] = orig_r_mas->node;
 			} else {
 				printk("Trying left\n");
 				// Put left into right.
@@ -2076,6 +2186,7 @@ static inline int mas_combine_separate(struct ma_state *mas,
 					b_end--;
 					break;
 				}
+				free[i++] = orig_l_mas->node;
 
 				mas_set_slot(orig_l_mas, 0);
 				orig_l_mas->index = orig_l_mas->min;
@@ -2086,6 +2197,8 @@ static inline int mas_combine_separate(struct ma_state *mas,
 						(mt_is_alloc(mas->tree) ? true : false));
 				// copy in prev.
 				mas_mab_cp(orig_l_mas, 0, end, b_node, 0);
+				l_mas.min = orig_l_mas->min;
+				free[i++] = orig_l_mas->node;
 				child += end + 1;
 				b_end += end + 1;
 				printk("b_end is %u\n", b_end);
@@ -2098,30 +2211,71 @@ static inline int mas_combine_separate(struct ma_state *mas,
 	}
 
 	l_mas.node = mt_mk_node(ma_mnode_ptr(mas_next_alloc(mas)),
-				    mte_node_type(orig_l_mas->node));
+				mte_node_type(orig_l_mas->node));
 	orig_l_mas->depth++;
 
 	printk("All done making replacement at %p\n", mas_mn(&l_mas));
+	printk("orig is %p\n", mas_mn(orig_l_mas));
+	printk("orig r is %p\n", mas_mn(orig_r_mas));
+	//mas_consume(orig_l_mas, orig_r_mas, free, &i, destroy, &d);
 	mab_mas_cp(b_node, 0, b_end, &l_mas, 0);
-	printk("\t\t==>Setting parent of %p to %p[%u]\n", l, mas_mn(&l_mas),
+	printk("\t\t==>Setting parent of %p to %p[%u]\n", left, mas_mn(&l_mas),
 	       child);
 
-	mte_set_parent(l, l_mas.node, child);
-	if (m)
-		mte_set_parent(m, l_mas.node, ++child);
+	mte_set_parent(left, l_mas.node, child);
+	if (middle)
+		mte_set_parent(middle, l_mas.node, ++child);
 
-	if (r)
-		mte_set_parent(r, l_mas.node, ++child);
+	if (right)
+		mte_set_parent(right, l_mas.node, ++child);
 
 	if (!b_end) {
 		mas_mn(&l_mas)->parent = ma_parent_ptr(
-				((unsigned long)mas->tree | MA_ROOT_PARENT));
+						       ((unsigned long)mas->tree | MA_ROOT_PARENT));
 	} else {
 		mas_mn(&l_mas)->parent = mas_mn(orig_l_mas)->parent;
 	}
 
 	mas_dup_state(orig_l_mas, &l_mas);
 	mas->depth = orig_l_mas->depth;
+done:
+
+	printk("Going to insert %p\n", mas_mn(&l_mas));
+	// Set up mas for insertion.
+	printk("Killing %p (%p)\n", mas_mn(mas), mas->node);
+	mas_set_node_dead(mas);
+	mas_dup_state(mas, orig_l_mas);
+	smp_wmb();
+
+	// Insert new sub-tree
+	_mas_replace(mas, false, false, false);
+
+
+
+	mt_dump(mas->tree);
+	if (!mte_is_leaf(mas->node))
+		mas_descend_adopt(mas);
+
+	do {
+		printk("%s: Free %u %p (%p)\n", __func__, i - 1,
+		       mte_to_node(free[i - 1]), free[i - 1]);
+		mte_free(free[--i]);
+	} while (i);
+
+	if (d) {
+		do {
+			printk("%s: destroy %u %p (%p)\n", __func__, d - 1,
+			       mte_to_node(destroy[d - 1]), destroy[d-1]);
+			mte_destroy_walk(destroy[--d], mas->tree);
+		} while (d);
+	}
+
+	mt_dump(mas->tree);
+	if (mte_is_leaf(mas->node))
+	    return b_end;
+
+	if (mt_is_alloc(mas->tree))
+		mas_update_gap(mas, false);
 
 	return b_end;
 }
@@ -2134,6 +2288,7 @@ static inline int mas_rebalance(struct ma_state *mas,
 	unsigned char b_end = 0;
 
 	printk("%s: rebalance %u levels\n", __func__, empty);
+	mt_dump(mas->tree);
 
 	mas_node_cnt(mas, 1 + empty * 2);
 	if (mas_is_err(mas))
@@ -2155,11 +2310,6 @@ static inline int mas_rebalance(struct ma_state *mas,
 	} else {
 		unsigned char shift;
 		mas_prev_sibling(&l_mas);
-		if (mas_is_none(&l_mas)) {
-			printk("%s: Drop a level..\n", __func__);
-			b_end = 0;
-			goto new_root_node;
-		}
 		shift = mas_data_end(&l_mas) + 1;
 		mab_shift_right(b_node, new_end, shift,
 				(mt_is_alloc(mas->tree) ? true : false));
@@ -2169,37 +2319,7 @@ static inline int mas_rebalance(struct ma_state *mas,
 		l_mas.index = l_mas.last = l_mas.min;
 	}
 
-	b_end = mas_combine_separate(mas, &l_mas, &r_mas, b_node, b_end, empty);
-
-new_root_node:
-	if (!b_end) {
-		printk("new root\n");
-	}
-
-	// Set up mas for insertion.
-	mas_set_node_dead(mas);
-	mas_dup_state(mas, &l_mas);
-	smp_wmb();
-
-	// FIXME: Mark modified nodes as dead.
-
-	// Insert new sub-tree
-	// FIXME: If mas isn't being replaced, kill the right node and add it to
-	// to the list.
-	_mas_replace(mas, false, false, false);
-
-	if (mte_is_leaf(mas->node))
-		return 1;
-
-	mas_descend_adopt(mas);
-
-	if (mt_is_alloc(mas->tree))
-		mas_update_gap(mas, false);
-
-	// Free list of collected nodes.
-	// FIXME
-
-	return 1;
+	return mas_combine_separate(mas, &l_mas, &r_mas, b_node, b_end, empty);
 }
 
 static inline int mas_split(struct ma_state *mas,
@@ -2209,17 +2329,16 @@ static inline int mas_split(struct ma_state *mas,
 
 	struct maple_enode *ancestor = MAS_NONE;
 	unsigned char split = 0;
-	int height = 0;
-	int i;
+	int j, i = 0, height = 0;
+	struct maple_enode *list[100];
 
 	printk("%s\n", __func__);
-	MA_STATE(orig_mas, mas->tree, mas->index, mas->last);
+	mt_dump(mas->tree);
 	MA_STATE(l_mas, mas->tree, mas->index, mas->last);
 	MA_STATE(r_mas, mas->tree, mas->index, mas->last);
 	MA_STATE(orig_l_mas, mas->tree, mas->index, mas->last);
 	MA_STATE(orig_r_mas, mas->tree, mas->index, mas->last);
 
-	mas_dup_state(&orig_mas, mas);
 
 	// Allocation failures will happen early.
 	// TODO: Increase this by one when optimizing with a rebalance.
@@ -2237,6 +2356,8 @@ static inline int mas_split(struct ma_state *mas,
 	 * into the highest common ancestor necessary to be modified (which may
 	 * be a new root).
 	 */
+	printk("list %u -> %p\n", i, mas_mn(mas));
+	list[i++] = mas->node;
 	printk("full cnt = %u\n", mas->full_cnt);
 	while (height++ <= mas->full_cnt) {
 		struct maple_node *l, *r;
@@ -2260,17 +2381,17 @@ static inline int mas_split(struct ma_state *mas,
 			 */
 			ancestor = mt_mk_node(ma_mnode_ptr(mas_next_alloc(mas)),
 					      type);
-			for (i = 0; i < mt_slots[type]; i++) {
-				if (!b_node->pivot[i])
+			for (j = 0; j < mt_slots[type]; j++) {
+				if (!b_node->pivot[j])
 					break;
-				mte_set_rcu_slot(ancestor, i, b_node->slot[i]);
-				if (i < mt_pivots[type])
-					mte_set_pivot(ancestor, i,
-						      b_node->pivot[i]);
+				mte_set_rcu_slot(ancestor, j, b_node->slot[j]);
+				if (j < mt_pivots[type])
+					mte_set_pivot(ancestor, j,
+						      b_node->pivot[j]);
 				if (mt_is_alloc(mas->tree)) {
-					printk("Set gap %i to %lu\n", i, b_node->gap[i]);
-					mte_set_gap(ancestor, i,
-						    b_node->gap[i]);
+					printk("Set gap %i to %lu\n", j, b_node->gap[j]);
+					mte_set_gap(ancestor, j,
+						    b_node->gap[j]);
 				}
 			}
 			// Set the parent for the children.
@@ -2302,12 +2423,12 @@ static inline int mas_split(struct ma_state *mas,
 						mas->min, type);
 			printk("Splitting at %u\n", split);
 			if (split < slot_cnt)
-				i = mab_mas_cp(b_node, 0, split, &l_mas, 0);
+				j = mab_mas_cp(b_node, 0, split, &l_mas, 0);
 			else
-				i = mab_mas_cp(b_node, 0, slot_cnt, &l_mas, 0);
+				j = mab_mas_cp(b_node, 0, slot_cnt, &l_mas, 0);
 
 			mte_set_pivot(r_mas.node, 0, r_mas.max);
-			mab_mas_cp(b_node, i, new_end, &r_mas, 0);
+			mab_mas_cp(b_node, j, new_end, &r_mas, 0);
 
 			printk("l_max is %lu\n", l_mas.max);
 			mas_set_slot(&l_mas, mte_parent_slot(mas->node));
@@ -2333,11 +2454,11 @@ static inline int mas_split(struct ma_state *mas,
 			split = slot_cnt / 2;
 			printk("split is %u %u %u\n", split, mas_get_slot(&l_mas),
 			       mas_get_slot(&r_mas));
-			i = mab_mas_cp(b_node, 0, split, &l_mas, 0);
-			l_mas.max = b_node->pivot[i - 1];
+			j = mab_mas_cp(b_node, 0, split, &l_mas, 0);
+			l_mas.max = b_node->pivot[j - 1];
 			printk("l_mas.max %lu\n", l_mas.max);
 			r_mas.min = l_mas.max + 1;
-			mab_mas_cp(b_node, i, slot_cnt, &r_mas, 0);
+			mab_mas_cp(b_node, j, slot_cnt, &r_mas, 0);
 			p_slot = mas_get_slot(&orig_l_mas);
 			printk("setting parents of %p and %p to %p and %p\n",
 			       mas_mn(&orig_l_mas), mas_mn(&orig_r_mas),
@@ -2365,44 +2486,46 @@ static inline int mas_split(struct ma_state *mas,
 			}
 		}
 
-		i = 0;
+		j = 0;
 		memset(b_node, 0, sizeof(struct maple_big_node));
 		if (mte_is_root(mas->node))
 			cp = false;
 		else {
 			mas_ascend(mas);
 			mas_set_slot(mas, mte_parent_slot(mas->node));
+			printk("list %u -> %p\n", i, mas_mn(mas));
+			list[i++] = mas->node;
 		}
 
 		if (cp && mas_get_slot(&l_mas))
-			i = mas_mab_cp(mas, 0, mas_get_slot(&l_mas) - 1,
+			j = mas_mab_cp(mas, 0, mas_get_slot(&l_mas) - 1,
 				       b_node, 0);
 
-		split = i;
+		split = j;
 		printk("Split is %u\n", split);
-		b_node->slot[i] = l_mas.node;
-		printk("l_mas %p[%u] piv %lu\n", mas_mn(&l_mas), i, l_mas.max);
-		b_node->pivot[i] = l_mas.max;
-		mas_set_slot(&l_mas, i);
+		b_node->slot[j] = l_mas.node;
+		printk("l_mas %p[%u] piv %lu\n", mas_mn(&l_mas), j, l_mas.max);
+		b_node->pivot[j] = l_mas.max;
+		mas_set_slot(&l_mas, j);
 		if (mt_is_alloc(mas->tree)) {
-			b_node->gap[i] = mas_find_gap(&l_mas);
-			b_node->gap[i + 1] = mas_find_gap(&r_mas);
+			b_node->gap[j] = mas_find_gap(&l_mas);
+			b_node->gap[j + 1] = mas_find_gap(&r_mas);
 			printk("Setting %p gap to %lu\n",
-			       mas_mn(&l_mas), b_node->gap[i]);
+			       mas_mn(&l_mas), b_node->gap[j]);
 			printk("Setting %p gap to %lu\n",
-			       mas_mn(&r_mas), b_node->gap[i+1]);
+			       mas_mn(&r_mas), b_node->gap[j+1]);
 		}
-		b_node->slot[++i] = r_mas.node;
-		b_node->pivot[i] = r_mas.max;
-		mas_set_slot(&r_mas, i);
-		printk("New right is in %u (%p)\n", i, b_node->slot[i]);
+		b_node->slot[++j] = r_mas.node;
+		b_node->pivot[j] = r_mas.max;
+		mas_set_slot(&r_mas, j);
+		printk("New right is in %u (%p)\n", j, b_node->slot[j]);
 		printk("piv right is %lu\n", r_mas.max);
 		mas_dup_state(&orig_l_mas, &l_mas);
 		mas_dup_state(&orig_r_mas, &r_mas);
 		if (cp)
-			i = mas_mab_cp(mas, split + 1,
-				       mt_slot_count(mas->node), b_node,
-				       i + 1);
+			j = mas_mab_cp(mas, split + 1,
+				       mt_slot_count(mas->node) - 1, b_node,
+				       j + 1);
 	}
 
 
@@ -2410,18 +2533,26 @@ static inline int mas_split(struct ma_state *mas,
 	printk("Using %p (%p)\n", ancestor, mte_to_node(ancestor));
 	BUG_ON(mas_is_none(mas));
 	// Set the original node as dead
-	printk("Killing orig_mas %p\n", mas_mn(&orig_mas));
-	mas_set_node_dead(&orig_mas);
+	printk("%s: Killing orig_mas %p\n", __func__, mte_to_node(list[0]));
+	mte_to_node(list[0])->parent = ma_parent_ptr(mte_to_node(list[0]));
 	smp_wmb();
 
 	// Insert the new data in the tree
 	_mas_replace(mas, false, false, false);
 
+	mas_descend_adopt(mas);
+	do {
+		printk("%s: Free %u %p (%p)\n", __func__, i - 1,
+		       mte_to_node(list[i - 1]), list[i-1]);
+		mte_free(list[--i]);
+		rcu_barrier();
+	} while (i);
+
 	if (mt_is_alloc(mas->tree))
 		mas_update_gap(mas, false);
 	//mt_dump(mas->tree);
 	printk("Start at %p\n", mas_mn(mas));
-	mas_descend_adopt(mas);
+	mt_dump(mas->tree);
 	return 1;
 }
 
@@ -2432,7 +2563,8 @@ static inline int mas_commit_b_node(struct ma_state *mas,
 	struct maple_enode *new_node;
 
 	printk("Commit\n");
-	if ((end < mt_min_slot_cnt(mas->node)) && !mte_is_root(mas->node)) {
+	if ((end < mt_min_slot_cnt(mas->node)) && !mte_is_root(mas->node) &&
+	    (mas->tree->ma_height > 1) ) {
 		printk("end %u min slot %u\n", end, mt_min_slot_cnt(mas->node));
 		return mas_rebalance(mas, b_node, end);
 	}
@@ -2663,7 +2795,6 @@ static inline bool mas_wr_walk(struct ma_state *mas, unsigned long *range_min,
 
 
 		end = mas_data_end(mas);
-		printk("End of %p is %u max of %u\n", mas_mn(mas), end, mt_slots[type]);
 		if (unlikely(!mas_node_walk(mas, type, range_min, range_max)))
 			return false;
 
@@ -2799,7 +2930,7 @@ static inline int mas_spanning_store(struct ma_state *mas, void *entry)
 {
 	unsigned long range_min, range_max;
 	unsigned char b_end = 0; // big_node end.
-	unsigned char count = 0;
+	unsigned char count = 0, left, right;
 	struct maple_big_node b_node;
 	int node_cnt = 0;
 
@@ -2827,7 +2958,17 @@ static inline int mas_spanning_store(struct ma_state *mas, void *entry)
 	if (mas_is_err(mas))
 		return 0;
 
+	//mt_dump(mas->tree);
 	mas_dup_state(&l_mas, mas);
+	mas->last = mas->index;
+	mas_node_walk(mas, mte_node_type(mas->node), &range_min, &range_max);
+	left = mas_get_slot(mas);
+	mas->index = mas->last = l_mas.last;
+	mas_node_walk(mas, mte_node_type(mas->node), &range_min, &range_max);
+	right = mas_get_slot(mas);
+	printk("-->Copy node %p and free %u-%u\n", mte_to_node(mas->node), left, right);
+
+	mas_dup_state(mas, &l_mas);
 
 	// Set up right side.
 	mas_dup_state(&r_mas, mas);
@@ -2838,6 +2979,7 @@ static inline int mas_spanning_store(struct ma_state *mas, void *entry)
 	r_mas.last = r_mas.index = mas->last;
 
 	// Set up left side.
+	mas_dup_state(&l_mas, mas);
 	mas_set_slot(&l_mas, 0);
 	__mas_walk(&l_mas, &range_min, &range_max);
 
@@ -2867,40 +3009,10 @@ static inline int mas_spanning_store(struct ma_state *mas, void *entry)
 	l_mas.depth = 0;
 	b_end = mas_combine_separate(mas, &l_mas, &r_mas, &b_node, b_end,
 			     count);
-	printk("Done looping over spanning mess at %p\n", mas_mn(&l_mas));
-	printk("\n\n");
-
-	if (!b_end) {
-		printk("new root\n");
-	}
-	printk("Going to insert %p\n", mas_mn(&l_mas));
-
-	// Set up mas for insertion.
-	mas_set_node_dead(mas);
-	mas_dup_state(mas, &l_mas);
-	smp_wmb();
-
-	// FIXME: Mark modified nodes as dead.
-
-	// Insert new sub-tree
-	// FIXME: If mas isn't being replaced, kill the right node and add it to
-	// to the list.
-	_mas_replace(mas, false, false, false);
 
-	if (mte_is_leaf(mas->node))
-		return 1;
-
-	mt_dump(mas->tree);
-
-	mas_descend_adopt(mas);
-
-	if (mt_is_alloc(mas->tree))
-		mas_update_gap(mas, false);
-
-	// Free list of collected nodes.
-	// FIXME
-
-	return 1;
+	//mt_dump(mas->tree);
+	printk("%s: complete\n", __func__);
+	return b_end;
 }
 
 static inline void *_mas_store(struct ma_state *mas, void *entry, bool overwrite)
@@ -3227,6 +3339,7 @@ static inline unsigned long mas_next_node(struct ma_state *mas,
 	unsigned long start_piv;
 
 restart_next_node:
+	//printk("%s: %d %p\n", __func__, __LINE__, mas_mn(mas));
 	level = 0;
 	while (1) {
 		unsigned char count;
@@ -3239,6 +3352,7 @@ restart_next_node:
 
 		mn = mas->node;
 		slot = mas_get_slot(mas);
+		//printk("slot is %u\n", slot);
 		start_piv = mas_get_safe_pivot(mas, slot);
 		level++;
 		mas_ascend(mas);
@@ -3269,8 +3383,12 @@ restart_next_node:
 			if (level == 1) {
 				mas_set_slot(mas, slot);
 				mas->node = mn;
-				if (mas_dead_node(mas, start_piv))
+				printk("Found node %p\n", mas_mn(mas));
+				if (mas_dead_node(mas, start_piv)) {
+					printk("it's dead\n");
+					MT_BUG_ON(mas->tree, 1);
 					goto restart_next_node;
+				}
 				return pivot;
 			}
 
@@ -5054,6 +5172,7 @@ static inline void mas_dfs_postorder(struct ma_state *mas, unsigned long max)
 	unsigned long p_min, p_max;
 
 	mas_set_slot(mas, mte_parent_slot(mas->node));
+	printk("next node from %p\n", mas_mn(mas));
 	mas_next_node(mas, max);
 	if (mas->node != MAS_NONE)
 		return;
@@ -5092,7 +5211,6 @@ void mt_validate(struct maple_tree *mt)
 	mas_start(&mas);
 	mas_first_entry(&mas, ULONG_MAX);
 	while (mas.node != MAS_NONE) {
-		printk("Checking %p\n", mas_mn(&mas));
 		if (!mte_is_root(mas.node)) {
 			end = mas_data_end(&mas);
 			if (end < mt_min_slot_cnt(mas.node)) {
diff --git a/lib/test_maple_tree.c b/lib/test_maple_tree.c
index c6adf2419a15..f374b070ab65 100644
--- a/lib/test_maple_tree.c
+++ b/lib/test_maple_tree.c
@@ -85,7 +85,7 @@ static noinline void check_load(struct maple_tree *mt, unsigned long index,
 				void *ptr)
 {
 	void *ret = mtree_test_load(mt, index);
-	printk("Load %lu returned %p expect %p\n", index, ret, ptr);
+//	printk("Load %lu returned %p expect %p\n", index, ret, ptr);
 	MT_BUG_ON(mt, ret != ptr);
 }
 
@@ -30139,7 +30139,7 @@ static noinline void check_alloc_rev_range(struct maple_tree *mt)
 
 	MA_STATE(mas, mt, 0, 0);
 
-#define DEBUG_REV_RANGE 1
+#define DEBUG_REV_RANGE 0
 	for (i = 0; i < range_cnt; i += 2) {
 		/* Inclusive, Inclusive (with the -1) */
 
@@ -30377,17 +30377,20 @@ static noinline void check_ranges(struct maple_tree *mt)
 	mt_set_non_kernel(10);
 	check_store_range(mt, r[10], r[11], xa_mk_value(r[10]), 0);
 	mtree_destroy(mt);
-//	rcu_barrier();
 
 	// Create tree of 1-200
 	check_seq(mt, 200, false);
 	// Store 45-168
 	check_store_range(mt, r[10], r[11], xa_mk_value(r[10]), 0);
+	mt_dump(mt);
+	printk("Destroying after tc\n");
 	mtree_destroy(mt);
 
 	check_seq(mt, 30, false);
 	check_store_range(mt, 6, 18, xa_mk_value(6), 0);
 	mtree_destroy(mt);
+	//rcu_barrier();
+	//return;
 
 	// Overwrite across multiple levels.
 	// Create tree of 1-400
@@ -30457,6 +30460,7 @@ static noinline void check_ranges(struct maple_tree *mt)
 	mt_set_non_kernel(0);
 	mtree_destroy(mt);
 
+
 	mt_set_non_kernel(50);
 	check_seq(mt, 400, false);
 	check_store_range(mt, 362, 367, xa_mk_value(362), 0);
@@ -30526,6 +30530,19 @@ static noinline void check_ranges(struct maple_tree *mt)
 	mtree_erase(mt, 230);
 	mt_set_non_kernel(0);
 	mtree_destroy(mt);
+
+	mtree_init(mt, MAPLE_ALLOC_RANGE);
+	for (i = 0; i <= 500; i++) {
+		int val = i*10;
+		int val2 = (i+1)*10;
+		check_store_range(mt, val, val2, xa_mk_value(val), 0);
+	}
+	mt_dump(mt);
+	check_store_range(mt, 4600, 4959, xa_mk_value(1), 0);
+	mt_dump(mt);
+	mt_validate(mt);
+	mtree_destroy(mt);
+
 }
 
 static noinline void check_next_entry(struct maple_tree *mt)
-- 
2.50.1