Refactor CTE level handling in insert_select_planner.c to shift every level (1→0, 2→1, 3→2, etc.) all in one pass

src/backend/distributed/planner/insert_select_planner.c

@@ -108,22 +108,17 @@ static void ProcessEntryPair(TargetEntry *insertEntry, TargetEntry *selectEntry,
 							 Form_pg_attribute attr, int targetEntryIndex,
 							 List **projectedEntries, List **nonProjectedEntries);
 
-/*
+typedef struct ShiftAllLevelsContext
- * DecrementCteLevelWalkerContext
- *
- *   'offset' is how much we shift ctelevelsup by (e.g. -1),
- *   'level' tracks the current query nesting level,
- *   so we know if RTE_CTE is referencing this level.
- */
-typedef struct DecrementCteLevelWalkerContext
 {
-	int oldLevel;
+	int baseLevel;  /* we shift ctelevelsup if it’s >= this number */
-	int newLevel;
+} ShiftAllLevelsContext;
-} DecrementCteLevelWalkerContext;
 							 
 
-static void DecrementCteLevelForQuery(Query *query, int oldLevel, int newLevel);
+static bool
-static bool DecrementCteLevelWalker(Node *node, DecrementCteLevelWalkerContext *context);
+ShiftAllLevelsWalker(Node *node, ShiftAllLevelsContext *ctx);
+static void
+ShiftAllCteLevelsInQuery(Query *subquery, int baseLevel);
+
 
 /* depth of current insert/select planner. */
 static int insertSelectPlannerLevel = 0;
@@ -560,79 +555,78 @@ PrepareInsertSelectForCitusPlanner(Query *insertSelectQuery)
 						copyObject(insertSelectQuery->cteList));
 		insertSelectQuery->cteList = NIL;
 
-		DecrementCteLevelForQuery(selectRte->subquery, 1, 0);
+		/* Suppose we physically appended the top-level cteList into the subquery, 
-		elog(DEBUG1, "Done shifting ctelevelsup 1->0 for subquery references");
+   		so references are at ctelevelsup=1, 2, etc. We want them all to shift by -1. */
-	}
-}
 
+		   ShiftAllCteLevelsInQuery(selectRte->subquery,
+			1       /* baseLevel => anything >= 1 is shifted */);
+
+		elog(DEBUG1, "Done shifting ctelevelsup X->X-1 for subquery references");
+	}
+}
 
 /*
- * DecrementCteLevelWalker
+ * ShiftAllLevelsWalker:
- *
+ *   Recursively visits each Query using QTW_EXAMINE_RTES_BEFORE so it
+ *   receives RangeTblEntry nodes. If an RTE_CTE has ctelevelsup >= baseLevel,
+ *   do ctelevelsup -= 1.
  */
 static bool
-DecrementCteLevelWalker(Node *node, DecrementCteLevelWalkerContext *context)
+ShiftAllLevelsWalker(Node *node, ShiftAllLevelsContext *ctx)
 {
     if (node == NULL)
-	{
         return false;
-	}
 
     if (IsA(node, RangeTblEntry))
     {
         RangeTblEntry *rte = (RangeTblEntry *) node;
 
-		if (rte->rtekind == RTE_CTE && rte->ctelevelsup == context->oldLevel)
+        if (rte->rtekind == RTE_CTE &&
+            rte->ctelevelsup >= ctx->baseLevel)
         {
-			rte->ctelevelsup = context->newLevel;
+            elog(DEBUG2, "Shifting ctelevelsup from %d to %d",
+                 rte->ctelevelsup, rte->ctelevelsup - 1);
+            rte->ctelevelsup -= 1;
         }
 
         return false;
     }
     else if (IsA(node, Query))
     {
+        /* For each Query, we call query_tree_walker with QTW_EXAMINE_RTES_BEFORE */
         Query *query = (Query *) node;
 
-		/*
-		 * Use QTW_EXAMINE_RTES_BEFORE so that this walker is called
-		 * on each RangeTblEntry in query->rtable, giving us a chance
-		 * to adjust ctelevelsup before we do the rest of the query tree.
-		 */
         query_tree_walker(query,
-						  DecrementCteLevelWalker,
+                          ShiftAllLevelsWalker,
-						  (void *) context,
+                          (void *) ctx,
                           QTW_EXAMINE_RTES_BEFORE);
 
         return false;
     }
-	else
-	{
-		/* fallback for expression nodes, etc. */
-		return expression_tree_walker(node,
-									  DecrementCteLevelWalker,
-									  (void *) context);
-	}
-}
 
+    /* fallback for expressions, etc. */
+    return expression_tree_walker(node, ShiftAllLevelsWalker, (void *) ctx);
+}
 
 /*
- * DecrementCteLevelForQuery
+ * ShiftAllCteLevelsInQuery
+ *   A convenience wrapper to shift ctelevelsup by 1 for any references >= baseLevel
  */
 static void
-DecrementCteLevelForQuery(Query *query, int oldLevel, int newLevel)
+ShiftAllCteLevelsInQuery(Query *subquery, int baseLevel)
 {
-	DecrementCteLevelWalkerContext ctx;
+    ShiftAllLevelsContext ctx;
+    ctx.baseLevel = baseLevel;
 
-	ctx.oldLevel = oldLevel;
+    (void) query_tree_walker(subquery,
-	ctx.newLevel = newLevel;
+                             ShiftAllLevelsWalker,
-
-	query_tree_walker(query,
-					  DecrementCteLevelWalker,
                              (void *) &ctx,
                              QTW_EXAMINE_RTES_BEFORE);
 }
 
 
+
+
 /*
  * CreateCombineQueryForRouterPlan is used for creating a dummy combineQuery
  * for a router plan, since router plans normally don't have one.