[PATCH] maple_tree: Fix use of node for global range in mas_wr_spanning_store()

[Liam Howlett <liam.howlett@xxxxxxxxxx>]

When writing a range with a NULL which expands to 0 - ULONG_MAX, don't
use a node to store this value.  Instead, call mas_new_root() which will
set the tree pointer to NULL and free all the nodes.

Fix a comment for the allocations in mas_wr_spanning_store().

Add mas_node_count_gfp() and use it to clean up mas_preallocate().

Clean up mas_preallocate() and ensure the ma_state is safe on return.

Update maple_tree.h to set alloc = NULL.

Fixes: d0aac5e48048 (Maple Tree: add new data structure)
Signed-off-by: Liam R. Howlett <Liam.Howlett@xxxxxxxxxx>
---
 include/linux/maple_tree.h |  1 +
 lib/maple_tree.c           | 34 +++++++++++++++++++++++++++-------
 2 files changed, 28 insertions(+), 7 deletions(-)

diff --git a/include/linux/maple_tree.h b/include/linux/maple_tree.h
index 3c6642358904..2c9dede989c7 100644
--- a/include/linux/maple_tree.h
+++ b/include/linux/maple_tree.h
@@ -434,6 +434,7 @@ struct ma_wr_state {
 		.node = MAS_START,					\
 		.min = 0,						\
 		.max = ULONG_MAX,					\
+		.alloc = NULL,						\
 	}
 
 #define MA_WR_STATE(name, ma_state, wr_entry)				\
diff --git a/lib/maple_tree.c b/lib/maple_tree.c
index 2cccd8f2153f..79e07c7dc323 100644
--- a/lib/maple_tree.c
+++ b/lib/maple_tree.c
@@ -1293,17 +1293,31 @@ static inline void mas_free(struct ma_state *mas, struct maple_enode *used)
  * there is not enough nodes.
  * @mas: The maple state
  * @count: The number of nodes needed
+ * @gfp: the gfp flags
  */
-static void mas_node_count(struct ma_state *mas, int count)
+static void mas_node_count_gfp(struct ma_state *mas, int count, gfp_t gfp)
 {
 	unsigned long allocated = mas_allocated(mas);
 
 	if (allocated < count) {
 		mas_set_alloc_req(mas, count - allocated);
-		mas_alloc_nodes(mas, GFP_NOWAIT | __GFP_NOWARN);
+		mas_alloc_nodes(mas, gfp);
 	}
 }
 
+/*
+ * mas_node_count() - Check if enough nodes are allocated and request more if
+ * there is not enough nodes.
+ * @mas: The maple state
+ * @count: The number of nodes needed
+ *
+ * Note: Uses GFP_NOWAIT | __GFP_NOWARN for gfp flags.
+ */
+static void mas_node_count(struct ma_state *mas, int count)
+{
+	return mas_node_count_gfp(mas, count, GFP_NOWAIT | __GFP_NOWARN);
+}
+
 /*
  * mas_start() - Sets up maple state for operations.
  * @mas: The maple state.
@@ -3962,7 +3976,7 @@ static inline int mas_wr_spanning_store(struct ma_wr_state *wr_mas)
 	if (unlikely(!mas->index && mas->last == ULONG_MAX))
 		return mas_new_root(mas, wr_mas->entry);
 	/*
-	 * Node rebalancing may occur due to this store, so there may be two new
+	 * Node rebalancing may occur due to this store, so there may be three new
 	 * entries per level plus a new root.
 	 */
 	height = mas_mt_height(mas);
@@ -3995,6 +4009,12 @@ static inline int mas_wr_spanning_store(struct ma_wr_state *wr_mas)
 		mas->last = l_mas.last = r_mas.last;
 	}
 
+	/* expanding NULLs may make this cover the entire range */
+	if (!l_mas.index && r_mas.last == ULONG_MAX) {
+		mas_set_range(mas, 0, ULONG_MAX);
+		return mas_new_root(mas, wr_mas->entry);
+	}
+
 	memset(&b_node, 0, sizeof(struct maple_big_node));
 	/* Copy l_mas and store the value in b_node. */
 	mas_store_b_node(&l_wr_mas, &b_node, l_wr_mas.node_end);
@@ -5657,15 +5677,15 @@ int mas_preallocate(struct ma_state *mas, void *entry, gfp_t gfp)
 {
 	int ret;
 
-	mas_set_alloc_req(mas, 1 + mas_mt_height(mas) * 3);
-	mas_alloc_nodes(mas, gfp);
+	mas_node_count_gfp(mas, 1 + mas_mt_height(mas) * 3, gfp);
 	if (likely(!mas_is_err(mas)))
 		return 0;
 
 	mas_set_alloc_req(mas, 0);
-	mas_destroy(mas);
 	ret = xa_err(mas->node);
-	mas->node = MAS_START;
+	mas_reset(mas);
+	mas_destroy(mas);
+	mas_reset(mas);
 	return ret;
 }
 
-- 
2.35.1