From 6b21ee2beaca497366a38b562a42a249e85c363c Mon Sep 17 00:00:00 2001
From: "Liam R. Howlett" <howlett@gmail.com>
Date: Fri, 14 Feb 2025 15:42:56 -0500
Subject: [PATCH] working on split rebalance now

Signed-off-by: Liam R. Howlett <howlett@gmail.com>
---
 lib/maple_tree.c | 645 +++++++++++++++++++++++++++++++----------------
 1 file changed, 426 insertions(+), 219 deletions(-)

diff --git a/lib/maple_tree.c b/lib/maple_tree.c
index 7133cc8242356..6f6b748657ee8 100644
--- a/lib/maple_tree.c
+++ b/lib/maple_tree.c
@@ -1707,6 +1707,7 @@ static inline void mas_put_in_tree(struct ma_state *mas,
 	unsigned char offset;
 	void __rcu **slots;
 
+	printk("%p is root?\n", mas->node);
 	if (mte_is_root(mas->node)) {
 		mas_mn(mas)->parent = ma_parent_ptr(mas_tree_parent(mas));
 		rcu_assign_pointer(mas->tree->ma_root, mte_mk_root(mas->node));
@@ -1760,9 +1761,10 @@ static inline bool mas_find_child(struct ma_state *mas, struct ma_state *child)
 	end = ma_data_end(node, mt, pivots, mas->max);
 	for (offset = mas->offset; offset <= end; offset++) {
 		entry = mas_slot_locked(mas, slots, offset);
-		//printk("check %p[%u/%u] %p\n", mas->node, end, offset, entry);
+		printk("check %p[%u/%u] %p\n", mas->node, end, offset, entry);
+		printk("parent of %p is %p\n", entry, mte_parent(entry));
 		if (mte_parent(entry) == node) {
-			//printk(" found entry at %u\n", offset);
+			printk(" found entry at %u\n", offset);
 			*child = *mas;
 			mas->offset = offset + 1;
 			child->offset = offset;
@@ -3122,6 +3124,7 @@ struct ma_node_part {
 	unsigned long gaps[2];
 	bool unfinished;
 	unsigned char skip;
+	bool leaf;
 };
 
 struct ma_node_info {
@@ -3145,38 +3148,12 @@ struct ma_node_state {
 		struct ma_node_info *info;
 		struct ma_node_part *part;
 	};
-	unsigned char start;            /* The start offset */
+	unsigned char start;            /* The start offset into the source*/
 	unsigned char size;             /* The size to copy or insert */
 	struct ma_node_info *dst;       /* The destination node */
 	bool use_part;
 };
 
-/*
- * Unsafe from reader side
- */
-static inline void mni_set_end(struct ma_node_info *mni)
-{
-	unsigned char offset;
-
-	if (mni->type == maple_arange_64) {
-		mni->end = ma_meta_end(mni->node, mni->type);
-		return;
-	}
-
-	offset = mt_pivots[mni->type] - 1;
-	if (likely(!mni->pivots[offset])) {
-		mni->end = ma_meta_end(mni->node, mni->type);
-		return;
-	}
-
-	if (likely(mni->pivots[offset] == mni->max)) {
-		mni->end = offset;
-		return;
-	}
-
-	mni->end = mt_pivots[mni->type];
-}
-
 /*
  * mns_node_part_leaf_init() - Initialize what is being inserted, calculate how
  * many slots will be skipped.
@@ -3190,12 +3167,12 @@ void mns_node_part_leaf_init(struct ma_node_part *ma_part,
 {
 	ma_part->pos = 0;
 	ma_part->size = 0;
-	//printk("%s: %lx - %lx store %lx - %lx\n", __func__, wr_mas->r_min, wr_mas->r_max, wr_mas->mas->index, wr_mas->mas->last);
+	printk("%s: %lx - %lx store %lx - %lx\n", __func__, wr_mas->r_min, wr_mas->r_max, wr_mas->mas->index, wr_mas->mas->last);
 	if (wr_mas->r_min < wr_mas->mas->index) {
 		ma_part->pivots[0] = wr_mas->mas->index - 1;
 		ma_part->slots[0] = wr_mas->content;
 		ma_part->size++;
-		//printk("r_min remains\n");
+		printk("r_min remains\n");
 	}
 
 	ma_part->pivots[ma_part->size] = wr_mas->mas->last;
@@ -3206,18 +3183,22 @@ void mns_node_part_leaf_init(struct ma_node_part *ma_part,
 		ma_part->pivots[ma_part->size] = wr_mas->end_piv;
 		ma_part->slots[ma_part->size] = src->slots[wr_mas->offset_end];
 		ma_part->size++;
-		//printk("r_max remains\n");
+		printk("r_max remains\n");
 	}
 
 	ma_part->unfinished = false;
 	ma_part->dst_max_off = 255;
 	ma_part->skip = wr_mas->offset_end - wr_mas->mas->offset + 1;
+	printk("skip %u = %u - %u + 1\n", ma_part->skip, wr_mas->offset_end, wr_mas->mas->offset);
+	printk("ma_part size %u\n", ma_part->size);
+	ma_part->leaf = true;
 }
 
 static inline
 void mni_node_part_init(struct ma_node_part *ma_part,
 		struct ma_node_info *left, struct ma_node_info *right)
 {
+	printk("set slot 0 to %p\n", left->enode);
 	ma_part->slots[0] = left->enode;
 	ma_part->pivots[0] = left->max;
 	ma_part->gaps[0] = left->max_gap;
@@ -3300,9 +3281,49 @@ void mni_node_init(struct ma_node_info *mni, struct maple_node *node,
 {
 	mni->node = node;
 	mni->type = type;
+	mni->enode = mt_mk_node(node, type);
 	_mni_node_init(mni, node, type);
 }
 
+/*
+ * Unsafe from reader side
+ */
+static inline void mni_set_end(struct ma_node_info *ni)
+{
+       unsigned char offset;
+
+       if (ni->type == maple_arange_64) {
+               ni->end = ma_meta_end(ni->node, ni->type);
+               return;
+       }
+
+       offset = mt_pivots[ni->type] - 1;
+       if (likely(!ni->pivots[offset])) {
+               ni->end = ma_meta_end(ni->node, ni->type);
+               return;
+       }
+
+       if (likely(ni->pivots[offset] == ni->max)) {
+               ni->end = offset;
+               return;
+       }
+
+       ni->end = mt_pivots[ni->type];
+}
+
+static inline
+void mni_mas_init(struct ma_node_info *mni, struct ma_state *mas)
+{
+	mni->node = mte_to_node(mas->node);
+	mni->type = mte_node_type(mas->node);
+	_mni_node_init(mni, mni->node, mni->type);
+	mni->enode = mas->node;
+	mni->end = mas->end;
+	mni->max = mas->max;
+	mni->min = mas->min;
+	mni_set_end(mni);
+}
+
 static inline
 void mns_mas_init(struct ma_node_state *mns, struct ma_node_info *mni,
 		struct ma_state *mas)
@@ -3326,26 +3347,24 @@ static inline
 void mns_mni_init(struct ma_node_state *mns, struct ma_node_info *dst,
 		unsigned char start, unsigned char len)
 {
-	mns->dst = ni;
+	mns->dst = dst;
 	mns->start = start;
-	mns->size = len + 1;
+	mns->size = len;
+	printk("Dst %p <= [%u] + %u\n", dst, start, len);
 }
 
 static inline
-bool mns_ends_in_null(struct ma_node_state *mns)
+bool mns_ends_in_null(struct ma_node_state *ns)
 {
 	void __rcu **s_slots;
-	unsigned long *s_piv;
 
-	if (mns->use_part) {
+	if (ns->use_part) {
 		s_slots = ns->part->slots;
-		s_piv = ns->part->pivots;
 	} else {
 		s_slots = ns->info->slots;
-		s_piv = ns->info->pivots;
 	}
-	
-	if (s_slots[mns->start + mns->size - 1])
+
+	if (s_slots[ns->start + ns->size - 1])
 		return false;
 
 	return true;
@@ -3557,27 +3576,6 @@ static __always_inline unsigned char mas_wr_new_end(struct ma_wr_state *wr_mas)
 	return new_end;
 }
 
-static inline void mas_wr_converged(struct ma_node_info *src,
-		struct ma_node_info *dst, struct ma_node_part *ma_part,
-		struct ma_state *mas)
-{
-	mni_node_init(dst, mas_pop_node(mas), src->type);
-	printk("%s: %p -> %p\n", __func__, src->node, dst->node);
-
-	if (mas->offset)
-		mni_cp(src, dst, mas->offset);
-
-	mni_insert_part(ma_part, dst);
-	src->offset += ma_part->skip;
-
-	if (src->offset <= mas->end)
-		mni_cp(src, dst, mas->end - src->offset + 1);
-
-	dst->node->parent = src->node->parent;
-	mni_finalise(dst);
-	mas_set_height(mas);
-}
-
 static void mas_wr_split_no_null(struct ma_node_info *src,
 		struct ma_node_info *left, struct ma_node_info *right,
 		unsigned char total, struct ma_node_part *ma_part)
@@ -3617,50 +3615,6 @@ static void mas_wr_split_no_null(struct ma_node_info *src,
 	right->min = left->max + 1;
 }
 
-static inline void mni_in_left(struct ma_node_info *src,
-		struct ma_node_info *left, struct ma_node_info *right,
-		struct ma_state *mas, unsigned char split,
-		unsigned char new_end, struct ma_node_part *ma_part)
-{
-	ma_part->dst_max_off = split;
-	if (mas->offset)
-		mni_cp(src, left, mas->offset);
-
-	mni_insert_part(ma_part, left);
-	src->offset+= ma_part->skip;
-	//printk("\t\tsrc L offset adjusted to %u\n", src->offset);
-	if (left->offset <= split)
-		mni_cp(src, left, split - left->offset + 1);
-
-	mas_wr_split_no_null(src, left, right, new_end, ma_part);
-	if (ma_part->unfinished)
-		mni_insert_part(ma_part, right);
-
-	right->min = left->max + 1;
-	mni_cp(src, right, mas->end - src->offset + 1);
-}
-
-static inline void mni_in_right(struct ma_node_info *src,
-		struct ma_node_info *left, struct ma_node_info *right,
-		struct ma_state *mas, unsigned char split,
-		unsigned char new_end, struct ma_node_part *ma_part)
-{
-	unsigned char cp;
-
-	cp = mas->offset - split - 1;
-	mni_cp(src, left, split + 1);
-	mas_wr_split_no_null(src, left, right, new_end, ma_part);
-	right->min = left->max + 1;
-	if (cp)
-		mni_cp(src, right, cp);
-
-	mni_insert_part(ma_part, right);
-	src->offset+= ma_part->skip;
-	//printk("\t\tsrc offset adjusted to %u\n", src->offset);
-	if (src->offset <= mas->end)
-		mni_cp(src, right, mas->end - src->offset + 1);
-}
-
 /*
  * mas_wr_rebalance_calc() - Try to calculate a rebalance that will work
  * @data_size: The total data to be written
@@ -3700,8 +3654,11 @@ void mas_wr_ascend_init(struct ma_state *mas, struct ma_node_info *ni)
 	ni->node = mte_to_node(mas->node);
 	ni->type = mte_node_type(mas->node);
 	_mni_node_init(ni, ni->node, ni->type);
+	mni_set_end(ni);
+	ni->insert_off = mas->offset;
 }
 
+#if 0
 static void mas_wr_rebalance_nodes(
 		struct ma_node_info *l_src,
 		struct ma_node_info *r_src,
@@ -3836,67 +3793,145 @@ static void mas_wr_rebalance_nodes(
 		}
 	}
 }
+#endif
 
 /*
  * Assemble all the states into the dst within those states
  *
  */
 static inline
-void mns_assemble(struct ma_node_state **states, unsigned char len)
+void mns_assemble(struct ma_node_state *states, unsigned char len)
 {
 	unsigned char i;
 	unsigned long max;
 	struct ma_node_info *last_dst;
 
-	last_dst = states[0]->dst;
+	last_dst = states[0].dst;
+	max = 0;
 	for (i = 0; i < len; i++) {
-		struct ma_node_state *ns = states[i];
+		struct ma_node_state *ns = &states[i];
 		void __rcu **s_slots, **d_slots;
 		unsigned long *s_piv, *d_piv;
 		unsigned long *s_gap, *d_gap;
 		unsigned char size;
 		unsigned char piv_overflow = 0;
 
-		if (last_dst != ns->dst)
+		if (last_dst != ns->dst) {
+			WARN_ON_ONCE(!max);
 			ns->dst->min = max + 1;
+		}
 
 		d_slots = ns->dst->slots + ns->dst->offset;
+		printk("Store at %p\n", d_slots);
 		d_piv = ns->dst->pivots + ns->dst->offset;
-		d_gap = ns->dst->gaps + ns->dst->offset;
 		size = ns->size;
+		printk("%u: dst %p [%u] size %lu\n", i, ns->dst, ns->dst->offset, size);
 
 		if (ns->use_part) {
 			s_slots = ns->part->slots + ns->start;
 			s_piv = ns->part->pivots + ns->start;
 			s_gap = ns->part->gaps + ns->start;
 			max = s_piv[size - 1];
+			printk("max from offset %u\n", size -1);
+			if (!ns->part->leaf) {
+				int j = 0;
+
+				printk("not a leaf\n");
+				do {
+					struct maple_enode *child;
+
+					child = ma_enode_ptr(ns->part->slots[j]);
+					printk("set %p %p[%u] parent\n", child,
+					       ns->dst->enode, j);
+					mte_set_parent(child, ns->dst->enode,
+						       j + ns->dst->offset);
+				} while (++j < ns->part->size);
+			}
 		} else {
+			printk("Use info %p + %u\n", ns->info->slots, ns->start);
 			s_slots = ns->info->slots + ns->start;
 			s_piv = ns->info->pivots + ns->start;
 			s_gap = ns->info->gaps + ns->start;
-			if (ns->start + size >= mt_pivots[ns->info->type]) {
+			if (ns->start + size > mt_pivots[ns->info->type]) {
 				piv_overflow = 1;
 				max = ns->info->max;
-			} else
+				printk("Use max for pivot\n");
+			} else {
 				max = s_piv[size - 1];
+			}
 		}
 
+		printk("Cp %p + %u to dst %p + %lu\n",
+		       s_slots, size,
+		       ns->dst->slots, ns->dst->offset);
 		memcpy(d_slots, s_slots, size * sizeof(void __rcu *));
-		if (d_gap)
+		if (ns->dst->gaps) {
+			d_gap = ns->dst->gaps + ns->dst->offset;
 			memcpy(d_gap, s_gap, size * sizeof(unsigned long));
+		}
 
 		if (ns->dst->offset + size >= mt_pivots[ns->dst->type])
 			piv_overflow = 1;
-		else if (piv_overflow) /* Source overflow */
+		else if (piv_overflow) { /* Source overflow */
+			printk("source overflow, set pivot to %lu\n", max);
 			*(d_piv + size - 1) = max;
+		}
 
 		ns->dst->offset += size;
+		ns->dst->end = ns->dst->offset - 1;
+		printk("\t\t set %p end to %u\n", ns->dst->node, ns->dst->end);
 		size -= piv_overflow;
 		memcpy(d_piv, s_piv, size * sizeof(unsigned long));
 		ns->dst->max = max;
 	}
 }
 
+static inline void mas_wr_converged(struct ma_node_info *src,
+		struct ma_node_info *dst, struct ma_node_part *ma_part,
+		struct ma_state *mas)
+
+{
+	struct ma_node_state converged[3];
+	unsigned char i = 0;
+	unsigned char off = 0;
+
+	mni_node_init(dst, mas_pop_node(mas), src->type);
+	printk("dst node is %p\n", dst->node);
+	printk("%s: %p -> %p\n", __func__, src->node, dst->node);
+
+	if (src->insert_off) {
+		printk("%s: cp start data\n", __func__);
+		mns_mni_init(&converged[i], dst, 0, src->insert_off);
+		converged[i].info = src;
+		off = src->insert_off;
+		i++;
+	}
+
+	printk("%s: insert data\n", __func__);
+	mns_mni_init(&converged[i], dst, 0, ma_part->size);
+	converged[i].part = ma_part;
+	converged[i].use_part = true;
+	off += ma_part->skip;
+	i++;
+
+	if (src->end >= off) {
+		printk ("end %u %u + %u\n", src->end, ma_part->skip,
+			src->insert_off);
+		printk("%s: cp end of data\n", __func__);
+		unsigned char start = off;
+		unsigned char size = src->end - start + 1;
+		mns_mni_init(&converged[i], dst, start, size);
+		converged[i].info = src;
+		i++;
+		off += size - 1;
+	}
+
+	mns_assemble(converged, i);
+	dst->node->parent = src->node->parent;
+	printk("Set parent of %p to %p\n", dst->node, dst->node->parent);
+	mni_finalise(dst);
+}
+
 /*
  * mas_wr_try_rebalance() - Try to rebalance two nodes, this may not work out.
  * @src: The source node state
@@ -3910,14 +3945,12 @@ void mns_assemble(struct ma_node_state **states, unsigned char len)
 static bool mas_wr_try_rebalance(struct ma_state *mas,
 		struct ma_node_info *src, unsigned char new_end,
 		struct ma_node_info *left, struct ma_node_info *right,
-		struct ma_node_part *ma_part)
+		struct ma_node_part *ma_part, unsigned char insert_end)
 {
 	struct ma_state tmp_mas;
-	struct ma_node_info src2_info, parent, new_parent;
-	struct ma_node_info  *l_src, *r_src;
-	struct ma_node_state src2;
-	unsigned char split, max;
-	unsigned char p_end, p_off;
+	struct ma_node_info src2, parent, new_parent;
+	struct ma_node_state state[4];
+	unsigned char split, max, i;
 	bool left_store = false;
 
 	/*
@@ -3927,65 +3960,147 @@ static bool mas_wr_try_rebalance(struct ma_state *mas,
 
 	tmp_mas = *mas;
 	mas_wr_ascend_init(&tmp_mas, &parent);
-	p_off = tmp_mas.offset;
-	p_end = ma_data_end(parent.node, parent.type, parent.pivots,
-			    parent.max);
-	//printk("parent  %p has end %u %p\n", parent.node, p_end, parent.slots[p_end]);
+	//printk("parent  %p has end %u %p\n", parent.node, parent.end, parent.slots[parent.end]);
 	max = mt_slots[src->type] - 1;
 	if (ma_is_leaf(src->type))
 		max--;
 
-	if (!p_off)
-		goto try_right;
+	src->end = mas->end;
+	if (!parent.insert_off) /* No left sibling */
+		goto try_right; /* There must be a right sibling */
 
+	/* Check for space in left sibling */
 	tmp_mas.offset--;
 	mas_descend(&tmp_mas);
-	mns_mas_init(&src2, &src2_info, &tmp_mas);
-	src2_info.insert_off = 255;
-	src2_info.end = ma_data_end(src2.node, src2.type, src2.pivots,
-			src2.max);
+	mni_mas_init(&src2, &tmp_mas);
 	split = mas_wr_rebalance_calc(src2.end + new_end, src2.type);
 	if (split) {
-		p_off--;
-		l_src = &src2;
-		r_src = src;
-		r_src->end = mas->end;
+		parent.insert_off--;
+		/* Left will be src2, right will be src */
+		left->min = src2.min;
+		right->max = src->max;
 	} else {
-		if (p_end <= p_off)
+		if (parent.end <= parent.insert_off)
 			return false;
 
 		mas_ascend(&tmp_mas);
 try_right:
-		tmp_mas.offset = p_off + 1;
+		tmp_mas.offset = parent.insert_off + 1;
 		mas_descend(&tmp_mas);
-		mns_mas_init(&src2, &tmp_mas);
-		src2.insert = 255;
-		src2.end = ma_data_end(src2.node, src2.type,
-					 src2.pivots, src2.max);
-		src->end = mas->end;
+		mni_mas_init(&src2, &tmp_mas);
+		mni_set_end(&src2);
 		split = mas_wr_rebalance_calc(src2.end + new_end, src2.type);
 		if (!split)
 			return false;
 
 		split = src2.end + new_end - split;
-		l_src = src;
-		r_src = &src2;
 		left_store = true;
+		/* Left will be src, right will be src2 */
+		left->min = src->min;
+		right->max = src2.max;
 	}
 
 	/*
 	 * At this point, the rebalance operation will succeed.
 	 */
 
-	left->min = l_src->min;
-	mas_wr_rebalance_nodes(l_src, r_src, left, right, left_store, split,
-			       ma_part, mas->offset, new_end);
+	i = 0;
+	if (left_store)  {
+		unsigned char space;
+
+		/*
+		 * Left pushes data right.
+		 * Fill left up to split from l_src and ma_part - Func_1
+		 */
+		if (mas->offset) {
+			state[i].info = src;
+			mns_mni_init(&state[i++], left, 0, mas->offset);
+		}
+
+		state[i].part = ma_part;
+		state[i].use_part = true;
+		space = split - mas->offset + 1;
+		if (space >= ma_part->size) {
+			mns_mni_init(&state[i++], left, 0, ma_part->size);
+		} else {
+			/* The insert part spans the left and right */
+			mns_mni_init(&state[i], left, 0, space);
+			if (mns_ends_in_null(&state[i])) {
+				state[i].size--;
+				split--;
+			}
+
+			space = ma_part->size - state[i].size;
+			i++;
+		/*
+		 * fill right from ma_part and l_src - Func_1
+		 */
+			state[i].part = ma_part;
+			state[i].use_part = true;
+			mns_mni_init(&state[i], right, state[i - 1].size,
+				     space);
+			i++;
+		}
+
+		if (split > insert_end) {
+			state[i].info = src;
+			mns_mni_init(&state[i], left, insert_end + 1,
+					split - insert_end + 1);
+
+			if (mns_ends_in_null(&state[i])) {
+				state[i].size--;
+				split--;
+			}
+			i++;
+		}
+
+		/*
+		 * fill right with right. - Func_2
+		 */
+		state[i].info = &src2;
+		mns_mni_init(&state[i++], right, split + 1, mas->end - split + 1);
+	} else {
+		/*
+		 * Right pushes data left
+		 * Copy left to new left - Func_2
+		 */
+		state[i].info = &src2;
+		mns_mni_init(&state[i], left, 0, split);
+		if (mns_ends_in_null(&state[i])) {
+			state[i].size--;
+			split--;
+		}
+		i++;
+
+		/*
+		 * copy from right or ma_part to split to new left - Func_1
+		 * fill right from ma_part and r_src - Func_1
+		 */
+
+		state[i].info = &src2;
+		mns_mni_init(&state[i++], right, split + 1,
+			     src2.end + 1);
+
+		if (mas->offset) {
+			state[i].info = src;
+			mns_mni_init(&state[i++], right, 0, mas->offset);
+		}
+
+		state[i].part = ma_part;
+		mns_mni_init(&state[i], right, 0, ma_part->size);
+		state[i++].use_part = true;
+		if (insert_end > mas->end) {
+			state[i].info = src;
+			mns_mni_init(&state[i++], right, insert_end + 1,
+				     mas->end - insert_end + 1);
+		}
+	}
 
 	mni_finalise(left);
 	mni_finalise(right);
 	mas_ascend(mas);
-	mas->end = p_end;
-	mas->offset = p_off;
+	mas->end = parent.end;
+	mas->offset = parent.insert_off;
 	mni_node_part_init(ma_part, left, right);
 	ma_part->skip = 2;
 	mas_wr_converged(&parent, &new_parent, ma_part, mas);
@@ -3994,6 +4109,7 @@ try_right:
 	return true;
 }
 
+#if 0
 /*
  * There is insufficient data in the node after a store.
  * This rebalance will succeed, not like the split variant that will attempt
@@ -4036,7 +4152,6 @@ static void mas_wr_rebalance(struct ma_wr_state *wr_mas)
 		printk("Ascend to parent %p\n", parent.node);
 		printk("offset of child is %u\n", mas->offset);
 		mas->depth--;
-		mns_set_end(&parent);
 		parent.insert = mas->offset;
 		printk("parent insert is %u\n", parent.insert);
 
@@ -4069,6 +4184,18 @@ static void mas_wr_rebalance(struct ma_wr_state *wr_mas)
 		 * Two possibilities:
 		 * 1. keep two nodes if possible and limit ripple
 		 * 2. make one node if possible and limit memory use
+		 *
+		 * Case 1:
+		 *  Left takes data from right.
+		 *  Fill left up to split from l_src and ma_part - Func_1
+		 *  Fill left up from l_src remainder - Func_2
+		 *  Fill left up to split from right. - Func_2
+		 *  fill right with remainder of right. - Func_2
+		 *
+		 *  Right takes data from left
+		 *  Copy left to new left up to split - Func_2
+		 *  Fill right with remainder of left - Func_2
+		 *  Fill right from old right or ma_part - Func_1
 		 */
 		if ((total > 2 * mt_min_slots[l_src.type] ) ||
 		    ma_is_root(parent.node)) {
@@ -4101,7 +4228,6 @@ static void mas_wr_rebalance(struct ma_wr_state *wr_mas)
 			break;
 		}
 
-
 		printk("consume\n");
 		/* Reduce two nodes into one */
 		if (l_store) {
@@ -4143,6 +4269,7 @@ static void mas_wr_rebalance(struct ma_wr_state *wr_mas)
 	mtree_range_walk(mas);
 	mt_dump(mas->tree, mt_dump_dec);
 }
+#endif
 
 /*
  * There is not enough room to contain the store in one node.
@@ -4160,8 +4287,9 @@ static void mas_wr_split(struct ma_wr_state *wr_mas)
 
 	trace_ma_op(__func__, mas);
 
-	//printk("\t\t%s\n", __func__);
-	//mt_dump(mas->tree, mt_dump_hex);
+	printk("\t\t%s\n", __func__);
+	mt_dump(mas->tree, mt_dump_hex);
+	printk ("Store %lu - %lu => %p\n", mas->index, mas->last, wr_mas->entry);
 	height = mas_mt_height(mas);
 	mas->depth = height;
 
@@ -4172,126 +4300,193 @@ static void mas_wr_split(struct ma_wr_state *wr_mas)
 	mns_node_part_leaf_init(&ma_part, wr_mas, &src_info);
 	total = mas->end + ma_part.size - 1;
 	split = (total + 1) / 2;
+	printk("Splitting leaf at %u store at %u\n", split, mas->offset);
 
-	if (mt_is_alloc(mas->tree))
+	if (mt_is_alloc(mas->tree)) {
 		right.alloc = left.alloc = true;
+	}
 
 	if (height > 1 &&
-			mas_wr_try_rebalance(mas, &src, total, &left, &right, &ma_part))
+			mas_wr_try_rebalance(mas, &src_info, total, &left,
+					     &right, &ma_part, wr_mas->offset_end))
 		goto rebalanced;
 
 	left.min = mas->min;
 	right.max = mas->max;
 	if (split >= mas->offset) {
-		unsigned char space; /* size remaining in left */
+		unsigned char space;
 
+		printk("Store in left\n");
 		if (mas->offset) {
-			state[i]->info = &src;
-			mns_mni_init(state[i++], &left, 0, mas->offset);
+			state[i].info = &src_info;
+			printk("src is info %p\n", src_info.node);
+			mns_mni_init(&state[i++], &left, 0, mas->offset);
 		}
 
-		state[i]->part = &ma_part;
-		state[i]->part = true;
+		state[i].part = &ma_part;
+		state[i].use_part = true;
 		space = split - mas->offset + 1;
-		if (space >= ma_part.size)
-			mns_mni_init(state[i++], &left, 0, ma_part.size);
+		if (space >= ma_part.size) {
+			mns_mni_init(&state[i++], &left, 0, ma_part.size);
 		} else {
-			mns_mni_init(state[i++], &left, 0, space);
-			state[i]->part = &ma_part;
-			state[i]->part = true;
-			mns_mni_init(state[i++], &right, space - 1, ma_part.size);
+			/* The insert part spans the left and right */
+			mns_mni_init(&state[i], &left, 0, space);
+			if (mns_ends_in_null(&state[i])) {
+				state[i].size--;
+				split--;
+			}
 
+			space = ma_part.size - state[i].size;
+			i++;
+			state[i].part = &ma_part;
+			state[i].use_part = true;
+			mns_mni_init(&state[i], &right, state[i - 1].size,
+				     space);
+			i++;
 		}
 
 		if (split > wr_mas->offset_end) {
-			state[i]->info = &src;
-			mns_mni_init(state[i++], &left, wr_mas->offset_end + 1,
+			state[i].info = &src_info;
+			printk("src is info %p\n", src_info.node);
+			mns_mni_init(&state[i], &left, wr_mas->offset_end + 1,
 					split - wr_mas->offset_end + 1);
-		}
 
-		if (mns_ends_in_null(state[i - 1])) {
-			state[i - 1]->size--;
-			split--;
+			if (mns_ends_in_null(&state[i])) {
+				state[i - 1].size--;
+				split--;
+			}
+			i++;
 		}
 
-		state[i]->info = &src;
-		mns_min_init(state[i++], &right, split + 1, mas-end - split + 1);
+		state[i].info = &src_info;
+		printk("src is info %p\n", src_info.node);
+		mns_mni_init(&state[i++], &right, split + 1, mas->end - split + 1);
 	} else {
-		state[i]->info = &src;
-		mns_mni_init(state[i++], &left, 0, split);
-		if (mns_ends_in_null(state[i - 1])) {
-			state[i - 1]->size--;
+		printk("Store in right\n");
+		printk("src is info %p\n", src_info.node);
+		state[i].info = &src_info;
+		split++; /* it's a size */
+		mns_mni_init(&state[i], &left, 0, split);
+		if (mns_ends_in_null(&state[i])) {
+			printk("Ends in null, move split\n");
+			state[i].size--;
 			split--;
 		}
+		i++;
 
-		if (mas->offset > split + 1) {
-			state[i]->info = &src;
-			mns_mni_init(state[i++], &right, split + 1,
-					mas->offset - split + 1);
+		if (mas->offset >= split) {
+			state[i].info = &src_info;
+		printk("src is info %p\n", src_info.node);
+			mns_mni_init(&state[i++], &right, split,
+				     mas->offset - split);
 		}
-		state[i]->part = &ma_part;
-		mns_mni_init(state[i], &right, 0, ma_part.size);
-		state[i++]->part = true;
+		printk("Store part next 0 - %u\n", ma_part.size - 1);
+		state[i].part = &ma_part;
+		mns_mni_init(&state[i], &right, 0, ma_part.size);
+		state[i++].use_part = true;
 		if (wr_mas->offset_end > mas->end) {
-			state[i]->info = &src;
-			mns_mni_init(state[i++], &right, wr_mas->offset_end + 1,
-					mas->end - wr_mas->offset_end + 1);
+			state[i].info = &src_info;
+		printk("src is info %p\n", src_info.node);
+			mns_mni_init(&state[i++], &right, wr_mas->offset_end + 1,
+				     mas->end - wr_mas->offset_end + 1);
 		}
 	}
 
-	mns_assemble(&state, i);
+	printk("\n\nAssemble %u\n", i);
+	mns_assemble(state, i);
+	printk("\n\n");
 
 	mni_finalise(&left);
 	mni_finalise(&right);
 	mni_node_part_init(&ma_part, &left, &right);
+	ma_part.leaf = false;
 
 	if (height == 1) {
-		if (mt_is_alloc(mas->tree))
-			src.type = maple_arange_64;
-		else
-			src.type = maple_range_64;
-
+		printk("src parent is %p\n", src_info.node->parent);
+		printk("New root\n");
 		goto new_root;
 	}
 
 	//printk("%d height is %d\n", __LINE__, height);
 	while (--height) {
 		i = 0;
-		mas_wr_ascend_init(mas, &src);
-		mas->end = ma_data_end(src.node, src.type, src.pivots,
-				src.max);
+		mas_wr_ascend_init(mas, &src_info);
+		mas->end = src_info.end;
 		total = mas->end + 1;
-		if (mas->end + 1 < mt_slots[src.type])
+		if (mas->end + 1 < mt_slots[src_info.type])
 			goto converged;
 
 		//printk("\tConsume %p type %u\n", src.node, src.type);
-		mni_node_init(&left, mas_pop_node(mas), src.type);
-		mni_node_init(&right, mas_pop_node(mas), src.type);
+		mni_node_init(&left, mas_pop_node(mas), src_info.type);
+		mni_node_init(&right, mas_pop_node(mas), src_info.type);
 		if ((height > 1) &&
-				(mas_wr_try_rebalance(mas, &src, mas->end + 1, &left,
-						      &right, &ma_part)))
+		    (mas_wr_try_rebalance(mas, &src_info, mas->end + 1, &left,
+					  &right, &ma_part, src_info.insert_off + 1)))
 			goto rebalanced;
 
-		left.min = src.min;
-		right.max = src.max;
+		left.min = src_info.min;
+		right.max = src_info.max;
 		split = (total + 1) / 2;
 
 		if (split >= mas->offset) {
+			unsigned char space; /* size remaining in left */
+
 			if (mas->offset) {
-				state[i]->info = &src;
-				msn_mni_init(state[i++], &left, 0, mas->offset);
+				state[i].info = &src_info;
+				mns_mni_init(&state[i++], &left, 0, mas->offset);
 			}
-		}
 
+			state[i].part = &ma_part;
+			state[i].use_part = true;
+			space = split - mas->offset + 1;
+			if (space >= ma_part.size) {
+				mns_mni_init(&state[i++], &left, 0, ma_part.size);
+			} else {
+				mns_mni_init(&state[i++], &left, 0, space);
+				state[i].part = &ma_part;
+				state[i].use_part = true;
+				mns_mni_init(&state[i++], &right, space - 1, ma_part.size);
+
+			}
 
-		/* Old
-		if (split >= mas->offset)
-			mni_in_left(&src, &left, &right, mas, split, total, &ma_part);
-		else
-			mni_in_right(&src, &left, &right, mas, split, total, &ma_part);
-			*/
+			if (split > wr_mas->offset_end) {
+				state[i].info = &src_info;
+				mns_mni_init(&state[i++], &left, wr_mas->offset_end + 1,
+					     split - wr_mas->offset_end + 1);
+			}
+
+			if (mns_ends_in_null(&state[i - 1])) {
+				state[i - 1].size--;
+				split--;
+			}
+
+			state[i].info = &src_info;
+			mns_mni_init(&state[i++], &right, split + 1, mas->end - split + 1);
+		} else {
+			state[i].info = &src_info;
+			mns_mni_init(&state[i], &left, 0, split);
+			if (mns_ends_in_null(&state[i])) {
+				state[i].size--;
+				split--;
+			}
+			i++;
+
+			if (mas->offset > split + 1) {
+				state[i].info = &src_info;
+				mns_mni_init(&state[i++], &right, split + 1,
+					     mas->offset - split + 1);
+			}
+			state[i].part = &ma_part;
+			mns_mni_init(&state[i], &right, 0, ma_part.size);
+			state[i++].use_part = true;
+			if (wr_mas->offset_end > mas->end) {
+				state[i].info = &src_info;
+				mns_mni_init(&state[i++], &right, wr_mas->offset_end + 1,
+					     mas->end - wr_mas->offset_end + 1);
+			}
+		}
 
-		mns_assemble(&state, i);
+		mns_assemble(state, i);
 		mni_finalise(&left);
 		mni_finalise(&right);
 		mni_node_part_init(&ma_part, &left, &right);
@@ -4300,16 +4495,26 @@ static void mas_wr_split(struct ma_wr_state *wr_mas)
 new_root:
 	/* Converged on new root */
 	mas->depth++;
-	mas->offset = 0;
+	src_info.insert_off = mas->offset = 0;
 	mas->end = 0;
 	mas_set_height(mas);
+	if (mt_is_alloc(mas->tree))
+		src_info.type = maple_arange_64;
+	else
+		src_info.type = maple_range_64;
+	src_info.end = 0; /* Empty node */
+
+	/*
+	 * src can only really provide the parent
+	 * set the skip to high enough to avoid using any data
+	 */
 converged:
-	mas_wr_converged(&src, &parent, &ma_part, mas);
+	mas_wr_converged(&src_info, &parent, &ma_part, mas);
 	mas->node = parent.enode;
 rebalanced:
-	mas_wmb_replace(mas, src.enode);
+	mas_wmb_replace(mas, src_info.enode);
 	mtree_range_walk(mas);
-	//mt_dump(wr_mas->mas->tree, mt_dump_hex);
+	mt_dump(wr_mas->mas->tree, mt_dump_hex);
 }
 
 /*
@@ -5005,7 +5210,8 @@ static inline void mas_wr_store_entry(struct ma_wr_state *wr_mas)
 		mas_wr_split(wr_mas);
 		break;
 	case wr_rebalance:
-		mas_wr_rebalance(wr_mas);
+		mas_wr_bnode(wr_mas);
+		//mas_wr_rebalance(wr_mas);
 		break;
 	case wr_exact_fit:
 		rcu_assign_pointer(wr_mas->slots[mas->offset], wr_mas->entry);
@@ -6368,6 +6574,7 @@ int mas_store_gfp(struct ma_state *mas, void *entry, gfp_t gfp)
 	MA_WR_STATE(wr_mas, mas, entry);
 	int ret = 0;
 
+	printk("store %lu-%lu\n", index, last);
 retry:
 	mas_wr_preallocate(&wr_mas, entry);
 	if (unlikely(mas_nomem(mas, gfp))) {
-- 
2.50.1