From: Liam R. Howlett <Liam.Howlett@oracle.com>
Date: Wed, 16 Apr 2025 15:39:06 +0000 (-0400)
Subject: Introduce helpers for can_rebalance directions
X-Git-Url: https://www.infradead.org/git/?a=commitdiff_plain;h=5d463b097adec9e47afb1a6afd405b0169391d12;p=users%2Fjedix%2Flinux-maple.git

Introduce helpers for can_rebalance directions

Signed-off-by: Liam R. Howlett <Liam.Howlett@oracle.com>
---

diff --git a/lib/maple_tree.c b/lib/maple_tree.c
index 88467f48b68e..158b6a676787 100644
--- a/lib/maple_tree.c
+++ b/lib/maple_tree.c
@@ -3947,6 +3947,48 @@ static void mt_wr_split_data(struct ma_node_info *src,
 	} while (node_off <= src->end);
 }
 
+
+static inline
+bool can_rebalance_right(struct ma_state *tmp_mas, struct ma_node_info *parent,
+		struct ma_node_info *src2, struct split_data *sd)
+{
+	if (parent->insert_off >= parent->end) /* No right node */
+		return false;
+
+	tmp_mas->offset++;
+	mas_descend(tmp_mas);
+	mni_mas_init(src2, tmp_mas);
+	mni_set_end(src2);
+	sd->split = mas_wr_rebalance_calc(src2->end + sd->new_end, src2->type);
+	if (!sd->split) {
+		mas_ascend(tmp_mas);
+		tmp_mas->offset--;
+		return false;
+	}
+	return true;
+}
+
+static inline
+bool can_rebalance_left(struct ma_state *tmp_mas, struct ma_node_info *parent,
+		struct ma_node_info *src2, struct split_data *sd)
+{
+	if (!parent->insert_off) /* No left node */
+		return false;
+
+	tmp_mas->offset--;
+	mas_descend(tmp_mas);
+	mni_mas_init(src2, tmp_mas);
+	mni_set_end(src2);
+	sd->split = mas_wr_rebalance_calc(src2->end + sd->new_end, src2->type);
+	if (!sd->split) {
+		mas_ascend(tmp_mas);
+		tmp_mas->offset++;
+		return false;
+	}
+
+	parent->insert_off--;
+	return true;
+}
 /*
  * mas_wr_try_rebalance() - Try to rebalance two nodes, this may not work out.
  * @src: The source node state
@@ -3976,39 +4018,17 @@ bool mas_wr_try_rebalance(struct ma_state *mas, struct ma_node_info *src,
 	mas_wr_ascend_init(&tmp_mas, &parent);
 
 	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);
-	mni_mas_init(&src2, &tmp_mas);
-	sd->split = mas_wr_rebalance_calc(src2.end + sd->new_end, src2.type);
-
-	if (sd->split) {
-		parent.insert_off--;
-		/* Left will be src2, right will be src */
+	if (can_rebalance_left(&tmp_mas, &parent, &src2, sd)) {
 		left->min = src2.min;
 		right->max = src->max;
-	} else {
-		if (parent.end <= parent.insert_off)
-			return false;
-
-		mas_ascend(&tmp_mas);
-try_right:
-		tmp_mas.offset = parent.insert_off + 1;
-		mas_descend(&tmp_mas);
-		mni_mas_init(&src2, &tmp_mas);
-		mni_set_end(&src2);
-		sd->split = mas_wr_rebalance_calc(src2.end + sd->new_end, src2.type);
-		if (!sd->split)
-			return false;
-
+	} else if (can_rebalance_right(&tmp_mas, &parent, &src2, sd)) {
 		sd->split = src2.end + sd->new_end - sd->split;
 		left_store = true;
 		/* Left will be src, right will be src2 */
 		left->min = src->min;
 		right->max = src2.max;
+	} else {
+		return false;
 	}
 
 	/* The rebalance operation will succeed. */
@@ -4054,7 +4074,7 @@ try_right:
 	return true;
 }
 
-#if 1
+#if 0
 static void mas_wr_rebalance(struct ma_wr_state *wr_mas)
 {
 	struct ma_state *mas;
@@ -4080,6 +4100,24 @@ static void mas_wr_rebalance(struct ma_wr_state *wr_mas)
 	height = mas_mt_height(mas);
 	mas->depth = height;
 
+	mni_mas_init(&src, mas);
+	src.end = mas->end;
+
+	mni_node_init(&left, mas_pop_node(mas), wr_mas->type);
+	if (mt_is_alloc(mas->tree))
+		left.is_alloc = true;
+
+	tmp_mas = *mas;
+	mas_wr_ascend_init(&tmp_mas, &parent);
+	if (mas_wr_try_right(&tmp_mas, &parent, &src2)) {
+		sd.insert = mas->offset;
+	} else {
+		/* Won't fail here */
+		MAS_BUG_ON(try_rebalance_left(&tmp_mas, &parent, &src2));
+		sd.insert = mas->offset + src2.end + 1;
+	}
+
+
 	/*
 	 * allocate left
 	 * set alloc in left
@@ -4113,8 +4151,6 @@ static void mas_wr_rebalance(struct ma_wr_state *wr_mas)
 	 *
 	 */
 
-	mni_mas_init(&src, mas);
-	tmp_mas = *mas;
 	while (height--) {
 		printk("height %d\n", height);
 		mas_wr_ascend_init(&tmp_mas, &parent);
@@ -5198,8 +5234,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_bnode(wr_mas);
-		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);