/*-------------------------------------------------------------------------

 *

 * tsquery_rewrite.c

 *    Utilities for reconstructing tsquery

 *

 * Portions Copyright (c) 1996-2025, PostgreSQL Global Development Group

 *

 *

 * IDENTIFICATION

 *    src/backend/utils/adt/tsquery_rewrite.c

 *

 *-------------------------------------------------------------------------

 */

#include "postgres.h"

#include "catalog/pg_type.h"

#include "executor/spi.h"

#include "miscadmin.h"

#include "tsearch/ts_utils.h"

#include "utils/builtins.h"

/*

 * If "node" is equal to "ex", return a copy of "subs" instead.

 * If "ex" matches a subset of node's children, return a modified version

 * of "node" in which those children are replaced with a copy of "subs".

 * Otherwise return "node" unmodified.

 *

 * The QTN_NOCHANGE bit is set in successfully modified nodes, so that

 * we won't uselessly recurse into them.

 * Also, set *isfind true if we make a replacement.

 */

static QTNode *

findeq(QTNode *node, QTNode *ex, QTNode *subs, bool *isfind)

{

  /* Can't match unless signature matches and node type matches. */

  if ((node->sign & ex->sign) != ex->sign ||

    node->valnode->type != ex->valnode->type)

    return node;

  /* Ignore nodes marked NOCHANGE, too. */

  if (node->flags & QTN_NOCHANGE)

    return node;

  if (node->valnode->type == QI_OPR)

  {

    /* Must be same operator. */

    if (node->valnode->qoperator.oper != ex->valnode->qoperator.oper)

      return node;

    if (node->nchild == ex->nchild)

    {

      /*

       * Simple case: when same number of children, match if equal.

       * (This is reliable when the children were sorted earlier.)

       */

      if (QTNEq(node, ex))

      {

        /* Match; delete node and return a copy of subs instead. */

        QTNFree(node);

        if (subs)

        {

          node = QTNCopy(subs);

          node->flags |= QTN_NOCHANGE;

        }

        else

          node = NULL;

        *isfind = true;

      }

    }

    else if (node->nchild > ex->nchild && ex->nchild > 0)

    {

      /*

       * AND and OR are commutative/associative, so we should check if a

       * subset of the children match.  For example, if node is A|B|C,

       * and ex is B|C, we have a match after we notionally convert node

       * to A|(B|C).  This does not work for NOT or PHRASE nodes, but we

       * can't get here for those node types because they have a fixed

       * number of children.

       *

       * Because we expect that the children are sorted, it suffices to

       * make one pass through the two lists to find the matches.

       */

      bool     *matched;

      int     nmatched;

      int     i,

            j;

      /* Assert that the subset rule is OK */

      Assert(node->valnode->qoperator.oper == OP_AND ||

           node->valnode->qoperator.oper == OP_OR);

      /* matched[] will record which children of node matched */

      matched = (bool *) palloc0(node->nchild * sizeof(bool));

      nmatched = 0;

      i = j = 0;

      while (i < node->nchild && j < ex->nchild)

      {

        int     cmp = QTNodeCompare(node->child[i], ex->child[j]);

        if (cmp == 0)

        {

          /* match! */

          matched[i] = true;

          nmatched++;

          i++, j++;

        }

        else if (cmp < 0)

        {

          /* node->child[i] has no match, ignore it */

          i++;

        }

        else

        {

          /* ex->child[j] has no match; we can give up immediately */

          break;

        }

      }

      if (nmatched == ex->nchild)

      {

        /* collapse out the matched children of node */

        j = 0;

        for (i = 0; i < node->nchild; i++)

        {

          if (matched[i])

            QTNFree(node->child[i]);

          else

            node->child[j++] = node->child[i];

        }

        /* and instead insert a copy of subs */

        if (subs)

        {

          subs = QTNCopy(subs);

          subs->flags |= QTN_NOCHANGE;

          node->child[j++] = subs;

        }

        node->nchild = j;

        /*

         * At this point we might have a node with zero or one child,

         * which should be simplified.  But we leave it to our caller

         * (dofindsubquery) to take care of that.

         */

        /*

         * Re-sort the node to put new child in the right place.  This

         * is a bit bogus, because it won't matter for findsubquery's

         * remaining processing, and it's insufficient to prepare the

         * tree for another search (we would need to re-flatten as

         * well, and we don't want to do that because we'd lose the

         * QTN_NOCHANGE marking on the new child).  But it's needed to

         * keep the results the same as the regression tests expect.

         */

        QTNSort(node);

        *isfind = true;

      }

      pfree(matched);

    }

  }

  else

  {

    Assert(node->valnode->type == QI_VAL);

    if (node->valnode->qoperand.valcrc != ex->valnode->qoperand.valcrc)

      return node;

    else if (QTNEq(node, ex))

    {

      QTNFree(node);

      if (subs)

      {

        node = QTNCopy(subs);

        node->flags |= QTN_NOCHANGE;

      }

      else

      {

        node = NULL;

      }

      *isfind = true;

    }

  }

  return node;

}

/*

 * Recursive guts of findsubquery(): attempt to replace "ex" with "subs"

 * at the root node, and if we failed to do so, recursively match against

 * child nodes.

 *

 * Delete any void subtrees resulting from the replacement.

 * In the following example '5' is replaced by empty operand:

 *

 *    AND   ->    6

 *   /   \

 *  5  OR

 *    /  \

 *     6  5

 */

static QTNode *

dofindsubquery(QTNode *root, QTNode *ex, QTNode *subs, bool *isfind)

{

  /* since this function recurses, it could be driven to stack overflow. */

  check_stack_depth();

  /* also, since it's a bit expensive, let's check for query cancel. */

  CHECK_FOR_INTERRUPTS();

  /* match at the node itself */

  root = findeq(root, ex, subs, isfind);

  /* unless we matched here, consider matches at child nodes */

  if (root && (root->flags & QTN_NOCHANGE) == 0 &&

    root->valnode->type == QI_OPR)

  {

    int     i,

          j = 0;

    /*

     * Any subtrees that are replaced by NULL must be dropped from the

     * tree.

     */

    for (i = 0; i < root->nchild; i++)

    {

      root->child[j] = dofindsubquery(root->child[i], ex, subs, isfind);

      if (root->child[j])

        j++;

    }

    root->nchild = j;

    /*

     * If we have just zero or one remaining child node, simplify out this

     * operator node.

     */

    if (root->nchild == 0)

    {

      QTNFree(root);

      root = NULL;

    }

    else if (root->nchild == 1 && root->valnode->qoperator.oper != OP_NOT)

    {

      QTNode     *nroot = root->child[0];

      pfree(root);

      root = nroot;

    }

  }

  return root;

}

/*

 * Substitute "subs" for "ex" throughout the QTNode tree at root.

 *

 * If isfind isn't NULL, set *isfind to show whether we made any substitution.

 *

 * Both "root" and "ex" must have been through QTNTernary and QTNSort

 * to ensure reliable matching.

 */

QTNode *

findsubquery(QTNode *root, QTNode *ex, QTNode *subs, bool *isfind)

{

  bool    DidFind = false;

  root = dofindsubquery(root, ex, subs, &DidFind);

  if (isfind)

    *isfind = DidFind;

  return root;

}

Datum

tsquery_rewrite_query(PG_FUNCTION_ARGS)

{

  TSQuery   query = PG_GETARG_TSQUERY_COPY(0);

  text     *in = PG_GETARG_TEXT_PP(1);

  TSQuery   rewritten = query;

  MemoryContext outercontext = CurrentMemoryContext;

  MemoryContext oldcontext;

  QTNode     *tree;

  char     *buf;

  SPIPlanPtr  plan;

  Portal    portal;

  bool    isnull;

  if (query->size == 0)

  {

    PG_FREE_IF_COPY(in, 1);

    PG_RETURN_POINTER(rewritten);

  }

  tree = QT2QTN(GETQUERY(query), GETOPERAND(query));

  QTNTernary(tree);

  QTNSort(tree);

  buf = text_to_cstring(in);

  SPI_connect();

  if ((plan = SPI_prepare(buf, 0, NULL)) == NULL)

    elog(ERROR, "SPI_prepare(\"%s\") failed", buf);

  if ((portal = SPI_cursor_open(NULL, plan, NULL, NULL, true)) == NULL)

    elog(ERROR, "SPI_cursor_open(\"%s\") failed", buf);

  SPI_cursor_fetch(portal, true, 100);

  if (SPI_tuptable == NULL ||

    SPI_tuptable->tupdesc->natts != 2 ||

    SPI_gettypeid(SPI_tuptable->tupdesc, 1) != TSQUERYOID ||

    SPI_gettypeid(SPI_tuptable->tupdesc, 2) != TSQUERYOID)

    ereport(ERROR,

        (errcode(ERRCODE_INVALID_PARAMETER_VALUE),

         errmsg("ts_rewrite query must return two tsquery columns")));

  while (SPI_processed > 0 && tree)

  {

    uint64    i;

    for (i = 0; i < SPI_processed && tree; i++)

    {

      Datum   qdata = SPI_getbinval(SPI_tuptable->vals[i], SPI_tuptable->tupdesc, 1, &isnull);

      Datum   sdata;

      if (isnull)

        continue;

      sdata = SPI_getbinval(SPI_tuptable->vals[i], SPI_tuptable->tupdesc, 2, &isnull);

      if (!isnull)

      {

        TSQuery   qtex = DatumGetTSQuery(qdata);

        TSQuery   qtsubs = DatumGetTSQuery(sdata);

        QTNode     *qex,

               *qsubs = NULL;

        if (qtex->size == 0)

        {

          if (qtex != (TSQuery) DatumGetPointer(qdata))

            pfree(qtex);

          if (qtsubs != (TSQuery) DatumGetPointer(sdata))

            pfree(qtsubs);

          continue;

        }

        qex = QT2QTN(GETQUERY(qtex), GETOPERAND(qtex));

        QTNTernary(qex);

        QTNSort(qex);

        if (qtsubs->size)

          qsubs = QT2QTN(GETQUERY(qtsubs), GETOPERAND(qtsubs));

        oldcontext = MemoryContextSwitchTo(outercontext);

        tree = findsubquery(tree, qex, qsubs, NULL);

        MemoryContextSwitchTo(oldcontext);

        QTNFree(qex);

        if (qtex != (TSQuery) DatumGetPointer(qdata))

          pfree(qtex);

        QTNFree(qsubs);

        if (qtsubs != (TSQuery) DatumGetPointer(sdata))

          pfree(qtsubs);

        if (tree)

        {

          /* ready the tree for another pass */

          QTNClearFlags(tree, QTN_NOCHANGE);

          QTNTernary(tree);

          QTNSort(tree);

        }

      }

    }

    SPI_freetuptable(SPI_tuptable);

    SPI_cursor_fetch(portal, true, 100);

  }

  SPI_freetuptable(SPI_tuptable);

  SPI_cursor_close(portal);

  SPI_freeplan(plan);

  SPI_finish();

  if (tree)

  {

    QTNBinary(tree);

    rewritten = QTN2QT(tree);

    QTNFree(tree);

    PG_FREE_IF_COPY(query, 0);

  }

  else

  {

    SET_VARSIZE(rewritten, HDRSIZETQ);

    rewritten->size = 0;

  }

  pfree(buf);

  PG_FREE_IF_COPY(in, 1);

  PG_RETURN_POINTER(rewritten);

}

Datum

tsquery_rewrite(PG_FUNCTION_ARGS)

{

  TSQuery   query = PG_GETARG_TSQUERY_COPY(0);

  TSQuery   ex = PG_GETARG_TSQUERY(1);

  TSQuery   subst = PG_GETARG_TSQUERY(2);

  TSQuery   rewritten = query;

  QTNode     *tree,

         *qex,

         *subs = NULL;

  if (query->size == 0 || ex->size == 0)

  {

    PG_FREE_IF_COPY(ex, 1);

    PG_FREE_IF_COPY(subst, 2);

    PG_RETURN_POINTER(rewritten);

  }

  tree = QT2QTN(GETQUERY(query), GETOPERAND(query));

  QTNTernary(tree);

  QTNSort(tree);

  qex = QT2QTN(GETQUERY(ex), GETOPERAND(ex));

  QTNTernary(qex);

  QTNSort(qex);

  if (subst->size)

    subs = QT2QTN(GETQUERY(subst), GETOPERAND(subst));

  tree = findsubquery(tree, qex, subs, NULL);

  QTNFree(qex);

  QTNFree(subs);

  if (!tree)

  {

    SET_VARSIZE(rewritten, HDRSIZETQ);

    rewritten->size = 0;

    PG_FREE_IF_COPY(ex, 1);

    PG_FREE_IF_COPY(subst, 2);

    PG_RETURN_POINTER(rewritten);

  }

  else

  {

    QTNBinary(tree);

    rewritten = QTN2QT(tree);

    QTNFree(tree);

  }

  PG_FREE_IF_COPY(query, 0);

  PG_FREE_IF_COPY(ex, 1);

  PG_FREE_IF_COPY(subst, 2);

  PG_RETURN_POINTER(rewritten);

}