[PATCH 1/3] maple_tree: introduce mas_prealloc_calc()

[Sidhartha Kumar]

Abstract the calculation used to determine the number of nodes needed for
a store operation into a separate function: mas_prealloc_calc().

Signed-off-by: Sidhartha Kumar <sidhartha.kumar@xxxxxxxxxx>
---
 lib/maple_tree.c | 85 ++++++++++++++++++++++++++++--------------------
 1 file changed, 50 insertions(+), 35 deletions(-)

diff --git a/lib/maple_tree.c b/lib/maple_tree.c
index 0e00a84e8e8f..e239197a57fc 100644
--- a/lib/maple_tree.c
+++ b/lib/maple_tree.c
@@ -5418,6 +5418,54 @@ void *mas_store(struct ma_state *mas, void *entry)
 }
 EXPORT_SYMBOL_GPL(mas_store);
 
+/**
+ * mas_prealloc_calc() - Calculate number of nodes needed for a
+ * store operation.
+ * @wr_mas: The maple write state
+ *
+ * Return: Number of nodes required for preallocation.
+ */
+int mas_prealloc_calc(struct ma_wr_state *wr_mas)
+{
+	struct ma_state *mas = wr_mas->mas;
+	unsigned char node_size;
+
+	if (unlikely(!mas->index && mas->last == ULONG_MAX))
+		return 1;
+
+	/* Root expand */
+	if (unlikely(mas_is_none(mas) || mas_is_ptr(mas)))
+		return 1;
+
+	if (unlikely(!mas_wr_walk(wr_mas))) {
+		/* Spanning store, use worst case for now */
+		return 1 + mas_mt_height(mas) * 3;
+	}
+
+	/* At this point, we are at the leaf node that needs to be altered. */
+	/* Exact fit, no nodes needed. */
+	if (wr_mas->r_min == mas->index && wr_mas->r_max == mas->last)
+		return 0;
+
+	mas_wr_end_piv(wr_mas);
+	node_size = mas_wr_new_end(wr_mas);
+	if (node_size >= mt_slots[wr_mas->type]) {
+		/* Split, worst case for now. */
+		return 1 + mas_mt_height(mas) * 2;
+	}
+
+	/* New root needs a singe node */
+	if (unlikely(mte_is_root(mas->node)))
+		return 1;
+
+	/* Potential spanning rebalance collapsing a node, use worst-case */
+	if (node_size  - 1 <= mt_min_slots[wr_mas->type])
+		return mas_mt_height(mas) * 2 - 1;
+
+	/* node store, slot store needs one node */
+	return 1;
+}
+
 /**
  * mas_store_gfp() - Store a value into the tree.
  * @mas: The maple state
@@ -5474,49 +5522,16 @@ EXPORT_SYMBOL_GPL(mas_store_prealloc);
 int mas_preallocate(struct ma_state *mas, void *entry, gfp_t gfp)
 {
 	MA_WR_STATE(wr_mas, mas, entry);
-	unsigned char node_size;
 	int request = 1;
 	int ret;
 
-
-	if (unlikely(!mas->index && mas->last == ULONG_MAX))
-		goto ask_now;
-
 	mas_wr_store_setup(&wr_mas);
 	wr_mas.content = mas_start(mas);
-	/* Root expand */
-	if (unlikely(mas_is_none(mas) || mas_is_ptr(mas)))
-		goto ask_now;
 
-	if (unlikely(!mas_wr_walk(&wr_mas))) {
-		/* Spanning store, use worst case for now */
-		request = 1 + mas_mt_height(mas) * 3;
-		goto ask_now;
-	}
-
-	/* At this point, we are at the leaf node that needs to be altered. */
-	/* Exact fit, no nodes needed. */
-	if (wr_mas.r_min == mas->index && wr_mas.r_max == mas->last)
+	request = mas_prealloc_calc(&wr_mas);
+	if (!request)
 		return 0;
 
-	mas_wr_end_piv(&wr_mas);
-	node_size = mas_wr_new_end(&wr_mas);
-	if (node_size >= mt_slots[wr_mas.type]) {
-		/* Split, worst case for now. */
-		request = 1 + mas_mt_height(mas) * 2;
-		goto ask_now;
-	}
-
-	/* New root needs a singe node */
-	if (unlikely(mte_is_root(mas->node)))
-		goto ask_now;
-
-	/* Potential spanning rebalance collapsing a node, use worst-case */
-	if (node_size  - 1 <= mt_min_slots[wr_mas.type])
-		request = mas_mt_height(mas) * 2 - 1;
-
-	/* node store, slot store needs one node */
-ask_now:
 	mas_node_count_gfp(mas, request, gfp);
 	mas->mas_flags |= MA_STATE_PREALLOC;
 	if (likely(!mas_is_err(mas)))
-- 
2.41.0