From e86f3d1a2de3e6aa5356145f1c1d63cdb38386bb Mon Sep 17 00:00:00 2001
From: "Liam R. Howlett" <Liam.Howlett@Oracle.com>
Date: Tue, 15 Sep 2020 12:03:11 -0400
Subject: [PATCH] maple_tree: Avoid case statements with ma_gaps and ma_pivots

Get the arrays and use them directly instead of switching on type every
time a value is needed.

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

diff --git a/lib/maple_tree.c b/lib/maple_tree.c
index 660629c06ea4..045b47c10bb8 100644
--- a/lib/maple_tree.c
+++ b/lib/maple_tree.c
@@ -477,13 +477,13 @@ static inline void mas_set_offset(struct ma_state *mas, int offset)
 }
 
 /*
- * ma_get_pivots() - Get the pivot of a node.
+ * ma_pivots() - Get the pivot of a node.
  * @node - the maple node.
  * @type - the node type.
  *
  * Returns: The value of the @piv in the @node.
  */
-static inline unsigned long *ma_get_pivots(struct maple_node *node,
+static inline unsigned long *ma_pivots(struct maple_node *node,
 					   enum maple_type type)
 {
 	switch (type) {
@@ -498,8 +498,8 @@ static inline unsigned long *ma_get_pivots(struct maple_node *node,
 	}
 }
 
-static inline unsigned long *ma_get_gaps(struct maple_node *node,
-					 enum maple_type type)
+static inline unsigned long *ma_gaps(struct maple_node *node,
+				     enum maple_type type)
 {
 	switch (type) {
 	case maple_arange_64:
@@ -512,12 +512,12 @@ static inline unsigned long *ma_get_gaps(struct maple_node *node,
 	}
 }
 
-static inline unsigned long _mte_get_pivot(const struct maple_enode *mn,
-				 unsigned char piv, enum maple_type type)
+static inline unsigned long mte_pivot(const struct maple_enode *mn,
+				 unsigned char piv)
 {
 	struct maple_node *node = mte_to_node(mn);
 
-	switch (type) {
+	switch (mte_node_type(mn)) {
 	case maple_arange_64:
 		return node->ma64.pivot[piv];
 	case maple_range_64:
@@ -529,19 +529,14 @@ static inline unsigned long _mte_get_pivot(const struct maple_enode *mn,
 	}
 }
 
-static inline unsigned long mte_get_pivot(const struct maple_enode *mn,
-					 unsigned char piv)
-{
-	return _mte_get_pivot(mn, piv, mte_node_type(mn));
-}
-
 static inline unsigned long _mas_safe_pivot(const struct ma_state *mas,
-				unsigned char piv, enum maple_type type)
+					    unsigned long *pivots,
+					    unsigned char piv, enum maple_type type)
 {
 	if (piv >= mt_pivots[type])
 		return mas->max;
 
-	return _mte_get_pivot(mas->node, piv, type);
+	return pivots[piv];
 }
 
 /*
@@ -555,17 +550,19 @@ static inline unsigned long mas_safe_pivot(const struct ma_state *mas,
 					 unsigned char piv)
 {
 	enum maple_type type = mte_node_type(mas->node);
+	unsigned long *pivots = ma_pivots(mas_mn(mas), type);
 
-	return _mas_safe_pivot(mas, piv, type);
+	return _mas_safe_pivot(mas, pivots, piv, type);
 }
 
 static inline unsigned long mas_safe_min(struct ma_state *mas,
+					 unsigned long *pivots,
 					 unsigned char piv)
 {
 	if (!piv)
 		return mas->min;
 
-	return mas_safe_pivot(mas, piv - 1) + 1;
+	return pivots[piv - 1] + 1;
 }
 
 static inline void ma_set_pivot(struct maple_node *mn, unsigned char piv,
@@ -592,10 +589,9 @@ static inline void mte_set_pivot(struct maple_enode *mn, unsigned char piv,
 	return ma_set_pivot(mte_to_node(mn), piv, mte_node_type(mn), val);
 }
 
-static inline void __rcu **ma_get_slots(struct maple_node *mn,
-					enum maple_type type)
+static inline void __rcu **ma_slots(struct maple_node *mn, enum maple_type mt)
 {
-	switch (type) {
+	switch (mt) {
 	default:
 	case maple_arange_64:
 		return mn->ma64.slot;
@@ -607,17 +603,22 @@ static inline void __rcu **ma_get_slots(struct maple_node *mn,
 	}
 }
 
-static inline struct maple_enode *mas_get_slot(struct ma_state *mas,
-		unsigned char offset)
+static inline void *ma_slot(struct ma_state *mas, void **slots,
+			    unsigned char offset)
 {
-	void **slots = ma_get_slots(mte_to_node(mas->node),
-				    mte_node_type(mas->node));
 	if (mt_in_rcu(mas->tree))
 		return rcu_dereference(slots[offset]);
 	else
 		return slots[offset];
 }
 
+static inline struct maple_enode *mas_get_slot(struct ma_state *mas,
+		unsigned char offset)
+{
+	return ma_slot(mas, ma_slots(mas_mn(mas), mte_node_type(mas->node)),
+		       offset);
+}
+
 /*
  * ma_set_slot() - Set a nodes rcu slot.
  *
@@ -733,23 +734,6 @@ static inline void mas_descend(struct ma_state *mas)
 	mas->node = mas_get_slot(mas, mas_offset(mas));
 }
 
-static inline unsigned long ma_get_gap(const struct maple_node *mn,
-				 unsigned char gap, enum maple_type type)
-{
-	switch (type) {
-	case maple_arange_64:
-		return mn->ma64.gap[gap];
-	default:
-		return 0;
-	}
-}
-
-static inline unsigned long mte_get_gap(const struct maple_enode *mn,
-				 unsigned char gap)
-{
-	return ma_get_gap(mte_to_node(mn), gap, mte_node_type(mn));
-}
-
 static inline void mte_set_gap(const struct maple_enode *mn,
 				 unsigned char gap, unsigned long val)
 {
@@ -792,12 +776,12 @@ ascend:
 	a_enode = mt_mk_node(a_node, a_type);
 	if (!set_min && a_slot) {
 		set_min = true;
-		min = mte_get_pivot(a_enode, a_slot - 1) + 1;
+		min = mte_pivot(a_enode, a_slot - 1) + 1;
 	}
 
 	if (!set_max && a_slot < mt_pivots[a_type]) {
 		set_max = true;
-		max = mte_get_pivot(a_enode, a_slot);
+		max = mte_pivot(a_enode, a_slot);
 	}
 
 no_parent:
@@ -1027,18 +1011,19 @@ done:
  */
 static inline unsigned char mas_data_end(const struct ma_state *mas)
 {
-	unsigned char slot = 0;
+	unsigned char offset = 0;
 	enum maple_type type = mte_node_type(mas->node);
+	unsigned long *pivots = ma_pivots(mas_mn(mas), type);
 	unsigned long piv = mas->min;
 
-	while (slot < mt_slots[type]) {
-		piv = _mas_safe_pivot(mas, slot, type);
+	while (offset < mt_slots[type]) {
+		piv = _mas_safe_pivot(mas, pivots, offset, type);
 		if (piv >= mas->max)
 			break;
-		slot++;
+		offset++;
 	}
 
-	return slot;
+	return offset;
 }
 
 /*
@@ -1054,7 +1039,7 @@ static inline unsigned long mas_leaf_max_gap(struct ma_state *mas)
 	unsigned long gap = 0;
 	void *entry = NULL;
 	int i;
-	void **slots = ma_get_slots(mte_to_node(mas->node), mt);
+	void **slots = ma_slots(mte_to_node(mas->node), mt);
 
 	if (ma_is_dense(mt)) {
 		for (i = 0; i < mt_slot_count(mas->node); i++) {
@@ -1102,13 +1087,15 @@ done:
  */
 static inline unsigned long mas_max_gap(struct ma_state *mas)
 {
+	enum maple_type mt = mte_node_type(mas->node);
+	unsigned long *gaps = ma_gaps(mas_mn(mas), mt);
 	unsigned long max_gap = 0;
 	unsigned char i;
 
-	for (i = 0; i < mt_slot_count(mas->node); i++) {
+	for (i = 0; i < mt_slots[mt]; i++) {
 		unsigned long gap;
 
-		gap = mte_get_gap(mas->node, i);
+		gap = gaps[i];
 		if (gap >  max_gap)
 			max_gap = gap;
 	}
@@ -1172,8 +1159,8 @@ static inline void mas_update_gap(struct ma_state *mas)
 
 	/* Get the gap reported in the parent */
 	pslot = mte_parent_slot(mas->node);
-	p_gap = ma_get_gap(mte_parent(mas->node), pslot,
-			mas_parent_enum(mas, mas->node));
+	p_gap = ma_gaps(mte_parent(mas->node),
+			mas_parent_enum(mas, mas->node))[pslot];
 
 	if (p_gap != max_gap)
 		mas_parent_gap(mas, pslot, max_gap);
@@ -1194,13 +1181,13 @@ static inline unsigned long mas_first_entry(struct ma_state *mas,
 	unsigned long pivot = mas->min;
 
 	while (!mte_is_leaf(mas->node)) {
-		mas->max = mte_get_pivot(mas->node, 0);
-		slots = ma_get_slots(mte_to_node(mas->node),
+		mas->max = mte_pivot(mas->node, 0);
+		slots = ma_slots(mte_to_node(mas->node),
 				     mte_node_type(mas->node));
 		mas->node = slots[0];
 	}
 
-	slots = ma_get_slots(mte_to_node(mas->node), mte_node_type(mas->node));
+	slots = ma_slots(mte_to_node(mas->node), mte_node_type(mas->node));
 
 	while ((pivot < limit) && (offset < mt_slot_count(mas->node))) {
 		pivot = mas_safe_pivot(mas, offset);
@@ -1228,7 +1215,7 @@ static inline void mas_adopt_children(struct ma_state *mas,
 {
 
 	enum maple_type type = mte_node_type(parent);
-	void **slots = ma_get_slots(mte_to_node(mas->node), type);
+	void **slots = ma_slots(mte_to_node(mas->node), type);
 	struct maple_enode *child;
 	unsigned char offset;
 
@@ -1263,7 +1250,7 @@ static inline void mas_replace(struct ma_state *mas, bool advanced)
 		parent = mte_parent(mas->node);
 		eparent = mt_mk_node(parent, ptype);
 		offset = mte_parent_slot(mas->node);
-		slots = ma_get_slots(parent, ptype);
+		slots = ma_slots(parent, ptype);
 		prev = slots[offset];
 	}
 
@@ -1297,7 +1284,7 @@ static inline bool mas_new_child(struct ma_state *mas, struct ma_state *child)
 	enum maple_type mt = mte_node_type(mas->node);
 	unsigned char offset;
 	struct maple_enode *entry;
-	void **slots = ma_get_slots(mte_to_node(mas->node), mt);
+	void **slots = ma_slots(mte_to_node(mas->node), mt);
 
 	for (offset= mas_offset(mas); offset < mt_slots[mt]; offset++) {
 		entry = slots[offset];
@@ -1432,18 +1419,23 @@ static inline void mas_mab_cp(struct ma_state *mas, unsigned char mas_start,
 			unsigned char mas_end, struct maple_big_node *b_node,
 			unsigned char mab_start)
 {
+	enum maple_type mt = mte_node_type(mas->node);
+	struct maple_node *node = mte_to_node(mas->node);
+	void **slots = ma_slots(node, mt);
+	unsigned long *gaps = NULL;
 	int i, j;
-	void **slots = ma_get_slots(mte_to_node(mas->node),
-				    mte_node_type(mas->node));
+
+	if (!mte_is_leaf(mas->node) && mt_is_alloc(mas->tree))
+		gaps = ma_gaps(node, mt);
 
 	for (i = mas_start, j = mab_start; i <= mas_end; i++, j++) {
 		b_node->slot[j] = rcu_dereference_protected(slots[i],
 					lockdep_is_held(&mas->tree->ma_lock));
 
-		if (!mte_is_leaf(mas->node) && mt_is_alloc(mas->tree))
-			b_node->gap[j] = mte_get_gap(mas->node, i);
+		if (gaps)
+			b_node->gap[j] = gaps[i];
 
-		if (i < mt_pivot_count(mas->node))
+		if (i < mt_pivots[mt])
 			b_node->pivot[j] = mas_safe_pivot(mas, i);
 		else
 			b_node->pivot[j] = mas->max;
@@ -1690,7 +1682,7 @@ static inline void mast_topiary(struct maple_subtree_state *mast)
 	l_off = mas_offset(mast->orig_l);
 	r_off = mas_offset(mast->orig_r);
 	if (mast->orig_l->node == mast->orig_r->node) {
-		slots = ma_get_slots(mte_to_node(mast->orig_l->node),
+		slots = ma_slots(mte_to_node(mast->orig_l->node),
 				     mte_node_type(mast->orig_l->node));
 		for (offset = l_off + 1; offset < r_off; offset++)
 			mat_add(mast->destroy, slots[offset]);
@@ -1702,7 +1694,7 @@ static inline void mast_topiary(struct maple_subtree_state *mast)
 
 	/* Now destroy l_off + 1 -> end and 0 -> r_off - 1 */
 	offset = l_off + 1;
-	slots = ma_get_slots(mte_to_node(mast->orig_l->node),
+	slots = ma_slots(mte_to_node(mast->orig_l->node),
 			     mte_node_type(mast->orig_l->node));
 	while (offset < mt_slot_count(mast->orig_l->node)) {
 		child = slots[offset++];
@@ -1711,7 +1703,7 @@ static inline void mast_topiary(struct maple_subtree_state *mast)
 		mat_add(mast->destroy, child);
 	}
 
-	slots = ma_get_slots(mte_to_node(mast->orig_r->node),
+	slots = ma_slots(mte_to_node(mast->orig_r->node),
 			     mte_node_type(mast->orig_r->node));
 	for (offset = 0; offset < r_off; offset++)
 		mat_add(mast->destroy, slots[offset]);
@@ -2737,12 +2729,14 @@ static inline bool mas_node_walk(struct ma_state *mas, enum maple_type type,
 {
 	unsigned char i;
 	unsigned long min = mas->min, pivot = 0;
+	unsigned long *pivots = ma_pivots(mas_mn(mas), type);
 	bool ret = true;
 
+
 	switch (type) {
 	default:
 		for (i = mas_offset(mas); i < mt_slots[type]; i++) {
-			pivot = _mas_safe_pivot(mas, i, type);
+			pivot = _mas_safe_pivot(mas, pivots, i, type);
 
 			if (!pivot && i) {
 				if (mas->max < mas->index) {
@@ -2792,12 +2786,13 @@ static inline void mas_cnt_empty(struct ma_state *mas)
 }
 
 /*
- *
- * mas_wr_walk(): Walk the tree for a write.  Tracks extra information which
- * is used in special cases of a write.
+ * mas_wr_walk(): Walk the tree for a write.
  * @range_min - pointer that will be set to the minimum of the slot range
  * @range_max - pointer that will be set to the maximum of the slot range
  * @entry - the value that will be written.
+ * Returns: True if found, false otherwise.
+ *
+ * Tracks extra information which is used in special cases of a write.
  */
 static inline bool mas_wr_walk(struct ma_state *mas, unsigned long *range_min,
 		unsigned long *range_max, void *entry)
@@ -2854,7 +2849,7 @@ static inline void mas_extend_null(struct ma_state *l_mas,
 	unsigned char cp_r_slot = r_slot;
 	unsigned long range_max = mas_safe_pivot(r_mas, r_slot);
 	unsigned long range_min = l_mas->min;
-	void **slots = ma_get_slots(mte_to_node(l_mas->node),
+	void **slots = ma_slots(mte_to_node(l_mas->node),
 				    mte_node_type(l_mas->node));
 	void *content = slots[l_slot];
 
@@ -2873,7 +2868,7 @@ static inline void mas_extend_null(struct ma_state *l_mas,
 		mas_set_offset(l_mas, l_slot - 1);
 	}
 
-	slots = ma_get_slots(mte_to_node(r_mas->node),
+	slots = ma_slots(mte_to_node(r_mas->node),
 			     mte_node_type(r_mas->node));
 	if (!slots[r_slot]) {
 		if (r_mas->last < range_max)
@@ -3161,7 +3156,7 @@ restart_prev_node:
 
 	offset--;
 	mt = mte_node_type(mas->node);
-	slots = ma_get_slots(mas_mn(mas), mt);
+	slots = ma_slots(mas_mn(mas), mt);
 	mas->max = mas_safe_pivot(mas, offset);
 	while (level > 1) {
 		level--;
@@ -3169,14 +3164,16 @@ restart_prev_node:
 		if (mas_dead_node(mas, start_piv))
 			goto restart_prev_node;
 		mt = mte_node_type(mas->node);
-		slots = ma_get_slots(mas_mn(mas), mt);
+		slots = ma_slots(mas_mn(mas), mt);
 		offset = mt_slots[mt];
 		do {} while(!mas_get_slot(mas, --offset));
 		mas->max = mas_safe_pivot(mas, offset);
 	}
 
 	mas_set_offset(mas, offset);
-	mas->min = mas_safe_min(mas, offset);
+	mas->min = mas_safe_min(mas, ma_pivots(mas_mn(mas),
+					       mte_node_type(mas->node)),
+				offset);
 	mas->node = rcu_dereference(slots[offset]);
 	if (mas_dead_node(mas, start_piv))
 		goto restart_prev_node;
@@ -3226,7 +3223,7 @@ restart_next_node:
 			goto restart_next_node;
 
 		mt = mte_node_type(mas->node);
-		slots = ma_get_slots(mas_mn(mas), mt);
+		slots = ma_slots(mas_mn(mas), mt);
 		prev_piv = mas_safe_pivot(mas, offset);
 		if (prev_piv > max)
 			goto no_entry;
@@ -3240,7 +3237,7 @@ restart_next_node:
 		level--;
 		mas->node = rcu_dereference(slots[offset]);
 		mt = mte_node_type(mas->node);
-		slots = ma_get_slots(mas_mn(mas), mt);
+		slots = ma_slots(mas_mn(mas), mt);
 		offset = 0;
 		pivot = mas_safe_pivot(mas, offset);
 	}
@@ -3295,17 +3292,20 @@ static inline bool mas_prev_nentry(struct ma_state *mas, unsigned long limit,
 static inline bool mas_next_nentry(struct ma_state *mas, unsigned long max,
 		unsigned long *range_start)
 {
+	enum maple_type type = mte_node_type(mas->node);
 	unsigned long pivot = mas->min;
 	unsigned long r_start = mas->min;
-	unsigned char slot = mas_offset(mas);
-	unsigned char count = mt_slot_count(mas->node);
+	unsigned char offset = mas_offset(mas);
+	unsigned char count = mt_slots[type];
+	unsigned long *pivots = ma_pivots(mas_mn(mas), type);
+	void **slots = ma_slots(mas_mn(mas), type);
 	void *entry;
 
-	r_start = mas_safe_min(mas, slot);
+	r_start = mas_safe_min(mas, pivots, offset);
 
-	while (slot < count) {
-		pivot = mas_safe_pivot(mas, slot);
-		if (!pivot && slot)
+	while (offset < count) {
+		pivot = _mas_safe_pivot(mas, pivots, offset, type);
+		if (!pivot && offset)
 			goto no_entry;
 
 		if (r_start > max)
@@ -3314,7 +3314,7 @@ static inline bool mas_next_nentry(struct ma_state *mas, unsigned long max,
 		if (r_start > mas->max)
 			goto no_entry;
 
-		entry = mas_get_slot(mas, slot);
+		entry = ma_slot(mas, slots, offset);
 		if (entry)
 			goto found;
 
@@ -3323,7 +3323,7 @@ static inline bool mas_next_nentry(struct ma_state *mas, unsigned long max,
 			goto no_entry;
 
 		r_start = pivot + 1;
-		slot++;
+		offset++;
 	}
 
 no_entry:
@@ -3333,7 +3333,7 @@ no_entry:
 found:
 	mas->last = pivot;
 	*range_start = r_start;
-	mas_set_offset(mas, slot);
+	mas_set_offset(mas, offset);
 	return true;
 }
 
@@ -3362,7 +3362,7 @@ static inline void *mas_last_entry(struct ma_state *mas,
 			entry = mas_get_slot(mas, slot - 1);
 			if (mte_is_leaf(mas->node)) {
 				mas->index = range_start - 1;
-				mas->index = mte_get_pivot(mas->node, slot - 1);
+				mas->index = mte_pivot(mas->node, slot - 1);
 				return entry;
 			}
 
@@ -3446,7 +3446,7 @@ next_node:
 static inline void *_mas_prev(struct ma_state *mas, unsigned long limit)
 {
 	unsigned long max = mas->max;
-	unsigned char slot;
+	unsigned char offset;
 
 	while (!mas_is_none(mas)) {
 		if (mas_prev_nentry(mas, limit, &max))
@@ -3462,8 +3462,10 @@ static inline void *_mas_prev(struct ma_state *mas, unsigned long limit)
 	}
 
 	mas->last = max;
-	slot = mas_offset(mas);
-	mas->index = mas_safe_min(mas, slot);
+	offset = mas_offset(mas);
+	mas->index = mas_safe_min(mas, ma_pivots(mas_mn(mas),
+						 mte_node_type(mas->node)),
+				  offset);
 	return mas_get_slot(mas, mas_offset(mas));
 }
 
@@ -3515,10 +3517,10 @@ bool _mas_rev_awalk(struct ma_state *mas, unsigned long size)
 		return true;
 	}
 
-	pivots = ma_get_pivots(node, type);
-	slots = ma_get_slots(node, type);
+	pivots = ma_pivots(node, type);
+	slots = ma_slots(node, type);
 	if (!ma_is_leaf(type))
-		gaps = ma_get_gaps(node, type);
+		gaps = ma_gaps(node, type);
 
 	if (offset == mt_pivots[type]) {
 		// Initial start, walk until the end of the data.
@@ -3531,10 +3533,10 @@ bool _mas_rev_awalk(struct ma_state *mas, unsigned long size)
 		}
 	}
 
-	min = _mas_safe_pivot(mas, offset, type) + 1;
+	min = _mas_safe_pivot(mas, pivots, offset, type) + 1;
 	do {
 		max = min - 1;
-		min = mas_safe_min(mas, offset);
+		min = mas_safe_min(mas, pivots, offset);
 		if (mas->last < min)
 			continue;
 
@@ -3545,8 +3547,7 @@ bool _mas_rev_awalk(struct ma_state *mas, unsigned long size)
 
 		if (!ma_is_leaf(type))
 			gap = gaps[offset];
-		else if (rcu_dereference_protected(slots[offset],
-				lockdep_is_held(&mas->tree->ma_lock)))
+		else if (rcu_dereference(slots[offset]))
 			continue; // no gap in leaf.
 		else
 			gap = max - min + 1;
@@ -3572,8 +3573,7 @@ bool _mas_rev_awalk(struct ma_state *mas, unsigned long size)
 	}
 
 	//descend
-	mas->node = rcu_dereference_protected(slots[offset],
-			      lockdep_is_held(&mas->tree->ma_lock));
+	mas->node = rcu_dereference(slots[offset]);
 	mas->min = min;
 	mas->max = max;
 	mas_set_offset(mas, mt_pivot_count(mas->node));
@@ -3591,64 +3591,63 @@ static inline bool _mas_awalk(struct ma_state *mas, unsigned long size)
 {
 	enum maple_type type = mte_node_type(mas->node);
 	unsigned long pivot, min, gap = 0;
-	unsigned char slot = 0, pivot_cnt = mt_pivots[type];
+	unsigned char offset = 0, pivot_cnt = mt_pivots[type];
+	unsigned long *gaps = NULL, *pivots = ma_pivots(mas_mn(mas), type);
 	bool found = false;
 
-	switch (type) {
-	default:
-		slot = mas_offset(mas);
-		fallthrough;
-	case maple_leaf_64:
-		min = mas_safe_min(mas, slot);
-		for (; slot <= pivot_cnt; slot++) {
-			pivot = _mas_safe_pivot(mas, slot, type);
-			if (slot && !pivot)
-				break;
+	if (ma_is_dense(type)) {
+		mas_set_offset(mas, mas->index - mas->min);
+		return true;
+	}
 
-			/* Not within lower bounds */
-			if (mas->index > pivot)
-				goto next_slot;
+	if (!ma_is_leaf(type)) {
+		offset = mas_offset(mas);
+		gaps = ma_gaps(mte_to_node(mas->node), type);
+	}
 
-			if (ma_is_leaf(type)) {
-				gap = 0;
-				if (!mas_get_slot(mas, slot))
-					gap = min(pivot, mas->last) -
-						max(mas->index, min) + 1;
-			} else {
-				gap = mte_get_gap(mas->node, slot);
-			}
+	min = mas_safe_min(mas, pivots, offset);
+	for (; offset <= pivot_cnt; offset++) {
+		pivot = _mas_safe_pivot(mas, pivots, offset, type);
+		if (offset && !pivot)
+			break;
 
-next_slot:
-			if (gap >= size) {
-				if (ma_is_leaf(type)) {
-					found = true;
-					break;
-				} else if (mas->index <= pivot) {
-					mas->node = mas_get_slot(mas, slot);
-					mas->min = min;
-					mas->max = pivot;
-					slot = 0;
-					break;
-				}
+		/* Not within lower bounds */
+		if (mas->index > pivot)
+			goto next_slot;
+
+		if (gaps)
+			gap = gaps[offset];
+		else if (!mas_get_slot(mas, offset))
+			gap = min(pivot, mas->last) -
+				max(mas->index, min) + 1;
+		else
+			goto next_slot;
+
+		if (gap >= size) {
+			if (ma_is_leaf(type)) {
+				found = true;
+				goto done;
 			}
-			min = pivot + 1;
-			if (mas->last < min) {
-				mas_set_err(mas, -EBUSY);
-				return true;
+			if (mas->index <= pivot) {
+				mas->node = mas_get_slot(mas, offset);
+				mas->min = min;
+				mas->max = pivot;
+				offset = 0;
+				break;
 			}
 		}
-		break;
-
-	case maple_dense: // placeholder.
-		slot = mas->index - mas->min;
-		found = true;
-		break;
+next_slot:
+		min = pivot + 1;
+		if (mas->last < min) {
+			mas_set_err(mas, -EBUSY);
+			return true;
+		}
 	}
 
 	if (mte_is_root(mas->node))
 		found = true;
-
-	mas_set_offset(mas, slot);
+done:
+	mas_set_offset(mas, offset);
 	return found;
 }
 
@@ -3768,7 +3767,7 @@ static inline bool mas_rewind_node(struct ma_state *mas)
 	return true;
 }
 
-static inline void mas_rev_awalk(struct ma_state *mas, unsigned long size)
+void mas_rev_awalk(struct ma_state *mas, unsigned long size)
 {
 	struct maple_enode *last = NULL;
 
@@ -3808,10 +3807,10 @@ static inline bool mas_skip_node(struct ma_state *mas)
 
 	mas_set_offset(mas, ++slot);
 	if (slot > 0)
-		mas->min = mte_get_pivot(mas->node, slot - 1) + 1;
+		mas->min = mte_pivot(mas->node, slot - 1) + 1;
 
 	if (slot < mt_pivot_count(mas->node))
-		mas->max = mte_get_pivot(mas->node, slot);
+		mas->max = mte_pivot(mas->node, slot);
 	return true;
 }
 
@@ -3845,6 +3844,8 @@ static int mas_fill_gap(struct ma_state *mas, void *entry, unsigned char slot,
 {
 	unsigned char pslot = mte_parent_slot(mas->node);
 	struct maple_enode *mn = mas->node;
+	unsigned long *pivots;
+	enum maple_type ptype;
 	/* mas->index is the start address for the search
 	 *  which may no longer be needed.
 	 * mas->last is the end address for the search
@@ -3858,8 +3859,10 @@ static int mas_fill_gap(struct ma_state *mas, void *entry, unsigned char slot,
 	 * calculation, so fix the ma_state here
 	 */
 	mas_ascend(mas);
-	mas->max = mas_safe_pivot(mas, pslot);
-	mas->min = mas_safe_min(mas, pslot);
+	ptype = mte_node_type(mas->node);
+	pivots = ma_pivots(mas_mn(mas), ptype);
+	mas->max = _mas_safe_pivot(mas, pivots, pslot, ptype);
+	mas->min = mas_safe_min(mas, pivots, pslot);
 	mas->node = mn;
 	mas_set_offset(mas, slot);
 	_mas_store(mas, entry, false);
@@ -3873,7 +3876,7 @@ void mas_set_fwd_index(struct ma_state *mas, unsigned long size)
 	// At this point, mas->node points to the right node and we have a
 	// slot that has a sufficient gap.
 	if (slot)
-		min = mte_get_pivot(mas->node, slot - 1) + 1;
+		min = mte_pivot(mas->node, slot - 1) + 1;
 
 	mas->min = min;
 	mas->max = mas_safe_pivot(mas, slot);
@@ -3994,8 +3997,8 @@ static inline int mas_alloc(struct ma_state *mas, void *entry,
 			return xa_err(mas->node);
 
 		if (!mas->index)
-			return mte_get_pivot(mas->node, 0);
-		return mte_get_pivot(mas->node, 1);
+			return mte_pivot(mas->node, 0);
+		return mte_pivot(mas->node, 1);
 	}
 
 	mas_awalk(mas, size); // Must be walking a tree.
@@ -4010,7 +4013,7 @@ static inline int mas_alloc(struct ma_state *mas, void *entry,
 	// slot that has a sufficient gap.
 	min = mas->min;
 	if (slot)
-		min = mte_get_pivot(mas->node, slot - 1) + 1;
+		min = mte_pivot(mas->node, slot - 1) + 1;
 
 	if (mas->index < min)
 		mas->index = min;
@@ -4341,7 +4344,7 @@ static inline void mas_dup_children(struct ma_state *mas, int *node_cnt)
 	struct maple_enode *oldchild, *echild;
 	unsigned char offset, end = mt_slot_count(mas->node);
 	int allocated = mas_alloc_cnt(mas);
-	void **slots = ma_get_slots(mte_to_node(mas->node),
+	void **slots = ma_slots(mte_to_node(mas->node),
 				    mte_node_type(mas->node));
 
 	if (allocated < end) {
@@ -4493,11 +4496,11 @@ static inline unsigned char mas_dead_leaves(struct ma_state *mas, void **slots)
 
 void **mas_destroy_descend(struct ma_state *mas)
 {
-	void **slots = ma_get_slots(mte_to_node(mas->node),
+	void **slots = ma_slots(mte_to_node(mas->node),
 				    mte_node_type(mas->node));
 	while (!mte_is_leaf(slots[0])) {
 		mas->node = slots[0];
-		slots = ma_get_slots(mte_to_node(mas->node),
+		slots = ma_slots(mte_to_node(mas->node),
 				     mte_node_type(mas->node));
 	}
 	return slots;
@@ -4535,7 +4538,7 @@ void mt_destroy_walk(struct rcu_head *head)
 		type = mas_parent_enum(&mas, mas.node);
 		slot = mte_parent_slot(mas.node);
 		mas.node = mt_mk_node(mte_parent(mas.node), type);
-		slots = ma_get_slots(mte_to_node(mas.node), type);
+		slots = ma_slots(mte_to_node(mas.node), type);
 
 		if ((slot == mt_slots[type] - 1) || !slots[slot + 1])
 			continue;
@@ -4723,20 +4726,6 @@ retry:
 	return ret;
 }
 
-int mtree_next(struct maple_tree *mt, unsigned long index, unsigned long *next)
-{
-	int ret = -ENOENT;
-
-	MA_STATE(mas, mt, index, index);
-	rcu_read_lock();
-	//mas_walk_next(&mas);
-	rcu_read_unlock();
-
-	if (mas.node)
-		return 0;
-	return ret;
-}
-
 void *mtree_erase(struct maple_tree *mt, unsigned long index)
 {
 	void *entry = NULL;
@@ -4947,6 +4936,7 @@ void mas_validate_gaps(struct ma_state *mas)
 	unsigned long gap = 0, max_gap = 0;
 	unsigned long p_end, p_start = mas->min;
 	unsigned char p_slot;
+	unsigned long *gaps = NULL;
 	int i;
 
 	if (ma_is_dense(mte_node_type(mte))) {
@@ -4962,12 +4952,16 @@ void mas_validate_gaps(struct ma_state *mas)
 		goto counted;
 	}
 
+	if (!mte_is_leaf(mte))
+		gaps = ma_gaps(mte_to_node(mte), mte_node_type(mte));
+
+
 	for (i = 0; i < mt_slot_count(mte); i++) {
 		p_end = mas_safe_pivot(mas, i);
 		if (!p_end && i)
 			p_end = mas->max;
 
-		if (mte_is_leaf(mte)) {
+		if (!gaps) {
 			if (mas_get_slot(mas, i)) {
 				gap = 0;
 				goto not_empty;
@@ -4977,7 +4971,7 @@ void mas_validate_gaps(struct ma_state *mas)
 		} else {
 			void *entry = mas_get_slot(mas, i);
 
-			gap = mte_get_gap(mte, i);
+			gap = gaps[i];
 			if (!entry) {
 				if (gap != p_end - p_start + 1) {
 					pr_err(MA_PTR"[%u] -> "MA_PTR" %lu != %lu - %lu + 1\n",
@@ -5016,14 +5010,13 @@ counted:
 	p_slot = mte_parent_slot(mas->node);
 	p_mn = mte_parent(mte);
 	MT_BUG_ON(mas->tree, max_gap > mas->max);
-	if (ma_get_gap(p_mn, p_slot, mas_parent_enum(mas, mte)) != max_gap) {
+	if (ma_gaps(p_mn, mas_parent_enum(mas, mte))[p_slot] != max_gap) {
 		pr_err("gap "MA_PTR"[%u] != %lu\n", p_mn, p_slot, max_gap);
 		mt_dump(mas->tree);
 	}
 
 	MT_BUG_ON(mas->tree,
-		ma_get_gap(p_mn, p_slot, mas_parent_enum(mas, mte)) !=
-		max_gap);
+		  ma_gaps(p_mn, mas_parent_enum(mas, mte))[p_slot] != max_gap);
 }
 
 void mas_validate_parent_slot(struct ma_state *mas)
@@ -5039,7 +5032,7 @@ void mas_validate_parent_slot(struct ma_state *mas)
 		return;
 
 	parent = mte_parent(mas->node);
-	slots = ma_get_slots(parent, p_type);
+	slots = ma_slots(parent, p_type);
 	MT_BUG_ON(mas->tree, mas_mn(mas) == parent);
 
 	// Check prev/next parent slot for duplicate node entry
@@ -5062,7 +5055,7 @@ void mas_validate_parent_slot(struct ma_state *mas)
 void mas_validate_child_slot(struct ma_state *mas)
 {
 	enum maple_type type = mte_node_type(mas->node);
-	void **slots = ma_get_slots(mte_to_node(mas->node), type);
+	void **slots = ma_slots(mte_to_node(mas->node), type);
 	struct maple_enode *child;
 	unsigned char i;
 
-- 
2.50.1