Support multiple iterators in the Red-Black Tree implementation.

While we don't need multiple iterators at the moment, the interface is
nicer and less dangerous this way.

Aleksander Alekseev, with some changes by me.

src/backend/access/gin/ginbulk.c

@@ -255,7 +255,7 @@ qsortCompareItemPointers(const void *a, const void *b)
 void
 ginBeginBAScan(BuildAccumulator *accum)
 {
-	rb_begin_iterate(accum->tree, LeftRightWalk);
+	rb_begin_iterate(accum->tree, LeftRightWalk, &accum->tree_walk);
 }
 
 /*
@@ -271,7 +271,7 @@ ginGetBAEntry(BuildAccumulator *accum,
 	GinEntryAccumulator *entry;
 	ItemPointerData *list;
 
-	entry = (GinEntryAccumulator *) rb_iterate(accum->tree);
+	entry = (GinEntryAccumulator *) rb_iterate(&accum->tree_walk);
 
 	if (entry == NULL)
 		return NULL;			/* no more entries */

src/backend/lib/rbtree.c

@@ -29,17 +29,6 @@
 #include "lib/rbtree.h"
 
 
-/*
- * Values of RBNode.iteratorState
- *
- * Note that iteratorState has an undefined value except in nodes that are
- * currently being visited by an active iteration.
- */
-#define InitialState	(0)
-#define FirstStepDone	(1)
-#define SecondStepDone	(2)
-#define ThirdStepDone	(3)
-
 /*
  * Colors of nodes (values of RBNode.color)
  */
@@ -53,10 +42,6 @@ struct RBTree
 {
 	RBNode	   *root;			/* root node, or RBNIL if tree is empty */
 
-	/* Iteration state */
-	RBNode	   *cur;			/* current iteration node */
-	RBNode	   *(*iterate) (RBTree *rb);
-
 	/* Remaining fields are constant after rb_create */
 
 	Size		node_size;		/* actual size of tree nodes */
@@ -75,8 +60,19 @@ struct RBTree
  */
 #define RBNIL (&sentinel)
 
-static RBNode sentinel = {InitialState, RBBLACK, RBNIL, RBNIL, NULL};
+static RBNode sentinel = {RBBLACK, RBNIL, RBNIL, NULL};
 
+/*
+ * Values used in the RBTreeIterator.next_state field, with an
+ * InvertedWalk iterator.
+ */
+typedef enum InvertedWalkNextStep
+{
+	NextStepBegin,
+	NextStepUp,
+	NextStepLeft,
+	NextStepRight
+} InvertedWalkNextStep;
 
 /*
  * rb_create: create an empty RBTree
@@ -123,8 +119,6 @@ rb_create(Size node_size,
 	Assert(node_size > sizeof(RBNode));
 
 	tree->root = RBNIL;
-	tree->cur = RBNIL;
-	tree->iterate = NULL;
 	tree->node_size = node_size;
 	tree->comparator = comparator;
 	tree->combiner = combiner;
@@ -437,7 +431,6 @@ rb_insert(RBTree *rb, const RBNode *data, bool *isNew)
 
 	x = rb->allocfunc (rb->arg);
 
-	x->iteratorState = InitialState;
 	x->color = RBRED;
 
 	x->left = RBNIL;
@@ -653,171 +646,194 @@ rb_delete(RBTree *rb, RBNode *node)
  *						  Traverse									  *
  **********************************************************************/
 
-/*
- * The iterator routines were originally coded in tail-recursion style,
- * which is nice to look at, but is trouble if your compiler isn't smart
- * enough to optimize it.  Now we just use looping.
- */
-#define descend(next_node) \
-	do { \
-		(next_node)->iteratorState = InitialState; \
-		node = rb->cur = (next_node); \
-		goto restart; \
-	} while (0)
+static RBNode *
+rb_left_right_iterator(RBTreeIterator *iter)
+{
+	if (iter->last_visited == NULL)
+	{
+		iter->last_visited = iter->rb->root;
+		while (iter->last_visited->left != RBNIL)
+			iter->last_visited = iter->last_visited->left;
 
-#define ascend(next_node) \
-	do { \
-		node = rb->cur = (next_node); \
-		goto restart; \
-	} while (0)
+		return iter->last_visited;
+	}
 
+	if (iter->last_visited->right != RBNIL)
+	{
+		iter->last_visited = iter->last_visited->right;
+		while (iter->last_visited->left != RBNIL)
+			iter->last_visited = iter->last_visited->left;
 
-static RBNode *
-rb_left_right_iterator(RBTree *rb)
-{
-	RBNode	   *node = rb->cur;
+		return iter->last_visited;
+	}
 
-restart:
-	switch (node->iteratorState)
+	for (;;)
 	{
-		case InitialState:
-			if (node->left != RBNIL)
-			{
-				node->iteratorState = FirstStepDone;
-				descend(node->left);
-			}
-			/* FALL THROUGH */
-		case FirstStepDone:
-			node->iteratorState = SecondStepDone;
-			return node;
-		case SecondStepDone:
-			if (node->right != RBNIL)
-			{
-				node->iteratorState = ThirdStepDone;
-				descend(node->right);
-			}
-			/* FALL THROUGH */
-		case ThirdStepDone:
-			if (node->parent)
-				ascend(node->parent);
+		RBNode	   *came_from = iter->last_visited;
+
+		iter->last_visited = iter->last_visited->parent;
+		if (iter->last_visited == NULL)
+		{
+			iter->is_over = true;
 			break;
-		default:
-			elog(ERROR, "unrecognized rbtree node state: %d",
-				 node->iteratorState);
+		}
+
+		if (iter->last_visited->left == came_from)
+			break;				/* came from left sub-tree, return current
+								 * node */
+
+		/* else - came from right sub-tree, continue to move up */
 	}
 
-	return NULL;
+	return iter->last_visited;
 }
 
 static RBNode *
-rb_right_left_iterator(RBTree *rb)
+rb_right_left_iterator(RBTreeIterator *iter)
 {
-	RBNode	   *node = rb->cur;
+	if (iter->last_visited == NULL)
+	{
+		iter->last_visited = iter->rb->root;
+		while (iter->last_visited->right != RBNIL)
+			iter->last_visited = iter->last_visited->right;
+
+		return iter->last_visited;
+	}
 
-restart:
-	switch (node->iteratorState)
+	if (iter->last_visited->left != RBNIL)
 	{
-		case InitialState:
-			if (node->right != RBNIL)
-			{
-				node->iteratorState = FirstStepDone;
-				descend(node->right);
-			}
-			/* FALL THROUGH */
-		case FirstStepDone:
-			node->iteratorState = SecondStepDone;
-			return node;
-		case SecondStepDone:
-			if (node->left != RBNIL)
-			{
-				node->iteratorState = ThirdStepDone;
-				descend(node->left);
-			}
-			/* FALL THROUGH */
-		case ThirdStepDone:
-			if (node->parent)
-				ascend(node->parent);
+		iter->last_visited = iter->last_visited->left;
+		while (iter->last_visited->right != RBNIL)
+			iter->last_visited = iter->last_visited->right;
+
+		return iter->last_visited;
+	}
+
+	for (;;)
+	{
+		RBNode	   *came_from = iter->last_visited;
+
+		iter->last_visited = iter->last_visited->parent;
+		if (iter->last_visited == NULL)
+		{
+			iter->is_over = true;
 			break;
-		default:
-			elog(ERROR, "unrecognized rbtree node state: %d",
-				 node->iteratorState);
+		}
+
+		if (iter->last_visited->right == came_from)
+			break;				/* came from right sub-tree, return current
+								 * node */
+
+		/* else - came from left sub-tree, continue to move up */
 	}
 
-	return NULL;
+	return iter->last_visited;
 }
 
 static RBNode *
-rb_direct_iterator(RBTree *rb)
+rb_direct_iterator(RBTreeIterator *iter)
 {
-	RBNode	   *node = rb->cur;
+	if (iter->last_visited == NULL)
+	{
+		iter->last_visited = iter->rb->root;
+		return iter->last_visited;
+	}
 
-restart:
-	switch (node->iteratorState)
+	if (iter->last_visited->left != RBNIL)
 	{
-		case InitialState:
-			node->iteratorState = FirstStepDone;
-			return node;
-		case FirstStepDone:
-			if (node->left != RBNIL)
+		iter->last_visited = iter->last_visited->left;
+		return iter->last_visited;
+	}
+
+	do
+	{
+		if (iter->last_visited->right != RBNIL)
+		{
+			iter->last_visited = iter->last_visited->right;
+			break;
+		}
+
+		/* go up and one step right */
+		for (;;)
+		{
+			RBNode	   *came_from = iter->last_visited;
+
+			iter->last_visited = iter->last_visited->parent;
+			if (iter->last_visited == NULL)
 			{
-				node->iteratorState = SecondStepDone;
-				descend(node->left);
+				iter->is_over = true;
+				break;
 			}
-			/* FALL THROUGH */
-		case SecondStepDone:
-			if (node->right != RBNIL)
+
+			if ((iter->last_visited->right != came_from) && (iter->last_visited->right != RBNIL))
 			{
-				node->iteratorState = ThirdStepDone;
-				descend(node->right);
+				iter->last_visited = iter->last_visited->right;
+				return iter->last_visited;
 			}
-			/* FALL THROUGH */
-		case ThirdStepDone:
-			if (node->parent)
-				ascend(node->parent);
-			break;
-		default:
-			elog(ERROR, "unrecognized rbtree node state: %d",
-				 node->iteratorState);
+		}
 	}
+	while (iter->last_visited != NULL);
 
-	return NULL;
+	return iter->last_visited;
 }
 
 static RBNode *
-rb_inverted_iterator(RBTree *rb)
+rb_inverted_iterator(RBTreeIterator *iter)
 {
-	RBNode	   *node = rb->cur;
+	RBNode	   *came_from;
 
-restart:
-	switch (node->iteratorState)
+loop:
+	switch ((InvertedWalkNextStep) iter->next_step)
 	{
-		case InitialState:
-			if (node->left != RBNIL)
+			/* First call, begin from root */
+		case NextStepBegin:
+			iter->last_visited = iter->rb->root;
+			iter->next_step = NextStepLeft;
+			goto loop;
+
+		case NextStepLeft:
+			while (iter->last_visited->left != RBNIL)
+				iter->last_visited = iter->last_visited->left;
+
+			iter->next_step = NextStepRight;
+			goto loop;
+
+		case NextStepRight:
+			if (iter->last_visited->right != RBNIL)
 			{
-				node->iteratorState = FirstStepDone;
-				descend(node->left);
+				iter->last_visited = iter->last_visited->right;
+				iter->next_step = NextStepLeft;
+				goto loop;
 			}
-			/* FALL THROUGH */
-		case FirstStepDone:
-			if (node->right != RBNIL)
+			else	/* not moved - return current, then go up */
+				iter->next_step = NextStepUp;
+			break;
+
+		case NextStepUp:
+			for (;;)
 			{
-				node->iteratorState = SecondStepDone;
-				descend(node->right);
+				came_from = iter->last_visited;
+				iter->last_visited = iter->last_visited->parent;
+				if (iter->last_visited == NULL)
+				{
+					iter->is_over = true;
+					break;		/* end of iteration */
+				}
+
+				if (came_from == iter->last_visited->right)
+				{
+					/* return current, then continue to go up */
+					break;
+				}
+
+				/* otherwise we came from the left */
+				iter->next_step = NextStepRight;
+				goto loop;
 			}
-			/* FALL THROUGH */
-		case SecondStepDone:
-			node->iteratorState = ThirdStepDone;
-			return node;
-		case ThirdStepDone:
-			if (node->parent)
-				ascend(node->parent);
 			break;
-		default:
-			elog(ERROR, "unrecognized rbtree node state: %d",
-				 node->iteratorState);
 	}
 
-	return NULL;
+	return iter->last_visited;
 }
 
 /*
@@ -827,33 +843,34 @@ restart:
  * returns NULL or the traversal stops being of interest.
  *
  * If the tree is changed during traversal, results of further calls to
- * rb_iterate are unspecified.
+ * rb_iterate are unspecified.  Multiple concurrent iterators on the same
+ * tree are allowed.
  *
- * Note: this used to return a separately palloc'd iterator control struct,
- * but that's a bit pointless since the data structure is incapable of
- * supporting multiple concurrent traversals.  Now we just keep the state
- * in RBTree.
+ * The iterator state is stored in the 'iter' struct.  The caller should
+ * treat it as opaque struct.
  */
 void
-rb_begin_iterate(RBTree *rb, RBOrderControl ctrl)
+rb_begin_iterate(RBTree *rb, RBOrderControl ctrl, RBTreeIterator *iter)
 {
-	rb->cur = rb->root;
-	if (rb->cur != RBNIL)
-		rb->cur->iteratorState = InitialState;
+	/* Common initialization for all traversal orders */
+	iter->rb = rb;
+	iter->last_visited = NULL;
+	iter->is_over = (rb->root == RBNIL);
 
 	switch (ctrl)
 	{
 		case LeftRightWalk:		/* visit left, then self, then right */
-			rb->iterate = rb_left_right_iterator;
+			iter->iterate = rb_left_right_iterator;
 			break;
 		case RightLeftWalk:		/* visit right, then self, then left */
-			rb->iterate = rb_right_left_iterator;
+			iter->iterate = rb_right_left_iterator;
 			break;
 		case DirectWalk:		/* visit self, then left, then right */
-			rb->iterate = rb_direct_iterator;
+			iter->iterate = rb_direct_iterator;
 			break;
 		case InvertedWalk:		/* visit left, then right, then self */
-			rb->iterate = rb_inverted_iterator;
+			iter->iterate = rb_inverted_iterator;
+			iter->next_step = NextStepBegin;
 			break;
 		default:
 			elog(ERROR, "unrecognized rbtree iteration order: %d", ctrl);
@@ -864,10 +881,10 @@ rb_begin_iterate(RBTree *rb, RBOrderControl ctrl)
  * rb_iterate: return the next node in traversal order, or NULL if no more
  */
 RBNode *
-rb_iterate(RBTree *rb)
+rb_iterate(RBTreeIterator *iter)
 {
-	if (rb->cur == RBNIL)
+	if (iter->is_over)
 		return NULL;
 
-	return rb->iterate(rb);
+	return iter->iterate(iter);
 }

src/include/access/gin_private.h

@@ -919,6 +919,7 @@ typedef struct
 	GinEntryAccumulator *entryallocator;
 	uint32		eas_used;
 	RBTree	   *tree;
+	RBTreeIterator tree_walk;
 } BuildAccumulator;
 
 extern void ginInitBA(BuildAccumulator *accum);

src/include/lib/rbtree.h

@@ -22,7 +22,6 @@
  */
 typedef struct RBNode
 {
-	char		iteratorState;	/* workspace for iterating through tree */
 	char color;					/* node's current color, red or black */
 	struct RBNode *left;		/* left child, or RBNIL if none */
 	struct RBNode *right;		/* right child, or RBNIL if none */
@@ -41,6 +40,22 @@ typedef enum RBOrderControl
 	InvertedWalk				/* postorder: left child, right child, node */
 } RBOrderControl;
 
+/*
+ * RBTreeIterator holds state while traversing a tree.  This is declared
+ * here so that callers can stack-allocate this, but must otherwise be
+ * treated as an opaque struct.
+ */
+typedef struct RBTreeIterator RBTreeIterator;
+
+struct RBTreeIterator
+{
+	RBTree	   *rb;
+	RBNode	   *(*iterate) (RBTreeIterator *iter);
+	RBNode	   *last_visited;
+	char		next_step;
+	bool		is_over;
+};
+
 /* Support functions to be provided by caller */
 typedef int (*rb_comparator) (const RBNode *a, const RBNode *b, void *arg);
 typedef void (*rb_combiner) (RBNode *existing, const RBNode *newdata, void *arg);
@@ -60,7 +75,8 @@ extern RBNode *rb_leftmost(RBTree *rb);
 extern RBNode *rb_insert(RBTree *rb, const RBNode *data, bool *isNew);
 extern void rb_delete(RBTree *rb, RBNode *node);
 
-extern void rb_begin_iterate(RBTree *rb, RBOrderControl ctrl);
-extern RBNode *rb_iterate(RBTree *rb);
+extern void rb_begin_iterate(RBTree *rb, RBOrderControl ctrl,
+				 RBTreeIterator *iter);
+extern RBNode *rb_iterate(RBTreeIterator *iter);
 
 #endif   /* RBTREE_H */