v31-0007-Review-radix-tree.patch

From 2c280fb3697501c70e4ce43808e3a5175bbc5eb2 Mon Sep 17 00:00:00 2001
From: Masahiko Sawada <sawada.mshk@gmail.com>
Date: Mon, 20 Feb 2023 11:28:50 +0900
Subject: [PATCH v31 07/14] Review radix tree.

Mainly improve the iteration codes and comments.
---
 src/include/lib/radixtree.h                   | 169 +++++++++---------
 src/include/lib/radixtree_iter_impl.h         |  85 ++++-----
 .../expected/test_radixtree.out               |   6 +-
 .../modules/test_radixtree/test_radixtree.c   | 103 +++++++----
 4 files changed, 197 insertions(+), 166 deletions(-)

diff --git a/src/include/lib/radixtree.h b/src/include/lib/radixtree.h
index e546bd705c..8bea606c62 100644
--- a/src/include/lib/radixtree.h
+++ b/src/include/lib/radixtree.h
@@ -83,7 +83,7 @@
  * RT_SET			- Set a key-value pair
  * RT_BEGIN_ITERATE	- Begin iterating through all key-value pairs
  * RT_ITERATE_NEXT	- Return next key-value pair, if any
- * RT_END_ITER		- End iteration
+ * RT_END_ITERATE	- End iteration
  * RT_MEMORY_USAGE	- Get the memory usage
  *
  * Interface for Shared Memory
@@ -152,8 +152,8 @@
 #define RT_INIT_NODE RT_MAKE_NAME(init_node)
 #define RT_FREE_NODE RT_MAKE_NAME(free_node)
 #define RT_FREE_RECURSE RT_MAKE_NAME(free_recurse)
-#define RT_EXTEND RT_MAKE_NAME(extend)
-#define RT_SET_EXTEND RT_MAKE_NAME(set_extend)
+#define RT_EXTEND_UP RT_MAKE_NAME(extend_up)
+#define RT_EXTEND_DOWN RT_MAKE_NAME(extend_down)
 #define RT_SWITCH_NODE_KIND RT_MAKE_NAME(grow_node_kind)
 #define RT_COPY_NODE RT_MAKE_NAME(copy_node)
 #define RT_REPLACE_NODE RT_MAKE_NAME(replace_node)
@@ -191,7 +191,7 @@
 #define RT_NODE_INSERT_LEAF RT_MAKE_NAME(node_insert_leaf)
 #define RT_NODE_INNER_ITERATE_NEXT RT_MAKE_NAME(node_inner_iterate_next)
 #define RT_NODE_LEAF_ITERATE_NEXT RT_MAKE_NAME(node_leaf_iterate_next)
-#define RT_UPDATE_ITER_STACK RT_MAKE_NAME(update_iter_stack)
+#define RT_ITER_SET_NODE_FROM RT_MAKE_NAME(iter_set_node_from)
 #define RT_ITER_UPDATE_KEY RT_MAKE_NAME(iter_update_key)
 #define RT_VERIFY_NODE RT_MAKE_NAME(verify_node)
 
@@ -612,7 +612,6 @@ static const RT_SIZE_CLASS_ELEM RT_SIZE_CLASS_INFO[] = {
 #endif
 
 /* Contains the actual tree and ancillary info */
-// WIP: this name is a bit strange
 typedef struct RT_RADIX_TREE_CONTROL
 {
 #ifdef RT_SHMEM
@@ -651,36 +650,40 @@ typedef struct RT_RADIX_TREE
  * Iteration support.
  *
  * Iterating the radix tree returns each pair of key and value in the ascending
- * order of the key. To support this, the we iterate nodes of each level.
+ * order of the key.
  *
- * RT_NODE_ITER struct is used to track the iteration within a node.
+ * RT_NODE_ITER is the struct for iteration of one radix tree node.
  *
  * RT_ITER is the struct for iteration of the radix tree, and uses RT_NODE_ITER
- * in order to track the iteration of each level. During iteration, we also
- * construct the key whenever updating the node iteration information, e.g., when
- * advancing the current index within the node or when moving to the next node
- * at the same level.
- *
- * XXX: Currently we allow only one process to do iteration. Therefore, rt_node_iter
- * has the local pointers to nodes, rather than RT_PTR_ALLOC.
- * We need either a safeguard to disallow other processes to begin the iteration
- * while one process is doing or to allow multiple processes to do the iteration.
+ * for each level to track the iteration within the node.
  */
 typedef struct RT_NODE_ITER
 {
-	RT_PTR_LOCAL node;			/* current node being iterated */
-	int			current_idx;	/* current position. -1 for initial value */
+	/*
+	 * Local pointer to the node we are iterating over.
+	 *
+	 * Since the radix tree doesn't support the shared iteration among multiple
+	 * processes, we use RT_PTR_LOCAL rather than RT_PTR_ALLOC.
+	 */
+	RT_PTR_LOCAL node;
+
+	/*
+	 * The next index of the chunk array in RT_NODE_KIND_3 and
+	 * RT_NODE_KIND_32 nodes, or the next chunk in RT_NODE_KIND_125 and
+	 * RT_NODE_KIND_256 nodes. 0 for the initial value.
+	 */
+	int		idx;
 } RT_NODE_ITER;
 
 typedef struct RT_ITER
 {
 	RT_RADIX_TREE *tree;
 
-	/* Track the iteration on nodes of each level */
-	RT_NODE_ITER stack[RT_MAX_LEVEL];
-	int			stack_len;
+	/* Track the nodes for each level. level = 0 is for a leaf node */
+	RT_NODE_ITER node_iters[RT_MAX_LEVEL];
+	int			top_level;
 
-	/* The key is constructed during iteration */
+	/* The key constructed during the iteration */
 	uint64		key;
 } RT_ITER;
 
@@ -1243,7 +1246,7 @@ RT_REPLACE_NODE(RT_RADIX_TREE *tree, RT_PTR_LOCAL parent,
  * it can store the key.
  */
 static pg_noinline void
-RT_EXTEND(RT_RADIX_TREE *tree, uint64 key)
+RT_EXTEND_UP(RT_RADIX_TREE *tree, uint64 key)
 {
 	int			target_shift;
 	RT_PTR_LOCAL root = RT_PTR_GET_LOCAL(tree, tree->ctl->root);
@@ -1282,7 +1285,7 @@ RT_EXTEND(RT_RADIX_TREE *tree, uint64 key)
  * Insert inner and leaf nodes from 'node' to bottom.
  */
 static pg_noinline void
-RT_SET_EXTEND(RT_RADIX_TREE *tree, uint64 key, RT_VALUE_TYPE *value_p, RT_PTR_LOCAL parent,
+RT_EXTEND_DOWN(RT_RADIX_TREE *tree, uint64 key, RT_VALUE_TYPE *value_p, RT_PTR_LOCAL parent,
 			  RT_PTR_ALLOC stored_node, RT_PTR_LOCAL node)
 {
 	int			shift = node->shift;
@@ -1613,7 +1616,7 @@ RT_SET(RT_RADIX_TREE *tree, uint64 key, RT_VALUE_TYPE *value_p)
 
 	/* Extend the tree if necessary */
 	if (key > tree->ctl->max_val)
-		RT_EXTEND(tree, key);
+		RT_EXTEND_UP(tree, key);
 
 	stored_child = tree->ctl->root;
 	parent = RT_PTR_GET_LOCAL(tree, stored_child);
@@ -1631,7 +1634,7 @@ RT_SET(RT_RADIX_TREE *tree, uint64 key, RT_VALUE_TYPE *value_p)
 
 		if (!RT_NODE_SEARCH_INNER(child, key, &new_child))
 		{
-			RT_SET_EXTEND(tree, key, value_p, parent, stored_child, child);
+			RT_EXTEND_DOWN(tree, key, value_p, parent, stored_child, child);
 			RT_UNLOCK(tree);
 			return false;
 		}
@@ -1805,16 +1808,9 @@ RT_DELETE(RT_RADIX_TREE *tree, uint64 key)
 }
 #endif
 
-static inline void
-RT_ITER_UPDATE_KEY(RT_ITER *iter, uint8 chunk, uint8 shift)
-{
-	iter->key &= ~(((uint64) RT_CHUNK_MASK) << shift);
-	iter->key |= (((uint64) chunk) << shift);
-}
-
 /*
- * Advance the slot in the inner node. Return the child if exists, otherwise
- * null.
+ * Scan the inner node and return the next child node if exist, otherwise
+ * return NULL.
  */
 static inline RT_PTR_LOCAL
 RT_NODE_INNER_ITERATE_NEXT(RT_ITER *iter, RT_NODE_ITER *node_iter)
@@ -1825,8 +1821,8 @@ RT_NODE_INNER_ITERATE_NEXT(RT_ITER *iter, RT_NODE_ITER *node_iter)
 }
 
 /*
- * Advance the slot in the leaf node. On success, return true and the value
- * is set to value_p, otherwise return false.
+ * Scan the leaf node, and return true and the next value is set to value_p
+ * if exists. Otherwise return false.
  */
 static inline bool
 RT_NODE_LEAF_ITERATE_NEXT(RT_ITER *iter, RT_NODE_ITER *node_iter,
@@ -1838,29 +1834,50 @@ RT_NODE_LEAF_ITERATE_NEXT(RT_ITER *iter, RT_NODE_ITER *node_iter,
 }
 
 /*
- * Update each node_iter for inner nodes in the iterator node stack.
+ * While descending the radix tree from the 'from' node to the bottom, we
+ * set the next node to iterate for each level.
  */
 static void
-RT_UPDATE_ITER_STACK(RT_ITER *iter, RT_PTR_LOCAL from_node, int from)
+RT_ITER_SET_NODE_FROM(RT_ITER *iter, RT_PTR_LOCAL from)
 {
-	int			level = from;
-	RT_PTR_LOCAL node = from_node;
+	int			level = from->shift / RT_NODE_SPAN;
+	RT_PTR_LOCAL node = from;
 
 	for (;;)
 	{
-		RT_NODE_ITER *node_iter = &(iter->stack[level--]);
+		RT_NODE_ITER *node_iter = &(iter->node_iters[level--]);
+
+#ifdef USE_ASSERT_CHECKING
+		if (node_iter->node)
+		{
+			/* We must have finished the iteration on the previous node */
+			if (RT_NODE_IS_LEAF(node_iter->node))
+			{
+				uint64 dummy;
+				Assert(!RT_NODE_LEAF_ITERATE_NEXT(iter, node_iter, &dummy));
+			}
+			else
+				Assert(!RT_NODE_INNER_ITERATE_NEXT(iter, node_iter));
+		}
+#endif
 
+		/* Set the node to the node iterator of this level */
 		node_iter->node = node;
-		node_iter->current_idx = -1;
+		node_iter->idx = 0;
 
-		/* We don't advance the leaf node iterator here */
 		if (RT_NODE_IS_LEAF(node))
-			return;
+		{
+			/* We will visit the leaf node when RT_ITERATE_NEXT() */
+			break;
+		}
 
-		/* Advance to the next slot in the inner node */
+		/*
+		 * Get the first child node from the node, which corresponds to the
+		 * lowest chunk within the node.
+		 */
 		node = RT_NODE_INNER_ITERATE_NEXT(iter, node_iter);
 
-		/* We must find the first children in the node */
+		/* The first child must be found */
 		Assert(node);
 	}
 }
@@ -1874,14 +1891,11 @@ RT_UPDATE_ITER_STACK(RT_ITER *iter, RT_PTR_LOCAL from_node, int from)
 RT_SCOPE RT_ITER *
 RT_BEGIN_ITERATE(RT_RADIX_TREE *tree)
 {
-	MemoryContext old_ctx;
 	RT_ITER    *iter;
 	RT_PTR_LOCAL root;
-	int			top_level;
 
-	old_ctx = MemoryContextSwitchTo(tree->context);
-
-	iter = (RT_ITER *) palloc0(sizeof(RT_ITER));
+	iter = (RT_ITER *) MemoryContextAllocZero(tree->context,
+											  sizeof(RT_ITER));
 	iter->tree = tree;
 
 	RT_LOCK_SHARED(tree);
@@ -1891,16 +1905,13 @@ RT_BEGIN_ITERATE(RT_RADIX_TREE *tree)
 		return iter;
 
 	root = RT_PTR_GET_LOCAL(tree, iter->tree->ctl->root);
-	top_level = root->shift / RT_NODE_SPAN;
-	iter->stack_len = top_level;
+	iter->top_level = root->shift / RT_NODE_SPAN;
 
 	/*
-	 * Descend to the left most leaf node from the root. The key is being
-	 * constructed while descending to the leaf.
+	 * Set the next node to iterate for each level from the level of the
+	 * root node.
 	 */
-	RT_UPDATE_ITER_STACK(iter, root, top_level);
-
-	MemoryContextSwitchTo(old_ctx);
+	RT_ITER_SET_NODE_FROM(iter, root);
 
 	return iter;
 }
@@ -1912,6 +1923,8 @@ RT_BEGIN_ITERATE(RT_RADIX_TREE *tree)
 RT_SCOPE bool
 RT_ITERATE_NEXT(RT_ITER *iter, uint64 *key_p, RT_VALUE_TYPE *value_p)
 {
+	Assert(value_p != NULL);
+
 	/* Empty tree */
 	if (!iter->tree->ctl->root)
 		return false;
@@ -1919,43 +1932,38 @@ RT_ITERATE_NEXT(RT_ITER *iter, uint64 *key_p, RT_VALUE_TYPE *value_p)
 	for (;;)
 	{
 		RT_PTR_LOCAL child = NULL;
-		RT_VALUE_TYPE value;
-		int			level;
-		bool		found;
-
-		/* Advance the leaf node iterator to get next key-value pair */
-		found = RT_NODE_LEAF_ITERATE_NEXT(iter, &(iter->stack[0]), &value);
 
-		if (found)
+		/* Get the next chunk of the leaf node */
+		if (RT_NODE_LEAF_ITERATE_NEXT(iter, &(iter->node_iters[0]), value_p))
 		{
 			*key_p = iter->key;
-			*value_p = value;
 			return true;
 		}
 
 		/*
-		 * We've visited all values in the leaf node, so advance inner node
-		 * iterators from the level=1 until we find the next child node.
+		 * We've visited all values in the leaf node, so advance all inner node
+		 * iterators by visiting inner nodes from the level = 1 until we find the
+		 * next inner node that has a child node.
 		 */
-		for (level = 1; level <= iter->stack_len; level++)
+		for (int level = 1; level <= iter->top_level; level++)
 		{
-			child = RT_NODE_INNER_ITERATE_NEXT(iter, &(iter->stack[level]));
+			child = RT_NODE_INNER_ITERATE_NEXT(iter, &(iter->node_iters[level]));
 
 			if (child)
 				break;
 		}
 
-		/* the iteration finished */
+		/* We've visited all nodes, so the iteration finished */
 		if (!child)
-			return false;
+			break;
 
 		/*
-		 * Set the node to the node iterator and update the iterator stack
-		 * from this node.
+		 * Found the new child node. We update the next node to iterate for each
+		 * level from the level of this child node.
 		 */
-		RT_UPDATE_ITER_STACK(iter, child, level - 1);
+		RT_ITER_SET_NODE_FROM(iter, child);
 
-		/* Node iterators are updated, so try again from the leaf */
+		/* Find key-value from the leaf node again */
 	}
 
 	return false;
@@ -2470,8 +2478,8 @@ RT_DUMP(RT_RADIX_TREE *tree)
 #undef RT_INIT_NODE
 #undef RT_FREE_NODE
 #undef RT_FREE_RECURSE
-#undef RT_EXTEND
-#undef RT_SET_EXTEND
+#undef RT_EXTEND_UP
+#undef RT_EXTEND_DOWN
 #undef RT_SWITCH_NODE_KIND
 #undef RT_COPY_NODE
 #undef RT_REPLACE_NODE
@@ -2509,8 +2517,7 @@ RT_DUMP(RT_RADIX_TREE *tree)
 #undef RT_NODE_INSERT_LEAF
 #undef RT_NODE_INNER_ITERATE_NEXT
 #undef RT_NODE_LEAF_ITERATE_NEXT
-#undef RT_UPDATE_ITER_STACK
-#undef RT_ITER_UPDATE_KEY
+#undef RT_RT_ITER_SET_NODE_FROM
 #undef RT_VERIFY_NODE
 
 #undef RT_DEBUG
diff --git a/src/include/lib/radixtree_iter_impl.h b/src/include/lib/radixtree_iter_impl.h
index 98c78eb237..5c1034768e 100644
--- a/src/include/lib/radixtree_iter_impl.h
+++ b/src/include/lib/radixtree_iter_impl.h
@@ -27,12 +27,10 @@
 #error node level must be either inner or leaf
 #endif
 
-	bool		found = false;
-	uint8		key_chunk;
+	uint8		key_chunk = 0;
 
 #ifdef RT_NODE_LEVEL_LEAF
-	RT_VALUE_TYPE		value;
-
+	Assert(value_p != NULL);
 	Assert(RT_NODE_IS_LEAF(node_iter->node));
 #else
 	RT_PTR_LOCAL child = NULL;
@@ -50,99 +48,92 @@
 			{
 				RT_NODE3_TYPE *n3 = (RT_NODE3_TYPE *) node_iter->node;
 
-				node_iter->current_idx++;
-				if (node_iter->current_idx >= n3->base.n.count)
-					break;
+				if (node_iter->idx >= n3->base.n.count)
+					return false;
+
 #ifdef RT_NODE_LEVEL_LEAF
-				value = n3->values[node_iter->current_idx];
+				*value_p = n3->values[node_iter->idx];
 #else
-				child = RT_PTR_GET_LOCAL(iter->tree, n3->children[node_iter->current_idx]);
+				child = RT_PTR_GET_LOCAL(iter->tree, n3->children[node_iter->idx]);
 #endif
-				key_chunk = n3->base.chunks[node_iter->current_idx];
-				found = true;
+				key_chunk = n3->base.chunks[node_iter->idx];
+				node_iter->idx++;
 				break;
 			}
 		case RT_NODE_KIND_32:
 			{
 				RT_NODE32_TYPE *n32 = (RT_NODE32_TYPE *) node_iter->node;
 
-				node_iter->current_idx++;
-				if (node_iter->current_idx >= n32->base.n.count)
-					break;
+				if (node_iter->idx >= n32->base.n.count)
+					return false;
 
 #ifdef RT_NODE_LEVEL_LEAF
-				value = n32->values[node_iter->current_idx];
+				*value_p = n32->values[node_iter->idx];
 #else
-				child = RT_PTR_GET_LOCAL(iter->tree, n32->children[node_iter->current_idx]);
+				child = RT_PTR_GET_LOCAL(iter->tree, n32->children[node_iter->idx]);
 #endif
-				key_chunk = n32->base.chunks[node_iter->current_idx];
-				found = true;
+				key_chunk = n32->base.chunks[node_iter->idx];
+				node_iter->idx++;
 				break;
 			}
 		case RT_NODE_KIND_125:
 			{
 				RT_NODE125_TYPE *n125 = (RT_NODE125_TYPE *) node_iter->node;
-				int			i;
+				int			chunk;
 
-				for (i = node_iter->current_idx + 1; i < RT_NODE_MAX_SLOTS; i++)
+				for (chunk = node_iter->idx; chunk < RT_NODE_MAX_SLOTS; chunk++)
 				{
-					if (RT_NODE_125_IS_CHUNK_USED((RT_NODE_BASE_125 *) n125, i))
+					if (RT_NODE_125_IS_CHUNK_USED((RT_NODE_BASE_125 *) n125, chunk))
 						break;
 				}
 
-				if (i >= RT_NODE_MAX_SLOTS)
-					break;
+				if (chunk >= RT_NODE_MAX_SLOTS)
+					return false;
 
-				node_iter->current_idx = i;
 #ifdef RT_NODE_LEVEL_LEAF
-				value = RT_NODE_LEAF_125_GET_VALUE(n125, i);
+				*value_p = RT_NODE_LEAF_125_GET_VALUE(n125, chunk);
 #else
-				child = RT_PTR_GET_LOCAL(iter->tree, RT_NODE_INNER_125_GET_CHILD(n125, i));
+				child = RT_PTR_GET_LOCAL(iter->tree, RT_NODE_INNER_125_GET_CHILD(n125, chunk));
 #endif
-				key_chunk = i;
-				found = true;
+				key_chunk = chunk;
+				node_iter->idx = chunk + 1;
 				break;
 			}
 		case RT_NODE_KIND_256:
 			{
 				RT_NODE256_TYPE *n256 = (RT_NODE256_TYPE *) node_iter->node;
-				int			i;
+				int			chunk;
 
-				for (i = node_iter->current_idx + 1; i < RT_NODE_MAX_SLOTS; i++)
+				for (chunk = node_iter->idx; chunk < RT_NODE_MAX_SLOTS; chunk++)
 				{
 #ifdef RT_NODE_LEVEL_LEAF
-					if (RT_NODE_LEAF_256_IS_CHUNK_USED(n256, i))
+					if (RT_NODE_LEAF_256_IS_CHUNK_USED(n256, chunk))
 #else
-					if (RT_NODE_INNER_256_IS_CHUNK_USED(n256, i))
+					if (RT_NODE_INNER_256_IS_CHUNK_USED(n256, chunk))
 #endif
 						break;
 				}
 
-				if (i >= RT_NODE_MAX_SLOTS)
-					break;
+				if (chunk >= RT_NODE_MAX_SLOTS)
+					return false;
 
-				node_iter->current_idx = i;
 #ifdef RT_NODE_LEVEL_LEAF
-				value = RT_NODE_LEAF_256_GET_VALUE(n256, i);
+				*value_p = RT_NODE_LEAF_256_GET_VALUE(n256, chunk);
 #else
-				child = RT_PTR_GET_LOCAL(iter->tree, RT_NODE_INNER_256_GET_CHILD(n256, i));
+				child = RT_PTR_GET_LOCAL(iter->tree, RT_NODE_INNER_256_GET_CHILD(n256, chunk));
 #endif
-				key_chunk = i;
-				found = true;
+				key_chunk = chunk;
+				node_iter->idx = chunk + 1;
 				break;
 			}
 	}
 
-	if (found)
-	{
-		RT_ITER_UPDATE_KEY(iter, key_chunk, node_iter->node->shift);
-#ifdef RT_NODE_LEVEL_LEAF
-		*value_p = value;
-#endif
-	}
+	/* Update the part of the key */
+	iter->key &= ~(((uint64) RT_CHUNK_MASK) << node_iter->node->shift);
+	iter->key |= (((uint64) key_chunk) << node_iter->node->shift);
 
 #ifdef RT_NODE_LEVEL_LEAF
-	return found;
+	return true;
 #else
 	return child;
 #endif
diff --git a/src/test/modules/test_radixtree/expected/test_radixtree.out b/src/test/modules/test_radixtree/expected/test_radixtree.out
index ce645cb8b5..7ad1ce3605 100644
--- a/src/test/modules/test_radixtree/expected/test_radixtree.out
+++ b/src/test/modules/test_radixtree/expected/test_radixtree.out
@@ -4,8 +4,10 @@ CREATE EXTENSION test_radixtree;
 -- an error if something fails.
 --
 SELECT test_radixtree();
-NOTICE:  testing basic operations with leaf node 4
-NOTICE:  testing basic operations with inner node 4
+NOTICE:  testing basic operations with leaf node 3
+NOTICE:  testing basic operations with inner node 3
+NOTICE:  testing basic operations with leaf node 15
+NOTICE:  testing basic operations with inner node 15
 NOTICE:  testing basic operations with leaf node 32
 NOTICE:  testing basic operations with inner node 32
 NOTICE:  testing basic operations with leaf node 125
diff --git a/src/test/modules/test_radixtree/test_radixtree.c b/src/test/modules/test_radixtree/test_radixtree.c
index afe53382f3..5a169854d9 100644
--- a/src/test/modules/test_radixtree/test_radixtree.c
+++ b/src/test/modules/test_radixtree/test_radixtree.c
@@ -43,12 +43,15 @@ typedef uint64 TestValueType;
  */
 static const bool rt_test_stats = false;
 
-static int	rt_node_kind_fanouts[] = {
-	0,
-	4,							/* RT_NODE_KIND_4 */
-	32,							/* RT_NODE_KIND_32 */
-	125,						/* RT_NODE_KIND_125 */
-	256							/* RT_NODE_KIND_256 */
+/*
+ * XXX: should we expose and use RT_SIZE_CLASS and RT_SIZE_CLASS_INFO?
+ */
+static int	rt_node_class_fanouts[] = {
+	3,		/* RT_CLASS_3 */
+	15,		/* RT_CLASS_32_MIN */
+	32, 	/* RT_CLASS_32_MAX */
+	125,	/* RT_CLASS_125 */
+	256		/* RT_CLASS_256 */
 };
 /*
  * A struct to define a pattern of integers, for use with the test_pattern()
@@ -260,10 +263,9 @@ test_basic(int children, bool test_inner)
  * Check if keys from start to end with the shift exist in the tree.
  */
 static void
-check_search_on_node(rt_radix_tree *radixtree, uint8 shift, int start, int end,
-					 int incr)
+check_search_on_node(rt_radix_tree *radixtree, uint8 shift, int start, int end)
 {
-	for (int i = start; i < end; i++)
+	for (int i = start; i <= end; i++)
 	{
 		uint64		key = ((uint64) i << shift);
 		TestValueType		val;
@@ -277,22 +279,26 @@ check_search_on_node(rt_radix_tree *radixtree, uint8 shift, int start, int end,
 	}
 }
 
+/*
+ * Insert 256 key-value pairs, and check if keys are properly inserted on each
+ * node class.
+ */
+/* Test keys [0, 256) */
+#define NODE_TYPE_TEST_KEY_MIN 0
+#define NODE_TYPE_TEST_KEY_MAX 256
 static void
-test_node_types_insert(rt_radix_tree *radixtree, uint8 shift, bool insert_asc)
+test_node_types_insert_asc(rt_radix_tree *radixtree, uint8 shift)
 {
-	uint64		num_entries;
-	int		ninserted = 0;
-	int		start = insert_asc ? 0 : 256;
-	int 	incr = insert_asc ? 1 : -1;
-	int		end = insert_asc ? 256 : 0;
-	int		node_kind_idx = 1;
+	uint64 num_entries;
+	int node_class_idx = 0;
+	uint64 key_checked = 0;
 
-	for (int i = start; i != end; i += incr)
+	for (int i = NODE_TYPE_TEST_KEY_MIN; i < NODE_TYPE_TEST_KEY_MAX; i++)
 	{
 		uint64		key = ((uint64) i << shift);
 		bool		found;
 
-		found = rt_set(radixtree, key, (TestValueType*) &key);
+		found = rt_set(radixtree, key, (TestValueType *) &key);
 		if (found)
 			elog(ERROR, "newly inserted key 0x" UINT64_HEX_FORMAT " is found", key);
 
@@ -300,24 +306,49 @@ test_node_types_insert(rt_radix_tree *radixtree, uint8 shift, bool insert_asc)
 		 * After filling all slots in each node type, check if the values
 		 * are stored properly.
 		 */
-		if (ninserted == rt_node_kind_fanouts[node_kind_idx] - 1)
+		if ((i + 1) == rt_node_class_fanouts[node_class_idx])
 		{
-			int check_start = insert_asc
-				? rt_node_kind_fanouts[node_kind_idx - 1]
-				: rt_node_kind_fanouts[node_kind_idx];
-			int check_end = insert_asc
-				? rt_node_kind_fanouts[node_kind_idx]
-				: rt_node_kind_fanouts[node_kind_idx - 1];
-
-			check_search_on_node(radixtree, shift, check_start, check_end, incr);
-			node_kind_idx++;
+			check_search_on_node(radixtree, shift, key_checked, i);
+			key_checked = i;
+			node_class_idx++;
 		}
-
-		ninserted++;
 	}
 
 	num_entries = rt_num_entries(radixtree);
+	if (num_entries != 256)
+		elog(ERROR,
+			 "rt_num_entries returned " UINT64_FORMAT ", expected " UINT64_FORMAT,
+			 num_entries, UINT64CONST(256));
+}
+
+/*
+ * Similar to test_node_types_insert_asc(), but inserts keys in descending order.
+ */
+static void
+test_node_types_insert_desc(rt_radix_tree *radixtree, uint8 shift)
+{
+	uint64 num_entries;
+	int node_class_idx = 0;
+	uint64 key_checked = NODE_TYPE_TEST_KEY_MAX - 1;
+
+	for (int i = NODE_TYPE_TEST_KEY_MAX - 1; i >= NODE_TYPE_TEST_KEY_MIN; i--)
+	{
+		uint64		key = ((uint64) i << shift);
+		bool		found;
+
+		found = rt_set(radixtree, key, (TestValueType *) &key);
+		if (found)
+			elog(ERROR, "newly inserted key 0x" UINT64_HEX_FORMAT " is found", key);
 
+		if ((i + 1) == rt_node_class_fanouts[node_class_idx])
+		{
+			check_search_on_node(radixtree, shift, i, key_checked);
+			key_checked = i;
+			node_class_idx++;
+		}
+	}
+
+	num_entries = rt_num_entries(radixtree);
 	if (num_entries != 256)
 		elog(ERROR,
 			 "rt_num_entries returned " UINT64_FORMAT ", expected " UINT64_FORMAT,
@@ -329,7 +360,7 @@ test_node_types_delete(rt_radix_tree *radixtree, uint8 shift)
 {
 	uint64		num_entries;
 
-	for (int i = 0; i < 256; i++)
+	for (int i = NODE_TYPE_TEST_KEY_MIN; i < NODE_TYPE_TEST_KEY_MAX; i++)
 	{
 		uint64		key = ((uint64) i << shift);
 		bool		found;
@@ -379,9 +410,9 @@ test_node_types(uint8 shift)
 	 * then delete all entries to make it empty, and insert and search entries
 	 * again.
 	 */
-	test_node_types_insert(radixtree, shift, true);
+	test_node_types_insert_asc(radixtree, shift);
 	test_node_types_delete(radixtree, shift);
-	test_node_types_insert(radixtree, shift, false);
+	test_node_types_insert_desc(radixtree, shift);
 
 	rt_free(radixtree);
 #ifdef RT_SHMEM
@@ -664,10 +695,10 @@ test_radixtree(PG_FUNCTION_ARGS)
 {
 	test_empty();
 
-	for (int i = 1; i < lengthof(rt_node_kind_fanouts); i++)
+	for (int i = 0; i < lengthof(rt_node_class_fanouts); i++)
 	{
-		test_basic(rt_node_kind_fanouts[i], false);
-		test_basic(rt_node_kind_fanouts[i], true);
+		test_basic(rt_node_class_fanouts[i], false);
+		test_basic(rt_node_class_fanouts[i], true);
 	}
 
 	for (int shift = 0; shift <= (64 - 8); shift += 8)
-- 
2.31.1