/*
 * compat/utils/walkers.c
 *
 * Ported from Apache Cloudberry (src/backend/optimizer/util/walkers.c).
 * Adapted for PostgreSQL 18 single-node mode: CBDB-only plan node types
 * (Motion, ShareInputScan, Sequence, PartitionSelector, Dynamic* scans,
 * SplitUpdate, AssertOp, Flow, TupleSplit, DQAExpr, RuntimeFilter,
 * TableFunctionScan, TableValueExpr) are omitted — they are never
 * generated by pg_orca and do not exist in PG18's NodeTag enum.
 *
 * Abstract base tags T_Plan, T_Scan, T_Join also do not exist in PG18
 * and are omitted accordingly.
 *
 * Ported functions:
 *   - extract_nodes_plan() and helpers
 *   - find_nodes()
 *   - check_collation() and helpers
 */

#include "postgres.h"

#include "miscadmin.h"
#include "nodes/nodeFuncs.h"
#include "nodes/plannodes.h"
#include "optimizer/optimizer.h"
#include "catalog/pg_collation.h"
#include "catalog/pg_type.h"

/* -----------------------------------------------------------------------
 * Internal types
 * ----------------------------------------------------------------------- */

/*
 * Minimal plan_tree_base_prefix: first field of every walker context.
 * extract_nodes_walker passes NULL as the base (no SubPlan lookup needed).
 */
typedef struct plan_tree_base_prefix
{
	Node	   *node;			/* PlannerInfo* or PlannedStmt*, may be NULL */
} plan_tree_base_prefix;

typedef struct extract_context
{
	plan_tree_base_prefix base; /* Required first field */
	List	   *nodes;
	NodeTag		nodeTag;
	bool		descendIntoSubqueries;
} extract_context;

/* -----------------------------------------------------------------------
 * Typed walker callback — avoids "no prototype" warnings
 * ----------------------------------------------------------------------- */
typedef bool (*node_walker_fn) (Node *node, void *context);

/* -----------------------------------------------------------------------
 * Forward declarations
 * ----------------------------------------------------------------------- */

static bool plan_tree_walker(Node *node, node_walker_fn walker,
							 void *context, bool recurse_into_subplans);
static bool extract_nodes_walker(Node *node, extract_context *context);

/* -----------------------------------------------------------------------
 * walk_plan_node_fields — walk the common Plan fields
 * ----------------------------------------------------------------------- */

static bool
walk_plan_node_fields(Plan *plan, node_walker_fn walker, void *context)
{
	if (walker((Node *) plan->targetlist, context))
		return true;
	if (walker((Node *) plan->qual, context))
		return true;
	if (walker((Node *) plan->lefttree, context))
		return true;
	if (walker((Node *) plan->righttree, context))
		return true;
	if (walker((Node *) plan->initPlan, context))
		return true;
	return false;
}

static bool
walk_scan_node_fields(Scan *scan, node_walker_fn walker, void *context)
{
	return walk_plan_node_fields((Plan *) scan, walker, context);
}

static bool
walk_join_node_fields(Join *join, node_walker_fn walker, void *context)
{
	if (walk_plan_node_fields((Plan *) join, walker, context))
		return true;
	return walker((Node *) join->joinqual, context);
}

/* -----------------------------------------------------------------------
 * plan_tree_walker — general Plan-tree walker for PG18
 *
 * CBDB-only node types and abstract base tags (T_Plan, T_Scan, T_Join)
 * are omitted; unknown tags fall through to expression_tree_walker.
 * ----------------------------------------------------------------------- */

static bool
plan_tree_walker(Node *node,
				 node_walker_fn walker,
				 void *context,
				 bool recurse_into_subplans)
{
	if (node == NULL)
		return false;

	check_stack_depth();

	switch (nodeTag(node))
	{
		case T_Result:
			if (walk_plan_node_fields((Plan *) node, walker, context))
				return true;
			if (walker((Node *) ((Result *) node)->resconstantqual, context))
				return true;
			break;

		case T_ProjectSet:
			if (walk_plan_node_fields((Plan *) node, walker, context))
				return true;
			break;

		case T_Append:
			if (walk_plan_node_fields((Plan *) node, walker, context))
				return true;
			if (walker((Node *) ((Append *) node)->appendplans, context))
				return true;
			break;

		case T_MergeAppend:
			if (walk_plan_node_fields((Plan *) node, walker, context))
				return true;
			if (walker((Node *) ((MergeAppend *) node)->mergeplans, context))
				return true;
			break;

		case T_RecursiveUnion:
			if (walk_plan_node_fields((Plan *) node, walker, context))
				return true;
			break;

		case T_BitmapAnd:
			if (walk_plan_node_fields((Plan *) node, walker, context))
				return true;
			if (walker((Node *) ((BitmapAnd *) node)->bitmapplans, context))
				return true;
			break;

		case T_BitmapOr:
			if (walk_plan_node_fields((Plan *) node, walker, context))
				return true;
			if (walker((Node *) ((BitmapOr *) node)->bitmapplans, context))
				return true;
			break;

		case T_SeqScan:
		case T_SampleScan:
		case T_BitmapHeapScan:
		case T_WorkTableScan:
		case T_NamedTuplestoreScan:
			if (walk_scan_node_fields((Scan *) node, walker, context))
				return true;
			break;

		case T_ForeignScan:
			if (walk_scan_node_fields((Scan *) node, walker, context))
				return true;
			if (walker((Node *) ((ForeignScan *) node)->fdw_exprs, context))
				return true;
			break;

		case T_CustomScan:
			if (walk_scan_node_fields((Scan *) node, walker, context))
				return true;
			if (walker((Node *) ((CustomScan *) node)->custom_plans, context))
				return true;
			if (walker((Node *) ((CustomScan *) node)->custom_exprs, context))
				return true;
			if (walker((Node *) ((CustomScan *) node)->custom_scan_tlist, context))
				return true;
			break;

		case T_ValuesScan:
			if (walker((Node *) ((ValuesScan *) node)->values_lists, context))
				return true;
			if (walk_scan_node_fields((Scan *) node, walker, context))
				return true;
			break;

		case T_TableFuncScan:
			if (walker((Node *) ((TableFuncScan *) node)->tablefunc, context))
				return true;
			if (walk_scan_node_fields((Scan *) node, walker, context))
				return true;
			break;

		case T_FunctionScan:
			if (walker((Node *) ((FunctionScan *) node)->functions, context))
				return true;
			if (walk_scan_node_fields((Scan *) node, walker, context))
				return true;
			break;

		case T_IndexScan:
			if (walk_scan_node_fields((Scan *) node, walker, context))
				return true;
			if (walker((Node *) ((IndexScan *) node)->indexqual, context))
				return true;
			break;

		case T_IndexOnlyScan:
			if (walk_scan_node_fields((Scan *) node, walker, context))
				return true;
			if (walker((Node *) ((IndexOnlyScan *) node)->indexqual, context))
				return true;
			break;

		case T_BitmapIndexScan:
			if (walk_scan_node_fields((Scan *) node, walker, context))
				return true;
			if (walker((Node *) ((BitmapIndexScan *) node)->indexqual, context))
				return true;
			break;

		case T_TidScan:
			if (walk_scan_node_fields((Scan *) node, walker, context))
				return true;
			if (walker((Node *) ((TidScan *) node)->tidquals, context))
				return true;
			break;

		case T_TidRangeScan:
			if (walk_scan_node_fields((Scan *) node, walker, context))
				return true;
			if (walker((Node *) ((TidRangeScan *) node)->tidrangequals, context))
				return true;
			break;

		case T_SubqueryScan:
			if (walk_scan_node_fields((Scan *) node, walker, context))
				return true;
			if (walker((Node *) ((SubqueryScan *) node)->subplan, context))
				return true;
			break;

		case T_NestLoop:
			if (walk_join_node_fields((Join *) node, walker, context))
				return true;
			break;

		case T_MergeJoin:
			if (walk_join_node_fields((Join *) node, walker, context))
				return true;
			if (walker((Node *) ((MergeJoin *) node)->mergeclauses, context))
				return true;
			break;

		case T_HashJoin:
			if (walk_join_node_fields((Join *) node, walker, context))
				return true;
			if (walker((Node *) ((HashJoin *) node)->hashclauses, context))
				return true;
			break;

		case T_Material:
			if (walk_plan_node_fields((Plan *) node, walker, context))
				return true;
			break;

		case T_Sort:
		case T_IncrementalSort:
			if (walk_plan_node_fields((Plan *) node, walker, context))
				return true;
			break;

		case T_Agg:
			if (walk_plan_node_fields((Plan *) node, walker, context))
				return true;
			break;

		case T_WindowAgg:
			if (walk_plan_node_fields((Plan *) node, walker, context))
				return true;
			if (walker(((WindowAgg *) node)->startOffset, context))
				return true;
			if (walker(((WindowAgg *) node)->endOffset, context))
				return true;
			break;

		case T_Unique:
			if (walk_plan_node_fields((Plan *) node, walker, context))
				return true;
			break;

		case T_Gather:
		case T_GatherMerge:
			if (walk_plan_node_fields((Plan *) node, walker, context))
				return true;
			break;

		case T_Hash:
			if (walk_plan_node_fields((Plan *) node, walker, context))
				return true;
			break;

		case T_SetOp:
			if (walk_plan_node_fields((Plan *) node, walker, context))
				return true;
			break;

		case T_Limit:
			if (walk_plan_node_fields((Plan *) node, walker, context))
				return true;
			if (walker((Node *) ((Limit *) node)->limitCount, context))
				return true;
			if (walker((Node *) ((Limit *) node)->limitOffset, context))
				return true;
			break;

		case T_LockRows:
			if (walk_plan_node_fields((Plan *) node, walker, context))
				return true;
			break;

		case T_ModifyTable:
			if (walk_plan_node_fields((Plan *) node, walker, context))
				return true;
			if (walker((Node *) ((ModifyTable *) node)->withCheckOptionLists, context))
				return true;
			if (walker((Node *) ((ModifyTable *) node)->onConflictSet, context))
				return true;
			if (walker((Node *) ((ModifyTable *) node)->onConflictWhere, context))
				return true;
			if (walker((Node *) ((ModifyTable *) node)->returningLists, context))
				return true;
			if (walker((Node *) ((ModifyTable *) node)->mergeActionLists, context))
				return true;
			break;

		case T_Memoize:
			if (walk_plan_node_fields((Plan *) node, walker, context))
				return true;
			if (walker((Node *) ((Memoize *) node)->param_exprs, context))
				return true;
			break;

		case T_CteScan:
			if (walk_scan_node_fields((Scan *) node, walker, context))
				return true;
			break;

		case T_SubPlan:
			{
				SubPlan    *subplan = (SubPlan *) node;

				if (expression_tree_walker((Node *) subplan->testexpr,
										   walker, context))
					return true;
				if (expression_tree_walker((Node *) subplan->args,
										   walker, context))
					return true;
				/*
				 * Do not recurse into the subplan's Plan tree here —
				 * extract_nodes_walker intentionally stops at SubPlan
				 * boundaries (matches Cloudberry behavior, see MPP-17168).
				 */
			}
			break;

		case T_IntList:
		case T_OidList:
			break;

		case T_Query:
			return query_tree_walker((Query *) node, walker, context, 0);

		default:
			return expression_tree_walker(node, walker, context);
	}

	return false;
}

/* -----------------------------------------------------------------------
 * extract_nodes_walker — internal recursive walker
 * ----------------------------------------------------------------------- */

static bool
extract_nodes_walker(Node *node, extract_context *context)
{
	if (node == NULL)
		return false;

	if (nodeTag(node) == context->nodeTag)
		context->nodes = lappend(context->nodes, node);

	if (nodeTag(node) == T_SubPlan)
	{
		SubPlan    *subplan = (SubPlan *) node;

		/*
		 * Walk SubPlan expressions even when !descendIntoSubqueries.
		 * Do not descend into the subplan's Plan tree (see MPP-17168).
		 */
		if (extract_nodes_walker((Node *) subplan->testexpr, context))
			return true;
		if (expression_tree_walker((Node *) subplan->args,
								   (node_walker_fn) extract_nodes_walker,
								   context))
			return true;
		return false;
	}

	if (nodeTag(node) == T_SubqueryScan && !context->descendIntoSubqueries)
		return false;

	return plan_tree_walker(node,
							(node_walker_fn) extract_nodes_walker,
							(void *) context,
							true);
}

/* -----------------------------------------------------------------------
 * find_nodes — public entry point
 *
 * Walk an expression/query tree rooted at 'node' and return the 0-based
 * index of the first node whose NodeTag matches any tag in 'nodeTags', or
 * -1 if no match is found.
 *
 * Ported from Cloudberry src/backend/optimizer/util/walkers.c.
 * ----------------------------------------------------------------------- */

typedef struct find_nodes_context
{
	List	   *nodeTags;
	int			foundNode;
} find_nodes_context;

static bool find_nodes_walker(Node *node, find_nodes_context *context);

int
find_nodes(Node *node, List *nodeTags)
{
	find_nodes_context context;

	Assert(NULL != node);
	context.nodeTags = nodeTags;
	context.foundNode = -1;
	find_nodes_walker(node, &context);
	return context.foundNode;
}

static bool
find_nodes_walker(Node *node, find_nodes_context *context)
{
	if (NULL == node)
		return false;

	if (IsA(node, Query))
	{
		/* Recurse into subselects */
		return query_tree_walker((Query *) node, find_nodes_walker,
								 (void *) context, 0 /* flags */);
	}

	{
		ListCell   *lc;
		int			i = 0;

		foreach(lc, context->nodeTags)
		{
			NodeTag		tag = (NodeTag) lfirst_int(lc);

			if (nodeTag(node) == tag)
			{
				context->foundNode = i;
				return true;
			}
			i++;
		}
	}

	return expression_tree_walker(node, find_nodes_walker, (void *) context);
}

/* -----------------------------------------------------------------------
 * extract_nodes_plan — public entry point
 *
 * Walk a Plan tree rooted at 'pl' and collect all nodes whose NodeTag
 * matches 'nodeTag'.  If 'descendIntoSubqueries' is false, SubqueryScan
 * subtrees are skipped.
 * ----------------------------------------------------------------------- */

List *
extract_nodes_plan(Plan *pl, int nodeTag, bool descendIntoSubqueries)
{
	extract_context context;

	Assert(pl);
	context.base.node = NULL;
	context.nodes = NIL;
	context.nodeTag = (NodeTag) nodeTag;
	context.descendIntoSubqueries = descendIntoSubqueries;
	extract_nodes_walker((Node *) pl, &context);
	return context.nodes;
}

/* -----------------------------------------------------------------------
 * check_collation — detect non-default collations in an expression tree
 *
 * Returns:
 *   -1  no nodes examined (node was NULL — should not happen in practice)
 *    0  all collations are default / invalid
 *    1  at least one non-default collation found
 *
 * Ported from Cloudberry src/backend/optimizer/util/walkers.c.
 * CBDB-only node types (T_TableValueExpr, T_DMLActionExpr) are omitted.
 * ----------------------------------------------------------------------- */

typedef struct check_collation_context
{
	int			foundNonDefaultCollation;
} check_collation_context;

static bool check_collation_walker(Node *node, check_collation_context *context);

static void
check_collation_in_list(List *colllist, check_collation_context *context)
{
	ListCell   *lc;

	foreach(lc, colllist)
	{
		Oid			coll = lfirst_oid(lc);

		if (InvalidOid != coll && DEFAULT_COLLATION_OID != coll)
		{
			context->foundNonDefaultCollation = 1;
			break;
		}
	}
}

static bool
check_collation_walker(Node *node, check_collation_context *context)
{
	Oid			collation,
				inputCollation,
				type;

	if (NULL == node)
		return false;

	if (IsA(node, Query))
	{
		/* Recurse into subselects */
		return query_tree_walker((Query *) node, check_collation_walker,
								 (void *) context, 0 /* flags */);
	}

	switch (nodeTag(node))
	{
		case T_Var:
		case T_Const:
		case T_OpExpr:
			type = exprType(node);
			collation = exprCollation(node);
			if (type == NAMEOID || type == NAMEARRAYOID)
			{
				if (collation != C_COLLATION_OID)
					context->foundNonDefaultCollation = 1;
			}
			else if (InvalidOid != collation && DEFAULT_COLLATION_OID != collation)
			{
				context->foundNonDefaultCollation = 1;
			}
			break;

		case T_ScalarArrayOpExpr:
		case T_DistinctExpr:
		case T_BoolExpr:
		case T_BooleanTest:
		case T_CaseExpr:
		case T_CaseTestExpr:
		case T_CoalesceExpr:
		case T_MinMaxExpr:
		case T_FuncExpr:
		case T_Aggref:
		case T_WindowFunc:
		case T_NullTest:
		case T_NullIfExpr:
		case T_RelabelType:
		case T_CoerceToDomain:
		case T_CoerceViaIO:
		case T_ArrayCoerceExpr:
		case T_SubLink:
		case T_ArrayExpr:
		case T_SubscriptingRef:
		case T_RowExpr:
		case T_RowCompareExpr:
		case T_FieldSelect:
		case T_FieldStore:
		case T_CoerceToDomainValue:
		case T_CurrentOfExpr:
		case T_NamedArgExpr:
		case T_ConvertRowtypeExpr:
		case T_CollateExpr:
		case T_XmlExpr:
		case T_SetToDefault:
		case T_PlaceHolderVar:
		case T_Param:
		case T_SubPlan:
		case T_AlternativeSubPlan:
		case T_GroupingFunc:
			collation = exprCollation(node);
			inputCollation = exprInputCollation(node);
			if ((InvalidOid != collation && DEFAULT_COLLATION_OID != collation) ||
				(InvalidOid != inputCollation && DEFAULT_COLLATION_OID != inputCollation))
			{
				context->foundNonDefaultCollation = 1;
			}
			break;

		case T_CollateClause:
			/* unsupported */
			context->foundNonDefaultCollation = 1;
			break;

		case T_RangeTblEntry:
			check_collation_in_list(((RangeTblEntry *) node)->colcollations, context);
			break;

		case T_RangeTblFunction:
			check_collation_in_list(((RangeTblFunction *) node)->funccolcollations, context);
			break;

		case T_CommonTableExpr:
			check_collation_in_list(((CommonTableExpr *) node)->ctecolcollations, context);
			break;

		case T_SetOperationStmt:
			check_collation_in_list(((SetOperationStmt *) node)->colCollations, context);
			break;

		default:
			/* make compiler happy */
			break;
	}

	if (context->foundNonDefaultCollation == 1)
		return true;			/* end recursion */

	return expression_tree_walker(node, check_collation_walker, (void *) context);
}

int
check_collation(Node *node)
{
	check_collation_context context;

	Assert(NULL != node);
	context.foundNonDefaultCollation = -1;
	check_collation_walker(node, &context);
	return context.foundNonDefaultCollation;
}