From 761e6c2583b37eb9d45d64de954d65d953277040 Mon Sep 17 00:00:00 2001
From: amitlan <amitlangote09@gmail.com>
Date: Thu, 3 Mar 2022 16:04:13 +0900
Subject: [PATCH v7 3/4] Add a plan_tree_walker()

Like planstate_tree_walker() but for uninitialized plan trees.
---
 src/backend/nodes/nodeFuncs.c | 116 ++++++++++++++++++++++++++++++++++
 src/include/nodes/nodeFuncs.h |   3 +
 2 files changed, 119 insertions(+)

diff --git a/src/backend/nodes/nodeFuncs.c b/src/backend/nodes/nodeFuncs.c
index 25cf282aab..5e5158ea0e 100644
--- a/src/backend/nodes/nodeFuncs.c
+++ b/src/backend/nodes/nodeFuncs.c
@@ -31,6 +31,10 @@ static bool planstate_walk_subplans(List *plans, bool (*walker) (),
 									void *context);
 static bool planstate_walk_members(PlanState **planstates, int nplans,
 								   bool (*walker) (), void *context);
+static bool plan_walk_subplans(List *plans,
+				   bool (*walker) (),
+				   void *context);
+static bool plan_walk_members(List *plans, bool (*walker) (), void *context);
 
 
 /*
@@ -4368,3 +4372,115 @@ planstate_walk_members(PlanState **planstates, int nplans,
 
 	return false;
 }
+
+/*
+ * plan_tree_walker --- walk plantrees
+ *
+ * The walker has already visited the current node, and so we need only
+ * recurse into any sub-nodes it has.
+ */
+bool
+plan_tree_walker(Plan *plan,
+				 bool (*walker) (),
+				 void *context)
+{
+	/* Guard against stack overflow due to overly complex plan trees */
+	check_stack_depth();
+
+	/* initPlan-s */
+	if (plan_walk_subplans(plan->initPlan, walker, context))
+		return true;
+
+	/* lefttree */
+	if (outerPlan(plan))
+	{
+		if (walker(outerPlan(plan), context))
+			return true;
+	}
+
+	/* righttree */
+	if (innerPlan(plan))
+	{
+		if (walker(innerPlan(plan), context))
+			return true;
+	}
+
+	/* special child plans */
+	switch (nodeTag(plan))
+	{
+		case T_Append:
+			if (plan_walk_members(((Append *) plan)->appendplans,
+								  walker, context))
+				return true;
+			break;
+		case T_MergeAppend:
+			if (plan_walk_members(((MergeAppend *) plan)->mergeplans,
+								  walker, context))
+				return true;
+			break;
+		case T_BitmapAnd:
+			if (plan_walk_members(((BitmapAnd *) plan)->bitmapplans,
+								  walker, context))
+				return true;
+			break;
+		case T_BitmapOr:
+			if (plan_walk_members(((BitmapOr *) plan)->bitmapplans,
+								  walker, context))
+				return true;
+			break;
+		case T_CustomScan:
+			if (plan_walk_members(((CustomScan *) plan)->custom_plans,
+								  walker, context))
+				return true;
+			break;
+		case T_SubqueryScan:
+			if (walker(((SubqueryScan *) plan)->subplan, context))
+				return true;
+			break;
+		default:
+			break;
+	}
+
+	return false;
+}
+
+/*
+ * Walk a list of SubPlans (or initPlans, which also use SubPlan nodes).
+ */
+static bool
+plan_walk_subplans(List *plans,
+				   bool (*walker) (),
+				   void *context)
+{
+	ListCell   *lc;
+	PlannedStmt *plannedstmt = (PlannedStmt *) context;
+
+	foreach(lc, plans)
+	{
+		SubPlan *sp = lfirst_node(SubPlan, lc);
+		Plan *p = list_nth(plannedstmt->subplans, sp->plan_id - 1);
+
+		if (walker(p, context))
+			return true;
+	}
+
+	return false;
+}
+
+/*
+ * Walk the constituent plans of a ModifyTable, Append, MergeAppend,
+ * BitmapAnd, or BitmapOr node.
+ */
+static bool
+plan_walk_members(List *plans, bool (*walker) (), void *context)
+{
+	ListCell *lc;
+
+	foreach(lc, plans)
+	{
+		if (walker(lfirst(lc), context))
+			return true;
+	}
+
+	return false;
+}
diff --git a/src/include/nodes/nodeFuncs.h b/src/include/nodes/nodeFuncs.h
index 93c60bde66..fca107ad65 100644
--- a/src/include/nodes/nodeFuncs.h
+++ b/src/include/nodes/nodeFuncs.h
@@ -158,5 +158,8 @@ extern bool raw_expression_tree_walker(Node *node, bool (*walker) (),
 struct PlanState;
 extern bool planstate_tree_walker(struct PlanState *planstate, bool (*walker) (),
 								  void *context);
+struct Plan;
+extern bool plan_tree_walker(struct Plan *plan, bool (*walker) (),
+				 void *context);
 
 #endif							/* NODEFUNCS_H */
-- 
2.24.1