From a33fb50678c196aa37088a1a37f85b4dd69548b3 Mon Sep 17 00:00:00 2001
From: "Liam R. Howlett" <Liam.Howlett@Oracle.com>
Date: Mon, 4 Nov 2019 16:00:27 -0500
Subject: [PATCH] maple_tree: Add validation functions and rework gap
 calculations

Add validation functions to verify gaps and limits.  This exposed places
where limits were not set correctly, so fix them.  Fixing the limits
exposed some other issues which resulted in rewriting the mas_next logic
and mas_find to better detect dead nodes.

Signed-off-by: Liam R. Howlett <Liam.Howlett@Oracle.com>
---
 lib/maple_tree.c | 359 ++++++++++++++++++++++++++++++-----------------
 1 file changed, 234 insertions(+), 125 deletions(-)

diff --git a/lib/maple_tree.c b/lib/maple_tree.c
index 6b01226b90f0..c85bbb6bb4a1 100644
--- a/lib/maple_tree.c
+++ b/lib/maple_tree.c
@@ -117,12 +117,15 @@ static inline enum maple_type mte_node_type(const struct maple_enode *entry)
 	return ((unsigned long)entry >> 3) & 15;
 }
 
-
 static inline bool ma_is_dense(const enum maple_type type)
 {
 	return type < maple_sparse_6;
 }
 
+static inline bool mte_is_dense(const struct maple_enode *entry)
+{
+	return ma_is_dense(mte_node_type(entry));
+}
 static inline bool ma_is_leaf(const enum maple_type type)
 {
 	return type < maple_range_16;
@@ -1256,9 +1259,9 @@ static inline unsigned long mas_leaf_max_gap(struct ma_state *mas)
 	unsigned long gap = 0;
 	unsigned long pstart, pend;
 
-	if (mt == maple_dense) {
+	if (ma_is_dense(mt)) {
 		for (int i = 0; i < mt_slot_count(mas->node); i++) {
-			if (mte_get_rcu_slot(mas->node, i)) {
+			if (!mt_is_empty(mte_get_rcu_slot(mas->node, i))) {
 				if (gap > max_gap)
 					max_gap = gap;
 				gap = 0;
@@ -1274,22 +1277,11 @@ static inline unsigned long mas_leaf_max_gap(struct ma_state *mas)
 	pstart = mas->min;
 	for (int i = 0; i < mt_slots[mt]; i++) {
 		pend = mas_get_safe_pivot(mas, i);
+		gap = pend - pstart + 1;
 
-		if (i && !pend)
-			break;
-
-		if (pend) {
-			gap = pend - pstart;
-		} else {
-			gap = mas->max - pstart;
-		}
-
-		if (mte_get_rcu_slot(mas->node, i))
+		if (!mt_is_empty(mte_get_rcu_slot(mas->node, i)))
 			goto next;
 
-		if (!gap)
-			gap++;
-
 		if (gap > max_gap)
 			max_gap = gap;
 
@@ -1370,30 +1362,39 @@ static inline void mas_update_gap(struct ma_state *mas)
  * mas->max if no node is found.  Node is returned as mas->node which may be
  * MAS_NONE.
  *
+ * Note, the slot for the first node is not valid if it is a leaf.
+ *
  * @mas: maple state
  * @max: The maximum index to consider valid.
  */
 static inline unsigned long mas_first_node(struct ma_state *mas,
-		unsigned long max)
+		unsigned long limit)
 {
-	unsigned char slot = mas_get_slot(mas) - 1;
+	int slot = mas_get_slot(mas) - 1;
 	unsigned char count = mt_slot_count(mas->node);
+	unsigned long min = mas->min;
 
 	while (++slot < count) {
 		struct maple_enode *mn;
 		unsigned long pivot;
 
 		pivot = mas_get_safe_pivot(mas, slot);
-		if (pivot > max)
+		if (pivot > limit)
 			goto no_entry;
 
 		mn = mte_get_rcu_slot(mas->node, slot);
 
-		if (mt_is_empty(mn))
+		if (mt_is_empty(mn)) {
+			min = pivot + 1;
 			continue;
+		}
 
-		if (!mte_is_leaf(mas->node))
+
+		if (!mte_is_leaf(mas->node)) {
+			mas->max = pivot;
+			mas->min = min;
 			mas->node = mn;
+		}
 		mas_set_slot(mas, slot);
 		return pivot;
 	}
@@ -1417,8 +1418,12 @@ static inline unsigned long mas_first_entry(struct ma_state *mas,
 		if (mas->node == MAS_NONE)
 			return pivot;
 
-		if (mte_is_leaf(mas->node))
+		if (mte_is_leaf(mas->node)) {
+			// Get the leaf slot.
+			mas_set_slot(mas, 0);
+			pivot = mas_first_node(mas, max);
 			return pivot;
+		}
 
 		mas_set_slot(mas, 0);
 	}
@@ -2338,7 +2343,7 @@ no_entry:
  * mas_prev_node() - Find the prev non-null entry at the same level in the
  * tree.  The prev value will be mas->node[mas_get_slot(mas)] or MAS_NONE.
  */
-static inline void mas_prev_node(struct ma_state *mas, unsigned long min)
+static inline void mas_prev_node(struct ma_state *mas, unsigned long limit)
 {
 	int level;
 	unsigned char slot;
@@ -2353,6 +2358,7 @@ restart_prev_node:
 		goto no_entry;
 
 	while (1) {
+		unsigned long min;
 		mas_encoded_parent(mas);
 		level++;
 
@@ -2372,7 +2378,12 @@ restart_prev_node:
 			unsigned long pivot = mas_get_safe_pivot(mas, slot);
 			unsigned char coalesce;
 
-			if (pivot < min)
+			if (slot)
+				min = mas_get_safe_pivot(mas, slot - 1);
+			else
+				min = mas->min;
+
+			if (pivot < limit)
 				goto no_entry;
 
 			if (slot != 0 && pivot == 0)
@@ -2385,6 +2396,8 @@ restart_prev_node:
 			if (level == 1) {
 				mas_set_slot(mas, slot);
 				mas->node = mn;
+				mas->max = pivot;
+				mas->min = min;
 				if (mas_dead_node(mas, start_piv))
 					goto restart_prev_node;
 				return;
@@ -2553,7 +2566,6 @@ static inline bool mas_next_nentry(struct ma_state *mas, unsigned long max,
 			goto no_entry;
 
 		*piv = pivot;
-
 		entry = mte_get_rcu_slot(mas->node, slot);
 		if (!mt_is_empty(entry))
 			goto found;
@@ -2574,49 +2586,50 @@ found:
 }
 
 
-static inline unsigned long mas_next_entry(struct ma_state *mas,
-		unsigned long max)
+static inline void *_mas_next(struct ma_state *mas, unsigned long max,
+		unsigned long *pivot)
 {
-	unsigned long pivot = max;
+	void *entry = NULL;
+	unsigned long index = mas->index;
 
-	if (mas->node == MAS_NONE)
-		return max;
+	do {
+		if (mas->node == MAS_NONE)
+			return NULL;
 
-	if (!mte_is_leaf(mas->node))
-		pivot = mas_first_entry(mas, max);
-	else {
-		while (mas->node != MAS_NONE) {
-			unsigned char p_slot;
+		if (!mte_is_leaf(mas->node)) {
+			*pivot = mas_first_entry(mas, max);
+		} else {
+			while (mas->node != MAS_NONE) {
+				unsigned char p_slot;
 
-			if (mas_next_nentry(mas, max, &pivot))
-				break;
+				if (mas_next_nentry(mas, max, pivot))
+					break;
 
-			if (mte_is_root(mas->node)) {
-				mas->node = MAS_NONE;
-				break;
-			}
+				if (mte_is_root(mas->node)) {
+					mas->node = MAS_NONE;
+					break;
+				}
 
-			p_slot = mte_parent_slot(mas->node);
-			mas_set_slot(mas, p_slot);
-			mas_next_node(mas, max);
-			if (mas->node != MAS_NONE)
-				mas_set_slot(mas, 0);
+				p_slot = mte_parent_slot(mas->node);
+				mas_set_slot(mas, p_slot);
+				mas_next_node(mas, max);
+				if (mas->node != MAS_NONE) {
+					mas_set_slot(mas, 0);
+				}
+			}
 		}
-	}
-	return pivot;
+		if (mas_is_none(mas))
+			return NULL;
+		entry = mte_get_rcu_slot(mas->node, mas_get_slot(mas));
+	} while (mas_dead_node(mas, index));
+
+	return entry;
 }
-static inline unsigned long mas_safe_next_entry(struct ma_state *mas,
-		unsigned long max)
+static inline void *mas_next(struct ma_state *mas, unsigned long max)
 {
 	unsigned long piv;
 
-retry:
-	piv = mas_next_entry(mas, max);
-	if (mas->node == MAS_NONE)
-		return piv;
-	if (mas_dead_node(mas, mas->index))
-		goto retry;
-	return piv;
+	return _mas_next(mas, max, &piv);
 }
 
 static inline int mas_coalesce_node(struct ma_state *mas, unsigned char end,
@@ -2882,12 +2895,15 @@ static inline void mas_coalesce_root(struct ma_state *mas)
 	struct maple_enode *this_enode;
 	enum maple_type this_type;
 	unsigned long piv;
+	unsigned long min, max;
 
 	this_enode = mas->node;
 	this_type = mte_node_type(this_enode);
 	end = mas_data_end(mas, this_type, &piv, &coalesce);
 	hard_data = ma_hard_data(end, coalesce);
 	if (end == 1 && coalesce == 1) {
+		min = mas->min;
+		max = mas->max;
 		mas_first_node(mas, ULONG_MAX);
 		if (mas_is_none(mas)) {
 			// Allocation failures
@@ -2897,6 +2913,8 @@ static inline void mas_coalesce_root(struct ma_state *mas)
 			return;
 		}
 		mas->node = this_enode;
+		mas->min = min;
+		mas->max = max;
 	}
 
 	if ((end <= coalesce)) {
@@ -3499,11 +3517,9 @@ static inline int mas_dead_node(struct ma_state *mas, unsigned long index)
 }
 
 /**
- * mas_find - Find the first entry that overlaps the range.
- *   If mas->node is root, finds mas->index and returns the value.  If the
- *   value is null, returns the next non-null range.
- *   If mas->node is not root, finds the next range after the range that
- *   contains mas->index.
+ * mas_find: If mas->node == MAS_START, find the first
+ * non-NULL entry >= mas->index.
+ * Otherwise, find the first non-NULL entry > mas->index
  *
  * If an entry exists, last and index are updated accordingly.
  *
@@ -3512,64 +3528,46 @@ static inline int mas_dead_node(struct ma_state *mas, unsigned long index)
 void *mas_find(struct ma_state *mas, unsigned long max)
 {
 	void *entry = NULL;
-	unsigned long last_piv;
+	unsigned long range_end;
 	unsigned char slot = 0;
 	unsigned char start_index = mas->index;
+	bool mas_start = false;
 
 	if (!mas_searchable(mas))
 		goto not_found;
 
-	do {
-		if (mas_is_start(mas)) {
-			_mas_walk(mas);
-			if (mas_is_ptr(mas))
-				entry = mas->tree->ma_root;
-
-			if (mas_is_none(mas))
-				goto not_found;
-
-			slot = mas_get_slot(mas);
-
-			if (slot != MAPLE_NODE_SLOTS)
-				entry = mte_get_rcu_slot(mas->node, slot);
 
-			if (mt_is_empty(entry)) {
-				last_piv = mas_safe_next_entry(mas, max);
-				if (mas_is_none(mas))
-					goto not_found;
-
-				if (!mte_is_leaf(mas->node)) {
-					mas_set_slot(mas, 0);
-					last_piv = mas_first_entry(mas, max);
-				}
-				entry = mte_get_rcu_slot(mas->node, slot);
-			} else {
-				last_piv = mas_get_safe_pivot(mas, slot);
-			}
-
-			mas->index = mas_get_prev_pivot(mas, slot) + 1;
-			if (last_piv)
-				mas->index--;
-		} else {
-
-			last_piv = mas_safe_next_entry(mas, max);
-			if (mas_is_none(mas))
-				goto not_found;
+	if (mas_is_start(mas)) {
+		mas_start = true;
+	}
+	_mas_walk(mas);
+retry:	
+	if (mas_start) {
+		if (mas_is_ptr(mas) && !mas->index)
+			entry = mas->tree->ma_root;
 
+		if (mas_is_none(mas))
+			goto not_found;
 
-			slot = mas_get_slot(mas);
+		slot = mas_get_slot(mas);
+		if (slot != MAPLE_NODE_SLOTS) {
 			entry = mte_get_rcu_slot(mas->node, slot);
+			range_end = mas_get_safe_pivot(mas, slot);
+		}
+	}
 
-			if (!entry)
-				goto not_found;
+	if (mt_is_empty(entry)) {
+		entry = _mas_next(mas, max, &range_end);
+	}
 
+	if (mas_dead_node(mas, start_index))
+		goto retry;
 
-			mas->index = mas_get_prev_pivot(mas, slot) + 1;
-		}
-	} while (mas_dead_node(mas, start_index));
+	if (entry) {
+		mas->index = range_end + 1;
+		mas->last = mas->index;
+	}
 
-	mas->last = last_piv;
-	mas_set_slot(mas, slot + 1);
 	return entry;
 
 not_found:
@@ -3611,9 +3609,10 @@ void *mt_find(struct maple_tree *mt, unsigned long start,
 	void *entry = NULL;
 
 	rcu_read_lock();
-	_mas_walk(&mas);
-
 	do {
+		mas.index = start;
+		_mas_walk(&mas);
+
 		if (mas_is_none(&mas))
 			goto done;
 
@@ -3627,27 +3626,15 @@ void *mt_find(struct maple_tree *mt, unsigned long start,
 		if (slot != MAPLE_NODE_SLOTS)
 			entry = mte_get_rcu_slot(mas.node, slot);
 
-retry:
 		if (xa_is_zero(entry))
 			entry = NULL;
 
-		if (mt_is_advanced(entry)) {
-			mas_safe_next_entry(&mas, max);
-			if (mas_is_none(&mas)) {
-				entry = NULL;
-			} else if (!mte_is_leaf(mas.node)) {
-				mas_set_slot(&mas, 0);
-				mas_first_entry(&mas, max);
-			}
+		if (mt_is_empty(entry))
+			entry = mas_next(&mas, max);
+	} while (mas_dead_node(&mas, start));
 
-			slot = mas_get_slot(&mas);
-			entry = mte_get_rcu_slot(mas.node, slot);
-			goto retry;
-		}
-	} while (mas_dead_node(&mas, mas.index));
 done:
 	rcu_read_unlock();
-
 	return entry;
 }
 EXPORT_SYMBOL(mt_find);
@@ -3674,7 +3661,7 @@ void *mt_find_after(struct maple_tree *mt, unsigned long *index,
 
 retry:
 	mas_set_slot(&mas, mas_get_slot(&mas) + 1);
-	mas_safe_next_entry(&mas, max);
+	mas_next(&mas, max);
 	if (mas_is_none(&mas))
 		goto done;
 
@@ -4614,4 +4601,126 @@ unsigned long mt_get_alloc_size(void)
 	return kmem_cache_get_alloc(maple_node_cache);
 }
 
-#endif
+// Tree validations
+
+/** Private
+ *
+ * Calculate the maximum gap in a leaf and check if that's what is reported in
+ * the parent.
+ */
+void mas_validate_gaps(struct ma_state *mas)
+{
+	struct maple_enode *mte = mas->node;
+	struct maple_node *p_mn;
+	unsigned long gap = 0, max_gap = 0;
+	unsigned long p_end, p_start = mas->min;
+	unsigned char p_slot;
+
+	if (mte_is_dense(mte)) {
+		for (int i = 0; i < mt_slot_count(mte); i++) {
+			if (!mt_is_empty(mte_get_rcu_slot(mas->node, i))) {
+				if (gap > max_gap)
+					max_gap = gap;
+				gap = 0;
+				continue;
+			}
+			gap++;
+		}
+		goto counted;
+	}
+
+	for (int i = 0; i < mt_slot_count(mte); i++) {
+		p_end = mas_get_safe_pivot(mas, i);
+		if (!p_end && i != 0)
+			break;
+
+		if (mte_is_leaf(mte)) {
+			gap = p_end - p_start + 1;
+			if (!mt_is_empty(mte_get_rcu_slot(mas->node, i)))
+				goto not_empty;
+		} else {
+			gap = mte_get_gap(mte, i);
+			if (mt_is_empty(mte_get_rcu_slot(mas->node, i)))
+				BUG_ON(gap != p_end - p_start + 1);
+			else
+				BUG_ON(gap >= p_end - p_start + 1);
+		}
+
+		if (gap > max_gap)
+			max_gap = gap;
+not_empty:
+		p_start = p_end + 1;
+		if (p_end >= mas->max)
+			break;
+	}
+
+counted:
+	if (mte_is_root(mte))
+		return;
+
+	p_slot = mte_parent_slot(mas->node);
+	p_mn = mte_parent(mte);
+	BUG_ON(_ma_get_gap(p_mn, p_slot, mas_parent_enum(mas, mte))
+			!= max_gap);
+}
+
+void mas_validate_limits(struct ma_state *mas) {
+
+	if (mte_is_root(mas->node))
+		return; // all limits are fine here.
+
+	for (int i = 0; i < mt_slot_count(mas->node); i++) {
+		unsigned long piv = mas_get_safe_pivot(mas, i);
+		if (!piv)
+			break;
+		BUG_ON(piv < mas->min);
+		BUG_ON(piv > mas->max);
+	}
+}
+
+/* Depth first search, post-order */
+static inline void mas_dfs_postorder(struct ma_state *mas, unsigned long max)
+{
+
+	struct maple_enode *p = MAS_NONE, *mn = mas->node;
+	unsigned long p_min, p_max;
+
+	mas_next_node(mas, max);
+	if (mas->node != MAS_NONE)
+		return;
+
+	if (mte_is_root(mn))
+		return;
+
+	mas->node = mn;
+	mas_encoded_parent(mas);
+	while (mas->node != MAS_NONE) {
+		p = mas->node;
+		p_min = mas->min;
+		p_max = mas->max;
+		mas_prev_node(mas, 0);
+	}
+
+	if (p == MAS_NONE)
+		return;
+
+	mas->node = p;
+	mas->max = p_max;
+	mas->min = p_min;
+}
+void mt_validate(struct maple_tree *mt)
+{
+	MA_STATE(mas, mt, 0, 0);
+	rcu_read_lock();
+	mas_start(&mas);
+	mas_first_entry(&mas, ULONG_MAX);
+	while (mas.node != MAS_NONE) {
+		mas_validate_limits(&mas);
+		if (mt_is_alloc(mt))
+			mas_validate_gaps(&mas);
+		mas_dfs_postorder(&mas, ULONG_MAX);
+	}
+	rcu_read_unlock();
+
+}
+#endif /* MT_DEBUG */
-- 
2.50.1