provsql.c

Go to the documentation of this file.

/**
 * @file provsql.c
 * @brief PostgreSQL planner hook for transparent provenance tracking.
 *
 * This file installs a @c planner_hook that intercepts every SELECT query
 * and rewrites it to propagate a provenance circuit token (UUID) alongside
 * normal result tuples.  The rewriting proceeds in three conceptual phases:
 *
 *  -# **Discovery** – scan the range table for relations/subqueries that
 *     already carry a @c provsql UUID column (@c get_provenance_attributes).
 *  -# **Expression building** – combine the discovered tokens according
 *     to the semiring operation that corresponds to the SQL operator in use
 *     (⊗ for joins, ⊕ for duplicate elimination, ⊖ for EXCEPT) and wrap
 *     aggregations (@c make_provenance_expression,
 *     @c make_aggregation_expression).
 *  -# **Splice** – append the resulting provenance expression to the target
 *     list and replace any explicit @c provenance() call in the query with
 *     the computed expression (@c add_to_select,
 *     @c replace_provenance_function_by_expression).
 */
#include "postgres.h"
#include "fmgr.h"
#include "miscadmin.h"
#include "pg_config.h"
#include "access/htup_details.h"
#include "access/sysattr.h"
#include "catalog/pg_aggregate.h"
#include "catalog/pg_class.h"           /* RELKIND_VIEW */
#include "catalog/pg_collation.h"
#include "catalog/pg_operator.h"
#include "catalog/pg_proc.h"
#include "catalog/pg_type.h"
#include "nodes/makefuncs.h"
#include "nodes/nodeFuncs.h"
#include "nodes/print.h"
#include "executor/executor.h"
#include "optimizer/planner.h"
#include "parser/parse_coerce.h"
#include "parser/parse_oper.h"
#include "utils/builtins.h"
#include "parser/parsetree.h"
#include "storage/lwlock.h"
#include "storage/shmem.h"
#include "utils/fmgroids.h"
#include "utils/guc.h"
#include "utils/lsyscache.h"
#include "utils/ruleutils.h"
#include "utils/syscache.h"
#include "catalog/namespace.h"
#include "catalog/pg_cast.h"
#include "commands/createas.h"
#include "executor/spi.h"
#include "tcop/utility.h"
#include <time.h>
 
#include "classify_query.h"
#include "provsql_mmap.h"
#include "provsql_shmem.h"
#include "provsql_utils.h"
#include "safe_query.h"
 
#if PG_VERSION_NUM < 100000
#error "ProvSQL requires PostgreSQL version 10 or later"
#endif
 
#include "compatibility.h"
 
PG_MODULE_MAGIC; ///< Required PostgreSQL extension magic block
 
/* -------------------------------------------------------------------------
 * Global state & forward declarations
 * ------------------------------------------------------------------------- */
 
bool provsql_interrupted = false;
static bool provsql_active = true; ///< @c true while ProvSQL query rewriting is enabled
bool provsql_where_provenance = false;
static bool provsql_update_provenance = false; ///< @c true when provenance tracking for DML is enabled
int provsql_verbose = 100; ///< Verbosity level; controlled by the @c provsql.verbose_level GUC
bool provsql_aggtoken_text_as_uuid = false; ///< When @c true, @c agg_token::text emits the underlying provenance UUID instead of @c "value (*)"
char *provsql_tool_search_path = NULL; ///< Colon-separated directory list prepended to @c PATH when invoking external tools (d4, c2d, minic2d, dsharp, weightmc, graph-easy); controlled by the @c provsql.tool_search_path GUC
int provsql_monte_carlo_seed = -1; ///< Seed for the Monte Carlo sampler; -1 means non-deterministic (std::random_device); controlled by the @c provsql.monte_carlo_seed GUC
int provsql_rv_mc_samples = 10000; ///< Default sample count for analytical-evaluator MC fallbacks; 0 disables fallback (callers raise instead); controlled by the @c provsql.rv_mc_samples GUC
bool provsql_simplify_on_load = true; ///< Run universal cmp-resolution passes when @c getGenericCircuit returns; controlled by the @c provsql.simplify_on_load GUC
bool provsql_hybrid_evaluation = true; ///< Run the hybrid-evaluator simplifier inside @c probability_evaluate; controlled by the @c provsql.hybrid_evaluation GUC
bool provsql_cmp_probability_evaluation = true; ///< Run closed-form / analytic probability evaluators for @c gate_cmps inside @c probability_evaluate (currently the Poisson-binomial pre-pass for HAVING-COUNT; future MIN / MAX / SUM evaluators will gate on the same GUC); controlled by the @c provsql.cmp_probability_evaluation GUC
bool provsql_boolean_provenance = false; ///< Opt-in safe-query optimisation: when @c true, rewrites hierarchical conjunctive queries to a read-once form whose probability is computable in linear time. The resulting circuit is tagged so that semiring evaluations admitting no homomorphism from Boolean functions refuse to run on it. Controlled by the @c provsql.boolean_provenance GUC.
 
 
extern void _PG_init(void);
extern void _PG_fini(void);
 
static planner_hook_type prev_planner = NULL; ///< Previous planner hook (chained)
 
static Query *process_query(const constants_t *constants, Query *q,
                            bool **removed, bool wrap_root);
static Expr *wrap_in_assume_boolean(const constants_t *constants, Expr *expr);
 
/* -------------------------------------------------------------------------
 * Provenance attribute construction
 * ------------------------------------------------------------------------- */
 
/**
 * @brief Build a Var node that references the provenance column of a relation.
 *
 * Creates a @c Var pointing to attribute @p attid of range-table entry
 * @p relid, typed as UUID, and marks the column as selected in the
 * permission bitmap so PostgreSQL grants access correctly.
 *
 * @param constants  Extension OID cache.
 * @param q          Owning query (needed to update permission info on PG 16+).
 * @param r          Range-table entry that owns the provenance column.
 * @param relid      1-based index of @p r in @p q->rtable.
 * @param attid      1-based attribute number of the provenance column in @p r.
 * @return  A freshly allocated @c Var node.
 */
static Var *make_provenance_attribute(const constants_t *constants, Query *q,
                                      RangeTblEntry *r, Index relid,
                                      AttrNumber attid) {
  Var *v = makeNode(Var);
 
  v->varno = relid;
  v->varattno = attid;
 
#if PG_VERSION_NUM >= 130000
  v->varnosyn = relid;
  v->varattnosyn = attid;
#else
  v->varnoold = relid;
  v->varoattno = attid;
#endif
 
  v->vartype = constants->OID_TYPE_UUID;
  v->varcollid = InvalidOid;
  v->vartypmod = -1;
  v->location = -1;
 
#if PG_VERSION_NUM >= 160000
  if (r->perminfoindex != 0) {
    RTEPermissionInfo *rpi =
      list_nth_node(RTEPermissionInfo, q->rteperminfos, r->perminfoindex - 1);
    rpi->selectedCols = bms_add_member(
      rpi->selectedCols, attid - FirstLowInvalidHeapAttributeNumber);
  }
#else
  r->selectedCols = bms_add_member(r->selectedCols,
                                   attid - FirstLowInvalidHeapAttributeNumber);
#endif
 
  return v;
}
 
/* -------------------------------------------------------------------------
 * Helper mutators: attribute-number fixup and type patching
 * ------------------------------------------------------------------------- */
 
/** @brief Context for the @c reduce_varattno_mutator tree walker. */
typedef struct reduce_varattno_mutator_context {
  Index varno;  ///< Range-table entry whose attribute numbers are being adjusted
  int *offset;  ///< Per-attribute cumulative shift to apply
} reduce_varattno_mutator_context;
 
/**
 * @brief Tree-mutator callback that adjusts Var attribute numbers.
 * @param node  Current expression tree node.
 * @param ctx   Pointer to a @c reduce_varattno_mutator_context.
 * @return      Possibly modified node.
 */
static Node *reduce_varattno_mutator(Node *node, void *ctx) {
  reduce_varattno_mutator_context *context = (reduce_varattno_mutator_context *)ctx;
  if (node == NULL)
    return NULL;
 
  if (IsA(node, Var)) {
    Var *v = (Var *)node;
 
    if (v->varno == context->varno) {
      v->varattno += context->offset[v->varattno - 1];
    }
  }
 
  return expression_tree_mutator(node, reduce_varattno_mutator, ctx);
}
 
/**
 * @brief Adjust Var attribute numbers in @p targetList after columns are removed.
 *
 * When provenance columns are stripped from a subquery's target list, the
 * remaining columns shift left.  This function applies a pre-computed
 * @p offset array (one entry per original column) to correct all @c Var
 * nodes that reference range-table entry @p varno.
 *
 * @param targetList  Target list of the outer query to patch.
 * @param varno       Range-table entry whose attribute numbers need fixing.
 * @param offset      Cumulative shift per original attribute (negative or zero).
 */
static void reduce_varattno_by_offset(List *targetList, Index varno,
                                      int *offset) {
  ListCell *lc;
  reduce_varattno_mutator_context context = {varno, offset};
 
  foreach (lc, targetList) {
    Node *te = lfirst(lc);
    expression_tree_mutator(te, reduce_varattno_mutator, &context);
  }
}
 
/** @brief Context for the @c aggregation_type_mutator tree walker. */
typedef struct aggregation_type_mutator_context {
  Index varno;                ///< Range-table entry index of the aggregate var
  Index varattno;             ///< Attribute number of the aggregate column
  const constants_t *constants; ///< Extension OID cache
} aggregation_type_mutator_context;
 
/**
 * @brief Check if a Var matches the target aggregate column.
 */
static bool is_target_agg_var(Node *node,
                              aggregation_type_mutator_context *context) {
  if (IsA(node, Var)) {
    Var *v = (Var *)node;
    return v->varno == context->varno && v->varattno == context->varattno;
  }
  return false;
}
 
/**
 * @brief Tree-mutator that retypes a specific Var to @c agg_token.
 *
 * When the target Var is inside a cast FuncExpr, replaces the cast
 * function with the equivalent agg_token→target cast from pg_cast.
 * When the Var appears bare (e.g. in a TargetEntry for display), it is
 * retyped to agg_token directly.  In all other contexts (arithmetic,
 * window functions, etc.), wraps the Var in an explicit agg_token→original
 * cast so that parent nodes receive the expected type.
 *
 * @param node  Current expression tree node.
 * @param ctx   Pointer to an @c aggregation_type_mutator_context (varno,
 *              varattno, and constants).
 * @return      Possibly modified node.
 */
static Node *
aggregation_type_mutator(Node *node, void *ctx) {
  aggregation_type_mutator_context *context = (aggregation_type_mutator_context *)ctx;
  if (node == NULL)
    return NULL;
 
  if (IsA(node, FuncExpr)) {
    FuncExpr *f = (FuncExpr *)node;
 
    /* Check if this is a cast wrapping our target Var */
    if (list_length(f->args) == 1 &&
        is_target_agg_var(linitial(f->args), context)) {
      /* Look up the cast from agg_token to the target type */
      HeapTuple castTuple = SearchSysCache2(CASTSOURCETARGET,
                                            ObjectIdGetDatum(context->constants->OID_TYPE_AGG_TOKEN),
                                            ObjectIdGetDatum(f->funcresulttype));
 
      if (HeapTupleIsValid(castTuple)) {
        Form_pg_cast castForm = (Form_pg_cast) GETSTRUCT(castTuple);
        if (OidIsValid(castForm->castfunc)) {
          f->funcid = castForm->castfunc;
        }
        ReleaseSysCache(castTuple);
      }
 
      /* Retype the Var inside */
      ((Var *)linitial(f->args))->vartype =
        context->constants->OID_TYPE_AGG_TOKEN;
 
      return (Node *)f;
    }
  }
 
  if (IsA(node, Var)) {
    Var *v = (Var *)node;
 
    if (v->varno == context->varno && v->varattno == context->varattno) {
      v->vartype = context->constants->OID_TYPE_AGG_TOKEN;
    }
  }
  return expression_tree_mutator(node, aggregation_type_mutator, ctx);
}
 
/**
 * @brief Retypes aggregation-result Vars in @p q from UUID to @c agg_token.
 *
 * After a subquery that contains @c provenance_aggregate is processed, its
 * result type is @c agg_token rather than plain UUID.  This mutator walks
 * the outer query and updates the type of every @c Var referencing that
 * result column so that subsequent type-checking passes correctly.
 *
 * @param constants   Extension OID cache.
 * @param q           Outer query to patch.
 * @param rteid       Range-table index of the subquery in @p q.
 * @param targetList  Target list of the subquery (to locate provenance_aggregate columns).
 */
static void fix_type_of_aggregation_result(const constants_t *constants,
                                           Query *q, Index rteid,
                                           List *targetList) {
  ListCell *lc;
  aggregation_type_mutator_context context = {0, 0, constants};
  Index attno = 1;
 
  foreach (lc, targetList) {
    TargetEntry *te = (TargetEntry *)lfirst(lc);
    if (IsA(te->expr, FuncExpr)) {
      FuncExpr *f = (FuncExpr *)te->expr;
 
      if (f->funcid == constants->OID_FUNCTION_PROVENANCE_AGGREGATE) {
        context.varno = rteid;
        context.varattno = attno;
        query_tree_mutator(q, aggregation_type_mutator, &context,
                           QTW_DONT_COPY_QUERY | QTW_IGNORE_RC_SUBQUERIES);
 
        /* Check if the retyped column is used in ORDER BY or GROUP BY */
        {
          ListCell *lc2;
          foreach (lc2, q->targetList) {
            TargetEntry *outer_te = (TargetEntry *)lfirst(lc2);
            if (IsA(outer_te->expr, Var)) {
              Var *v = (Var *)outer_te->expr;
              if (v->varno == rteid && v->varattno == attno &&
                  outer_te->ressortgroupref > 0) {
                ListCell *lc3;
                foreach (lc3, q->sortClause) {
                  SortGroupClause *sgc = (SortGroupClause *)lfirst(lc3);
                  if (sgc->tleSortGroupRef == outer_te->ressortgroupref)
                    provsql_error("ORDER BY on aggregate results from "
                                  "a subquery not supported");
                }
                foreach (lc3, q->groupClause) {
                  SortGroupClause *sgc = (SortGroupClause *)lfirst(lc3);
                  if (sgc->tleSortGroupRef == outer_te->ressortgroupref)
                    provsql_error("GROUP BY on aggregate results from "
                                  "a subquery not supported");
                }
              }
            }
          }
        }
      }
    }
    ++attno;
  }
}
 
/**
 * @brief Inline CTE references as subqueries within a query.
 *
 * Replaces each non-recursive RTE_CTE entry in @p rtable with an
 * RTE_SUBQUERY containing a copy of the CTE's query, looking up
 * definitions in @p cteList.  Recurses into newly inlined subqueries
 * to handle nested CTE references (ctelevelsup > 0).
 *
 * @param rtable   Range table to scan for RTE_CTE entries.
 * @param cteList  CTE definitions to look up names in.
 */
static void inline_ctes_in_rtable(List *rtable, List *cteList) {
  ListCell *lc;
  foreach (lc, rtable) {
    RangeTblEntry *r = (RangeTblEntry *)lfirst(lc);
    if (r->rtekind == RTE_CTE) {
      ListCell *lc2;
      foreach (lc2, cteList) {
        CommonTableExpr *cte = (CommonTableExpr *)lfirst(lc2);
        if (strcmp(cte->ctename, r->ctename) == 0) {
          if (cte->cterecursive) {
            provsql_error("Recursive CTEs not supported");
          } else {
            r->rtekind = RTE_SUBQUERY;
            r->subquery = copyObject((Query *)cte->ctequery);
            r->ctename = NULL;
            r->ctelevelsup = 0;
            /* Recurse: the inlined subquery may reference other CTEs
             * from the same cteList */
            inline_ctes_in_rtable(r->subquery->rtable, cteList);
          }
          break;
        }
      }
    } else if (r->rtekind == RTE_SUBQUERY && r->subquery != NULL) {
      /* Recurse into existing subqueries (e.g., UNION branches) to
       * inline CTE references they may contain */
      inline_ctes_in_rtable(r->subquery->rtable, cteList);
    }
  }
}
 
/**
 * @brief Inline all CTE references in @p q as subqueries.
 */
static void inline_ctes(Query *q) {
  if (q->cteList == NIL)
    return;
  inline_ctes_in_rtable(q->rtable, q->cteList);
  q->cteList = NIL;
}
 
/**
 * @brief Collect all provenance Var nodes reachable from @p q's range table.
 *
 * Walks every RTE in @p q->rtable:
 * - @c RTE_RELATION: looks for a column named @c provsql of type UUID.
 * - @c RTE_SUBQUERY: recursively calls @c process_query and splices the
 *   resulting provenance column back into the parent's column list, also
 *   patching outer Var attribute numbers if inner columns were removed.
 * - @c RTE_CTE: non-recursive CTEs are inlined as @c RTE_SUBQUERY before
 *   the main loop, then processed as above.  Recursive CTEs raise an error.
 * - @c RTE_FUNCTION: handled when the function returns a single UUID column
 *   named @c provsql.
 * - @c RTE_JOIN / @c RTE_VALUES / @c RTE_GROUP: handled passively (the
 *   underlying base-table RTEs supply the tokens).
 *
 * @param constants  Extension OID cache.
 * @param q          Query whose range table is scanned (subquery RTEs are
 *                   modified in place by the recursive call).
 * @return  List of @c Var nodes, one per provenance source; @c NIL if the
 *          query has no provenance-bearing relation.
 */
static List *get_provenance_attributes(const constants_t *constants, Query *q) {
  List *prov_atts = NIL;
 
  for(Index rteid = 1; rteid <= q->rtable->length; ++rteid) {
    RangeTblEntry *r = list_nth_node(RangeTblEntry, q->rtable, rteid-1);
 
    if (r->rtekind == RTE_RELATION) {
      ListCell *lc;
      AttrNumber attid = 1;
 
      /* PG 14 and 15 leave the OLD/NEW rule-placeholder RTEs (relkind
       * = RELKIND_VIEW, inFromCl = false) in the rewritten range table
       * for any view body.  PG 16+ removes them.  They are never
       * scanned and the planner does not build a RelOptInfo for them,
       * so any Var we point at them later fails find_base_rel().
       * Filter them out here; any post-rewrite RTE_RELATION whose
       * relkind is still a view is one of these artifacts. */
      if (r->relkind == RELKIND_VIEW)
        continue;
 
      foreach (lc, r->eref->colnames) {
        const char *v = strVal(lfirst(lc));
 
        if (!strcmp(v, PROVSQL_COLUMN_NAME) &&
            get_atttype(r->relid, attid) == constants->OID_TYPE_UUID) {
          prov_atts =
            lappend(prov_atts,
                    make_provenance_attribute(constants, q, r, rteid, attid));
        }
 
        ++attid;
      }
    } else if (r->rtekind == RTE_SUBQUERY) {
      bool *inner_removed = NULL;
      int old_targetlist_length =
        r->subquery->targetList ? r->subquery->targetList->length : 0;
      Query *new_subquery =
        process_query(constants, r->subquery, &inner_removed, false);
      if (new_subquery != NULL) {
        int i = 0;
        int *offset = (int *)palloc(old_targetlist_length * sizeof(int));
        unsigned varattnoprovsql;
        ListCell *cell, *prev;
 
        r->subquery = new_subquery;
 
        if (inner_removed != NULL) {
          for (cell = list_head(r->eref->colnames), prev = NULL;
               cell != NULL;) {
            if (inner_removed[i]) {
              r->eref->colnames =
                my_list_delete_cell(r->eref->colnames, cell, prev);
              if (prev)
                cell = my_lnext(r->eref->colnames, prev);
              else
                cell = list_head(r->eref->colnames);
            } else {
              prev = cell;
              cell = my_lnext(r->eref->colnames, cell);
            }
            ++i;
          }
          for (i = 0; i < old_targetlist_length; ++i) {
            offset[i] =
              (i == 0 ? 0 : offset[i - 1]) - (inner_removed[i] ? 1 : 0);
          }
 
          reduce_varattno_by_offset(q->targetList, rteid, offset);
        }
 
        varattnoprovsql = 0;
        for (cell = list_head(new_subquery->targetList); cell != NULL;
             cell = my_lnext(new_subquery->targetList, cell)) {
          TargetEntry *te = (TargetEntry *)lfirst(cell);
          ++varattnoprovsql;
          if (te->resname && !strcmp(te->resname, PROVSQL_COLUMN_NAME))
            break;
        }
 
        if (cell != NULL) {
          r->eref->colnames = list_insert_nth(r->eref->colnames, varattnoprovsql-1,
                                              makeString(pstrdup(PROVSQL_COLUMN_NAME)));
          prov_atts =
            lappend(prov_atts, make_provenance_attribute(
                      constants, q, r, rteid, varattnoprovsql));
        }
        fix_type_of_aggregation_result(constants, q, rteid,
                                       r->subquery->targetList);
      }
    } else if (r->rtekind == RTE_JOIN) {
      if (r->jointype == JOIN_INNER || r->jointype == JOIN_LEFT ||
          r->jointype == JOIN_FULL || r->jointype == JOIN_RIGHT) {
        // Nothing to do, there will also be RTE entries for the tables
        // that are part of the join, from which we will extract the
        // provenance information
      } else { // Semijoin (should be feasible, but check whether the second
               // provenance information is available) Antijoin (feasible with
               // negation)
        provsql_error("JOIN type not supported");
      }
    } else if (r->rtekind == RTE_FUNCTION) {
      ListCell *lc;
      AttrNumber attid = 1;
 
      foreach (lc, r->functions) {
        RangeTblFunction *func = (RangeTblFunction *)lfirst(lc);
 
        if (func->funccolcount == 1) {
          FuncExpr *expr = (FuncExpr *)func->funcexpr;
          if (expr->funcresulttype == constants->OID_TYPE_UUID &&
              !strcmp(get_rte_attribute_name(r, attid), PROVSQL_COLUMN_NAME)) {
            prov_atts = lappend(prov_atts, make_provenance_attribute(
                                  constants, q, r, rteid, attid));
          }
        } else {
          provsql_error("FROM function with multiple output "
                        "attributes not supported");
        }
 
        attid += func->funccolcount;
      }
    } else if (r->rtekind == RTE_VALUES) {
      // Nothing to do, no provenance attribute in literal values
#if PG_VERSION_NUM >= 180000
    } else if (r->rtekind == RTE_GROUP) {
      // Introduced in PostgreSQL 18, we already handle group by from
      // groupClause
#endif
    } else {
      provsql_error("FROM clause not supported");
    }
  }
 
  return prov_atts;
}
 
/* -------------------------------------------------------------------------
 * Target-list surgery
 * ------------------------------------------------------------------------- */
 
/**
 * @brief Strip provenance UUID columns from @p q's SELECT list.
 *
 * Scans the target list and removes every @c Var entry whose column name is
 * @c provsql and whose type is UUID.  The remaining entries have their
 * @c resno values decremented to fill the gaps.
 *
 * @param constants  Extension OID cache.
 * @param q          Query to modify in place.
 * @param removed    Out-param: allocated boolean array (length =
 *                   original target list length) where @c true means the
 *                   corresponding entry was removed.  The caller must
 *                   @c pfree this array when done.
 * @return  Bitmapset of @c ressortgroupref values whose entries were
 *          removed (so the caller can clean up GROUP BY / ORDER BY).
 */
static Bitmapset *
remove_provenance_attributes_select(const constants_t *constants, Query *q,
                                    bool **removed) {
  int nbRemoved = 0;
  int i = 0;
  Bitmapset *ressortgrouprefs = NULL;
  ListCell *cell, *prev;
  *removed = (bool *)palloc(q->targetList->length * sizeof(bool));
 
  for (cell = list_head(q->targetList), prev = NULL; cell != NULL;) {
    TargetEntry *rt = (TargetEntry *)lfirst(cell);
    (*removed)[i] = false;
 
    if (rt->expr->type == T_Var) {
      Var *v = (Var *)rt->expr;
 
      if (v->vartype == constants->OID_TYPE_UUID) {
        const char *colname;
 
        if (rt->resname)
          colname = rt->resname;
        else {
          /* This case occurs, for example, when grouping by a column
           * that is projected out */
          RangeTblEntry *r = (RangeTblEntry *)list_nth(q->rtable, v->varno - 1);
          colname = strVal(list_nth(r->eref->colnames, v->varattno - 1));
        }
 
        if (!strcmp(colname, PROVSQL_COLUMN_NAME)) {
          q->targetList = my_list_delete_cell(q->targetList, cell, prev);
 
          (*removed)[i] = true;
          ++nbRemoved;
 
          if (rt->ressortgroupref > 0)
            ressortgrouprefs =
              bms_add_member(ressortgrouprefs, rt->ressortgroupref);
        }
      }
    }
 
    if ((*removed)[i]) {
      if (prev) {
        cell = my_lnext(q->targetList, prev);
      } else {
        cell = list_head(q->targetList);
      }
    } else {
      rt->resno -= nbRemoved;
      prev = cell;
      cell = my_lnext(q->targetList, cell);
    }
 
    ++i;
  }
 
  return ressortgrouprefs;
}
 
/**
 * @brief Semiring operation used to combine provenance tokens.
 *
 * @c SR_TIMES corresponds to the multiplicative operation (joins, Cartesian
 * products), @c SR_PLUS to the additive operation (duplicate elimination), and
 * @c SR_MONUS to the monus / set-difference operation (EXCEPT).
 *
 * @see https://provsql.org/lean-docs/Provenance/QueryRewriting.html
 *      Lean 4 formalization of rewriting rules (R1)--(R5) and correctness
 *      theorem @c Query.rewriting_valid.
 */
typedef enum {
  SR_PLUS,  ///< Semiring addition (UNION, SELECT DISTINCT)
  SR_MONUS, ///< Semiring monus / set difference (EXCEPT)
  SR_TIMES  ///< Semiring multiplication (JOIN, Cartesian product)
} semiring_operation;
 
/* -------------------------------------------------------------------------
 * Semiring expression builders
 * ------------------------------------------------------------------------- */
 
/**
 * @brief Wrap @p toExpr in a @c provenance_eq gate if @p fromOpExpr is an
 *        equality between two tracked columns.
 *
 * Used for where-provenance: each equijoin condition (and some WHERE
 * equalities) introduces an @c eq gate that records which attribute positions
 * were compared.  Because this function is also called for WHERE predicates,
 * it applies extra guards and silently returns @p toExpr unchanged when the
 * expression does not match the expected shape (both sides must be @c Var
 * nodes, possibly wrapped in a @c RelabelType).
 *
 * @param constants    Extension OID cache.
 * @param fromOpExpr   The equality @c OpExpr to inspect.
 * @param toExpr       Existing provenance expression to wrap.
 * @param columns      Per-RTE column-numbering array.  EQ gate positions
 *                     carry the same sequential-number caveat as PROJECT
 *                     gate positions (see @c build_column_map()); they are
 *                     only correct when each operand's RTE is either a join
 *                     RTE or a subquery, not a bare provenance-tracked base
 *                     table.
 * @return  @p toExpr wrapped in @c provenance_eq(toExpr, col1, col2), or
 *          @p toExpr unchanged if the shape is unsupported.
 */
static Expr *add_eq_from_OpExpr_to_Expr(const constants_t *constants,
                                        OpExpr *fromOpExpr, Expr *toExpr,
                                        int **columns) {
  Datum first_arg;
  Datum second_arg;
  FuncExpr *fc;
  Const *c1;
  Const *c2;
  Var *v1;
  Var *v2;
 
  if (my_lnext(fromOpExpr->args, list_head(fromOpExpr->args))) {
    /* Sometimes Var is nested within a RelabelType */
    if (IsA(linitial(fromOpExpr->args), Var)) {
      v1 = linitial(fromOpExpr->args);
    } else if (IsA(linitial(fromOpExpr->args), RelabelType)) {
      /* In the WHERE case it can be a Const */
      RelabelType *rt1 = linitial(fromOpExpr->args);
      if (IsA(rt1->arg, Var)) { /* Can be Param in the WHERE case */
        v1 = (Var *)rt1->arg;
      } else
        return toExpr;
    } else
      return toExpr;
    if (!columns[v1->varno - 1])
      return toExpr;
    first_arg = Int16GetDatum(columns[v1->varno - 1][v1->varattno - 1]);
 
    if (IsA(lsecond(fromOpExpr->args), Var)) {
      v2 = lsecond(fromOpExpr->args);
    } else if (IsA(lsecond(fromOpExpr->args), RelabelType)) {
      /* In the WHERE case it can be a Const */
      RelabelType *rt2 = lsecond(fromOpExpr->args);
      if (IsA(rt2->arg, Var)) { /* Can be Param in the WHERE case */
        v2 = (Var *)rt2->arg;
      } else
        return toExpr;
    } else
      return toExpr;
    if (!columns[v2->varno - 1])
      return toExpr;
    second_arg = Int16GetDatum(columns[v2->varno - 1][v2->varattno - 1]);
 
    fc = makeNode(FuncExpr);
    fc->funcid = constants->OID_FUNCTION_PROVENANCE_EQ;
    fc->funcvariadic = false;
    fc->funcresulttype = constants->OID_TYPE_UUID;
    fc->location = -1;
 
    c1 = makeConst(constants->OID_TYPE_INT, -1, InvalidOid, sizeof(int16),
                   first_arg, false, true);
 
    c2 = makeConst(constants->OID_TYPE_INT, -1, InvalidOid, sizeof(int16),
                   second_arg, false, true);
 
    fc->args = list_make3(toExpr, c1, c2);
    return (Expr *)fc;
  }
  return toExpr;
}
 
/**
 * @brief Walk a join-condition or WHERE quals node and add @c eq gates for
 *        every equality it contains.
 *
 * Dispatches to @c add_eq_from_OpExpr_to_Expr for simple @c OpExpr nodes
 * and iterates over the arguments of an AND @c BoolExpr.  OR/NOT inside a
 * join ON clause are rejected with an error.
 *
 * @param constants  Extension OID cache.
 * @param quals      Root of the quals tree (@c OpExpr or @c BoolExpr), or
 *                   @c NULL (in which case @p result is returned unchanged).
 * @param result     Provenance expression to wrap.
 * @param columns    Per-RTE column-numbering array.
 * @return  Updated provenance expression with zero or more @c eq gates added.
 */
static Expr *add_eq_from_Quals_to_Expr(const constants_t *constants,
                                       Node *quals, Expr *result,
                                       int **columns) {
  OpExpr *oe;
 
  if (!quals)
    return result;
 
  if (IsA(quals, OpExpr)) {
    oe = (OpExpr *)quals;
    result = add_eq_from_OpExpr_to_Expr(constants, oe, result, columns);
  } /* Sometimes OpExpr is nested within a BoolExpr */
  else if (IsA(quals, BoolExpr)) {
    BoolExpr *be = (BoolExpr *)quals;
    /* In some cases, there can be an OR or a NOT specified with ON clause */
    if (be->boolop == OR_EXPR || be->boolop == NOT_EXPR) {
      provsql_error("Boolean operators OR and NOT in a join...on "
                    "clause are not supported");
    } else {
      ListCell *lc2;
      foreach (lc2, be->args) {
        if (IsA(lfirst(lc2), OpExpr)) {
          oe = (OpExpr *)lfirst(lc2);
          result = add_eq_from_OpExpr_to_Expr(constants, oe, result, columns);
        }
      }
    }
  } else { /* Handle other cases */
  }
  return result;
}
 
/**
 * @brief Build the per-row provenance token for an aggregate rewrite.
 *
 * Used by both @c make_aggregation_expression (for the agg_token /
 * @c provenance_semimod path) and @c make_rv_aggregate_expression (for
 * the inline RV-aggregate path).  Combines @p prov_atts via
 * @c provenance_times (under @c SR_TIMES) or @c provenance_monus
 * (under @c SR_MONUS); a single @c prov_att is returned as-is.
 *
 * @return  An @c Expr returning UUID; never @c NULL.
 */
static Expr *combine_prov_atts(const constants_t *constants,
                               List *prov_atts, semiring_operation op) {
  FuncExpr *combine;
 
  if (my_lnext(prov_atts, list_head(prov_atts)) == NULL)
    return (Expr *)linitial(prov_atts);
 
  combine = makeNode(FuncExpr);
  if (op == SR_TIMES) {
    ArrayExpr *array = makeNode(ArrayExpr);
 
    combine->funcid = constants->OID_FUNCTION_PROVENANCE_TIMES;
    combine->funcvariadic = true;
 
    array->array_typeid = constants->OID_TYPE_UUID_ARRAY;
    array->element_typeid = constants->OID_TYPE_UUID;
    array->elements = prov_atts;
    array->location = -1;
 
    combine->args = list_make1(array);
  } else { // SR_MONUS
    combine->funcid = constants->OID_FUNCTION_PROVENANCE_MONUS;
    combine->args = prov_atts;
  }
  combine->funcresulttype = constants->OID_TYPE_UUID;
  combine->location = -1;
  return (Expr *)combine;
}
 
/**
 * @brief Inline rewrite of an RV-returning aggregate into the same
 *        aggregate over provenance-wrapped per-row arguments.
 *
 * Originally Phase 1 of the SUM-over-RV story (see @c aggregation-of-rvs.md);
 * extended to any aggregate whose result type is @c random_variable
 * (e.g. @c provsql.sum, @c provsql.avg).  Replaces @c agg(@c x) with an
 * Aggref whose per-row argument is lifted through @c rv_aggregate_semimod
 * to attach the row's provenance: each row contributes
 * @c mixture(prov_token, X_i, as_random(0)).  The aggregate itself
 * (@c aggfnoid) is preserved verbatim, so its SFUNC / FFUNC decide what
 * gate shape to build from the per-row mixtures.  In particular:
 *  - @c sum(random_variable) collects the mixtures into a single
 *    @c gate_arith @c PLUS root, realising
 *    @f$\mathrm{SUM}(x) = \sum_i \mathbf{1}\{\varphi_i\} \cdot X_i@f$;
 *  - @c avg(random_variable) walks each mixture to recover @c prov_i
 *    and emits @c gate_arith(DIV, @c sum(mixture(p_i,x_i,0)),
 *    @c sum(mixture(p_i,1,0))), the natural lift of "AVG = SUM / COUNT"
 *    into the @c random_variable algebra.
 *
 * Routing happens at @c make_aggregation_expression on
 * @c agg_ref->aggtype @c == @c OID_TYPE_RANDOM_VARIABLE, so any future
 * RV-returning aggregate inherits the same per-row provenance wrap
 * without further C-side dispatch.
 *
 * SR_PLUS (UNION outer level) is handled by the caller; this builder
 * never runs for SR_PLUS.
 */
static Expr *make_rv_aggregate_expression(const constants_t *constants,
                                          Aggref *agg_ref, List *prov_atts,
                                          semiring_operation op) {
  Expr *prov_expr = combine_prov_atts(constants, prov_atts, op);
  Expr *rv_arg = ((TargetEntry *)linitial(agg_ref->args))->expr;
  FuncExpr *wrap;
  Aggref *new_agg;
  TargetEntry *te;
 
  /* Wrap the per-row RV in mixture(prov, rv, as_random(0)) via the
   * rv_aggregate_semimod SQL helper.  Going through the helper avoids
   * having to look up the specific (uuid, random_variable,
   * random_variable) overload of mixture and the (double precision)
   * overload of as_random at OID-cache time. */
  wrap = makeNode(FuncExpr);
  wrap->funcid = constants->OID_FUNCTION_RV_AGGREGATE_SEMIMOD;
  wrap->funcresulttype = constants->OID_TYPE_RANDOM_VARIABLE;
  wrap->args = list_make2(prov_expr, rv_arg);
  wrap->location = -1;
 
  /* Rebuild an Aggref calling the SAME aggregate (sum_rv, avg_rv, ...),
   * but with the argument now wrapped.  Inherit the original Aggref's
   * clause positioning; aggargtypes / aggtype stay random_variable. */
  te = makeNode(TargetEntry);
  te->resno = 1;
  te->expr = (Expr *)wrap;
 
  new_agg = makeNode(Aggref);
  new_agg->aggfnoid = agg_ref->aggfnoid;
  new_agg->aggtype = constants->OID_TYPE_RANDOM_VARIABLE;
  new_agg->aggargtypes = list_make1_oid(constants->OID_TYPE_RANDOM_VARIABLE);
  new_agg->aggkind = AGGKIND_NORMAL;
  new_agg->aggtranstype = InvalidOid;
  new_agg->args = list_make1(te);
  new_agg->location = agg_ref->location;
#if PG_VERSION_NUM >= 140000
  new_agg->aggno = new_agg->aggtransno = -1;
#endif
 
  return (Expr *)new_agg;
}
 
/**
 * @brief Build the provenance expression for a single aggregate function.
 *
 * For @c SR_PLUS (union context) returns the first provenance attribute
 * directly.  For @c SR_TIMES or @c SR_MONUS, constructs:
 * @code
 *   provenance_aggregate(fn_oid, result_type,
 *                        original_aggref,
 *                        array_agg(provenance_semimod(arg, times_or_monus_token)))
 * @endcode
 * COUNT(*) and COUNT(expr) are remapped to SUM so that the semimodule
 * semantics (scalar × token → token) work correctly.
 *
 * @param constants  Extension OID cache.
 * @param agg_ref    The original @c Aggref node from the query.
 * @param prov_atts  List of provenance @c Var nodes.
 * @param op         Semiring operation (determines how tokens are combined).
 * @return  Provenance expression of type @c agg_token.
 */
static Expr *make_aggregation_expression(const constants_t *constants,
                                         Aggref *agg_ref, List *prov_atts,
                                         semiring_operation op) {
  Expr *result;
  FuncExpr *expr, *expr_s;
  Aggref *agg = makeNode(Aggref);
  FuncExpr *plus = makeNode(FuncExpr);
  TargetEntry *te_inner = makeNode(TargetEntry);
  Const *fn = makeNode(Const);
  Const *typ = makeNode(Const);
 
  if (op == SR_PLUS) {
    result = linitial(prov_atts);
  } else {
    Oid aggregation_function = agg_ref->aggfnoid;
 
    /* Aggregates that return random_variable (sum_rv, avg_rv, and any
     * future RV-returning aggregate) get a different rewrite: instead
     * of going through provenance_semimod (which builds a gate_value
     * from CAST(val AS VARCHAR), nonsensical for an RV), each per-row
     * argument is wrapped in mixture(prov, rv, as_random(0)) and the
     * original aggregate's SFUNC / FFUNC decide what gate shape to
     * build from the resulting mixtures.  See aggregation-of-rvs.md. */
    if (OidIsValid(constants->OID_TYPE_RANDOM_VARIABLE) &&
        agg_ref->aggtype == constants->OID_TYPE_RANDOM_VARIABLE) {
      return make_rv_aggregate_expression(constants, agg_ref, prov_atts, op);
    }
 
    if (my_lnext(prov_atts, list_head(prov_atts)) == NULL)
      expr = linitial(prov_atts);
    else {
      expr = makeNode(FuncExpr);
      if (op == SR_TIMES) {
        ArrayExpr *array = makeNode(ArrayExpr);
 
        expr->funcid = constants->OID_FUNCTION_PROVENANCE_TIMES;
        expr->funcvariadic = true;
 
        array->array_typeid = constants->OID_TYPE_UUID_ARRAY;
        array->element_typeid = constants->OID_TYPE_UUID;
        array->elements = prov_atts;
        array->location = -1;
 
        expr->args = list_make1(array);
      } else { // SR_MONUS
        expr->funcid = constants->OID_FUNCTION_PROVENANCE_MONUS;
        expr->args = prov_atts;
      }
      expr->funcresulttype = constants->OID_TYPE_UUID;
      expr->location = -1;
    }
 
    // semimodule function
    expr_s = makeNode(FuncExpr);
    expr_s->funcid = constants->OID_FUNCTION_PROVENANCE_SEMIMOD;
    expr_s->funcresulttype = constants->OID_TYPE_UUID;
 
    // check the particular case of count
    if (aggregation_function == F_COUNT_ ||
        aggregation_function == F_COUNT_ANY) // count(*) or count(arg)
    {
      Const *one = makeConst(constants->OID_TYPE_INT, -1, InvalidOid,
                             sizeof(int32), Int32GetDatum(1), false, true);
      expr_s->args = list_make2(one, expr);
      aggregation_function = F_SUM_INT4;
    } else {
      expr_s->args =
        list_make2(((TargetEntry *)linitial(agg_ref->args))->expr, expr);
    }
 
    expr_s->location = -1;
 
    // aggregating all semirings in an array
    te_inner->resno = 1;
    te_inner->expr = (Expr *)expr_s;
    agg->aggfnoid = constants->OID_FUNCTION_ARRAY_AGG;
    agg->aggtype = constants->OID_TYPE_UUID_ARRAY;
    agg->args = list_make1(te_inner);
    agg->aggkind = AGGKIND_NORMAL;
    agg->location = -1;
#if PG_VERSION_NUM >= 140000
    agg->aggno = agg->aggtransno = -1;
#endif
 
    agg->aggargtypes = list_make1_oid(constants->OID_TYPE_UUID);
 
    // final aggregation function
    plus->funcid = constants->OID_FUNCTION_PROVENANCE_AGGREGATE;
 
    fn = makeConst(constants->OID_TYPE_INT, -1, InvalidOid, sizeof(int32),
                   Int32GetDatum(aggregation_function), false, true);
 
    typ = makeConst(constants->OID_TYPE_INT, -1, InvalidOid, sizeof(int32),
                    Int32GetDatum(agg_ref->aggtype), false, true);
 
    plus->funcresulttype = constants->OID_TYPE_AGG_TOKEN;
    plus->args = list_make4(fn, typ, agg_ref, agg);
    plus->location = -1;
 
    result = (Expr *)plus;
  }
 
  return result;
}
 
/* -------------------------------------------------------------------------
 * HAVING / WHERE-on-aggregates rewriting
 * ------------------------------------------------------------------------- */
 
/* Forward declaration needed because having_BoolExpr_to_provenance and
 * having_Expr_to_provenance_cmp are mutually recursive. */
static FuncExpr *having_Expr_to_provenance_cmp(Expr *expr, const constants_t *constants, bool negated);
 
/* Forward declaration: defined alongside the other tree walkers
 * further down in the file. */
static bool needs_having_lift(Node *havingQual, const constants_t *constants);
 
/**
 * @brief Convert a comparison @c OpExpr on aggregate results into a
 *        @c provenance_cmp gate expression.
 *
 * Each argument of @p opExpr must be one of:
 * - A @c Var of type @c agg_token (or a @c FuncExpr implicit-cast wrapper
 *   around one) → cast to UUID via @c agg_token_to_uuid.
 * - A scalar @c Const → wrapped in @c provenance_semimod(const, gate_one()).
 *
 * If @p negated is true the operator OID is replaced by its negator so that
 * NOT(a < b) becomes a >= b at the provenance level.
 *
 * @param opExpr     The comparison expression from the HAVING clause.
 * @param constants  Extension OID cache.
 * @param negated    Whether the expression appears under a NOT.
 * @return  A @c provenance_cmp(lhs, op_oid, rhs) @c FuncExpr.
 */
static FuncExpr *having_OpExpr_to_provenance_cmp(OpExpr *opExpr, const constants_t *constants, bool negated) {
  FuncExpr *cmpExpr;
  Node *arguments[2];
  Const *oid;
  Oid opno = opExpr->opno;
 
  for (unsigned i = 0; i < 2; ++i) {
    Node *node = (Node *)lfirst(list_nth_cell(opExpr->args, i));
 
    if (IsA(node, FuncExpr)) {
      FuncExpr *fe = (FuncExpr *)node;
      if (fe->funcformat == COERCE_IMPLICIT_CAST ||
          fe->funcformat == COERCE_EXPLICIT_CAST) {
        if (fe->args->length == 1)
          node = lfirst(list_head(fe->args));
      }
    }
 
    if (IsA(node, FuncExpr)) {
      FuncExpr *fe = (FuncExpr *)node;
      if (fe->funcid == constants->OID_FUNCTION_PROVENANCE_AGGREGATE) {
        // We need to add an explicit cast to UUID
        FuncExpr *castToUUID = makeNode(FuncExpr);
 
        castToUUID->funcid = constants->OID_FUNCTION_AGG_TOKEN_UUID;
        castToUUID->funcresulttype = constants->OID_TYPE_UUID;
        castToUUID->args = list_make1(fe);
        castToUUID->location = -1;
 
        arguments[i] = (Node *)castToUUID;
      } else {
        provsql_error("cannot handle complex HAVING expressions");
      }
    } else if (IsA(node, Var)) {
      Var *v = (Var *)node;
 
      if (v->vartype == constants->OID_TYPE_AGG_TOKEN) {
        // We need to add an explicit cast to UUID
        FuncExpr *castToUUID = makeNode(FuncExpr);
 
        castToUUID->funcid = constants->OID_FUNCTION_AGG_TOKEN_UUID;
        castToUUID->funcresulttype = constants->OID_TYPE_UUID;
        castToUUID->args = list_make1(v);
        castToUUID->location = -1;
 
        arguments[i] = (Node *)castToUUID;
      } else {
        provsql_error("cannot handle complex HAVING expressions");
      }
    } else if (IsA(node, Const)) {
      Const *literal = (Const *)node;
      FuncExpr *oneExpr, *semimodExpr;
 
      // gate_one() expression
      oneExpr = makeNode(FuncExpr);
      oneExpr->funcid = constants->OID_FUNCTION_GATE_ONE;
      oneExpr->funcresulttype = constants->OID_TYPE_UUID;
      oneExpr->args = NIL;
      oneExpr->location = -1;
 
      // provenance_semimod(literal, gate_one())
      semimodExpr = makeNode(FuncExpr);
      semimodExpr->funcid = constants->OID_FUNCTION_PROVENANCE_SEMIMOD;
      semimodExpr->funcresulttype = constants->OID_TYPE_UUID;
      semimodExpr->args = list_make2((Expr *)literal, (Expr *)oneExpr);
      semimodExpr->location = -1;
 
      arguments[i] = (Node *)semimodExpr;
    } else {
      provsql_error("cannot handle complex HAVING expressions");
    }
  }
 
  if (negated) {
    opno = get_negator(opno);
    if (!opno)
      provsql_error("Missing negator");
  }
 
  oid = makeConst(constants->OID_TYPE_INT, -1, InvalidOid, sizeof(int32),
                  Int32GetDatum(opno), false, true);
 
  cmpExpr = makeNode(FuncExpr);
  cmpExpr->funcid = constants->OID_FUNCTION_PROVENANCE_CMP;
  cmpExpr->funcresulttype = constants->OID_TYPE_UUID;
  cmpExpr->args = list_make3(arguments[0], oid, arguments[1]);
  cmpExpr->location = opExpr->location;
 
  return cmpExpr;
}
 
/**
 * @brief Convert a Boolean combination of HAVING comparisons into a
 *        @c provenance_times / @c provenance_plus gate expression.
 *
 * Applies De Morgan duality when @p negated is true: AND becomes
 * @c provenance_plus (OR) and vice-versa.  NOT is handled by flipping
 * @p negated and delegating to @c having_Expr_to_provenance_cmp.
 *
 * @param be         Boolean expression from the HAVING clause.
 * @param constants  Extension OID cache.
 * @param negated    Whether the expression appears under a NOT.
 * @return  A @c FuncExpr combining the sub-expressions.
 */
static FuncExpr *having_BoolExpr_to_provenance(BoolExpr *be, const constants_t *constants, bool negated) {
  if(be->boolop == NOT_EXPR) {
    Expr *expr = (Expr *) lfirst(list_head(be->args));
    return having_Expr_to_provenance_cmp(expr, constants, !negated);
  } else {
    FuncExpr *result;
    List *l = NULL;
    ListCell *lc;
    ArrayExpr *array = makeNode(ArrayExpr);
 
    array->array_typeid = constants->OID_TYPE_UUID_ARRAY;
    array->element_typeid = constants->OID_TYPE_UUID;
    array->location = -1;
 
    result = makeNode(FuncExpr);
    result->funcresulttype = constants->OID_TYPE_UUID;
    result->funcvariadic = true;
    result->location = be->location;
    result->args = list_make1(array);
 
    if ((be->boolop == AND_EXPR && !negated) || (be->boolop == OR_EXPR && negated))
      result->funcid = constants->OID_FUNCTION_PROVENANCE_TIMES;
    else if ((be->boolop == AND_EXPR && negated) || (be->boolop == OR_EXPR && !negated))
      result->funcid = constants->OID_FUNCTION_PROVENANCE_PLUS;
    else
      provsql_error("Unknown Boolean operator");
 
    foreach (lc, be->args) {
      Expr *expr = (Expr *)lfirst(lc);
      FuncExpr *arg = having_Expr_to_provenance_cmp(expr, constants, negated);
      l = lappend(l, arg);
    }
 
    array->elements = l;
 
    return result;
  }
}
 
/**
 * @brief Dispatch a HAVING sub-expression to the appropriate converter.
 *
 * Entry point for the mutual recursion between
 * @c having_BoolExpr_to_provenance and @c having_OpExpr_to_provenance_cmp.
 *
 * @param expr       Sub-expression to convert (@c BoolExpr or @c OpExpr).
 * @param constants  Extension OID cache.
 * @param negated    Whether the expression appears under a NOT.
 * @return  Converted @c FuncExpr.
 */
static FuncExpr *having_Expr_to_provenance_cmp(Expr *expr, const constants_t *constants, bool negated)
{
  if (IsA(expr, BoolExpr))
    return having_BoolExpr_to_provenance((BoolExpr *)expr, constants, negated);
  else if (IsA(expr, OpExpr))
    return having_OpExpr_to_provenance_cmp((OpExpr *)expr, constants, negated);
  else
    provsql_error("Unknown structure within Boolean expression");
}
 
/* -------------------------------------------------------------------------
 * Random-variable WHERE-clause rewriting
 *
 * Mirror of the HAVING trio above.  An OpExpr whose @c opfuncid matches
 * one of the @c random_variable_{eq,ne,le,lt,ge,gt} procedures is an
 * RV comparison; the planner hook lifts it out of @c jointree->quals,
 * builds an equivalent @c provenance_cmp(left_uuid, op_oid, right_uuid)
 * @c FuncExpr, and conjoins the resulting UUID into the row's
 * provenance via @c provenance_times.  The lifted WHERE conjunct is
 * removed (or the whole WHERE replaced by @c NULL when only RV cmps
 * were present); what remains is purely Boolean and the executor
 * evaluates it in the usual way.  The RV-cmp operators themselves are
 * boolean placeholders -- their procedure raises if reached, which can
 * happen only when the planner hook is bypassed (e.g. provsql.active
 * off).
 * ------------------------------------------------------------------------- */
 
/**
 * @brief Test whether @p funcoid is one of the @c random_variable_*
 *        comparison procedures, and if so return its
 *        @c ComparisonOperator index.
 *
 * @param  constants  Extension OID cache.
 * @param  funcoid    Procedure OID to test (typically @c OpExpr->opfuncid).
 * @return Index in @c [0..6) on match, @c -1 otherwise.  Match indices
 *         line up with @c ComparisonOperator (EQ=0, NE=1, LE=2, LT=3,
 *         GE=4, GT=5).
 */
static int rv_cmp_index(const constants_t *constants, Oid funcoid)
{
  for (int i = 0; i < 6; ++i) {
    if (funcoid == constants->OID_FUNCTION_RV_CMP[i])
      return i;
  }
  return -1;
}
 
/**
 * @brief Wrap an expression returning @c random_variable in a
 *        binary-coercible cast to @c uuid.
 *
 * Operand of the comparison may be a Var, a constant lifted by an
 * implicit cast, or another OpExpr (e.g. <tt>a + b</tt>).
 * @c random_variable and @c uuid share the same byte layout, so we
 * emit a @c RelabelType node -- the planner sees a zero-cost type
 * relabel, the executor never dispatches through a runtime
 * conversion function.
 */
static Expr *
wrap_random_variable_uuid(Node *operand, const constants_t *constants)
{
  RelabelType *rt = makeNode(RelabelType);
  rt->arg = (Expr *) operand;
  rt->resulttype = constants->OID_TYPE_UUID;
  rt->resulttypmod = -1;
  rt->resultcollid = InvalidOid;
  rt->relabelformat = COERCE_IMPLICIT_CAST;
  rt->location = -1;
  return (Expr *) rt;
}
 
/* Forward declaration: the BoolExpr and Expr walkers below are mutually
 * recursive (BoolExpr recurses into Expr for each AND/OR child). */
static FuncExpr *rv_Expr_to_provenance(Expr *expr,
                                       const constants_t *constants,
                                       bool negated);
 
/**
 * @brief Convert a single RV-comparison @c OpExpr into a
 *        @c provenance_cmp() FuncExpr returning UUID.
 *
 * If @p negated is true the operator OID is replaced by its negator
 * (so <tt>NOT (a &gt; b)</tt> becomes <tt>a &le; b</tt> at the
 * provenance level), exactly as @c having_OpExpr_to_provenance_cmp
 * does.
 *
 * @param opExpr    The comparison expression from the WHERE clause.
 *                  Must satisfy @c rv_cmp_index(opExpr->opfuncid) &ge; 0;
 *                  callers are responsible for the type check.
 * @param constants Extension OID cache.
 * @param negated   Whether the expression appears under a NOT.
 */
static FuncExpr *
rv_OpExpr_to_provenance_cmp(OpExpr *opExpr, const constants_t *constants,
                            bool negated)
{
  FuncExpr *cmpExpr;
  Const    *oid_const;
  Oid       opno = opExpr->opno;
  Node     *left = (Node *)linitial(opExpr->args);
  Node     *right = (Node *)lsecond(opExpr->args);
 
  if (negated) {
    opno = get_negator(opno);
    if (!opno)
      provsql_error("Missing negator for random_variable comparison");
  }
 
  oid_const = makeConst(constants->OID_TYPE_INT, -1, InvalidOid,
                        sizeof(int32), Int32GetDatum(opno), false, true);
 
  cmpExpr = makeNode(FuncExpr);
  cmpExpr->funcid = constants->OID_FUNCTION_PROVENANCE_CMP;
  cmpExpr->funcresulttype = constants->OID_TYPE_UUID;
  cmpExpr->args = list_make3(
    wrap_random_variable_uuid(left, constants),
    oid_const,
    wrap_random_variable_uuid(right, constants));
  cmpExpr->location = opExpr->location;
 
  return cmpExpr;
}
 
/**
 * @brief Convert a Boolean combination of RV comparisons into a
 *        @c provenance_times / @c provenance_plus expression.
 *
 * Same De Morgan handling as @c having_BoolExpr_to_provenance: under
 * negation, AND ↔ OR (which means PROVENANCE_TIMES ↔ PROVENANCE_PLUS).
 * NOT flips @c negated and recurses.
 */
static FuncExpr *
rv_BoolExpr_to_provenance(BoolExpr *be, const constants_t *constants,
                          bool negated)
{
  FuncExpr   *result;
  ArrayExpr  *array;
  List       *l = NIL;
  ListCell   *lc;
 
  if (be->boolop == NOT_EXPR) {
    Expr *child = (Expr *)linitial(be->args);
    return rv_Expr_to_provenance(child, constants, !negated);
  }
 
  array = makeNode(ArrayExpr);
  array->array_typeid = constants->OID_TYPE_UUID_ARRAY;
  array->element_typeid = constants->OID_TYPE_UUID;
  array->location = -1;
 
  result = makeNode(FuncExpr);
  result->funcresulttype = constants->OID_TYPE_UUID;
  result->funcvariadic = true;
  result->location = be->location;
  result->args = list_make1(array);
 
  if ((be->boolop == AND_EXPR && !negated) ||
      (be->boolop == OR_EXPR  && negated))
    result->funcid = constants->OID_FUNCTION_PROVENANCE_TIMES;
  else if ((be->boolop == AND_EXPR && negated) ||
           (be->boolop == OR_EXPR  && !negated))
    result->funcid = constants->OID_FUNCTION_PROVENANCE_PLUS;
  else
    provsql_error("Unknown Boolean operator in random_variable WHERE clause");
 
  foreach (lc, be->args) {
    FuncExpr *arg = rv_Expr_to_provenance((Expr *)lfirst(lc),
                                          constants, negated);
    l = lappend(l, arg);
  }
  array->elements = l;
 
  return result;
}
 
/**
 * @brief Dispatch a WHERE sub-expression to the appropriate RV converter.
 *
 * Entry point for the mutual recursion between
 * @c rv_BoolExpr_to_provenance and @c rv_OpExpr_to_provenance_cmp.
 */
static FuncExpr *
rv_Expr_to_provenance(Expr *expr, const constants_t *constants, bool negated)
{
  if (IsA(expr, BoolExpr))
    return rv_BoolExpr_to_provenance((BoolExpr *)expr, constants, negated);
  if (IsA(expr, OpExpr)) {
    OpExpr *opExpr = (OpExpr *)expr;
    if (rv_cmp_index(constants, opExpr->opfuncid) >= 0)
      return rv_OpExpr_to_provenance_cmp(opExpr, constants, negated);
  }
  provsql_error("Unsupported sub-expression in random_variable WHERE clause "
                "(only Boolean combinations of RV comparisons are accepted)");
  return NULL; /* unreachable, silences -Wreturn-type */
}
 
/**
 * @brief Test whether an Expr (sub-)tree contains any RV comparison.
 *
 * Used by the WHERE-clause extractor to decide whether a top-level
 * conjunct mentions any random_variable comparator and therefore
 * needs lifting (or, if the conjunct mixes RV and non-RV operators
 * in a way we cannot rewrite, errors).
 */
static bool
expr_contains_rv_cmp(Node *node, const constants_t *constants)
{
  if (node == NULL)
    return false;
  if (IsA(node, OpExpr)) {
    OpExpr *opExpr = (OpExpr *)node;
    if (rv_cmp_index(constants, opExpr->opfuncid) >= 0)
      return true;
  }
  if (IsA(node, BoolExpr)) {
    BoolExpr *be = (BoolExpr *)node;
    ListCell *lc;
    foreach (lc, be->args) {
      if (expr_contains_rv_cmp(lfirst(lc), constants))
        return true;
    }
    return false;
  }
  return false;
}
 
/**
 * @brief Test whether @p expr is a Boolean combination of @em only
 *        random_variable comparisons (no other leaves allowed).
 *
 * Mirrors @c check_expr_on_aggregate / @c check_boolexpr_on_aggregate
 * for the agg_token WHERE-to-HAVING migration path.  Recursively
 * accepts:
 * - @c BoolExpr (AND/OR/NOT) all of whose children pass; and
 * - @c OpExpr matching one of the @c random_variable_* comparators.
 *
 * Anything else (a non-RV @c OpExpr, a @c Var, a @c Const, a non-cmp
 * @c FuncExpr) makes the expression mixed and unsupportable by the
 * RV-only walker, so the function returns @c false and the caller
 * raises a clear error.
 */
static bool
check_expr_on_rv(Expr *expr, const constants_t *constants)
{
  if (expr == NULL)
    return false;
  if (IsA(expr, OpExpr))
    return rv_cmp_index(constants, ((OpExpr *)expr)->opfuncid) >= 0;
  if (IsA(expr, BoolExpr)) {
    BoolExpr *be = (BoolExpr *)expr;
    ListCell *lc;
    foreach (lc, be->args) {
      if (!check_expr_on_rv((Expr *)lfirst(lc), constants))
        return false;
    }
    return true;
  }
  return false;
}
 
/* The earlier RV-only WHERE walker (@c extract_rv_cmps_from_quals)
 * has been folded into the unified classifier
 * @c migrate_probabilistic_quals further down in this file; both the
 * agg_token and the random_variable migration paths are now special
 * cases of one walk over @c q->jointree->quals.  See the comment on
 * @c qual_class for the routing matrix. */
 
/**
 * @brief Build the combined provenance expression to be added to the SELECT list.
 *
 * Combines the tokens in @p prov_atts according to @p op:
 * - @c SR_PLUS  → use the first token directly (union branch; the outer
 *                  @c array_agg / @c provenance_plus is added later if needed).
 * - @c SR_TIMES → wrap all tokens in @c provenance_times(...).
 * - @c SR_MONUS → wrap all tokens in @c provenance_monus(...).
 *
 * When @p aggregation or @p group_by_rewrite is true, wraps the result in
 * @c array_agg + @c provenance_plus to collapse groups.  A @c provenance_delta
 * gate is added for plain aggregations without a HAVING clause.
 *
 * If a HAVING clause is present it is removed from @p q->havingQual and
 * converted into a provenance expression via @c having_Expr_to_provenance_cmp.
 *
 * If @c provsql_where_provenance is enabled, equality gates (@c provenance_eq)
 * are prepended for join conditions and WHERE equalities, and a projection gate
 * is appended if the output columns form a proper subset of the input columns.
 *
 * @param constants        Extension OID cache.
 * @param q                Query being rewritten (HAVING is cleared if present).
 * @param prov_atts        List of provenance @c Var nodes.
 * @param aggregation      True if the query contains aggregate functions.
 * @param group_by_rewrite True if a GROUP BY requires the plus-aggregate wrapper.
 * @param op               Semiring operation to use for combining tokens.
 * @param columns          Per-RTE column-numbering array (for where-provenance).
 *                         For provenance-tracked @c RTE_RELATION entries, the
 *                         -1 sentinel is used to identify them; the PROJECT
 *                         gate positions for their columns use @c varattno
 *                         rather than the query-order-dependent sequential
 *                         numbers (see @c build_column_map() for the
 *                         rationale).
 * @param nbcols           Total number of non-provenance output columns.
 * @param wrap_assumed_boolean If true, wrap the result in
 *                         @c provenance_assumed_boolean so downstream
 *                         probability evaluators may treat it as Boolean.
 * @return  The provenance @c Expr to be appended to the target list.
 */
static Expr *make_provenance_expression(const constants_t *constants, Query *q,
                                        List *prov_atts, bool aggregation,
                                        bool group_by_rewrite,
                                        semiring_operation op, int **columns,
                                        int nbcols, bool wrap_assumed_boolean) {
  Expr *result;
  ListCell *lc_v;
 
  if (op == SR_PLUS) {
    result = linitial(prov_atts);
  } else {
    if (my_lnext(prov_atts, list_head(prov_atts)) == NULL) {
      result = linitial(prov_atts);
    } else {
      FuncExpr *expr = makeNode(FuncExpr);
      if (op == SR_TIMES) {
        ArrayExpr *array = makeNode(ArrayExpr);
 
        expr->funcid = constants->OID_FUNCTION_PROVENANCE_TIMES;
        expr->funcvariadic = true;
 
        array->array_typeid = constants->OID_TYPE_UUID_ARRAY;
        array->element_typeid = constants->OID_TYPE_UUID;
        array->elements = prov_atts;
        array->location = -1;
 
        expr->args = list_make1(array);
      } else { // SR_MONUS
        expr->funcid = constants->OID_FUNCTION_PROVENANCE_MONUS;
        expr->args = prov_atts;
      }
      expr->funcresulttype = constants->OID_TYPE_UUID;
      expr->location = -1;
 
      result = (Expr *)expr;
    }
 
    if (group_by_rewrite || aggregation) {
      Aggref *agg = makeNode(Aggref);
      FuncExpr *plus = makeNode(FuncExpr);
      TargetEntry *te_inner = makeNode(TargetEntry);
 
      q->hasAggs = true;
 
      te_inner->resno = 1;
      te_inner->expr = (Expr *)result;
 
      agg->aggfnoid = constants->OID_FUNCTION_ARRAY_AGG;
      agg->aggtype = constants->OID_TYPE_UUID_ARRAY;
      agg->args = list_make1(te_inner);
      agg->aggkind = AGGKIND_NORMAL;
      agg->location = -1;
#if PG_VERSION_NUM >= 140000
      agg->aggno = agg->aggtransno = -1;
#endif
 
      agg->aggargtypes = list_make1_oid(constants->OID_TYPE_UUID);
 
      plus->funcid = constants->OID_FUNCTION_PROVENANCE_PLUS;
      plus->args = list_make1(agg);
      plus->funcresulttype = constants->OID_TYPE_UUID;
      plus->location = -1;
 
      result = (Expr *)plus;
    }
 
    /* HAVING quals come in two flavours.  A qual that references an
     * agg_token Var or a provenance_aggregate() wrapper must be lifted
     * into a provenance_cmp gate so the per-group truth value is
     * carried by the provenance circuit (and the corresponding gate_agg
     * remains evaluable).  Anything else -- a deterministic scalar
     * predicate, or one over random_variable aggregates collapsed by
     * expected() / variance() / moment() to a plain double -- is left
     * in q->havingQual for PostgreSQL to evaluate natively, and the
     * per-group provenance still gets a delta wrapper. */
    {
      bool lift_having = q->havingQual != NULL &&
                         needs_having_lift((Node *) q->havingQual, constants);
 
      if (aggregation && !lift_having) {
        FuncExpr *deltaExpr = makeNode(FuncExpr);
 
        // adding the delta gate to the provenance circuit
        deltaExpr->funcid = constants->OID_FUNCTION_PROVENANCE_DELTA;
        deltaExpr->args = list_make1(result);
        deltaExpr->funcresulttype = constants->OID_TYPE_UUID;
        deltaExpr->location = -1;
 
        result = (Expr *)deltaExpr;
      }
 
      if (lift_having) {
        result = (Expr*) having_Expr_to_provenance_cmp((Expr*)q->havingQual, constants, false);
        q->havingQual = NULL;
      }
    }
  }
 
  /* Part to handle eq gates used for where-provenance.
   * Placed before projection gates because they need
   * to be deeper in the provenance tree. */
  if (provsql_where_provenance && q->jointree) {
    ListCell *lc;
    foreach (lc, q->jointree->fromlist) {
      if (IsA(lfirst(lc), JoinExpr)) {
        JoinExpr *je = (JoinExpr *)lfirst(lc);
        /* Study equalities coming from From clause */
        result =
          add_eq_from_Quals_to_Expr(constants, je->quals, result, columns);
      }
    }
    /* Study equalities coming from WHERE clause */
    result = add_eq_from_Quals_to_Expr(constants, q->jointree->quals, result,
                                       columns);
  }
 
  if (provsql_where_provenance) {
    ArrayExpr *array = makeNode(ArrayExpr);
    FuncExpr *fe = makeNode(FuncExpr);
    bool projection = false;
    int nb_column = 0;
    /* Cumulative offset of each RTE within the TIMES gate's concatenated
     * locator vector.  WhereCircuit::evaluate(TIMES) appends the locator
     * vector of each child input in q->rtable order, so a column at
     * varattno k of the i-th provenance-tracked base RTE lands at
     * prov_offset[i] + k in the concat.  varattno alone (the recent fix
     * documented at the top of this file) is correct only when there is a
     * single provenance-tracked input; for multi-input joins it omits the
     * preceding inputs' nb_user_cols and the project gate then reads from
     * the wrong table's locator slice.
     *
     * 1-indexed by rteid for direct indexing via Var->varno; entry 0 is
     * unused.  Length q->rtable->length + 1. */
    int *prov_offset = (int *)palloc0((q->rtable->length + 1) * sizeof(int));
    int cum = 0;
    Index r;
 
    fe->funcid = constants->OID_FUNCTION_PROVENANCE_PROJECT;
    fe->funcvariadic = true;
    fe->funcresulttype = constants->OID_TYPE_UUID;
    fe->location = -1;
 
    array->array_typeid = constants->OID_TYPE_INT_ARRAY;
    array->element_typeid = constants->OID_TYPE_INT;
    array->elements = NIL;
    array->location = -1;
 
    for (r = 1; r <= (Index)q->rtable->length; ++r) {
      prov_offset[r] = cum;
      if (columns[r-1]) {
        RangeTblEntry *rte_r = (RangeTblEntry *)list_nth(q->rtable, r-1);
        int ncols = list_length(rte_r->eref->colnames);
        bool is_prov = false;
        int nb_user = 0;
        int k;
        for (k = 0; k < ncols; ++k) {
          if (columns[r-1][k] == -1) is_prov = true;
          else if (columns[r-1][k] > 0) nb_user++;
        }
        if (is_prov) cum += nb_user;
      }
    }
 
    foreach (lc_v, q->targetList) {
      TargetEntry *te_v = (TargetEntry *)lfirst(lc_v);
      if (IsA(te_v->expr, Var)) {
        Var *vte_v = (Var *)te_v->expr;
        RangeTblEntry *rte_v =
          (RangeTblEntry *)lfirst(list_nth_cell(q->rtable, vte_v->varno - 1));
        int value_v;
#if PG_VERSION_NUM >= 180000
        if (rte_v->rtekind == RTE_GROUP) {
          Expr *ge = lfirst(list_nth_cell(rte_v->groupexprs, vte_v->varattno - 1));
          if(IsA(ge, Var)) {
            Var *v = (Var *) ge;
            value_v = columns[v->varno - 1] ?
                      columns[v->varno - 1][v->varattno - 1] : 0;
          } else {
            Const *ce = makeConst(constants->OID_TYPE_INT, -1, InvalidOid,
                                  sizeof(int32), Int32GetDatum(0), false, true);
 
            array->elements = lappend(array->elements, ce);
            value_v = 0;
          }
        } else
#endif
        if (rte_v->rtekind != RTE_JOIN) { // Normal RTE
          if (rte_v->rtekind == RTE_RELATION && columns[vte_v->varno - 1]) {
            /* Determine whether this base table is provenance-tracked by
             * scanning for the sentinel -1 entry that build_column_map()
             * assigns to the provsql column. */
            bool is_prov = false;
            int ncols_rte = list_length(rte_v->eref->colnames);
            for (int k = 0; k < ncols_rte; k++) {
              if (columns[vte_v->varno - 1][k] == -1) {
                is_prov = true;
                break;
              }
            }
            if (is_prov) {
              int raw = columns[vte_v->varno - 1][vte_v->varattno - 1];
              /* Local position within this table is `varattno` (the
               * provsql column is appended last by add_provenance(), so
               * user columns occupy 1..nb_user_cols exactly matching the
               * IN gate's Locator vector).  We then shift by
               * prov_offset[varno] to land in the right slice of the
               * TIMES gate's concatenated locator vector when the query
               * joins multiple provenance-tracked relations. */
              value_v = (raw == -1) ? -1
                                    : (int)vte_v->varattno
                        + prov_offset[vte_v->varno];
            } else {
              /* Non-provenance base table: no IN gate exists for it, so
               * the position would be out of range regardless.  Explicitly
               * record 0 so evaluate() returns an empty locator set and
               * the positions array stays in sync with the output column
               * count. */
              Const *ce =
                makeConst(constants->OID_TYPE_INT, -1, InvalidOid,
                          sizeof(int32), Int32GetDatum(0), false, true);
              array->elements = lappend(array->elements, ce);
              projection = true;
              continue;
            }
          } else {
            /* RTE_SUBQUERY and others: the sequential number equals the
            * column's 1-indexed position in the subquery's output list,
            * which matches what the child gate's evaluate() expects. */
            value_v = columns[vte_v->varno - 1] ?
                      columns[vte_v->varno - 1][vte_v->varattno - 1] : 0;
          }
        } else { // Join RTE
          Var *jav_v = (Var *)lfirst(
            list_nth_cell(rte_v->joinaliasvars, vte_v->varattno - 1));
          if (jav_v && IsA(jav_v, Var) && columns[jav_v->varno - 1]) {
            RangeTblEntry *jrte_v = (RangeTblEntry *)lfirst(
              list_nth_cell(q->rtable, jav_v->varno - 1));
            if (jrte_v->rtekind == RTE_RELATION) {
              /* Provenance-tracking check and varattno fix — same rationale
               * as the RTE_RELATION branch above. */
              bool is_prov = false;
              int ncols_jrte = list_length(jrte_v->eref->colnames);
              for (int k = 0; k < ncols_jrte; k++) {
                if (columns[jav_v->varno - 1][k] == -1) {
                  is_prov = true;
                  break;
                }
              }
              if (is_prov) {
                int raw = columns[jav_v->varno - 1][jav_v->varattno - 1];
                value_v = (raw == -1) ? -1
                                      : (int)jav_v->varattno
                          + prov_offset[jav_v->varno];
              } else {
                Const *ce =
                  makeConst(constants->OID_TYPE_INT, -1, InvalidOid,
                            sizeof(int32), Int32GetDatum(0), false, true);
                array->elements = lappend(array->elements, ce);
                projection = true;
                continue;
              }
            } else {
              value_v = columns[jav_v->varno - 1][jav_v->varattno - 1];
            }
          } else {
            value_v = 0;
          }
        }
 
        /* If this is a valid column */
        if (value_v > 0) {
          Const *ce =
            makeConst(constants->OID_TYPE_INT, -1, InvalidOid, sizeof(int32),
                      Int32GetDatum(value_v), false, true);
 
          array->elements = lappend(array->elements, ce);
 
          if (value_v != ++nb_column)
            projection = true;
        } else {
          if (value_v != -1)
            projection = true;
        }
      } else { // we have a function in target
        Const *ce = makeConst(constants->OID_TYPE_INT, -1, InvalidOid,
                              sizeof(int32), Int32GetDatum(0), false, true);
 
        array->elements = lappend(array->elements, ce);
        projection = true;
      }
    }
 
    if (nb_column != nbcols)
      projection = true;
 
    if (projection) {
      fe->args = list_make2(result, array);
      result = (Expr *)fe;
    } else {
      pfree(array);
      pfree(fe);
    }
  }
 
  /* Wrap the finished per-row root in a @c gate_assumed_boolean when
   * our caller (the safe-query rewrite path in @c process_query) asks
   * for it.  Wrapping here -- before @c add_to_select and
   * @c replace_provenance_function_by_expression -- means every
   * per-row root reference in the final target list carries the
   * marker uniformly.  Subqueries that this same Query body opens
   * (per-atom DISTINCT projections inserted by the rewriter) are
   * handled by their own deeper @c process_query / @c make_provenance_expression
   * calls with @c wrap_assumed_boolean = false, so the marker sits
   * only at the outermost root that surfaces as the user-visible
   * row provenance. */
  if (wrap_assumed_boolean &&
      OidIsValid(constants->OID_FUNCTION_ASSUME_BOOLEAN))
    result = wrap_in_assume_boolean(constants, result);
 
  return result;
}
 
/* -------------------------------------------------------------------------
 * Set-operation & DISTINCT rewriting
 * ------------------------------------------------------------------------- */
 
#if PG_VERSION_NUM >= 180000
typedef struct {
  Index group_rtindex;
  List *groupexprs;
} resolve_group_rte_ctx;
 
static Node *
resolve_group_rte_vars_mutator(Node *node, void *raw_ctx) {
  resolve_group_rte_ctx *ctx = (resolve_group_rte_ctx *)raw_ctx;
  if (node == NULL)
    return NULL;
  if (IsA(node, Var)) {
    Var *v = (Var *)node;
    if (v->varno == ctx->group_rtindex) {
      Node *resolved = copyObject(list_nth(ctx->groupexprs, v->varattno - 1));
      /* Clear varnullingrels: the group-step nulling bits reference the
       * group_rtindex RTE which does not exist in the fresh inner query.
       * Leaving them set causes the planner to access simple_rel_array at
       * group_rtindex (which has no RelOptInfo), triggering
       * "unrecognized RTE kind: 9". */
      if (IsA(resolved, Var))
        ((Var *)resolved)->varnullingrels = NULL;
      return resolved;
    }
  }
  return expression_tree_mutator(node, resolve_group_rte_vars_mutator, raw_ctx);
}
 
/**
 * @brief Strip PG 18's virtual @c RTE_GROUP entry from @p q in place.
 *
 * @c parseCheckAggregates() appends an @c RTE_GROUP entry at the end of
 * @c q->rtable whenever the query has a @c GROUP @c BY clause; references
 * to grouped columns in @c targetList and @c jointree->quals point at that
 * synthetic RTE rather than the underlying base tables.  ProvSQL's
 * rewriters need a flat range-table to do their own index arithmetic, so
 * we remove the @c RTE_GROUP and resolve every @c Var(@c group_rtindex,
 * @c i) back to its base-table expression before going further.
 *
 * Idempotent: when @c q->hasGroupRTE is already false, returns without
 * doing anything.
 */
void strip_group_rte_pg18(Query *q) {
  resolve_group_rte_ctx grp_ctx;
  bool found = false;
  ListCell *lc;
  Index idx = 1;
  int rte_len = 0;
 
  if (!q->hasGroupRTE)
    return;
 
  foreach (lc, q->rtable) {
    RangeTblEntry *r = (RangeTblEntry *) lfirst(lc);
    if (r->rtekind == RTE_GROUP) {
      grp_ctx.group_rtindex = idx;
      grp_ctx.groupexprs    = r->groupexprs;
      found    = true;
      rte_len  = idx - 1;
      break;
    }
    idx++;
  }
 
  if (!found)
    return;
 
  q->rtable      = list_truncate(q->rtable, rte_len);
  q->hasGroupRTE = false;
 
  foreach (lc, q->targetList) {
    TargetEntry *te = (TargetEntry *) lfirst(lc);
    te->expr = (Expr *) resolve_group_rte_vars_mutator(
      (Node *) te->expr, &grp_ctx);
  }
  if (q->jointree && q->jointree->quals)
    q->jointree->quals = resolve_group_rte_vars_mutator(
      q->jointree->quals, &grp_ctx);
}
#endif
 
/* Forward declaration — defined later but needed by rewrite_agg_distinct */
static bool provenance_function_walker(Node *node, void *data);
 
/**
 * @brief Build the inner GROUP-BY subquery for one @c AGG(DISTINCT key).
 *
 * Produces:
 * @code
 *   SELECT key_expr, gb_col1, gb_col2, ...
 *   FROM   <same tables as q>
 *   GROUP BY key_expr, gb_col1, gb_col2, ...
 * @endcode
 *
 * @param q           Original query (supplies FROM / WHERE).
 * @param key_expr    The DISTINCT argument expression.
 * @param groupby_tes Non-aggregate target entries that are GROUP BY columns.
 * @return  Fresh inner @c Query.
 */
static Query *build_inner_for_distinct_key(Query *q, Expr *key_expr,
                                           List *groupby_tes) {
  Query *inner;
  List *new_tl = NIL;
  List *new_gc = NIL;
  ListCell *lc;
  int resno = 1, sgref = 1;
 
  inner = copyObject(q);
 
  inner->hasAggs    = false;
  inner->sortClause = NIL;
  inner->limitCount = NULL;
  inner->limitOffset = NULL;
  inner->distinctClause = NIL;
  inner->hasDistinctOn = false;
  inner->havingQual  = NULL;
 
  /* First column: the DISTINCT key */
  {
    TargetEntry *kte = makeNode(TargetEntry);
    SortGroupClause *sgc = makeNode(SortGroupClause);
 
    kte->expr   = copyObject(key_expr);
    kte->resno  = resno++;
    kte->resname = "key";
    sgc->tleSortGroupRef = kte->ressortgroupref = sgref++;
    get_sort_group_operators(exprType((Node *)kte->expr), true, true, false,
                             &sgc->sortop, &sgc->eqop, NULL, &sgc->hashable);
    new_gc = list_make1(sgc);
    new_tl = list_make1(kte);
  }
 
  /* Remaining columns: GROUP BY columns from the original query */
  foreach (lc, groupby_tes) {
    TargetEntry *gyte = copyObject((TargetEntry *)lfirst(lc));
    SortGroupClause *sgc = makeNode(SortGroupClause);
 
    gyte->resno   = resno++;
    gyte->resjunk = false;
    sgc->tleSortGroupRef = gyte->ressortgroupref = sgref++;
    get_sort_group_operators(exprType((Node *)gyte->expr), true, true, false,
                             &sgc->sortop, &sgc->eqop, NULL, &sgc->hashable);
    new_gc = lappend(new_gc, sgc);
    new_tl = lappend(new_tl, gyte);
  }
 
  inner->targetList  = new_tl;
  inner->groupClause = new_gc;
  return inner;
}
 
/**
 * @brief Wrap @p inner in an outer query that applies the original aggregate.
 *
 * Produces:
 * @code
 *   SELECT AGG(key_col), gb_col1, gb_col2, ...
 *   FROM   inner
 *   GROUP BY gb_col1, gb_col2, ...
 * @endcode
 * The DISTINCT flag is cleared; @p inner provides exactly one row per
 * (key, group-by) combination, so the plain aggregate gives the right count.
 *
 * @param orig_agg_te  Original @c TargetEntry containing @c AGG(DISTINCT key).
 * @param inner        Inner query from @c build_inner_for_distinct_key.
 * @param n_gb         Number of GROUP BY columns (trailing entries in @p inner).
 * @param constants    Extension OID cache.
 * @return  Fresh outer @c Query.
 */
static Query *build_outer_for_distinct_key(TargetEntry *orig_agg_te,
                                           Query *inner, int n_gb,
                                           const constants_t *constants) {
  Query *outer = makeNode(Query);
  RangeTblEntry *rte = makeNode(RangeTblEntry);
  Alias *alias = makeNode(Alias), *eref = makeNode(Alias);
  RangeTblRef *rtr = makeNode(RangeTblRef);
  FromExpr *jt = makeNode(FromExpr);
  List *new_tl = NIL, *new_gc = NIL;
  ListCell *lc;
  int resno = 1, sgref = 1;
  int inner_len = list_length(inner->targetList);
  int attno;
 
  /* Wrap inner in a subquery RTE */
  alias->aliasname = eref->aliasname = "d";
  eref->colnames = NIL;
  foreach (lc, inner->targetList) {
    TargetEntry *te = lfirst(lc);
    eref->colnames = lappend(eref->colnames,
                             makeString(te->resname ? pstrdup(te->resname) : ""));
  }
  rte->alias   = alias;
  rte->eref    = eref;
  rte->rtekind = RTE_SUBQUERY;
  rte->subquery = inner;
  rte->inFromCl = true;
#if PG_VERSION_NUM < 160000
  rte->requiredPerms = ACL_SELECT;
#endif
 
  rtr->rtindex = 1;
  jt->fromlist = list_make1(rtr);
 
  outer->commandType = CMD_SELECT;
  outer->canSetTag   = true;
  outer->rtable      = list_make1(rte);
  outer->jointree    = jt;
  outer->hasAggs     = true;
 
  /* First output column: the aggregate over the key (col 1 of inner) */
  {
    TargetEntry *agg_te = copyObject(orig_agg_te);
    Aggref *ar = (Aggref *)agg_te->expr;
    Var *key_var = makeNode(Var);
    TargetEntry *arg_te = makeNode(TargetEntry);
 
    key_var->varno      = 1;
    key_var->varattno   = 1;    /* key is first column of inner */
    key_var->vartype    = linitial_oid(ar->aggargtypes);
    key_var->varcollid  = exprCollation((Node *)((TargetEntry *)linitial(ar->args))->expr);
    key_var->vartypmod  = -1;
    key_var->location   = -1;
    arg_te->resno = 1;
    arg_te->expr  = (Expr *)key_var;
 
    ar->args        = list_make1(arg_te);
    ar->aggdistinct = NIL;
    agg_te->resno   = resno++;
    new_tl = list_make1(agg_te);
  }
 
  /* Remaining output columns: GROUP BY cols (trailing cols of inner) */
  for (attno = inner_len - n_gb + 1; attno <= inner_len; attno++) {
    TargetEntry *inner_te = list_nth(inner->targetList, attno - 1);
    Var *gb_var = makeNode(Var);
    TargetEntry *gb_te = makeNode(TargetEntry);
    SortGroupClause *sgc = makeNode(SortGroupClause);
 
    gb_var->varno      = 1;
    gb_var->varattno   = attno;
    gb_var->vartype    = exprType((Node *)inner_te->expr);
    gb_var->varcollid  = exprCollation((Node *)inner_te->expr);
    gb_var->vartypmod  = -1;
    gb_var->location   = -1;
 
    gb_te->resno   = resno++;
    gb_te->expr    = (Expr *)gb_var;
    gb_te->resname = inner_te->resname;
 
    sgc->tleSortGroupRef = gb_te->ressortgroupref = sgref++;
    sgc->nulls_first = false;
    get_sort_group_operators(gb_var->vartype, true, true, false,
                             &sgc->sortop, &sgc->eqop, NULL, &sgc->hashable);
    new_gc = lappend(new_gc, sgc);
    new_tl = lappend(new_tl, gb_te);
  }
 
  outer->targetList  = new_tl;
  outer->groupClause = new_gc;
  return outer;
}
 
/**
 * @brief Rewrite every @c AGG(DISTINCT key) in @p q using independent subqueries.
 *
 * For a single DISTINCT aggregate, produces a subquery:
 * @code
 *   SELECT AGG(key), gb...  FROM (SELECT key, gb... FROM t GROUP BY key, gb...) GROUP BY gb...
 * @endcode
 * For multiple DISTINCT aggregates with different keys, produces an JOIN
 * of one such subquery per aggregate, joined on the GROUP BY columns.
 * Non-DISTINCT aggregates are left untouched.
 *
 * @param q            Query to inspect and possibly rewrite.
 * @param constants    Extension OID cache.
 * @return   Rewritten query, or @c NULL if no @c AGG(DISTINCT) was found.
 */
static Query *rewrite_agg_distinct(Query *q, const constants_t *constants) {
  List *distinct_agg_tes = NIL;
  List *groupby_tes = NIL;
  ListCell *lc;
 
#if PG_VERSION_NUM >= 180000
  /* In PostgreSQL 18, parseCheckAggregates() injects a virtual RTE_GROUP
   * entry at the END of the range table.  GROUP BY column Vars in the
   * SELECT list point to this entry (varno == group_rtindex) instead of
   * the underlying base-table RTE.
   *
   * Strip that entry now, before we do any index arithmetic (fll, rtr->rtindex,
   * agg_idx) or copy q->targetList into groupby_tes.  Once removed:
   *  - q->rtable contains only real RTEs, so appending outer-subquery RTEs
   *    lands at the correct indices.
   *  - groupby_tes will carry resolved (base-table) Var expressions, so
   *    the WHERE equalities and the inner-query target list are correct.
   * We also resolve the Var(group_rtindex) refs in q's own targetList and
   * WHERE clause so the final query doesn't reference the stripped entry. */
  strip_group_rte_pg18(q);
#endif
 
  /* Extract AGG(DISTINCT) and GROUP BY targets from the target list.
   * Regular AGG() aggregations and expressions containing provenance()
   * are left untouched. */
  foreach (lc, q->targetList) {
    TargetEntry *te = lfirst(lc);
    if (IsA(te->expr, Aggref)) {
      Aggref *ar = (Aggref *)te->expr;
      if (list_length(ar->aggdistinct) > 0)
        distinct_agg_tes = lappend(distinct_agg_tes, te);
    } else if (provenance_function_walker((Node *)te->expr,
                                          (void *)constants)) {
      /* Expression contains provenance() — skip it, it will be
       * handled later by the provenance rewriter */
    } else {
      /* Non-aggregate column — treat as GROUP BY key */
      TargetEntry *te_copy = copyObject(te);
      te_copy->resjunk = false;
      groupby_tes = lappend(groupby_tes, te_copy);
    }
  }
 
  if (distinct_agg_tes == NIL)
    return NULL;
 
  {
    int n_aggs = list_length(distinct_agg_tes);
    int n_gb   = list_length(groupby_tes);
    List *outer_queries = NIL;
 
    /* -----------------------------------------------------------------------
     * For each DISTINCT aggregate, build:
     *   inner_i: SELECT key_i, gb... FROM original... GROUP BY key_i, gb...
     *   outer_i: SELECT AGG(key_i) ASS agg_i, gb... FROM inner_i GROUP BY gb...
     *
     * Then produce a final query:
     *   SELECT gb..., agg_0, ..., agg_{N-1}
     *   FROM original... JOIN outer_0 ON gb... = gb... [JOIN ...]
     *  keeping the same order for the output columns.
     *
     * Column order in the final target list follows q->targetList:
     *   - DISTINCT agg i  → Var(n+i, 1)   (agg col of outer_i)
     * ----------------------------------------------------------------------- */
 
    /* Build one inner + one outer query per DISTINCT aggregate */
    foreach (lc, distinct_agg_tes) {
      TargetEntry *agg_te = lfirst(lc);
      Aggref *ar = (Aggref *)agg_te->expr;
      if(list_length(ar->args) != 1)
        provsql_error("AGG(DISTINCT) with more than one argument is not supported");
      else {
        Expr *key_expr = (Expr *)((TargetEntry *)linitial(ar->args))->expr;
        Query *inner = build_inner_for_distinct_key(q, key_expr, groupby_tes);
        Query *outer = build_outer_for_distinct_key(agg_te, inner, n_gb, constants);
        outer_queries = lappend(outer_queries, outer);
      }
    }
 
    {
      /* One subquery RTE per outer query */
      int i = 0;
      foreach (lc, outer_queries) {
        Query *oq = lfirst(lc);
        RangeTblEntry *rte = makeNode(RangeTblEntry);
        Alias *alias = makeNode(Alias), *eref = makeNode(Alias);
        ListCell *lc2;
        char buf[16];
 
        snprintf(buf, sizeof(buf), "d%d", i + 1);
        alias->aliasname = eref->aliasname = pstrdup(buf);
        eref->colnames = NIL;
        foreach (lc2, oq->targetList) {
          TargetEntry *te = lfirst(lc2);
          eref->colnames = lappend(eref->colnames,
                                   makeString(te->resname ? pstrdup(te->resname) : ""));
        }
        rte->alias    = alias;
        rte->eref     = eref;
        rte->rtekind  = RTE_SUBQUERY;
        rte->subquery = oq;
        rte->inFromCl = true;
#if PG_VERSION_NUM < 160000
        rte->requiredPerms = ACL_SELECT;
#endif
        q->rtable = lappend(q->rtable, rte);
        i++;
      }
 
      /* Build FROM list and WHERE conditions for the implicit join.
       * Use a simple FROM original..., outer_i, ... WHERE original.gb_j = outer_i.gb_j */
      {
        FromExpr *jt = q->jointree;
        List *from_list = jt->fromlist;
        unsigned fll = list_length(from_list);
        List *where_args = NIL;
 
        for (i = fll+1; i <= fll+n_aggs; i++) {
          RangeTblRef *rtr = makeNode(RangeTblRef);
          ListCell *lc2;
          unsigned j=0;
 
          rtr->rtindex = i;
          from_list = lappend(from_list, rtr);
 
          /* outer_0.gb_j = outer_i.gb_j for each GROUP BY column j */
          foreach(lc2, groupby_tes) {
            TargetEntry *gb_te = lfirst(lc2);
            int gb_attno = ++j + 1; /* col 1 = agg, cols 2+ = GB */
            Oid ytype  = exprType((Node *)gb_te->expr);
            Oid opno   = find_equality_operator(ytype, ytype);
            Operator opInfo = SearchSysCache1(OPEROID, ObjectIdGetDatum(opno));
            Form_pg_operator opform;
            OpExpr *oe = makeNode(OpExpr);
            Expr *le = copyObject(gb_te->expr);
            Var *rv = makeNode(Var);
            Oid collation=exprCollation((Node*) le);
 
            if (!HeapTupleIsValid(opInfo))
              provsql_error("could not find equality operator for type %u",
                            ytype);
            opform = (Form_pg_operator)GETSTRUCT(opInfo);
 
            oe->opno         = opno;
            oe->opfuncid     = opform->oprcode;
            oe->opresulttype = opform->oprresult;
            oe->opcollid     = InvalidOid;
            oe->inputcollid  = collation;
            oe->location     = -1;
            ReleaseSysCache(opInfo);
 
            rv->varno = i; rv->varattno = gb_attno;
            rv->vartype = ytype; rv->varcollid = collation;
            rv->vartypmod = -1; rv->location = -1;
 
            oe->args   = list_make2(le, rv);
            where_args = lappend(where_args, oe);
          }
        }
 
        if (list_length(where_args) == 0) {
          jt->quals = NULL;
        } else if (list_length(where_args) == 1) {
          jt->quals = linitial(where_args);
        } else {
          BoolExpr *be = makeNode(BoolExpr);
          be->boolop   = AND_EXPR;
          be->args     = where_args;
          be->location = -1;
          jt->quals    = (Node *)be;
        }
      }
 
      /* Build final target list in original column order.
       * DISTINCT agg i → Var(i+1, 1);  GROUP BY col j → Var(1, 2+j). */
      {
        int agg_idx   = list_length(q->jointree->fromlist) - n_aggs + 1;
        ListCell *lc2;
 
        foreach (lc2, q->targetList) {
          TargetEntry *te = lfirst(lc2);
 
          if (IsA(te->expr, Aggref) &&
              ((Aggref *)te->expr)->aggdistinct != NIL) {
            Var *v = makeNode(Var);
            v->varno = agg_idx++; /* outer_{agg_idx} RTE */
            v->varattno = 1;      /* agg result is col 1 of each outer */
            v->vartypmod = -1;
            v->location  = -1;
            te->expr = (Expr*)v;
          }
        }
      }
 
      return q;
    }
  }
}
 
 
/* -------------------------------------------------------------------------
 * Aggregation replacement mutator
 * ------------------------------------------------------------------------- */
 
/** @brief Context for the @c aggregation_mutator tree walker. */
typedef struct aggregation_mutator_context {
  List *prov_atts;              ///< List of provenance Var nodes
  semiring_operation op;        ///< Semiring operation for combining tokens
  const constants_t *constants; ///< Extension OID cache
} aggregation_mutator_context;
 
/**
 * @brief Tree-mutator that replaces Aggrefs with provenance-aware aggregates.
 * @param node  Current expression tree node.
 * @param ctx   Pointer to an @c aggregation_mutator_context (prov_atts,
 *              op, and constants).
 * @return      Possibly modified node.
 */
static Node *aggregation_mutator(Node *node, void *ctx) {
  aggregation_mutator_context *context = (aggregation_mutator_context *)ctx;
  if (node == NULL)
    return NULL;
 
  if (IsA(node, Aggref)) {
    Aggref *ar_v = (Aggref *)node;
    return (Node *)make_aggregation_expression(context->constants, ar_v,
                                               context->prov_atts, context->op);
  }
 
  return expression_tree_mutator(node, aggregation_mutator, ctx);
}
 
/**
 * @brief Wrap a @c provenance_aggregate FuncExpr with a cast to the
 *        original aggregate return type.
 *
 * @param prov_agg  The provenance_aggregate FuncExpr to wrap.
 * @param constants Extension OID cache.
 * @return          Cast FuncExpr wrapping @p prov_agg.
 */
static Node *wrap_agg_token_with_cast(FuncExpr *prov_agg,
                                      const constants_t *constants) {
  Const *typ_const = (Const *)lsecond(prov_agg->args);
  Oid target_type = DatumGetObjectId(typ_const->constvalue);
  CoercionPathType pathtype;
  Oid castfuncid;
 
  pathtype = find_coercion_pathway(target_type,
                                   constants->OID_TYPE_AGG_TOKEN,
                                   COERCION_EXPLICIT, &castfuncid);
  if (pathtype == COERCION_PATH_FUNC && OidIsValid(castfuncid)) {
    FuncExpr *cast = makeNode(FuncExpr);
    cast->funcid = castfuncid;
    cast->funcresulttype = target_type;
    cast->funcretset = false;
    cast->funcvariadic = false;
    cast->funcformat = COERCE_IMPLICIT_CAST;
    cast->args = list_make1(prov_agg);
    cast->location = -1;
    return (Node *)cast;
  }
 
  provsql_error("no cast from agg_token to %s for arithmetic on aggregate",
                format_type_be(target_type));
  return (Node *)prov_agg; /* unreachable */
}
 
/**
 * @brief Cast @c provenance_aggregate arguments of an operator or
 *        function when the formal parameter type requires it.
 *
 * For each argument in @p args that is a @c provenance_aggregate call,
 * check the corresponding formal parameter type of the parent function
 * @p parent_funcid.  If the formal type is polymorphic or @c agg_token
 * itself, the argument is left alone.  Otherwise a cast to the original
 * aggregate return type is inserted.
 *
 * @param args           Argument list to inspect (modified in place).
 * @param parent_funcid  OID of the parent function / operator implementor.
 * @param constants      Extension OID cache.
 */
static void maybe_cast_agg_token_args(List *args, Oid parent_funcid,
                                      const constants_t *constants) {
  HeapTuple tp;
  Form_pg_proc procForm;
  ListCell *lc;
  int i;
 
  tp = SearchSysCache1(PROCOID, ObjectIdGetDatum(parent_funcid));
  if (!HeapTupleIsValid(tp))
    return;
  procForm = (Form_pg_proc) GETSTRUCT(tp);
 
  i = 0;
  foreach(lc, args) {
    Node *arg = lfirst(lc);
 
    if (i < procForm->pronargs && IsA(arg, FuncExpr) &&
        ((FuncExpr *)arg)->funcid == constants->OID_FUNCTION_PROVENANCE_AGGREGATE) {
      Oid formal_type = procForm->proargtypes.values[i];
 
      if (formal_type != constants->OID_TYPE_AGG_TOKEN &&
          !IsPolymorphicType(formal_type)) {
        lfirst(lc) = wrap_agg_token_with_cast((FuncExpr *)arg, constants);
      }
    }
    i++;
  }
 
  ReleaseSysCache(tp);
}
 
/**
 * @brief Tree-mutator that casts @c provenance_aggregate results back
 *        to the original aggregate return type where needed.
 *
 * After the aggregation mutator replaces Aggrefs with
 * @c provenance_aggregate calls (returning @c agg_token), this
 * post-processing step inserts casts where the surrounding expression
 * expects a different type (e.g. @c SUM(id)+1).  Arguments to
 * functions that accept @c agg_token or polymorphic types are left
 * alone.
 *
 * @param node Current expression tree node.
 * @param ctx  Pointer to the @c constants_t OID cache.
 * @return     Possibly modified node.
 */
static Node *cast_agg_token_mutator(Node *node, void *ctx) {
  const constants_t *constants = (const constants_t *)ctx;
  Node *result;
 
  if (node == NULL)
    return NULL;
 
  /* Recurse first, then fix up arguments at this level. */
  result = expression_tree_mutator(node, cast_agg_token_mutator, ctx);
 
  if (IsA(result, OpExpr)) {
    OpExpr *op = (OpExpr *)result;
    set_opfuncid(op);
    maybe_cast_agg_token_args(op->args, op->opfuncid, constants);
  } else if (IsA(result, FuncExpr)) {
    FuncExpr *fe = (FuncExpr *)result;
    if (fe->funcid != constants->OID_FUNCTION_PROVENANCE_AGGREGATE)
      maybe_cast_agg_token_args(fe->args, fe->funcid, constants);
  }
 
  return result;
}
 
/**
 * @brief Replace every @c Aggref in @p q with a provenance-aware aggregate.
 *
 * Walks the query tree and substitutes each @c Aggref node with the result
 * of @c make_aggregation_expression, which wraps the original aggregate in
 * the semimodule machinery (@c provenance_semimod + @c array_agg +
 * @c provenance_aggregate).
 *
 * @param constants  Extension OID cache.
 * @param q          Query to mutate in place.
 * @param prov_atts  List of provenance @c Var nodes.
 * @param op         Semiring operation for combining tokens across rows.
 */
static void
replace_aggregations_by_provenance_aggregate(const constants_t *constants,
                                             Query *q, List *prov_atts,
                                             semiring_operation op) {
 
  aggregation_mutator_context context = {prov_atts, op, constants};
  ListCell *lc;
 
  query_tree_mutator(q, aggregation_mutator, &context,
                     QTW_DONT_COPY_QUERY | QTW_IGNORE_RT_SUBQUERIES);
 
  /* Post-processing: for target-list entries where a provenance_aggregate
   * result is nested inside an outer expression (e.g. SUM(id)+1),
   * insert a cast from agg_token back to the original aggregate return
   * type.  Standalone provenance_aggregate entries are left as agg_token
   * so they display as "value (*)". */
  foreach(lc, q->targetList) {
    TargetEntry *te = (TargetEntry *)lfirst(lc);
    if (te->expr == NULL)
      continue;
    /* Skip standalone provenance_aggregate calls */
    if (IsA(te->expr, FuncExpr) &&
        ((FuncExpr *)te->expr)->funcid == constants->OID_FUNCTION_PROVENANCE_AGGREGATE)
      continue;
    te->expr = (Expr *)cast_agg_token_mutator((Node *)te->expr,
                                              (void *)constants);
  }
}
 
/**
 * @brief Append the provenance expression to @p q's target list.
 *
 * Inserts a new @c TargetEntry named @c provsql immediately before any
 * @c resjunk entries (which must remain last) and adjusts the @c resno
 * of subsequent entries accordingly.
 *
 * @param q           Query to modify in place.
 * @param provenance  Expression to add (becomes the @c provsql output column).
 */
static void add_to_select(Query *q, Expr *provenance) {
  TargetEntry *newte = makeNode(TargetEntry);
  bool inserted = false;
  unsigned resno = 0;
 
  newte->expr = provenance;
  newte->resname = (char *)PROVSQL_COLUMN_NAME;
 
  if (IsA(provenance, Var)) {
    RangeTblEntry *rte = list_nth(q->rtable, ((Var *)provenance)->varno - 1);
    newte->resorigtbl = rte->relid;
    newte->resorigcol = ((Var *)provenance)->varattno;
  }
 
  /* Make sure to insert before all resjunk Target Entry */
  for (ListCell *cell = list_head(q->targetList); cell != NULL;) {
    TargetEntry *te = (TargetEntry *)lfirst(cell);
 
    if (!inserted)
      ++resno;
 
    if (te->resjunk) {
      if (!inserted) {
        newte->resno = resno;
        q->targetList = list_insert_nth(q->targetList, resno - 1, newte);
        cell = list_nth_cell(q->targetList, resno);
        te = (TargetEntry *)lfirst(cell);
        inserted = true;
      }
 
      ++te->resno;
    }
 
    cell = my_lnext(q->targetList, cell);
  }
 
  if (!inserted) {
    newte->resno = resno + 1;
    q->targetList = lappend(q->targetList, newte);
  }
}
 
/* -------------------------------------------------------------------------
 * Provenance function replacement
 * ------------------------------------------------------------------------- */
 
/** @brief Context for the @c provenance_mutator tree walker. */
typedef struct provenance_mutator_context {
  Expr *provsql;                ///< Provenance expression to substitute for provenance() calls
  const constants_t *constants; ///< Extension OID cache
  bool provsql_has_aggref;      ///< @c true when @c provsql contains an @c Aggref (set once by @c replace_provenance_function_by_expression).  When @c true, a @c provenance() substitution that lands inside another @c Aggref's argument tree would produce a nested same-level aggregate -- @c parse_agg.c forbids that shape, the planner's @c preprocess_aggrefs_walker does not recurse through @c Aggref boundaries, and the inner @c Aggref's @c aggno stays at the @c -1 sentinel and crashes @c ExecInterpExpr on @c ecxt_aggvalues[-1].
  bool inside_aggref;           ///< @c true while descending the argument tree of an @c Aggref node.
} provenance_mutator_context;
 
/**
 * @brief @c expression_tree_walker predicate: returns @c true on the first
 *        @c Aggref it encounters.
 *
 * Used to decide whether the provenance expression about to be substituted
 * would inject a nested aggregate when a @c provenance() call lives inside
 * another @c Aggref's argument tree.
 */
static bool
expr_contains_aggref_walker(Node *node, void *context) {
  if (node == NULL)
    return false;
  if (IsA(node, Aggref))
    return true;
  return expression_tree_walker(node, expr_contains_aggref_walker, context);
}
 
/**
 * @brief Tree-mutator that replaces provenance() calls with the actual provenance expression.
 * @param node  Current expression tree node.
 * @param ctx   Pointer to a @c provenance_mutator_context (provenance
 *              expression and constants).
 * @return      Possibly modified node.
 */
static Node *provenance_mutator(Node *node, void *ctx) {
  provenance_mutator_context *context = (provenance_mutator_context *)ctx;
  if (node == NULL)
    return NULL;
 
  if (IsA(node, Aggref)) {
    /* Descend into the Aggref's arguments with @c inside_aggref set so we
     * can refuse substitutions that would create a nested same-level
     * aggregate.  Save and restore the flag so sibling sub-expressions
     * outside this Aggref see the original value. */
    bool saved = context->inside_aggref;
    Node *result;
    context->inside_aggref = true;
    result = expression_tree_mutator(node, provenance_mutator, ctx);
    context->inside_aggref = saved;
    return result;
  }
 
  if (IsA(node, FuncExpr)) {
    FuncExpr *f = (FuncExpr *)node;
 
    if (f->funcid == context->constants->OID_FUNCTION_PROVENANCE) {
      if (context->inside_aggref && context->provsql_has_aggref) {
        provsql_error(
          "applying an SQL aggregate on top of a ProvSQL-introduced "
          "aggregation is not supported: the inner provenance() would "
          "be substituted with an expression containing an aggregate, "
          "producing a nested same-level aggregate that PostgreSQL "
          "rejects.  Evaluate the per-row provenance in a subquery "
          "and aggregate the resulting scalar outside, or drop the "
          "surrounding aggregate.");
      }
      return (Node *)copyObject(context->provsql);
    }
  } else if (IsA(node, RangeTblEntry) || IsA(node, RangeTblFunction)) {
    // A provenance() expression in a From (not within a subquery) is
    // non-sensical
    return node;
  }
 
  return expression_tree_mutator(node, provenance_mutator, ctx);
}
 
/**
 * @brief Replace every explicit @c provenance() call in @p q with @p provsql.
 *
 * Users can write @c provenance() in the target list or WHERE to refer to the
 * provenance token of the current tuple.  This mutator substitutes those calls
 * with the actual computed provenance expression.
 *
 * @param constants  Extension OID cache.
 * @param q          Query to mutate in place.
 * @param provsql    Provenance expression to substitute.
 */
static void
replace_provenance_function_by_expression(const constants_t *constants,
                                          Query *q, Expr *provsql) {
  provenance_mutator_context context;
 
  context.provsql = provsql;
  context.constants = constants;
  context.provsql_has_aggref =
    expr_contains_aggref_walker((Node *) provsql, NULL);
  context.inside_aggref = false;
 
  query_tree_mutator(q, provenance_mutator, &context,
                     QTW_DONT_COPY_QUERY | QTW_IGNORE_RT_SUBQUERIES);
}
 
/**
 * @brief Convert a SELECT DISTINCT into an equivalent GROUP BY.
 *
 * ProvSQL cannot handle DISTINCT directly (it would collapse provenance
 * tokens that should remain separate).  This function moves every entry
 * from @p q->distinctClause into @p q->groupClause (skipping any that are
 * already there) and clears @p q->distinctClause.
 *
 * @param q  Query to modify in place.
 */
static void transform_distinct_into_group_by(Query *q) {
  // First check which are already in the group by clause
  // Should be either none or all as "SELECT DISTINCT a, b ... GROUP BY a"
  // is invalid
  Bitmapset *already_in_group_by = NULL;
  ListCell *lc;
  foreach (lc, q->groupClause) {
    SortGroupClause *sgc = (SortGroupClause *)lfirst(lc);
    already_in_group_by =
      bms_add_member(already_in_group_by, sgc->tleSortGroupRef);
  }
 
  foreach (lc, q->distinctClause) {
    SortGroupClause *sgc = (SortGroupClause *)lfirst(lc);
    if (!bms_is_member(sgc->tleSortGroupRef, already_in_group_by)) {
      q->groupClause = lappend(q->groupClause, sgc);
    }
  }
 
  q->distinctClause = NULL;
}
 
/**
 * @brief Remove sort/group references that belonged to removed provenance columns.
 *
 * After @c remove_provenance_attributes_select strips provenance entries from
 * the target list, any GROUP BY, ORDER BY, or DISTINCT clause that referenced
 * them by @c tleSortGroupRef must be cleaned up.
 *
 * @param q                      Query to modify in place.
 * @param removed_sortgrouprefs  Bitmapset of @c ressortgroupref values to remove.
 */
static void
remove_provenance_attribute_groupref(Query *q,
                                     const Bitmapset *removed_sortgrouprefs) {
  List **lists[3] = {&q->groupClause, &q->distinctClause, &q->sortClause};
  int i = 0;
 
  for (i = 0; i < 3; ++i) {
    ListCell *cell, *prev;
 
    for (cell = list_head(*lists[i]), prev = NULL; cell != NULL;) {
      SortGroupClause *sgc = (SortGroupClause *)lfirst(cell);
      if (bms_is_member(sgc->tleSortGroupRef, removed_sortgrouprefs)) {
        *lists[i] = my_list_delete_cell(*lists[i], cell, prev);
 
        if (prev) {
          cell = my_lnext(*lists[i], prev);
        } else {
          cell = list_head(*lists[i]);
        }
      } else {
        prev = cell;
        cell = my_lnext(*lists[i], cell);
      }
    }
  }
}
 
/**
 * @brief Strip the provenance column's type info from a set-operation node.
 *
 * When a provenance column is removed from a UNION/EXCEPT query's target list,
 * the matching entries in the @c SetOperationStmt's @c colTypes, @c colTypmods,
 * and @c colCollations lists must also be removed.
 *
 * @param q        Query containing @c setOperations.
 * @param removed  Boolean array (from @c remove_provenance_attributes_select)
 *                 indicating which columns were removed.
 */
static void remove_provenance_attribute_setoperations(Query *q, bool *removed) {
  SetOperationStmt *so = (SetOperationStmt *)q->setOperations;
  List **lists[3] = {&so->colTypes, &so->colTypmods, &so->colCollations};
  int i = 0;
 
  for (i = 0; i < 3; ++i) {
    ListCell *cell, *prev;
    int j;
 
    for (cell = list_head(*lists[i]), prev = NULL, j = 0; cell != NULL; ++j) {
      if (removed[j]) {
        *lists[i] = my_list_delete_cell(*lists[i], cell, prev);
 
        if (prev) {
          cell = my_lnext(*lists[i], prev);
        } else {
          cell = list_head(*lists[i]);
        }
      } else {
        prev = cell;
        cell = my_lnext(*lists[i], cell);
      }
    }
  }
}
 
/**
 * @brief Wrap a non-ALL set operation in an outer GROUP BY query.
 *
 * UNION / EXCEPT (without ALL) would deduplicate tuples before ProvSQL can
 * attach provenance tokens.  To avoid this, the set operation is converted to
 * UNION ALL / EXCEPT ALL and a new outer query is built that groups the results
 * by all non-provenance columns, collecting tokens into an array for the
 * @c provenance_plus evaluation.
 *
 * After this rewrite the recursive call to @c process_query handles the
 * now-ALL inner set operation normally.
 *
 * @param q  Query whose @c setOperations is non-ALL (modified to ALL in place).
 * @return   New outer query that wraps @p q as a subquery RTE.
 */
static Query *rewrite_non_all_into_external_group_by(Query *q) {
  Query *new_query = makeNode(Query);
  RangeTblEntry *rte = makeNode(RangeTblEntry);
  FromExpr *jointree = makeNode(FromExpr);
  RangeTblRef *rtr = makeNode(RangeTblRef);
 
  SetOperationStmt *stmt = (SetOperationStmt *)q->setOperations;
 
  ListCell *lc;
  int sortgroupref = 0;
 
  stmt->all = true;
  // we might leave sub nodes of the SetOperationsStmt tree with all = false
  // but only for recursive trees of operators and only union can be recursive
  // https://doxygen.postgresql.org/prepunion_8c_source.html#l00479
  // we will set therefore set them later in process_set_operation_union
 
  rte->rtekind = RTE_SUBQUERY;
  rte->subquery = q;
  rte->eref = copyObject(((RangeTblEntry *)linitial(q->rtable))->eref);
  rte->inFromCl = true;
#if PG_VERSION_NUM < 160000
  // For PG_VERSION_NUM >= 160000, rte->perminfoindex==0 so no need to
  // care about permissions
  rte->requiredPerms = ACL_SELECT;
#endif
 
  rtr->rtindex = 1;
  jointree->fromlist = list_make1(rtr);
 
  new_query->commandType = CMD_SELECT;
  new_query->canSetTag = true;
  new_query->rtable = list_make1(rte);
  new_query->jointree = jointree;
  new_query->targetList = copyObject(q->targetList);
 
  if (new_query->targetList) {
    foreach (lc, new_query->targetList) {
      TargetEntry *te = (TargetEntry *)lfirst(lc);
      SortGroupClause *sgc = makeNode(SortGroupClause);
 
      sgc->tleSortGroupRef = te->ressortgroupref = ++sortgroupref;
 
      get_sort_group_operators(exprType((Node *)te->expr), false, true, false,
                               &sgc->sortop, &sgc->eqop, NULL, &sgc->hashable);
 
      new_query->groupClause = lappend(new_query->groupClause, sgc);
    }
  } else {
    GroupingSet *gs = makeNode(GroupingSet);
    gs->kind = GROUPING_SET_EMPTY;
    gs->content = 0;
    gs->location = -1;
    new_query->groupingSets = list_make1(gs);
  }
 
  return new_query;
}
 
/* -------------------------------------------------------------------------
 * Detection walkers
 * ------------------------------------------------------------------------- */
 
/**
 * @brief Tree walker that returns true if any @c provenance() call is found.
 *
 * Used to detect whether a query explicitly calls @c provenance(), which
 * triggers the substitution in @c replace_provenance_function_by_expression.
 * @param node  Current expression tree node.
 * @param data  Pointer to @c constants_t (cast from @c void*).
 * @return      @c true if a @c provenance() call is found anywhere in @p node.
 */
static bool provenance_function_walker(Node *node, void *data) {
  const constants_t *constants = (const constants_t *)data;
  if (node == NULL)
    return false;
 
  if (IsA(node, FuncExpr)) {
    FuncExpr *f = (FuncExpr *)node;
 
    if (f->funcid == constants->OID_FUNCTION_PROVENANCE)
      return true;
  }
 
  return expression_tree_walker(node, provenance_function_walker, data);
}
 
/**
 * @brief Check whether a @c provenance() call appears in the GROUP BY list.
 *
 * When the user writes @c GROUP BY provenance(), ProvSQL must not add its own
 * group-by wrapper (the query is already grouping on the token).
 *
 * @param constants  Extension OID cache.
 * @param q          Query to inspect.
 * @return  True if any GROUP BY key contains a @c provenance() call.
 */
static bool provenance_function_in_group_by(const constants_t *constants,
                                            Query *q) {
  ListCell *lc;
 
  /* Build the set of ressortgrouprefs that are actually in GROUP BY
   * (not ORDER BY or DISTINCT, which also set ressortgroupref). */
  Bitmapset *group_refs = NULL;
  foreach (lc, q->groupClause) {
    SortGroupClause *sgc = (SortGroupClause *)lfirst(lc);
    group_refs = bms_add_member(group_refs, sgc->tleSortGroupRef);
  }
 
  foreach (lc, q->targetList) {
    TargetEntry *te = (TargetEntry *)lfirst(lc);
    if (te->ressortgroupref > 0 &&
        bms_is_member(te->ressortgroupref, group_refs)) {
      if(expression_tree_walker((Node *)te, provenance_function_walker,
                                (void *)constants)) {
        return true;
      }
 
#if PG_VERSION_NUM >= 180000
      // Starting from PostgreSQL 18, the content of the GROUP BY is not
      // in the groupClause but in an associated RTE_GROUP RangeTblEntry
      if(IsA(te->expr, Var)) {
        Var *v = (Var *) te->expr;
        RangeTblEntry *r = (RangeTblEntry *)list_nth(q->rtable, v->varno - 1);
        if(r->rtekind == RTE_GROUP)
          if(expression_tree_walker((Node *) r->groupexprs, provenance_function_walker,
                                    (void *)constants)) {
            return true;
          }
      }
#endif
    }
  }
 
  return false;
}
 
/**
 * @brief Tree walker that detects any provenance-bearing relation or provenance() call.
 * @param node  Current expression tree node.
 * @param data  Pointer to @c constants_t (cast from @c void*).
 * @return      @c true if provenance rewriting is needed for this node.
 */
/**
 * @brief Recursive helper for @c has_provenance_walker that detects
 *        rv_cmp @c OpExpr and @c provenance() @c FuncExpr in
 *        expression subtrees.
 *
 * Stops at Query boundaries: @c SubLink subselects (used as
 * scalar/array subqueries in expressions) are not rewritten by the
 * outer planner_hook pass, so a tracked relation inside one must not
 * cause the OUTER query's gate to engage.  Only the @c testexpr of a
 * SubLink is followed (it lives in the outer's evaluation scope).
 */
static bool has_rv_or_provenance_call(Node *node, void *data) {
  const constants_t *constants = (const constants_t *)data;
  if (node == NULL)
    return false;
 
  if (IsA(node, OpExpr)) {
    OpExpr *op = (OpExpr *)node;
    if (rv_cmp_index(constants, op->opfuncid) >= 0)
      return true;
  }
 
  if (IsA(node, FuncExpr)) {
    FuncExpr *f = (FuncExpr *)node;
    if (f->funcid == constants->OID_FUNCTION_PROVENANCE)
      return true;
  }
 
  if (IsA(node, SubLink)) {
    SubLink *sl = (SubLink *)node;
    return has_rv_or_provenance_call((Node *)sl->testexpr, data);
  }
 
  /* Query nodes are opaque here; expression_tree_walker returns false
   * on them.  Explicit short-circuit just makes the intent obvious. */
  if (IsA(node, Query))
    return false;
 
  return expression_tree_walker(node, has_rv_or_provenance_call, data);
}
 
static bool has_provenance_walker(Node *node, void *data) {
  const constants_t *constants = (const constants_t *)data;
  if (node == NULL)
    return false;
 
  if (IsA(node, Query)) {
    Query *q = (Query *)node;
    ListCell *rc;
 
    /* Walk into CTE subqueries explicitly: they will be inlined as
     * subqueries by the rewriter, so a tracked-table inside one (or
     * an rv_cmp / provenance() call) matters for this query. */
    foreach (rc, q->cteList) {
      CommonTableExpr *cte = (CommonTableExpr *)lfirst(rc);
      if (has_provenance_walker((Node *)cte->ctequery, data))
        return true;
    }
 
    /* Walk this query's own expressions for rv_cmp OpExpr and
     * provenance() FuncExpr.  Use the SubLink-aware walker so we
     * don't descend into expression-context subqueries (they get
     * planned standalone; an rv_cmp inside one matters only to
     * that planning pass).
     *
     * This intentionally replaces a single query_tree_walker call:
     * that helper recurses with the passed walker into BOTH rtable
     * RTEs (RTE_SUBQUERY) and SubLink subselects, which would erase
     * the SubLink/RTE_SUBQUERY distinction we need. */
    if (has_rv_or_provenance_call((Node *)q->targetList, data))
      return true;
    if (has_rv_or_provenance_call((Node *)q->jointree, data))
      return true;
    if (has_rv_or_provenance_call((Node *)q->havingQual, data))
      return true;
    if (has_rv_or_provenance_call((Node *)q->returningList, data))
      return true;
 
    foreach (rc, q->rtable) {
      RangeTblEntry *r = (RangeTblEntry *)lfirst(rc);
      if (r->rtekind == RTE_RELATION) {
        ListCell *lc;
        AttrNumber attid = 1;
 
        foreach (lc, r->eref->colnames) {
          const char *v = strVal(lfirst(lc));
 
          if (!strcmp(v, PROVSQL_COLUMN_NAME) &&
              get_atttype(r->relid, attid) == constants->OID_TYPE_UUID) {
            return true;
          }
 
          ++attid;
        }
      } else if (r->rtekind == RTE_FUNCTION) {
        ListCell *lc;
        AttrNumber attid = 1;
 
        foreach (lc, r->functions) {
          RangeTblFunction *func = (RangeTblFunction *)lfirst(lc);
 
          if (func->funccolcount == 1) {
            FuncExpr *expr = (FuncExpr *)func->funcexpr;
            if (expr->funcresulttype == constants->OID_TYPE_UUID &&
                !strcmp(get_rte_attribute_name(r, attid),
                        PROVSQL_COLUMN_NAME)) {
              return true;
            }
          }
 
          attid += func->funccolcount;
        }
      } else if (r->rtekind == RTE_SUBQUERY && r->subquery != NULL) {
        /* A FROM-source subquery contributes its provenance to ours;
         * process_query recurses on it explicitly, so we must detect
         * tracked relations / rv_cmp / provenance() inside it. */
        if (has_provenance_walker((Node *)r->subquery, data))
          return true;
      }
    }
  }
 
  /* For non-Query nodes, use the expression-only walker.  It detects
   * rv_cmp OpExpr and provenance() FuncExpr inside arbitrary
   * sub-expressions (BoolExpr around an rv comparison, RV cmp under
   * IS-DISTINCT-FROM, ...) but stops at Query boundaries so a sibling
   * subquery's tracked rtable doesn't make THIS query's gate engage
   * (subqueries have their own planner_hook pass). */
  return has_rv_or_provenance_call(node, data);
}
 
/**
 * @brief Return true if @p q involves any provenance-bearing relation or
 *        contains an explicit @c provenance() call.
 *
 * This is the gate condition checked by @c provsql_planner before doing any
 * rewriting: if neither condition holds the query is passed through unchanged.
 *
 * @param constants  Extension OID cache.
 * @param q          Query to inspect.
 * @return  True if provenance rewriting is needed.
 */
static bool has_provenance(const constants_t *constants, Query *q) {
  return has_provenance_walker((Node *)q, (void *)constants);
}
 
/**
 * @brief Tree walker that detects any Var of type agg_token.
 * @param node  Current expression tree node.
 * @param data  Pointer to a @c constants_t (extension OID cache).
 * @return      @c true if an agg_token Var is found in @p node.
 */
static bool aggtoken_walker(Node *node, void *data) {
  const constants_t *constants = (const constants_t *) data;
  if (node == NULL)
    return false;
 
  if (IsA(node, Var)) {
    Var *v = (Var *) node;
    if(v->vartype == constants->OID_TYPE_AGG_TOKEN)
      return true;
  }
 
  return expression_tree_walker(node, aggtoken_walker, data);
}
 
/**
 * @brief Return true if @p node contains a @c Var of type @c agg_token.
 *
 * Used to detect whether a WHERE clause references an aggregate result
 * (which must be moved to HAVING).
 *
 * @param node       Expression tree to inspect.
 * @param constants  Extension OID cache.
 * @return  True if an @c agg_token @c Var is found anywhere in @p node.
 */
static bool has_aggtoken(Node *node, const constants_t *constants) {
  return expression_tree_walker(node, aggtoken_walker, (void*) constants);
}
 
/**
 * @brief Walker for @c needs_having_lift: detect any operand shape that
 *        the HAVING-lift rewriter (@c having_OpExpr_to_provenance_cmp)
 *        needs to handle specially.
 *
 * Returns @c true on:
 *   - a @c Var of type @c agg_token; or
 *   - a @c FuncExpr whose @c funcid is @c provenance_aggregate (the
 *     wrapper the planner-hook puts around aggregates over tracked
 *     non-RV columns -- yields @c agg_token).
 *
 * Anything else (deterministic scalars, plain @c Const, @c FuncExpr
 * over @c random_variable like @c expected / @c variance / @c moment,
 * comparisons of those) is left for PostgreSQL to evaluate natively;
 * the HAVING-lift never needs to touch it.
 */
static bool having_lift_walker(Node *node, void *data) {
  const constants_t *constants = (const constants_t *) data;
  if (node == NULL)
    return false;
 
  if (IsA(node, Var)) {
    Var *v = (Var *) node;
    if (v->vartype == constants->OID_TYPE_AGG_TOKEN)
      return true;
  }
 
  if (IsA(node, FuncExpr)) {
    FuncExpr *fe = (FuncExpr *) node;
    if (fe->funcid == constants->OID_FUNCTION_PROVENANCE_AGGREGATE)
      return true;
  }
 
  return expression_tree_walker(node, having_lift_walker, data);
}
 
/**
 * @brief Return true if @p havingQual contains anything the HAVING-lift
 *        path needs to handle (an @c agg_token Var or a
 *        @c provenance_aggregate wrapper).  A qual that returns @c false
 *        is left in place for PostgreSQL to evaluate, while the
 *        per-group provenance still gets a @c gate_delta wrapper.
 *
 * This is what lets a HAVING like @c expected(avg(rv)) > 20 work
 * directly: @c provsql.avg returns @c random_variable (not
 * @c agg_token), @c expected collapses to a scalar @c double, and the
 * surrounding comparison is a plain Boolean that PostgreSQL can filter
 * groups by without any provenance-side rewriting.
 */
static bool needs_having_lift(Node *havingQual, const constants_t *constants) {
  return expression_tree_walker(havingQual, having_lift_walker,
                                (void *) constants);
}
 
/**
 * @brief Rewrite an EXCEPT query into a LEFT JOIN with monus provenance.
 *
 * EXCEPT cannot be handled directly because it deduplicates.  This function
 * transforms:
 * @code
 *   SELECT … FROM A EXCEPT SELECT … FROM B
 * @endcode
 * into a LEFT JOIN of A and B on equality of all non-provenance columns,
 * clears @c setOperations, and leaves the monus token combination to
 * @c make_provenance_expression (which will see @c SR_MONUS).
 *
 * Only simple (non-chained) EXCEPT is supported; chained EXCEPT raises an
 * error.
 *
 * @param constants  Extension OID cache.
 * @param q          Query to rewrite in place.
 * @return  Always true (errors out on unsupported cases).
 */
static bool transform_except_into_join(const constants_t *constants, Query *q) {
  SetOperationStmt *setOps = (SetOperationStmt *)q->setOperations;
  RangeTblEntry *rte = makeNode(RangeTblEntry);
  FromExpr *fe = makeNode(FromExpr);
  JoinExpr *je = makeNode(JoinExpr);
  BoolExpr *expr = makeNode(BoolExpr);
  ListCell *lc;
  int attno = 1;
 
  if (!IsA(setOps->larg, RangeTblRef) || !IsA(setOps->rarg, RangeTblRef)) {
    provsql_error("Unsupported chain of EXCEPT operations");
  }
 
  expr->boolop = AND_EXPR;
  expr->location = -1;
  expr->args = NIL;
 
  foreach (lc, q->targetList) {
    TargetEntry *te = (TargetEntry *)lfirst(lc);
    Var *v;
 
    if (!IsA(te->expr, Var))
      provsql_error("EXCEPT query format not supported");
 
    v = (Var *)te->expr;
 
    if (v->vartype != constants->OID_TYPE_UUID) {
      OpExpr *oe = makeNode(OpExpr);
      Oid opno = find_equality_operator(v->vartype, v->vartype);
      Operator opInfo = SearchSysCache1(OPEROID, ObjectIdGetDatum(opno));
      Form_pg_operator opform;
      Var *leftArg, *rightArg;
 
      if (!HeapTupleIsValid(opInfo))
        provsql_error("could not find operator with OID %u to compare variables of type %u",
                      opno, v->vartype);
 
      opform = (Form_pg_operator)GETSTRUCT(opInfo);
      leftArg = makeNode(Var);
      rightArg = makeNode(Var);
 
      oe->opno = opno;
      oe->opfuncid = opform->oprcode;
      oe->opresulttype = opform->oprresult;
      oe->opcollid = InvalidOid;
      oe->inputcollid = DEFAULT_COLLATION_OID;
 
      leftArg->varno = ((RangeTblRef *)setOps->larg)->rtindex;
      rightArg->varno = ((RangeTblRef *)setOps->rarg)->rtindex;
      leftArg->varattno = rightArg->varattno = attno;
 
#if PG_VERSION_NUM >= 130000
      leftArg->varnosyn = rightArg->varnosyn = 0;
      leftArg->varattnosyn = rightArg->varattnosyn = 0;
#else
      leftArg->varnoold = leftArg->varno;
      rightArg->varnoold = rightArg->varno;
      leftArg->varoattno = rightArg->varoattno = attno;
#endif
 
      leftArg->vartype = rightArg->vartype = v->vartype;
      leftArg->varcollid = rightArg->varcollid = InvalidOid;
      leftArg->vartypmod = rightArg->vartypmod = -1;
      leftArg->location = rightArg->location = -1;
 
      oe->args = list_make2(leftArg, rightArg);
      oe->location = -1;
      expr->args = lappend(expr->args, oe);
 
      ReleaseSysCache(opInfo);
    }
 
    ++attno;
  }
 
  /* Populate the JOIN RTE's eref / joinaliasvars / joinleftcols /
   * joinrightcols by walking the larg and rarg subqueries' targetLists.
   * Execution doesn't need these (outer Vars reference the input RTEs
   * directly), but PostgreSQL's ruleutils deparser walks them when
   * pg_get_querydef / EXPLAIN VERBOSE traverse the rewritten tree and
   * segfaults on NULL eref. Non-USING LEFT JOIN: joinmergedcols = 0,
   * output is left columns followed by right columns. */
  {
    RangeTblRef *larg_ref = (RangeTblRef *)setOps->larg;
    RangeTblRef *rarg_ref = (RangeTblRef *)setOps->rarg;
    RangeTblEntry *larg_rte =
      (RangeTblEntry *)list_nth(q->rtable, larg_ref->rtindex - 1);
    RangeTblEntry *rarg_rte =
      (RangeTblEntry *)list_nth(q->rtable, rarg_ref->rtindex - 1);
    List *aliasvars = NIL;
    List *leftcols = NIL;
    List *rightcols = NIL;
    List *colnames = NIL;
    ListCell *lc_te;
    int colno;
 
    colno = 1;
    foreach (lc_te, larg_rte->subquery->targetList) {
      TargetEntry *te = (TargetEntry *)lfirst(lc_te);
      if (te->resjunk) {
        colno++;
        continue;
      }
      aliasvars = lappend(aliasvars,
                          makeVar(larg_ref->rtindex, colno,
                                  exprType((Node *)te->expr),
                                  exprTypmod((Node *)te->expr),
                                  exprCollation((Node *)te->expr),
                                  0));
      leftcols = lappend_int(leftcols, colno);
      rightcols = lappend_int(rightcols, 0);
      colnames = lappend(colnames,
                         makeString(pstrdup(te->resname ? te->resname
                                                        : "?column?")));
      colno++;
    }
    colno = 1;
    foreach (lc_te, rarg_rte->subquery->targetList) {
      TargetEntry *te = (TargetEntry *)lfirst(lc_te);
      if (te->resjunk) {
        colno++;
        continue;
      }
      aliasvars = lappend(aliasvars,
                          makeVar(rarg_ref->rtindex, colno,
                                  exprType((Node *)te->expr),
                                  exprTypmod((Node *)te->expr),
                                  exprCollation((Node *)te->expr),
                                  0));
      leftcols = lappend_int(leftcols, 0);
      rightcols = lappend_int(rightcols, colno);
      colnames = lappend(colnames,
                         makeString(pstrdup(te->resname ? te->resname
                                                        : "?column?")));
      colno++;
    }
 
    rte->alias = NULL;
    rte->eref = makeAlias("unnamed_join", colnames);
    rte->joinaliasvars = aliasvars;
#if PG_VERSION_NUM >= 130000
    rte->joinleftcols = leftcols;
    rte->joinrightcols = rightcols;
    rte->joinmergedcols = 0;
#else
    (void) leftcols;
    (void) rightcols;
#endif
  }
 
  rte->rtekind = RTE_JOIN;
  rte->jointype = JOIN_LEFT;
 
  q->rtable = lappend(q->rtable, rte);
 
  je->jointype = JOIN_LEFT;
 
  je->larg = setOps->larg;
  je->rarg = setOps->rarg;
  je->quals = (Node *)expr;
  je->rtindex = list_length(q->rtable);
 
  fe->fromlist = list_make1(je);
 
  q->jointree = fe;
 
  // TODO: Add group by in the right-side table
 
  q->setOperations = 0;
 
  return true;
}
 
/**
 * @brief Recursively annotate a UNION tree with the provenance UUID type.
 *
 * Walks the @c SetOperationStmt tree of a UNION and appends the UUID type
 * to @c colTypes / @c colTypmods / @c colCollations on every node, and sets
 * @c all = true so that PostgreSQL does not deduplicate the combined stream.
 * The non-ALL deduplication has already been moved to an outer GROUP BY by
 * @c rewrite_non_all_into_external_group_by before this is called.
 *
 * @param constants  Extension OID cache.
 * @param stmt       Root (or subtree) of the UNION @c SetOperationStmt.
 * @param q          Outer query (to look up subquery RTEs for agg_token type updates).
 */
static void process_set_operation_union(const constants_t *constants,
                                        SetOperationStmt *stmt,
                                        Query *q) {
  if (stmt->op != SETOP_UNION) {
    provsql_error("Unsupported mixed set operations");
  }
  if (IsA(stmt->larg, SetOperationStmt)) {
    process_set_operation_union(constants, (SetOperationStmt *)(stmt->larg), q);
  }
  if (IsA(stmt->rarg, SetOperationStmt)) {
    process_set_operation_union(constants, (SetOperationStmt *)(stmt->rarg), q);
  }
 
  /* Update colTypes for columns that became agg_token after rewriting.
   * Use the left branch's subquery to detect agg_token columns. */
  if (IsA(stmt->larg, RangeTblRef)) {
    Index rtindex = ((RangeTblRef *)stmt->larg)->rtindex;
    RangeTblEntry *rte = list_nth_node(RangeTblEntry, q->rtable, rtindex - 1);
    if (rte->rtekind == RTE_SUBQUERY && rte->subquery != NULL) {
      ListCell *lc_type = list_head(stmt->colTypes);
      ListCell *lc_te = list_head(rte->subquery->targetList);
      while (lc_type != NULL && lc_te != NULL) {
        TargetEntry *te = (TargetEntry *)lfirst(lc_te);
        if (exprType((Node *)te->expr) == constants->OID_TYPE_AGG_TOKEN) {
          lfirst_oid(lc_type) = constants->OID_TYPE_AGG_TOKEN;
        }
        lc_type = my_lnext(stmt->colTypes, lc_type);
        lc_te = my_lnext(rte->subquery->targetList, lc_te);
      }
    }
  }
 
  stmt->colTypes = lappend_oid(stmt->colTypes, constants->OID_TYPE_UUID);
  stmt->colTypmods = lappend_int(stmt->colTypmods, -1);
  stmt->colCollations = lappend_int(stmt->colCollations, 0);
  stmt->all = true;
}
 
/**
 * @brief Add a WHERE condition filtering out zero-provenance tuples.
 *
 * For EXCEPT queries, tuples whose provenance evaluates to zero (i.e., the
 * right-hand side fully subsumes the left-hand side) must be excluded from
 * the result.  This function appends @c provsql <> gate_zero() to
 * @p q->jointree->quals, ANDing with any existing WHERE condition.
 *
 * @param constants  Extension OID cache.
 * @param q          Query to modify in place.
 * @param provsql    Provenance expression that was added to the SELECT list.
 */
static void add_select_non_zero(const constants_t *constants, Query *q,
                                Expr *provsql) {
  FuncExpr *gate_zero = makeNode(FuncExpr);
  OpExpr *oe = makeNode(OpExpr);
 
  gate_zero->funcid = constants->OID_FUNCTION_GATE_ZERO;
  gate_zero->funcresulttype = constants->OID_TYPE_UUID;
 
  oe->opno = constants->OID_OPERATOR_NOT_EQUAL_UUID;
  oe->opfuncid = constants->OID_FUNCTION_NOT_EQUAL_UUID;
  oe->opresulttype = BOOLOID;
  oe->args = list_make2(provsql, gate_zero);
  oe->location = -1;
 
  if (q->jointree->quals != NULL) {
    BoolExpr *be = makeNode(BoolExpr);
 
    be->boolop = AND_EXPR;
    be->args = list_make2(oe, q->jointree->quals);
    be->location = -1;
 
    q->jointree->quals = (Node *)be;
  } else
    q->jointree->quals = (Node *)oe;
}
 
/**
 * @brief Append @p expr to @p havingQual with an AND, creating one if needed.
 *
 * If @p havingQual is NULL, returns @p expr directly.  If it is already an
 * AND @c BoolExpr, appends to its argument list.  Otherwise wraps both in a
 * new AND node.
 *
 * @param havingQual  Existing HAVING qualifier, or NULL.
 * @param expr        Expression to conjoin.
 * @return  The updated HAVING qualifier.
 */
static Node *add_to_havingQual(Node *havingQual, Expr *expr)
{
  if(!havingQual) {
    havingQual = (Node*) expr;
  } else if(IsA(havingQual, BoolExpr) && ((BoolExpr*)havingQual)->boolop==AND_EXPR) {
    BoolExpr *be = (BoolExpr*)havingQual;
    be->args = lappend(be->args, expr);
  } else if(IsA(havingQual, OpExpr) || IsA(havingQual, BoolExpr)) {
    /* BoolExpr that is not an AND (OR/NOT): wrap with a new AND node. */
    BoolExpr *be = makeNode(BoolExpr);
    be->boolop=AND_EXPR;
    be->location=-1;
    be->args = list_make2(havingQual, expr);
    havingQual = (Node*) be;
  } else
    provsql_error("Unknown structure within Boolean expression");
 
  return havingQual;
}
 
/**
 * @brief Check whether @p op is a supported comparison on an aggregate result.
 *
 * Returns true iff @p op is a two-argument operator where at least one
 * argument is a @c Var of type @c agg_token (or an implicit-cast wrapper
 * thereof) and the other is a @c Const (possibly cast).  This is the set
 * of WHERE-on-aggregate patterns that ProvSQL can safely move to a HAVING
 * clause.
 *
 * @param op         The @c OpExpr to inspect.
 * @param constants  Extension OID cache.
 * @return  True if the pattern is supported, false otherwise.
 */
static bool check_selection_on_aggregate(OpExpr *op, const constants_t *constants)
{
  bool ok=true;
  bool found_agg_token=false;
 
  if(op->args->length != 2)
    return false;
 
  for(unsigned i=0; i<2; ++i) {
    Node *arg = lfirst(list_nth_cell(op->args, i));
 
    // Check both arguments are either an aggtoken or a constant
    // (possibly after a cast)
    if((IsA(arg, Var) && ((Var*)arg)->vartype==constants->OID_TYPE_AGG_TOKEN)) {
      found_agg_token=true;
    } else if(IsA(arg, Const)) {
    } else if(IsA(arg, FuncExpr)) {
      FuncExpr *fe = (FuncExpr*) arg;
      if(fe->funcformat != COERCE_IMPLICIT_CAST && fe->funcformat != COERCE_EXPLICIT_CAST) {
        ok=false;
        break;
      }
      if(fe->args->length != 1) {
        ok=false;
        break;
      }
      if(!IsA(lfirst(list_head(fe->args)), Const)) {
        ok=false;
        break;
      }
    } else {
      ok=false;
      break;
    }
  }
 
  return ok && found_agg_token;
}
 
/**
 * @brief Check whether every leaf of a Boolean expression is a supported
 *        comparison on an aggregate result.
 *
 * Recursively validates @c OpExpr leaves via @c check_selection_on_aggregate
 * and descends into nested @c BoolExpr nodes.
 *
 * @param be         The Boolean expression to validate.
 * @param constants  Extension OID cache.
 * @return  True if all leaves are supported, false if any is not.
 */
static bool check_boolexpr_on_aggregate(BoolExpr *be, const constants_t *constants)
{
  ListCell *lc;
 
  foreach (lc, be->args) {
    Node *n=lfirst(lc);
    if(IsA(n, OpExpr)) {
      if(!check_selection_on_aggregate((OpExpr*) n, constants))
        return false;
    } else if(IsA(n, BoolExpr)) {
      if(!check_boolexpr_on_aggregate((BoolExpr*) n, constants))
        return false;
    } else
      return false;
  }
 
  return true;
}
 
/**
 * @brief Top-level dispatcher for supported WHERE-on-aggregate patterns.
 *
 * @param expr       Expression to validate (@c OpExpr or @c BoolExpr).
 * @param constants  Extension OID cache.
 * @return  True if ProvSQL can handle this expression.
 */
static bool check_expr_on_aggregate(Expr *expr, const constants_t *constants) {
  switch(expr->type) {
  case T_BoolExpr:
    return check_boolexpr_on_aggregate((BoolExpr*) expr, constants);
  case T_OpExpr:
    return check_selection_on_aggregate((OpExpr*) expr, constants);
  default:
    provsql_error("Unknown structure within Boolean expression");
  }
}
 
/* -------------------------------------------------------------------------
 * Main query transformation
 * ------------------------------------------------------------------------- */
 
/**
 * @brief Build the per-RTE column-numbering map used by where-provenance.
 *
 * Assigns a sequential position (1, 2, 3, …) to every non-provenance,
 * non-join, non-empty column across all RTEs in @p q->rtable.  The
 * @c provsql column is assigned -1 so callers can detect provenance-tracked
 * RTEs.  Join-RTE columns and empty-named columns (used for anonymous GROUP
 * BY keys) are assigned 0.
 *
 * @note For @c RTE_RELATION entries that are provenance-tracked, the
 *       sequential numbers produced here must @b not be used as PROJECT gate
 *       positions.  Because numbering is query-order-dependent, the sequential
 *       number for a column of a provenance table that is not the first RTE
 *       will exceed @c nb_columns of that table's IN gate, causing
 *       @c WhereCircuit::evaluate() to return an empty locator set.  Instead,
 *       callers should use @c varattno directly (see
 *       @c make_provenance_expression()).  The -1 sentinel is the reliable
 *       way to identify a provenance-tracked RTE.
 *
 * @param q         Query whose range table is mapped.
 * @param columns   Pre-allocated array of length @p q->rtable->length.
 *                  Each element is allocated and filled by this function.
 * @param nbcols    Out-param: total number of non-provenance output columns.
 */
static void build_column_map(Query *q, int **columns, int *nbcols) {
  unsigned i = 0;
  ListCell *l;
 
  *nbcols = 0;
 
  foreach (l, q->rtable) {
    RangeTblEntry *r = (RangeTblEntry *)lfirst(l);
    ListCell *lc;
 
    columns[i] = 0;
    if (r->eref) {
      unsigned j = 0;
 
      columns[i] = (int *)palloc(r->eref->colnames->length * sizeof(int));
 
      foreach (lc, r->eref->colnames) {
        if (!lfirst(lc)) {
          /* Column without name — used e.g. when grouping by a discarded column */
          columns[i][j] = ++(*nbcols);
        } else {
          const char *v = strVal(lfirst(lc));
 
          if (strcmp(v, "") && r->rtekind != RTE_JOIN) { /* join RTE columns ignored */
            if (!strcmp(v, PROVSQL_COLUMN_NAME))
              columns[i][j] = -1;
            else
              columns[i][j] = ++(*nbcols);
          } else {
            columns[i][j] = 0;
          }
        }
 
        ++j;
      }
    }
 
    ++i;
  }
}
 
/**
 * @brief Categorisation of a top-level WHERE conjunct.
 *
 * Drives the unified WHERE classifier that replaced the original pair
 * @c migrate_aggtoken_quals_to_having + @c extract_rv_cmps_from_quals.
 * Both probabilistic flavours (agg_token's "moved to HAVING" world and
 * random_variable's "lifted to provenance" world) are special cases of
 * "this conjunct involves a probabilistic value the executor cannot
 * evaluate as a Boolean directly, so the planner has to route it to a
 * different evaluation site".  The classifier reports which site, or
 * (for unsupported mixes) errors.
 */
typedef enum {
  QUAL_DETERMINISTIC,    /**< no probabilistic value; stays in WHERE        */
  QUAL_PURE_AGG,         /**< pure agg_token expression; route to HAVING    */
  QUAL_PURE_RV,          /**< pure random_variable expression; lift to provenance */
  QUAL_MIXED_AGG_DET,    /**< agg_token mixed with non-agg leaves; error    */
  QUAL_MIXED_RV_DET,     /**< random_variable mixed with non-RV leaves; error */
  QUAL_MIXED_AGG_RV      /**< agg_token and random_variable in the same expr; error */
} qual_class;
 
/**
 * @brief Classify @p expr along the @c qual_class axis.
 *
 * Decision table (the predicates @c has_aggtoken,
 * @c expr_contains_rv_cmp, @c check_expr_on_aggregate, and
 * @c check_expr_on_rv each return whether the expression "contains" or
 * "is purely" the corresponding flavour):
 *
 * | aggtoken | rv_cmp | check_agg | check_rv | classification        |
 * |----------|--------|-----------|----------|-----------------------|
 * | yes      | yes    |    -      |    -     | QUAL_MIXED_AGG_RV     |
 * | yes      | no     |   true    |    -     | QUAL_PURE_AGG         |
 * | yes      | no     |   false   |    -     | QUAL_MIXED_AGG_DET    |
 * | no       | yes    |    -      |   true   | QUAL_PURE_RV          |
 * | no       | yes    |    -      |   false  | QUAL_MIXED_RV_DET     |
 * | no       | no     |    -      |    -     | QUAL_DETERMINISTIC    |
 */
static qual_class classify_qual(Expr *expr, const constants_t *constants)
{
  bool has_agg = has_aggtoken((Node *)expr, constants);
  bool has_rv  = expr_contains_rv_cmp((Node *)expr, constants);
 
  if (has_agg && has_rv)
    return QUAL_MIXED_AGG_RV;
  if (has_agg) {
    if (check_expr_on_aggregate(expr, constants))
      return QUAL_PURE_AGG;
    return QUAL_MIXED_AGG_DET;
  }
  if (has_rv) {
    if (check_expr_on_rv(expr, constants))
      return QUAL_PURE_RV;
    return QUAL_MIXED_RV_DET;
  }
  return QUAL_DETERMINISTIC;
}
 
/** @brief Raise the user-facing error appropriate to a mixed @p c.
 *
 * Each @c provsql_error call is @c ereport(ERROR), which does not
 * return; the explicit @c break statements below are present only to
 * keep @c -Wimplicit-fallthrough happy (PostgreSQL's @c elog macro is
 * not marked @c noreturn for the compiler's flow analysis). */
static void error_for_mixed_qual(qual_class c)
{
  switch (c) {
    case QUAL_MIXED_AGG_DET:
      provsql_error("Complex selection on aggregation results not supported");
      break;
    case QUAL_MIXED_RV_DET:
      provsql_error("WHERE clause mixes random_variable comparisons with "
                    "other predicates inside the same Boolean expression; "
                    "split the non-RV part into its own AND conjunct");
      break;
    case QUAL_MIXED_AGG_RV:
      provsql_error("WHERE clause mixes agg_token (HAVING-style) and "
                    "random_variable (per-tuple) comparisons inside the "
                    "same Boolean expression; this combination is not "
                    "supported");
      break;
    default:
      /* QUAL_DETERMINISTIC / QUAL_PURE_AGG / QUAL_PURE_RV: not a mixed case. */
      break;
  }
}
 
/**
 * @brief Unified WHERE classifier &ndash; routes each top-level conjunct
 *        to the right evaluation site in a single pass.
 *
 * Replaces and consolidates the original
 * @c migrate_aggtoken_quals_to_having (agg-only) and
 * @c extract_rv_cmps_from_quals (rv-only).  The two old functions were
 * structurally isomorphic: each walked the WHERE clause, classified
 * each top-level conjunct, and routed pure-X conjuncts somewhere
 * semantic (HAVING vs the returned rv_cmps list); the deterministic
 * conjuncts stayed in WHERE.  Doing it in one pass means the rare
 * conjunct that mixes agg_token and random_variable (which neither old
 * function would have caught cleanly) gets a deterministic, useful
 * error message.
 *
 * Supported shapes mirror the union of the two predecessors:
 * - Whole WHERE is a single conjunct: classify and route or error.
 * - Top-level AND of conjuncts: classify each, route, and (after
 *   walking) collapse the AND if it has zero or one remaining children
 *   so downstream code does not see a degenerate Boolean node.
 * - Top-level OR / NOT containing both deterministic and probabilistic
 *   leaves: error.
 *
 * @param constants  Extension OID cache.
 * @param q          Query whose @c jointree->quals and @c havingQual
 *                   may both be mutated in place.
 * @return List of @c FuncExpr nodes (one per lifted RV conjunct), each
 *         producing a @c UUID.  The caller conjoins these into
 *         @c prov_atts before @c make_provenance_expression.
 */
static List *
migrate_probabilistic_quals(const constants_t *constants, Query *q)
{
  List *rv_cmps = NIL;
  Node *quals;
 
  if (!q->jointree || !q->jointree->quals)
    return NIL;
 
  quals = q->jointree->quals;
 
  /* Whole WHERE is one conjunct (single OpExpr, or non-AND BoolExpr
   * which we treat opaquely &ndash; the per-flavour pure checks
   * @c check_expr_on_aggregate / @c check_expr_on_rv recurse through
   * the BoolExpr structure themselves). */
  if (!IsA(quals, BoolExpr) || ((BoolExpr *)quals)->boolop != AND_EXPR) {
    qual_class c = classify_qual((Expr *)quals, constants);
    error_for_mixed_qual(c);
 
    switch (c) {
      case QUAL_PURE_AGG:
        q->havingQual = add_to_havingQual(q->havingQual, (Expr *)quals);
        q->jointree->quals = NULL;
        break;
      case QUAL_PURE_RV:
        rv_cmps = lappend(rv_cmps,
                          rv_Expr_to_provenance((Expr *)quals,
                                                constants, false));
        q->jointree->quals = NULL;
        break;
      case QUAL_DETERMINISTIC:
        /* Leave WHERE alone. */
        break;
      default:
        /* Errors handled by error_for_mixed_qual. */
        break;
    }
    return rv_cmps;
  }
 
  /* Top-level AND: walk conjuncts. */
  {
    BoolExpr *be = (BoolExpr *)quals;
    ListCell *cell, *prev;
 
    for (cell = list_head(be->args), prev = NULL; cell != NULL;) {
      Expr *conjunct = (Expr *)lfirst(cell);
      qual_class c = classify_qual(conjunct, constants);
 
      error_for_mixed_qual(c);
 
      switch (c) {
        case QUAL_PURE_AGG:
          q->havingQual = add_to_havingQual(q->havingQual, conjunct);
          be->args = my_list_delete_cell(be->args, cell, prev);
          if (prev)
            cell = my_lnext(be->args, prev);
          else
            cell = list_head(be->args);
          break;
        case QUAL_PURE_RV:
          rv_cmps = lappend(rv_cmps,
                            rv_Expr_to_provenance(conjunct,
                                                  constants, false));
          be->args = my_list_delete_cell(be->args, cell, prev);
          if (prev)
            cell = my_lnext(be->args, prev);
          else
            cell = list_head(be->args);
          break;
        case QUAL_DETERMINISTIC:
          prev = cell;
          cell = my_lnext(be->args, cell);
          break;
        default:
          /* Errors handled by error_for_mixed_qual. */
          break;
      }
    }
 
    /* Collapse degenerate ANDs so downstream code sees a tidy WHERE. */
    if (be->args == NIL)
      q->jointree->quals = NULL;
    else if (list_length(be->args) == 1)
      q->jointree->quals = (Node *)linitial(be->args);
  }
 
  return rv_cmps;
}
 
/** @brief Context for the @c insert_agg_token_casts_mutator. */
typedef struct insert_agg_token_casts_context {
  Query *query;               ///< Outer query (to look up subquery RTEs)
  const constants_t *constants; ///< Extension OID cache
} insert_agg_token_casts_context;
 
/**
 * @brief Look up the original aggregate return type for an agg_token Var.
 *
 * Navigates from the Var's varno/varattno to the subquery's target list,
 * finds the provenance_aggregate() FuncExpr, and extracts the type OID
 * from its second argument (aggtype).
 */
static Oid get_agg_token_orig_type(Var *v, insert_agg_token_casts_context *ctx) {
  RangeTblEntry *rte;
  TargetEntry *te;
 
  if (v->varno < 1 || v->varno > list_length(ctx->query->rtable))
    return InvalidOid;
 
  rte = list_nth_node(RangeTblEntry, ctx->query->rtable, v->varno - 1);
  if (rte->rtekind != RTE_SUBQUERY || rte->subquery == NULL)
    return InvalidOid;
 
  if (v->varattno < 1 || v->varattno > list_length(rte->subquery->targetList))
    return InvalidOid;
 
  te = list_nth_node(TargetEntry, rte->subquery->targetList, v->varattno - 1);
  if (IsA(te->expr, FuncExpr)) {
    FuncExpr *f = (FuncExpr *)te->expr;
    if (f->funcid == ctx->constants->OID_FUNCTION_PROVENANCE_AGGREGATE) {
      Const *aggtype_const = (Const *)lsecond(f->args);
      return DatumGetObjectId(aggtype_const->constvalue);
    }
  }
  return InvalidOid;
}
 
/**
 * @brief Wrap an agg_token Var in a cast to its original type, in place.
 */
static void cast_agg_token_in_list(ListCell *lc,
                                   insert_agg_token_casts_context *ctx) {
  Var *v = (Var *)lfirst(lc);
  Oid target = get_agg_token_orig_type(v, ctx);
  HeapTuple castTuple;
 
  if (!OidIsValid(target))
    return;
 
  castTuple = SearchSysCache2(CASTSOURCETARGET,
                              ObjectIdGetDatum(ctx->constants->OID_TYPE_AGG_TOKEN),
                              ObjectIdGetDatum(target));
  if (HeapTupleIsValid(castTuple)) {
    Form_pg_cast castForm = (Form_pg_cast)GETSTRUCT(castTuple);
    if (OidIsValid(castForm->castfunc)) {
      FuncExpr *fc = makeNode(FuncExpr);
      fc->funcid = castForm->castfunc;
      fc->funcresulttype = target;
      fc->funcretset = false;
      fc->funcvariadic = false;
      fc->funcformat = COERCE_IMPLICIT_CAST;
      fc->funccollid = InvalidOid;
      fc->inputcollid = InvalidOid;
      fc->args = list_make1(v);
      fc->location = -1;
      lfirst(lc) = fc;
    }
    ReleaseSysCache(castTuple);
  }
}
 
/**
 * @brief Wrap any agg_token Vars in an argument list.
 */
static void cast_agg_token_args(List *args,
                                insert_agg_token_casts_context *ctx) {
  ListCell *lc;
  foreach (lc, args) {
    if (IsA(lfirst(lc), Var) &&
        ((Var *)lfirst(lc))->vartype == ctx->constants->OID_TYPE_AGG_TOKEN)
      cast_agg_token_in_list(lc, ctx);
  }
}
 
/**
 * @brief Insert agg_token casts for Vars used in expressions.
 *
 * After the WHERE-to-HAVING migration, agg_token Vars remaining in
 * expression nodes (OpExpr, WindowFunc, CoalesceExpr, MinMaxExpr, etc.)
 * need explicit casts to their original type so that operators and
 * functions receive correct values.  The original type is looked up
 * from the provenance_aggregate() call in the subquery.
 */
static Node *
insert_agg_token_casts_mutator(Node *node, void *data) {
  insert_agg_token_casts_context *ctx = (insert_agg_token_casts_context *)data;
 
  if (node == NULL)
    return NULL;
 
  if (IsA(node, OpExpr)) {
    cast_agg_token_args(((OpExpr *)node)->args, ctx);
    return (Node *)node;
  }
  if (IsA(node, WindowFunc)) {
    cast_agg_token_args(((WindowFunc *)node)->args, ctx);
    return (Node *)node;
  }
  if (IsA(node, CoalesceExpr)) {
    cast_agg_token_args(((CoalesceExpr *)node)->args, ctx);
    return (Node *)node;
  }
  if (IsA(node, MinMaxExpr)) {
    cast_agg_token_args(((MinMaxExpr *)node)->args, ctx);
    return (Node *)node;
  }
  if (IsA(node, NullIfExpr)) {
    cast_agg_token_args(((NullIfExpr *)node)->args, ctx);
    return (Node *)node;
  }
 
  return expression_tree_mutator(node, insert_agg_token_casts_mutator, data);
}
 
/**
 * @brief Walk query and insert agg_token casts where needed.
 */
static void insert_agg_token_casts(const constants_t *constants, Query *q) {
  insert_agg_token_casts_context ctx = {q, constants};
  query_tree_mutator(q, insert_agg_token_casts_mutator, &ctx,
                     QTW_DONT_COPY_QUERY | QTW_IGNORE_RC_SUBQUERIES);
}
 
/**
 * @brief Wrap @p expr in a @c provsql.assume_boolean FuncExpr.
 *
 * Used by @c make_provenance_expression when its caller (the
 * safe-query rewrite path in @c process_query) flagged the result
 * as needing the @c gate_assumed_boolean structural marker.
 * Wrapping at expression-build time rather than at splice time
 * means @c add_to_select and
 * @c replace_provenance_function_by_expression both consume the
 * already-wrapped expression, so every per-row root occurrence in
 * the final target list -- the auto-added @c provsql column and
 * every substituted user-side @c provenance() call -- carries the
 * wrapper uniformly.
 *
 * @param constants  Extension OID cache.
 * @param expr       Provenance expression to wrap.
 * @return  A @c FuncExpr applying @c provsql.assume_boolean to @p expr.
 */
static Expr *wrap_in_assume_boolean(const constants_t *constants,
                                    Expr *expr) {
  FuncExpr *wrap = makeNode(FuncExpr);
  wrap->funcid = constants->OID_FUNCTION_ASSUME_BOOLEAN;
  wrap->funcresulttype = constants->OID_TYPE_UUID;
  wrap->funcretset = false;
  wrap->funcvariadic = false;
  wrap->funcformat = COERCE_EXPLICIT_CALL;
  wrap->funccollid = InvalidOid;
  wrap->inputcollid = InvalidOid;
  wrap->args = list_make1(expr);
  wrap->location = -1;
  return (Expr *) wrap;
}
 
/**
 * @brief Rewrite a single SELECT query to carry provenance.
 *
 * This is the recursive entry point for the provenance rewriter.  It is
 * called from @c provsql_planner for top-level queries and re-entered from
 * @c get_provenance_attributes for subqueries in FROM.
 *
 * High-level steps:
 *  1. Strip any @c provsql column propagated into this query's target list.
 *  2. Detect and rewrite structural forms requiring pre-processing:
 *     non-ALL set operations (wrap in outer GROUP BY), AGG DISTINCT (push
 *     into a subquery), DISTINCT (convert to GROUP BY).
 *  3. Collect provenance attributes via @c get_provenance_attributes.
 *  4. Build a column-numbering map for where-provenance (@c build_column_map).
 *  5. Handle aggregates, migrate WHERE-on-aggregate to HAVING, and set ops.
 *  6. Build and splice the combined provenance expression.
 *
 * @param constants  Extension OID cache.
 * @param q          Query to rewrite (modified in place).
 * @param removed    Out-param: boolean array indicating which original target
 *                   list entries were provenance columns and were removed.
 *                   May be @c NULL if the caller does not need this info.
 * @param wrap_root  If true, mark this query's provenance expression as a
 *                   safe-query root that must be wrapped in
 *                   @c provsql.assume_boolean before splicing.
 * @return  The (possibly restructured) rewritten query, or @c NULL if the
 *          query has no FROM clause and can be skipped.
 */
static Query *process_query(const constants_t *constants, Query *q,
                            bool **removed, bool wrap_root) {
  List *prov_atts;
  bool has_union = false;
  bool has_difference = false;
  bool supported = true;
  bool group_by_rewrite = false;
  int nbcols = 0;
  int **columns;
  unsigned i = 0;
  if (provsql_verbose >= 50)
    elog_node_display(NOTICE, "ProvSQL: Before query rewriting", q, true);
 
  if (q->rtable == NULL) {
    /* FROM-less SELECT: the rest of the rewriter indexes into
     * q->rtable, so it can't process anything tied to a base relation.
     * But a WHERE-on-RV is still meaningful in this shape (e.g.
     *   SELECT 1 WHERE normal(0,1) > 2)
     * since the comparison produces a pure-rv gate that's lifted into
     * a synthesised provsql column on the single result row.  Run only
     * the qual migration + targetList splice and return; everything
     * else this function does (column mapping, set-ops, aggregation
     * rewriting, ...) assumes a non-empty rtable. */
    List *rv_cmps = migrate_probabilistic_quals(constants, q);
    if (rv_cmps != NIL) {
      Expr *provenance;
      RangeTblEntry *values_rte;
      RangeTblRef *rtr;
      Var *v;
 
      if (list_length(rv_cmps) == 1) {
        provenance = (Expr *)linitial(rv_cmps);
      } else {
        /* Multiple rv conjuncts: combine via provenance_times. */
        FuncExpr *times = makeNode(FuncExpr);
        ArrayExpr *array = makeNode(ArrayExpr);
        times->funcid = constants->OID_FUNCTION_PROVENANCE_TIMES;
        times->funcresulttype = constants->OID_TYPE_UUID;
        times->funcvariadic = true;
        times->location = -1;
        array->array_typeid = constants->OID_TYPE_UUID_ARRAY;
        array->element_typeid = constants->OID_TYPE_UUID;
        array->elements = rv_cmps;
        array->location = -1;
        times->args = list_make1(array);
        provenance = (Expr *)times;
      }
 
      /* Bind the lifted expression to a single evaluation by wrapping
       * it in a synthesized FROM (VALUES (<expr>)) AS _prov_(provsql).
       * Without this, multiple references to the same provenance
       * expression in the outer targetList (the user's provenance()
       * call, plus the auto-added provsql column) each re-invoke any
       * rv constructor inside, producing distinct UUIDs per call
       * because uniform / normal / ... mint a fresh leaf gate each
       * time.  Wrapping in VALUES gives one evaluation site that all
       * outer references read from. */
      values_rte = makeNode(RangeTblEntry);
      values_rte->rtekind = RTE_VALUES;
      values_rte->values_lists = list_make1(list_make1(provenance));
      values_rte->coltypes = list_make1_oid(constants->OID_TYPE_UUID);
      values_rte->coltypmods = list_make1_int(-1);
      values_rte->colcollations = list_make1_oid(InvalidOid);
      values_rte->eref = makeAlias(
        "_prov_",
        list_make1(makeString(pstrdup(PROVSQL_COLUMN_NAME))));
      values_rte->inh = false;
      values_rte->inFromCl = true;
#if PG_VERSION_NUM < 160000
      values_rte->requiredPerms = 0;
#endif
      q->rtable = list_make1(values_rte);
 
      rtr = makeNode(RangeTblRef);
      rtr->rtindex = 1;
      if (q->jointree == NULL) {
        q->jointree = makeNode(FromExpr);
      }
      q->jointree->fromlist = list_make1(rtr);
 
      v = makeVar(1, 1, constants->OID_TYPE_UUID, -1, InvalidOid, 0);
 
      /* Substitute any provenance() FuncExpr in the targetList with
       * a reference to the bound expression. */
      replace_provenance_function_by_expression(constants, q, (Expr *)v);
 
      /* Append a provsql column reading the same Var so callers that
       * expect the auto-added column find it. */
      {
        TargetEntry *te = makeTargetEntry(
          (Expr *)copyObject(v),
          list_length(q->targetList) + 1,
          pstrdup(PROVSQL_COLUMN_NAME),
          false);
        q->targetList = lappend(q->targetList, te);
      }
    }
    return q;
  }
 
  /* Inline non-recursive CTE references as subqueries so we can track
   * provenance through them. Must happen before set operation handling
   * since UNION/EXCEPT branches may reference CTEs. */
  inline_ctes(q);
 
  {
    Bitmapset *removed_sortgrouprefs = NULL;
 
    if (q->targetList) {
      removed_sortgrouprefs =
        remove_provenance_attributes_select(constants, q, removed);
      if (removed_sortgrouprefs != NULL)
        remove_provenance_attribute_groupref(q, removed_sortgrouprefs);
      if (q->setOperations)
        remove_provenance_attribute_setoperations(q, *removed);
    }
  }
 
  if(provsql_active) {
    columns = (int **)palloc(q->rtable->length * sizeof(int *));
 
    if (q->setOperations) {
      // TODO: Nest set operations as subqueries in FROM,
      // so that we only do set operations on base tables
 
      SetOperationStmt *stmt = (SetOperationStmt *)q->setOperations;
      if (!stmt->all) {
        /* Check if any branch has aggregates — non-ALL set operations
         * on aggregate results are not supported because agg_token
         * lacks comparison operators for deduplication */
        ListCell *lc_rte;
        foreach (lc_rte, q->rtable) {
          RangeTblEntry *rte = (RangeTblEntry *)lfirst(lc_rte);
          if (rte->rtekind == RTE_SUBQUERY && rte->subquery &&
              rte->subquery->hasAggs)
            provsql_error("Non-ALL set operations (UNION, EXCEPT) on "
                          "aggregate results not supported");
        }
        q = rewrite_non_all_into_external_group_by(q);
        return process_query(constants, q, removed, wrap_root);
      }
    }
 
    if (q->hasAggs) {
      Query *rewritten = rewrite_agg_distinct(q, constants);
      if (rewritten)
        return process_query(constants, rewritten, removed, wrap_root);
    }
 
    /* Opt-in safe-query optimisation slot: when on, try to rewrite
     * hierarchical conjunctive queries to a read-once form whose
     * probability is computable in linear time via independent
     * evaluation.  See try_safe_query_rewrite().
     *
     * The rewriter is gated on the presence of the assume_boolean()
     * helper (installed by the 1.6.0 upgrade script).  Without it we
     * cannot wrap the per-row root in a gate_assumed_boolean, which is
     * what downstream evaluators inspect to refuse unsound evaluation,
     * so we refuse to rewrite on schemas that still predate the
     * helper. */
    if (provsql_boolean_provenance &&
        OidIsValid(constants->OID_FUNCTION_ASSUME_BOOLEAN)) {
      Query *rewritten = try_safe_query_rewrite(constants, q);
      if (rewritten)
        return process_query(constants, rewritten, removed, true);
    }
 
    // get_provenance_attributes will also recursively process subqueries
    // by calling process_query
    prov_atts = get_provenance_attributes(constants, q);
 
    if (prov_atts == NIL) {
      /* If the WHERE clause contains a random_variable comparison, we
       * still need to take the rewriting path so the result tuple
       * carries the comparator's gate_cmp UUID as its provenance.
       * Synthesize a single gate_one() prov_att; the combination
       * provenance_times(one, rv_cmp) collapses to rv_cmp downstream
       * because gate_one is the multiplicative identity. */
      if (q->jointree && q->jointree->quals &&
          expr_contains_rv_cmp(q->jointree->quals, constants)) {
        FuncExpr *one_expr = makeNode(FuncExpr);
        one_expr->funcid = constants->OID_FUNCTION_GATE_ONE;
        one_expr->funcresulttype = constants->OID_TYPE_UUID;
        one_expr->args = NIL;
        one_expr->location = -1;
        prov_atts = list_make1(one_expr);
      } else {
        return q;
      }
    }
 
    if (q->hasSubLinks) {
      provsql_error("Subqueries (EXISTS, IN, scalar subquery) not supported");
      supported = false;
    }
 
    if (supported && q->distinctClause) {
      if (q->hasDistinctOn) {
        provsql_error("DISTINCT ON not supported");
        supported = false;
      } else if (q->hasAggs) {
        provsql_error("DISTINCT on aggregate results not supported");
      } else if (list_length(q->distinctClause) < list_length(q->targetList)) {
        provsql_error("Inconsistent DISTINCT and GROUP BY clauses not "
                      "supported");
        supported = false;
      } else {
        transform_distinct_into_group_by(q);
      }
    }
 
    if (supported && q->setOperations) {
      SetOperationStmt *stmt = (SetOperationStmt *)q->setOperations;
 
      if (stmt->op == SETOP_UNION) {
        process_set_operation_union(constants, stmt, q);
        has_union = true;
      } else if (stmt->op == SETOP_EXCEPT) {
        if (!transform_except_into_join(constants, q))
          supported = false;
        has_difference = true;
      } else {
        provsql_error("Set operations other than UNION and EXCEPT not "
                      "supported");
        supported = false;
      }
    }
 
    if (supported && q->groupClause &&
        !provenance_function_in_group_by(constants, q)) {
      group_by_rewrite = true;
    }
 
    if (supported && q->groupingSets) {
      if (q->groupClause || list_length(q->groupingSets) > 1 ||
          ((GroupingSet *)linitial(q->groupingSets))->kind !=
          GROUPING_SET_EMPTY) {
        provsql_error("GROUPING SETS, CUBE, and ROLLUP not supported");
        supported = false;
      } else {
        // Simple GROUP BY ()
        group_by_rewrite = true;
      }
    }
 
    if (supported)
      build_column_map(q, columns, &nbcols);
 
    if (supported) {
      Expr *provenance;
      List *rv_cmps;
 
      /* Single unified pass over WHERE: each top-level conjunct is
       * routed to the right evaluation site (HAVING for agg_token,
       * the returned rv_cmps list for random_variable, left in WHERE
       * otherwise).  Mixed shapes raise a clear error.  Replaces the
       * historical pair migrate_aggtoken_quals_to_having +
       * extract_rv_cmps_from_quals; see the qual_class doc above for
       * the routing matrix.
       *
       * Must run before replace_aggregations_by_provenance_aggregate
       * so the lifted RV cmps factor into each row's contribution to
       * any surrounding agg_token: otherwise the cmp lands at group
       * level with row-typed Vars the executor cannot resolve, or
       * gets discarded by the HAVING-replaces-result branch of
       * make_provenance_expression.
       *
       * Skipped for SR_PLUS / SR_MONUS (UNION / EXCEPT outer level):
       * each branch is rewritten by its own recursive process_query
       * call, so an outer-level WHERE on RV here is exotic; the
       * fallback after make_provenance_expression handles it. */
      rv_cmps = migrate_probabilistic_quals(constants, q);
      if (rv_cmps != NIL && !has_union && !has_difference) {
        prov_atts = list_concat(prov_atts, rv_cmps);
        rv_cmps = NIL;
      }
 
      if (q->hasAggs) {
        ListCell *lc_sort;
 
        // Compute aggregation expressions
        replace_aggregations_by_provenance_aggregate(
          constants, q, prov_atts,
          has_union ? SR_PLUS : (has_difference ? SR_MONUS : SR_TIMES));
 
        // If there are any sort clauses on something whose type is now
        // aggregate token, we throw an error: sorting aggregation values
        // when provenance is captured is ill-defined
        foreach (lc_sort, q->sortClause) {
          SortGroupClause *sort = (SortGroupClause *)lfirst(lc_sort);
          ListCell *lc_te;
          foreach (lc_te, q->targetList) {
            TargetEntry *te = (TargetEntry *)lfirst(lc_te);
            if (sort->tleSortGroupRef == te->ressortgroupref) {
              if (exprType((Node *)te->expr) == constants->OID_TYPE_AGG_TOKEN)
                provsql_error("ORDER BY on the result of an aggregate function is "
                              "not supported");
              break;
            }
          }
        }
      }
 
      /* Insert casts for agg_token Vars used in arithmetic or window
       * functions, now that WHERE-to-HAVING migration is done */
      insert_agg_token_casts(constants, q);
 
      provenance = make_provenance_expression(
        constants, q, prov_atts, q->hasAggs, group_by_rewrite,
        has_union ? SR_PLUS : (has_difference ? SR_MONUS : SR_TIMES), columns,
        nbcols, wrap_root);
 
      /* Fallback for the rare set-op outer WHERE case: conjoin via
       * provenance_times after the aggregation wrappers.  Correct only
       * when no aggregation collapses rows above this point. */
      if (rv_cmps != NIL) {
        FuncExpr *times = makeNode(FuncExpr);
        ArrayExpr *array = makeNode(ArrayExpr);
        times->funcid = constants->OID_FUNCTION_PROVENANCE_TIMES;
        times->funcresulttype = constants->OID_TYPE_UUID;
        times->funcvariadic = true;
        times->location = -1;
        array->array_typeid = constants->OID_TYPE_UUID_ARRAY;
        array->element_typeid = constants->OID_TYPE_UUID;
        array->elements = lcons(provenance, rv_cmps);
        array->location = -1;
        times->args = list_make1(array);
        provenance = (Expr *)times;
      }
 
      add_to_select(q, provenance);
      replace_provenance_function_by_expression(constants, q, provenance);
 
      if (has_difference)
        add_select_non_zero(constants, q, provenance);
    }
 
    for (i = 0; i < q->rtable->length; ++i) {
      if (columns[i])
        pfree(columns[i]);
    }
  }
 
  if (provsql_verbose >= 50)
    elog_node_display(NOTICE, "ProvSQL: After query rewriting", q, true);
 
  return q;
}
 
/* -------------------------------------------------------------------------
 * INSERT ... SELECT provenance propagation
 * ------------------------------------------------------------------------- */
 
/**
 * @brief Propagate provenance through INSERT ... SELECT.
 *
 * If the source SELECT involves provenance-tracked tables and the target
 * table has a provsql column, rewrites the source SELECT to carry
 * provenance and maps its provsql output to the target's provsql column,
 * replacing the default uuid_generate_v4().
 *
 * If the target has no provsql column, emits a warning instead.
 */
static void process_insert_select(const constants_t *constants, Query *q) {
  ListCell *lc;
  Index src_rteid = 0;
  RangeTblEntry *src_rte = NULL;
  RangeTblEntry *tgt_rte;
  AttrNumber provsql_attno = 0;
  TargetEntry *provsql_te = NULL;
 
  /* Find the source SELECT subquery with provenance */
  foreach (lc, q->rtable) {
    RangeTblEntry *r = (RangeTblEntry *)lfirst(lc);
    ++src_rteid;
    if (r->rtekind == RTE_SUBQUERY && r->subquery &&
        has_provenance(constants, r->subquery)) {
      src_rte = r;
      break;
    }
  }
 
  if (src_rte == NULL)
    return;
 
  /* Check if the target table has a provsql column */
  tgt_rte = list_nth_node(RangeTblEntry, q->rtable, q->resultRelation - 1);
  if (tgt_rte->rtekind == RTE_RELATION) {
    AttrNumber attid = 1;
    foreach (lc, tgt_rte->eref->colnames) {
      if (!strcmp(strVal(lfirst(lc)), PROVSQL_COLUMN_NAME) &&
          get_atttype(tgt_rte->relid, attid) == constants->OID_TYPE_UUID)
        provsql_attno = attid;
      ++attid;
    }
  }
 
  if (provsql_attno == 0) {
    provsql_warning("INSERT ... SELECT on provenance-tracked "
                    "tables: source provenance is not propagated "
                    "to inserted rows");
    return;
  }
 
  /* Find the provsql target entry and verify it's a UUID default */
  foreach (lc, q->targetList) {
    TargetEntry *te = (TargetEntry *)lfirst(lc);
    if (te->resno == provsql_attno &&
        exprType((Node *)te->expr) == constants->OID_TYPE_UUID) {
      provsql_te = te;
      break;
    }
  }
 
  if (provsql_te == NULL) {
    /* The target's provsql column is not in the INSERT's targetList
     * (no DEFAULT on the column since 1.6.0; the user did not name
     * the column either).  Synthesise a TE here so we have something
     * to substitute the source provsql Var into below. */
    provsql_te = makeNode(TargetEntry);
    provsql_te->resno = provsql_attno;
    provsql_te->resname = pstrdup(PROVSQL_COLUMN_NAME);
    /* expr is set to the source Var by the substitution block below. */
    q->targetList = lappend(q->targetList, provsql_te);
  }
 
  /* Rewrite the source SELECT to carry provenance */
  {
    bool *removed = NULL;
    Query *new_subquery = process_query(constants, src_rte->subquery, &removed, false);
    AttrNumber src_provsql_attno = 0;
 
    if (new_subquery == NULL)
      return;
 
    src_rte->subquery = new_subquery;
 
    /* Find the provsql column in the rewritten subquery, verify its type */
    foreach (lc, new_subquery->targetList) {
      TargetEntry *te = (TargetEntry *)lfirst(lc);
      if (te->resname && !strcmp(te->resname, PROVSQL_COLUMN_NAME) &&
          exprType((Node *)te->expr) == constants->OID_TYPE_UUID) {
        src_provsql_attno = te->resno;
        break;
      }
    }
 
    if (src_provsql_attno == 0)
      return;
 
    /* Replace the target's provsql default with a Var from the source */
    {
      Var *v = makeNode(Var);
      v->varno = src_rteid;
      v->varattno = src_provsql_attno;
      v->vartype = constants->OID_TYPE_UUID;
      v->vartypmod = -1;
      v->varcollid = InvalidOid;
      v->location = -1;
      provsql_te->expr = (Expr *)v;
    }
 
    /* Update the subquery RTE's column names to include provsql */
    src_rte->eref->colnames = lappend(src_rte->eref->colnames,
                                      makeString(pstrdup(PROVSQL_COLUMN_NAME)));
  }
}
 
/* -------------------------------------------------------------------------
 * Planner hook & extension lifecycle
 * ------------------------------------------------------------------------- */
 
/**
 * @brief PostgreSQL planner hook — entry point for provenance rewriting.
 *
 * Replaces (or chains after) the standard planner.  For every CMD_SELECT
 * that involves at least one provenance-bearing relation or an explicit
 * @c provenance() call, rewrites the query via @c process_query before
 * handing the result to the standard planner.  Non-SELECT commands and
 * queries without provenance are passed through unchanged.
 * @param q              The query to plan.
 * @param cursorOptions  Cursor options bitmask.
 * @param boundParams    Pre-bound parameter values.
 * @return               The planned statement.
 */
/** @brief Executor nesting depth.
 *
 * Tracks how deep we are inside @c Executor invocations.  Incremented
 * in @c provsql_executor_start, decremented in @c provsql_executor_end.
 * The classifier @c NOTICE only fires when this is zero, which
 * corresponds to the user's outermost statement being planned (before
 * any executor entry).  Plans built for PL/pgSQL function bodies that
 * the rewriter inserts -- @c provenance_times, @c provenance_plus,
 * @c provenance_aggregate, ... -- happen during execution of the
 * user's plan, so they see depth >= 1 and skip the NOTICE. */
static int provsql_executor_depth = 0;
 
static PlannedStmt *provsql_planner(Query *q,
#if PG_VERSION_NUM >= 130000
                                    const char *query_string,
#endif
                                    int cursorOptions,
                                    ParamListInfo boundParams) {
  if (q->commandType == CMD_INSERT && q->rtable) {
    const constants_t constants = get_constants(false);
    if (constants.ok)
      process_insert_select(&constants, q);
  } else if (q->commandType == CMD_SELECT) {
    /* No rtable check here: a FROM-less SELECT (e.g.
     *   SELECT 1 WHERE normal(0,1) > 2)
     * still needs the hook to engage when the WHERE contains an
     * rv_cmp.  has_provenance walks the tree and returns false fast
     * on FROM-less queries that have neither rv_cmp nor provenance(),
     * so widening the gate costs nothing in the common case. */
    const constants_t constants = get_constants(false);
 
    /* Query-time TID / BID / OPAQUE classifier.  Emits a NOTICE for
     * the user's outermost SELECT when the GUC is on.  Runs on the
     * user's original Query before any provsql rewriting so the
     * reported kind reflects the SQL the user wrote.  Gating on
     * @c provsql_executor_depth @c == @c 0 skips the spurious extra
     * planning calls triggered by PL/pgSQL function bodies the
     * rewriter inserts (@c provenance_times, ...), whose internal
     * SELECTs go through the planner hook during execution of the
     * user's plan. */
    if (provsql_executor_depth == 0
        && provsql_classify_top_level
        && q->rtable != NIL) {
      ProvSQLClassification cls;
      provsql_classify_query(q, &cls);
      provsql_classify_emit_notice(&cls);
      list_free(cls.source_relids);
    }
 
    if (constants.ok && has_provenance(&constants, q)) {
      bool *removed = NULL;
      Query *new_query;
      clock_t begin = 0;
 
#if PG_VERSION_NUM >= 150000
      if (provsql_verbose >= 20)
        provsql_notice("Main query before query rewriting:\n%s\n",
                       pg_get_querydef(q, true));
#endif
 
      if (provsql_verbose >= 40)
        begin = clock();
 
      new_query = process_query(&constants, q, &removed, false);
 
      if (provsql_verbose >= 40)
        provsql_notice("planner time spent=%f",
                       (double)(clock() - begin) / CLOCKS_PER_SEC);
 
      if (new_query != NULL)
        q = new_query;
 
#if PG_VERSION_NUM >= 150000
      if (provsql_verbose >= 20)
        provsql_notice("Main query after query rewriting:\n%s\n",
                       pg_get_querydef(q, true));
#endif
    }
  }
 
  if (prev_planner)
    return prev_planner(q,
#if PG_VERSION_NUM >= 130000
                        query_string,
#endif
                        cursorOptions, boundParams);
  else
    return standard_planner(q,
#if PG_VERSION_NUM >= 130000
                            query_string,
#endif
                            cursorOptions, boundParams);
}
 
/* -------------------------------------------------------------------------
 * Executor hooks (depth tracking only)
 *
 * We install ExecutorStart / ExecutorEnd hooks solely to maintain
 * @c provsql_executor_depth, which the classifier in @c provsql_planner
 * consults to distinguish the user's outermost statement from nested
 * PL/pgSQL bodies the rewriter calls into.  No other behaviour changes.
 * ------------------------------------------------------------------------- */
static ExecutorStart_hook_type prev_ExecutorStart = NULL;
static ExecutorEnd_hook_type   prev_ExecutorEnd   = NULL;
 
static void provsql_executor_start(QueryDesc *queryDesc, int eflags) {
  provsql_executor_depth++;
  PG_TRY();
  {
    if (prev_ExecutorStart)
      prev_ExecutorStart(queryDesc, eflags);
    else
      standard_ExecutorStart(queryDesc, eflags);
  }
  PG_CATCH();
  {
    provsql_executor_depth--;
    PG_RE_THROW();
  }
  PG_END_TRY();
}
 
static void provsql_executor_end(QueryDesc *queryDesc) {
#if PG_VERSION_NUM >= 130000
  PG_TRY();
  {
    if (prev_ExecutorEnd)
      prev_ExecutorEnd(queryDesc);
    else
      standard_ExecutorEnd(queryDesc);
  }
  PG_FINALLY();
  {
    provsql_executor_depth--;
  }
  PG_END_TRY();
#else
  /* PG < 13 lacks PG_FINALLY: emulate by running the cleanup on the
   * error path (via PG_CATCH + PG_RE_THROW) and on the success path
   * (after PG_END_TRY).  Functionally equivalent. */
  PG_TRY();
  {
    if (prev_ExecutorEnd)
      prev_ExecutorEnd(queryDesc);
    else
      standard_ExecutorEnd(queryDesc);
  }
  PG_CATCH();
  {
    provsql_executor_depth--;
    PG_RE_THROW();
  }
  PG_END_TRY();
  provsql_executor_depth--;
#endif
}
 
/* -------------------------------------------------------------------------
 * ProcessUtility hook: CTAS lineage inheritance.
 *
 * When a @c CREATE @c TABLE @c AS (or @c CREATE @c MATERIALIZED @c VIEW,
 * or @c SELECT @c INTO -- PG's parser transforms all three into
 * @c CreateTableAsStmt) projects a @c provsql column lifted verbatim
 * from a tracked source, the resulting relation's atoms are not freshly
 * minted UUIDs but lineage tokens of one or more base @c
 * add_provenance / @c repair_key relations.  The hook intercepts the
 * utility statement, classifies the inner @c SELECT via
 * @c provsql_classify_query, lets PG run the CTAS, then populates
 * @c provsql_table_info (with the inherited @c kind / BID @c block_key)
 * and the ancestor registry (with the transitive union of source
 * ancestor sets) on the just-created relation.  A @c provenance_guard
 * trigger is installed on the new table so any subsequent INSERT /
 * UPDATE that supplies a non-NULL @c provsql still flips the table to
 * OPAQUE the standard way.
 *
 * The hook deliberately fires only when the inner @c SELECT projects
 * a @c provsql column from a tracked source -- otherwise the new
 * relation has no @c provsql column and the lineage metadata would be
 * operationally pointless.  Users who want a tracked CTAS-derived
 * table without inherited lineage still call @c add_provenance on it
 * afterwards (that path seeds @c {self} and overrides whatever this
 * hook may have recorded).
 * ------------------------------------------------------------------------- */
 
static ProcessUtility_hook_type prev_ProcessUtility = NULL;
 
/** @brief State captured by the pre-execution pass for the post-execution one. */
typedef struct ProvSQLCtasCapture {
  bool                fire;             ///< true when the post-pass should run
  Query              *inner_query;      ///< cloned for safety; freed by pfree on completion
  provsql_table_kind  inherited_kind;
  uint16              ancestor_n;
  Oid                 ancestors[PROVSQL_TABLE_INFO_MAX_ANCESTORS];
  Oid                 source_relid;     ///< Single source whose block_key we want to align (BID only)
  uint16              source_block_key_n;
  AttrNumber          source_block_key[PROVSQL_TABLE_INFO_MAX_BLOCK_KEY];
} ProvSQLCtasCapture;
 
/**
 * @brief Decide whether @p parsetree is a CTAS that should trigger
 *        the ancestry hook, and if so populate @p cap with the inner
 *        classification, the (single) source's block-key columns, and
 *        the transitive ancestor union.
 *
 * Fires only when the inner @c SELECT's target list projects a base-
 * level @c Var (possibly through @c RelabelType wrappers) that
 * resolves to the @c provsql column of an @c RTE_RELATION whose
 * metadata is non-OPAQUE.  Anything else (no @c provsql in the
 * projection, classifier says OPAQUE, the projected source is itself
 * OPAQUE) leaves @c cap->fire false and the post-pass becomes a
 * no-op.
 */
static void provsql_ProcessUtility_capture(Node *parsetree,
                                           ProvSQLCtasCapture *cap) {
  CreateTableAsStmt    *stmt;
  Query                *qry;
  ProvSQLClassification cls;
  ListCell             *lc;
  AttrNumber            prov_resno = InvalidAttrNumber;
  Oid                   source_relid = InvalidOid;
  ProvenanceTableInfo   source_info;
  Bitmapset            *ancestor_bms = NULL;
  int                   bms_member;
  uint16                ancestor_n;
 
  cap->fire = false;
  if (!provsql_active)
    return;
  if (parsetree == NULL || !IsA(parsetree, CreateTableAsStmt))
    return;
  stmt = (CreateTableAsStmt *) parsetree;
  if (stmt->query == NULL || !IsA(stmt->query, Query))
    return;
  qry = (Query *) stmt->query;
  if (qry->commandType != CMD_SELECT)
    return;
 
  provsql_classify_query(qry, &cls);
  if (cls.kind == PROVSQL_TABLE_OPAQUE) {
    list_free(cls.source_relids);
    return;
  }
 
  /* Walk the inner target list for a TLE whose Var resolves to the
   * provsql column of a tracked, non-OPAQUE source.  First match wins
   * (CTAS preserves the TLE's column name verbatim in the new
   * table, so a single provsql TLE is the normal case). */
  foreach (lc, qry->targetList) {
    TargetEntry   *te = (TargetEntry *) lfirst(lc);
    Node          *e  = (Node *) te->expr;
    Var           *v;
    RangeTblEntry *rte;
    AttrNumber     prov_attno;
 
    if (te->resjunk)
      continue;
    while (e != NULL && IsA(e, RelabelType))
      e = (Node *) ((RelabelType *) e)->arg;
    if (e == NULL || !IsA(e, Var))
      continue;
    v = (Var *) e;
    if (v->varlevelsup != 0)
      continue;
    if (v->varno < 1 || (int) v->varno > list_length(qry->rtable))
      continue;
    rte = (RangeTblEntry *) list_nth(qry->rtable, v->varno - 1);
    if (rte->rtekind != RTE_RELATION)
      continue;
    prov_attno = get_attnum(rte->relid, PROVSQL_COLUMN_NAME);
    if (prov_attno == InvalidAttrNumber || v->varattno != prov_attno)
      continue;
    if (!provsql_lookup_table_info(rte->relid, &source_info))
      continue;
    if (source_info.kind == PROVSQL_TABLE_OPAQUE)
      continue;
    prov_resno   = te->resno;
    source_relid = rte->relid;
    break;
  }
  if (prov_resno == InvalidAttrNumber) {
    list_free(cls.source_relids);
    return;
  }
 
  /* Transitive ancestor union: lookup each classifier-reported source's
   * registered ancestry, fall back to {source} when none recorded
   * (defensive: the SQL add_provenance / repair_key seed should always
   * give us a non-empty set).  Bitmapset dedupes; we then walk it in
   * ascending order to get the sorted Oid array the registry stores. */
  foreach (lc, cls.source_relids) {
    Oid src_relid = lfirst_oid(lc);
    uint16 src_n;
    Oid    src_ancestors[PROVSQL_TABLE_INFO_MAX_ANCESTORS];
    if (provsql_lookup_ancestry(src_relid, &src_n, src_ancestors)) {
      for (uint16 i = 0; i < src_n; ++i)
        ancestor_bms = bms_add_member(ancestor_bms, (int) src_ancestors[i]);
    } else {
      ancestor_bms = bms_add_member(ancestor_bms, (int) src_relid);
    }
  }
  list_free(cls.source_relids);
 
  ancestor_n = 0;
  bms_member = -1;
  while ((bms_member = bms_next_member(ancestor_bms, bms_member)) >= 0) {
    if (ancestor_n >= PROVSQL_TABLE_INFO_MAX_ANCESTORS) {
      /* Cap exceeded: refuse to fire rather than truncate the ancestor
       * set silently (a partial set would let the safe-query
       * disjointness check accept a join that shouldn't be safe). */
      bms_free(ancestor_bms);
      return;
    }
    cap->ancestors[ancestor_n++] = (Oid) bms_member;
  }
  bms_free(ancestor_bms);
 
  cap->fire               = true;
  cap->inner_query        = qry;
  cap->inherited_kind     = (provsql_table_kind) cls.kind;
  cap->ancestor_n         = ancestor_n;
  cap->source_relid       = source_relid;
  cap->source_block_key_n = source_info.block_key_n;
  memcpy(cap->source_block_key, source_info.block_key,
         source_info.block_key_n * sizeof(AttrNumber));
}
 
/** @brief Map @c provsql_table_kind to its textual label
 *  (@c set_table_info accepts text). */
static const char *provsql_ctas_kind_label(provsql_table_kind k) {
  switch (k) {
    case PROVSQL_TABLE_TID:    return "tid";
    case PROVSQL_TABLE_BID:    return "bid";
    case PROVSQL_TABLE_OPAQUE: return "opaque";
  }
  return "opaque";
}
 
/** @brief Forward declaration of the C SQL entry points. */
extern Datum set_table_info(PG_FUNCTION_ARGS);
extern Datum set_ancestors(PG_FUNCTION_ARGS);
 
/**
 * @brief Apply @p cap to the freshly-created relation @c stmt->into->rel.
 *
 * For BID sources: walks the inner query's target list to align each
 * source block-key column to its output @c resno.  If any block-key
 * column is missing from the projection (the CTAS dropped it), the
 * new relation cannot honour the BID invariant under that column --
 * the hook demotes to TID rather than asserting a now-stale block
 * key.
 *
 * Installs @c provenance_guard via SPI so subsequent INSERT /
 * UPDATE OF provsql on the new relation flip its kind to OPAQUE
 * through the standard guard path.
 */
static void provsql_ProcessUtility_apply(Node *parsetree,
                                         ProvSQLCtasCapture *cap) {
  CreateTableAsStmt *stmt;
  Oid                new_relid;
  AttrNumber         prov_attno;
  provsql_table_kind eff_kind;
  uint16             eff_block_key_n = 0;
  AttrNumber         eff_block_key[PROVSQL_TABLE_INFO_MAX_BLOCK_KEY];
  Datum              kind_datum;
  Datum              block_key_datum;
  Datum              ancestors_datum;
  Datum             *block_key_elems;
  Datum             *ancestor_elems;
  ArrayType         *block_key_arr;
  ArrayType         *ancestors_arr;
  const char        *nspname;
  const char        *relname;
  StringInfoData     trigger_sql;
 
  if (!cap->fire)
    return;
  stmt = (CreateTableAsStmt *) parsetree;
 
  new_relid = RangeVarGetRelid(stmt->into->rel, NoLock, true);
  if (new_relid == InvalidOid)
    return;
 
  /* Confirm the new relation actually has a @c provsql @c uuid column.
   * CTAS preserves TLE column names, so this is essentially the
   * post-execution verification of what the pre-pass already
   * required. */
  prov_attno = get_attnum(new_relid, PROVSQL_COLUMN_NAME);
  if (prov_attno == InvalidAttrNumber)
    return;
  if (get_atttype(new_relid, prov_attno) != UUIDOID)
    return;
 
  /* BID block-key alignment: each source block-key column must
   * survive in the inner-query target list.  When all do, the new
   * relation's effective block key is the corresponding output
   * resno.  When any is missing, demote to TID (a partial block
   * key would falsely advertise mutual exclusion the rows no
   * longer have). */
  eff_kind = cap->inherited_kind;
  if (eff_kind == PROVSQL_TABLE_BID) {
    bool ok = true;
    for (uint16 i = 0; i < cap->source_block_key_n; ++i) {
      AttrNumber src_attno = cap->source_block_key[i];
      ListCell  *lc;
      bool       found = false;
      foreach (lc, cap->inner_query->targetList) {
        TargetEntry   *te = (TargetEntry *) lfirst(lc);
        Node          *e  = (Node *) te->expr;
        Var           *v;
        RangeTblEntry *rte;
        if (te->resjunk)
          continue;
        while (e != NULL && IsA(e, RelabelType))
          e = (Node *) ((RelabelType *) e)->arg;
        if (e == NULL || !IsA(e, Var))
          continue;
        v = (Var *) e;
        if (v->varlevelsup != 0)
          continue;
        if (v->varno < 1
            || (int) v->varno > list_length(cap->inner_query->rtable))
          continue;
        rte = (RangeTblEntry *)
            list_nth(cap->inner_query->rtable, v->varno - 1);
        if (rte->rtekind != RTE_RELATION)
          continue;
        if (rte->relid == cap->source_relid
            && v->varattno == src_attno) {
          if (eff_block_key_n >= PROVSQL_TABLE_INFO_MAX_BLOCK_KEY) {
            ok = false;
            break;
          }
          eff_block_key[eff_block_key_n++] = te->resno;
          found = true;
          break;
        }
      }
      if (!found) {
        ok = false;
        break;
      }
    }
    if (!ok) {
      eff_kind        = PROVSQL_TABLE_TID;
      eff_block_key_n = 0;
    }
  }
 
  /* Marshal arguments and invoke the SQL-level helpers via
   * DirectFunctionCall: this reaches the worker through the same IPC
   * path the user-facing SQL functions use, including the relcache
   * invalidation broadcast on the way out. */
  kind_datum = CStringGetTextDatum(provsql_ctas_kind_label(eff_kind));
  if (eff_block_key_n == 0) {
    block_key_arr = construct_empty_array(INT2OID);
  } else {
    block_key_elems = palloc(eff_block_key_n * sizeof(Datum));
    for (uint16 i = 0; i < eff_block_key_n; ++i)
      block_key_elems[i] = Int16GetDatum(eff_block_key[i]);
    block_key_arr = construct_array(block_key_elems, eff_block_key_n,
                                    INT2OID, 2, true, 's');
    pfree(block_key_elems);
  }
  block_key_datum = PointerGetDatum(block_key_arr);
  DirectFunctionCall3(set_table_info,
                      ObjectIdGetDatum(new_relid),
                      kind_datum,
                      block_key_datum);
 
  if (cap->ancestor_n == 0) {
    ancestors_arr = construct_empty_array(OIDOID);
  } else {
    ancestor_elems = palloc(cap->ancestor_n * sizeof(Datum));
    for (uint16 i = 0; i < cap->ancestor_n; ++i)
      ancestor_elems[i] = ObjectIdGetDatum(cap->ancestors[i]);
    ancestors_arr = construct_array(ancestor_elems, cap->ancestor_n,
                                    OIDOID, sizeof(Oid), true, 'i');
    pfree(ancestor_elems);
  }
  ancestors_datum = PointerGetDatum(ancestors_arr);
  DirectFunctionCall2(set_ancestors,
                      ObjectIdGetDatum(new_relid),
                      ancestors_datum);
 
  /* Install the provenance_guard trigger via SPI.  Users who later
   * INSERT / UPDATE OF provsql with a non-NULL value will then
   * trigger the standard kind flip to OPAQUE; users who omit the
   * column on INSERT get a fresh @c uuid_generate_v4 leaf (which
   * already disconnects the row from the inherited lineage, but the
   * guard prevents the more dangerous shared-UUID aliasing path).
   *
   * Materialized views are exempt: PG forbids triggers on them, and
   * they cannot be modified through DML anyway (only @c REFRESH @c
   * MATERIALIZED @c VIEW changes the contents -- which re-runs the
   * inner SELECT and the freshly-projected rows continue to carry
   * lineage from the same sources). */
  /* PG 14 renamed CreateTableAsStmt.relkind -> objtype (same ObjectType,
   * same OBJECT_* values; pure field rename). */
#if PG_VERSION_NUM >= 140000
  if (stmt->objtype == OBJECT_MATVIEW)
#else
  if (stmt->relkind == OBJECT_MATVIEW)
#endif
    return;
 
  nspname = get_namespace_name(get_rel_namespace(new_relid));
  relname = get_rel_name(new_relid);
  if (nspname == NULL || relname == NULL)
    return;
  initStringInfo(&trigger_sql);
  appendStringInfo(&trigger_sql,
      "CREATE TRIGGER provenance_guard "
      "BEFORE INSERT OR UPDATE OF provsql ON %s.%s "
      /* "EXECUTE PROCEDURE" is the legacy form, kept as a valid synonym
       * of "EXECUTE FUNCTION" through PG 18 -- matches the rest of the
       * codebase and stays PG 10-compatible.  Promote when PG 10 drops
       * out of the support floor. */
      "FOR EACH ROW EXECUTE PROCEDURE provsql.provenance_guard()",
      quote_identifier(nspname), quote_identifier(relname));
  if (SPI_connect() != SPI_OK_CONNECT)
    provsql_error("CTAS lineage hook: SPI_connect failed");
  if (SPI_exec(trigger_sql.data, 0) != SPI_OK_UTILITY)
    provsql_error("CTAS lineage hook: failed to install provenance_guard "
                  "on %s.%s", nspname, relname);
  SPI_finish();
  pfree(trigger_sql.data);
}
 
static void provsql_ProcessUtility(
    PlannedStmt *pstmt,
    const char *queryString,
#if PG_VERSION_NUM >= 140000
    bool readOnlyTree,
#endif
    ProcessUtilityContext context,
    ParamListInfo params,
    QueryEnvironment *queryEnv,
    DestReceiver *dest,
#if PG_VERSION_NUM >= 130000
    QueryCompletion *qc
#else
    char *completionTag
#endif
    ) {
  Node              *parsetree = pstmt ? pstmt->utilityStmt : NULL;
  ProvSQLCtasCapture cap = {0};
 
  provsql_ProcessUtility_capture(parsetree, &cap);
 
  if (prev_ProcessUtility)
    prev_ProcessUtility(pstmt, queryString,
#if PG_VERSION_NUM >= 140000
                        readOnlyTree,
#endif
                        context, params, queryEnv, dest,
#if PG_VERSION_NUM >= 130000
                        qc
#else
                        completionTag
#endif
                        );
  else
    standard_ProcessUtility(pstmt, queryString,
#if PG_VERSION_NUM >= 140000
                            readOnlyTree,
#endif
                            context, params, queryEnv, dest,
#if PG_VERSION_NUM >= 130000
                            qc
#else
                            completionTag
#endif
                            );
 
  provsql_ProcessUtility_apply(parsetree, &cap);
}
 
/**
 * @brief Extension initialization – called once when the shared library is loaded.
 *
 * Registers the GUC variables (@c provsql.active, @c where_provenance,
 * @c update_provenance, @c verbose_level, @c aggtoken_text_as_uuid,
 * @c tool_search_path), installs the planner hook and shared-memory hooks,
 * and launches the background MMap worker.
 *
 * Must be loaded via @c shared_preload_libraries; raises an error otherwise.
 */
void _PG_init(void) {
  if (!process_shared_preload_libraries_in_progress)
    provsql_error("provsql needs to be added to the shared_preload_libraries "
                  "configuration variable");
 
  DefineCustomBoolVariable("provsql.active",
                           "Should ProvSQL track provenance?",
                           "1 is standard ProvSQL behavior, 0 means provsql attributes will be dropped.",
                           &provsql_active,
                           true,
                           PGC_USERSET,
                           0,
                           NULL,
                           NULL,
                           NULL);
  DefineCustomBoolVariable("provsql.where_provenance",
                           "Should ProvSQL track where-provenance?",
                           "1 turns where-provenance on, 0 off.",
                           &provsql_where_provenance,
                           false,
                           PGC_USERSET,
                           0,
                           NULL,
                           NULL,
                           NULL);
  DefineCustomBoolVariable("provsql.update_provenance",
                           "Should ProvSQL track update provenance?",
                           "1 turns update provenance on, 0 off.",
                           &provsql_update_provenance,
                           false,
                           PGC_USERSET,
                           0,
                           NULL,
                           NULL,
                           NULL);
  DefineCustomBoolVariable("provsql.aggtoken_text_as_uuid",
                           "Output agg_token cells as the underlying UUID "
                           "instead of \"value (*)\".",
                           "Off by default for psql-friendly output. UI "
                           "layers (notably ProvSQL Studio) flip this on "
                           "per session so aggregate cells expose the "
                           "circuit root UUID for click-through; the "
                           "display value is recovered via "
                           "provsql.agg_token_value_text(uuid).",
                           &provsql_aggtoken_text_as_uuid,
                           false,
                           PGC_USERSET,
                           0,
                           NULL,
                           NULL,
                           NULL);
  DefineCustomIntVariable("provsql.verbose_level",
                          "Level of verbosity for ProvSQL informational and debug messages",
                          "0 for quiet (default), 1-9 for informational messages, 10-100 for debug information.",
                          &provsql_verbose,
                          0,
                          0,
                          100,
                          PGC_USERSET,
                          1,
                          NULL,
                          NULL,
                          NULL);
  DefineCustomStringVariable("provsql.tool_search_path",
                             "Directories prepended to PATH when ProvSQL spawns external tools.",
                             "Colon-separated list of directories searched before the server's PATH "
                             "when locating d4, c2d, minic2d, dsharp, weightmc, or graph-easy. "
                             "Empty (default) means rely on the server's PATH alone.",
                             &provsql_tool_search_path,
                             "",
                             PGC_USERSET,
                             0,
                             NULL,
                             NULL,
                             NULL);
  DefineCustomBoolVariable("provsql.simplify_on_load",
                           "Apply universal cmp-resolution passes when "
                           "loading a provenance circuit.",
                           "When on (default), every GenericCircuit returned "
                           "by getGenericCircuit goes through RangeCheck "
                           "(and any future universal pass): comparators "
                           "decidable to certain Boolean values become "
                           "Bernoulli gate_input gates with probability 0 "
                           "or 1, transparent to every downstream consumer "
                           "(semiring evaluators, MC, view_circuit, PROV "
                           "export). Set off to inspect raw circuit "
                           "structure (e.g. when debugging gate-creation "
                           "paths).",
                           &provsql_simplify_on_load,
                           true,
                           PGC_USERSET,
                           0,
                           NULL,
                           NULL,
                           NULL);
  /* Debug-only: hidden from SHOW ALL and postgresql.conf.sample.
   * On is strictly better for end users (analytic answers where
   * possible, lower MC variance, more methods usable on continuous
   * circuits); off only serves developer A/B against pure MC and as
   * a bisection escape valve if a closure rule misbehaves. */
  DefineCustomBoolVariable("provsql.hybrid_evaluation",
                           "Run the hybrid-evaluator simplifier and "
                           "island decomposer inside probability_evaluate. "
                           "Debug only.",
                           "When on (default), probability_evaluate runs "
                           "the HybridEvaluator peephole simplifier "
                           "between RangeCheck and AnalyticEvaluator and "
                           "the per-cmp MC island decomposer after "
                           "AnalyticEvaluator. Off bypasses both and lets "
                           "unresolved comparators fall through to "
                           "whole-circuit MC. End users have no reason "
                           "to flip this; it exists for developer A/B "
                           "testing against the unfolded path and as a "
                           "bisection knob if a closure rule turns out "
                           "to be unsound on some workload.",
                           &provsql_hybrid_evaluation,
                           true,
                           PGC_USERSET,
                           GUC_NO_SHOW_ALL | GUC_NOT_IN_SAMPLE,
                           NULL,
                           NULL,
                           NULL);
  /* Debug-only: hidden from SHOW ALL and postgresql.conf.sample.
   * Umbrella for closed-form / analytic resolution of gate_cmp
   * probabilities in probability_evaluate (currently the
   * Poisson-binomial HAVING-COUNT pre-pass; future MIN / MAX / SUM
   * pre-passes gate on the same flag).  On is strictly better for
   * end users (each resolver replaces an exponential DNF
   * construction with O(N) or O(N x C) arithmetic); off only serves
   * developer A/B against the unoptimised enumerate_valid_worlds
   * path and as a bisection escape valve. */
  DefineCustomBoolVariable("provsql.cmp_probability_evaluation",
                           "Run closed-form / analytic probability "
                           "evaluators for gate_cmps inside "
                           "probability_evaluate. Debug only.",
                           "When on (default), probability_evaluate "
                           "runs pre-passes that recognise specific "
                           "gate_cmp shapes (currently HAVING COUNT(*) "
                           "op C over distinct gate_input leaves) and "
                           "replace each cmp with a Bernoulli "
                           "gate_input carrying the closed-form "
                           "probability, bypassing the DNF that "
                           "provsql_having's enumerate_valid_worlds "
                           "would otherwise emit. Off forces the cmp "
                           "to fall through to that enumeration path. "
                           "Future MIN / MAX / SUM probability "
                           "evaluators will gate on the same flag. "
                           "End users have no reason to flip this; it "
                           "exists for developer A/B testing and as a "
                           "bisection escape valve.",
                           &provsql_cmp_probability_evaluation,
                           true,
                           PGC_USERSET,
                           GUC_NO_SHOW_ALL | GUC_NOT_IN_SAMPLE,
                           NULL,
                           NULL,
                           NULL);
  DefineCustomBoolVariable("provsql.boolean_provenance",
                           "Opt-in safe-query optimisation for hierarchical conjunctive queries.",
                           "When on, the planner rewrites self-join-free "
                           "hierarchical conjunctive queries (and independent "
                           "UCQs) over TID/BID tables to a read-once form "
                           "whose probability is computable in linear time "
                           "via the existing BooleanCircuit independent "
                           "evaluator. The rewrite preserves the Boolean "
                           "polynomial of the existential lineage, so any "
                           "evaluation that factors through a homomorphism "
                           "from Boolean functions remains sound (notably "
                           "probability and Shapley / Banzhaf). The "
                           "resulting root gate is tagged; semiring "
                           "evaluators are individually marked at compile "
                           "time as compatible or not with this rewrite, "
                           "and incompatible evaluators refuse to run on a "
                           "tagged circuit. Off by default because the "
                           "rewrite changes the multiset of result rows and "
                           "is therefore unsound for per-row provenance "
                           "interrogations and aggregation queries.",
                           &provsql_boolean_provenance,
                           false,
                           PGC_USERSET,
                           0,
                           NULL,
                           NULL,
                           NULL);
  DefineCustomBoolVariable("provsql.classify_top_level",
                           "Emit a NOTICE classifying each top-level SELECT.",
                           "When on, every top-level SELECT that "
                           "touches a relation triggers a NOTICE of "
                           "the form `ProvSQL: query result is "
                           "<KIND> (sources: ...)` where <KIND> is "
                           "TID, BID, or OPAQUE under the existing "
                           "provsql_table_kind taxonomy and the "
                           "sources list names the provenance-"
                           "tracked base relations the query touches. "
                           "Read-only : the classifier does not "
                           "rewrite the query.  Studio reads the "
                           "NOTICE to label query results with their "
                           "certified kind.",
                           &provsql_classify_top_level,
                           false,
                           PGC_USERSET,
                           0,
                           NULL,
                           NULL,
                           NULL);
  DefineCustomIntVariable("provsql.monte_carlo_seed",
                          "Seed for the Monte Carlo sampler.",
                          "-1 (default) seeds from std::random_device for "
                          "non-deterministic sampling. Any other value "
                          "(including 0) is used as a literal seed for "
                          "std::mt19937_64, making "
                          "probability_evaluate(..., 'monte-carlo', n) "
                          "reproducible across runs and across the Bernoulli "
                          "and continuous (gate_rv) sampling paths.",
                          &provsql_monte_carlo_seed,
                          -1,
                          -1,
                          INT_MAX,
                          PGC_USERSET,
                          0,
                          NULL,
                          NULL,
                          NULL);
  DefineCustomIntVariable("provsql.rv_mc_samples",
                          "Default sample count for analytical-evaluator MC fallbacks.",
                          "Used when an analytical evaluator (Expectation, "
                          "future hybrid evaluator, etc.) cannot decompose a "
                          "sub-circuit and needs to fall back to Monte Carlo. "
                          "Default 10000. Set to 0 to disable the fallback "
                          "entirely: callers raise an exception rather than "
                          "sampling, which is useful when only analytical "
                          "answers are acceptable. Unrelated to "
                          "probability_evaluate(..., 'monte-carlo', n) where "
                          "the sample count is an explicit argument.",
                          &provsql_rv_mc_samples,
                          10000,
                          0,
                          INT_MAX,
                          PGC_USERSET,
                          0,
                          NULL,
                          NULL,
                          NULL);
 
  // Emit warnings for undeclared provsql.* configuration parameters
  EmitWarningsOnPlaceholders("provsql");
 
  prev_planner = planner_hook;
  prev_shmem_startup = shmem_startup_hook;
  prev_ExecutorStart = ExecutorStart_hook;
  prev_ExecutorEnd   = ExecutorEnd_hook;
  prev_ProcessUtility = ProcessUtility_hook;
#if (PG_VERSION_NUM >= 150000)
  prev_shmem_request = shmem_request_hook;
  shmem_request_hook = provsql_shmem_request;
#else
  provsql_shmem_request();
#endif
 
  planner_hook        = provsql_planner;
  shmem_startup_hook  = provsql_shmem_startup;
  ExecutorStart_hook  = provsql_executor_start;
  ExecutorEnd_hook    = provsql_executor_end;
  ProcessUtility_hook = provsql_ProcessUtility;
 
  RegisterProvSQLMMapWorker();
}
 
/**
 * @brief Extension teardown — restores the planner and shmem hooks.
 */
void _PG_fini(void) {
  planner_hook        = prev_planner;
  shmem_startup_hook  = prev_shmem_startup;
  ExecutorStart_hook  = prev_ExecutorStart;
  ExecutorEnd_hook    = prev_ExecutorEnd;
  ProcessUtility_hook = prev_ProcessUtility;
}