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

 *

 * nodeLimit.c

 *    Routines to handle limiting of query results where appropriate

 *

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

 * Portions Copyright (c) 1994, Regents of the University of California

 *

 *

 * IDENTIFICATION

 *    src/backend/executor/nodeLimit.c

 *

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

 */

/*

 * INTERFACE ROUTINES

 *    ExecLimit   - extract a limited range of tuples

 *    ExecInitLimit - initialize node and subnodes..

 *    ExecEndLimit  - shutdown node and subnodes

 */

#include "postgres.h"

#include "executor/executor.h"

#include "executor/nodeLimit.h"

#include "miscadmin.h"

static void recompute_limits(LimitState *node);

static int64 compute_tuples_needed(LimitState *node);

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

 *    ExecLimit

 *

 *    This is a very simple node which just performs LIMIT/OFFSET

 *    filtering on the stream of tuples returned by a subplan.

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

 */

static TupleTableSlot *     /* return: a tuple or NULL */

ExecLimit(PlanState *pstate)

{

  LimitState *node = castNode(LimitState, pstate);

  ExprContext *econtext = node->ps.ps_ExprContext;

  ScanDirection direction;

  TupleTableSlot *slot;

  PlanState  *outerPlan;

  CHECK_FOR_INTERRUPTS();

  /*

   * get information from the node

   */

  direction = node->ps.state->es_direction;

  outerPlan = outerPlanState(node);

  /*

   * The main logic is a simple state machine.

   */

  switch (node->lstate)

  {

    case LIMIT_INITIAL:

      /*

       * First call for this node, so compute limit/offset. (We can't do

       * this any earlier, because parameters from upper nodes will not

       * be set during ExecInitLimit.)  This also sets position = 0 and

       * changes the state to LIMIT_RESCAN.

       */

      recompute_limits(node);

      /* FALL THRU */

    case LIMIT_RESCAN:

      /*

       * If backwards scan, just return NULL without changing state.

       */

      if (!ScanDirectionIsForward(direction))

        return NULL;

      /*

       * Check for empty window; if so, treat like empty subplan.

       */

      if (node->count <= 0 && !node->noCount)

      {

        node->lstate = LIMIT_EMPTY;

        return NULL;

      }

      /*

       * Fetch rows from subplan until we reach position > offset.

       */

      for (;;)

      {

        slot = ExecProcNode(outerPlan);

        if (TupIsNull(slot))

        {

          /*

           * The subplan returns too few tuples for us to produce

           * any output at all.

           */

          node->lstate = LIMIT_EMPTY;

          return NULL;

        }

        /*

         * Tuple at limit is needed for comparison in subsequent

         * execution to detect ties.

         */

        if (node->limitOption == LIMIT_OPTION_WITH_TIES &&

          node->position - node->offset == node->count - 1)

        {

          ExecCopySlot(node->last_slot, slot);

        }

        node->subSlot = slot;

        if (++node->position > node->offset)

          break;

      }

      /*

       * Okay, we have the first tuple of the window.

       */

      node->lstate = LIMIT_INWINDOW;

      break;

    case LIMIT_EMPTY:

      /*

       * The subplan is known to return no tuples (or not more than

       * OFFSET tuples, in general).  So we return no tuples.

       */

      return NULL;

    case LIMIT_INWINDOW:

      if (ScanDirectionIsForward(direction))

      {

        /*

         * Forwards scan, so check for stepping off end of window.  At

         * the end of the window, the behavior depends on whether WITH

         * TIES was specified: if so, we need to change the state

         * machine to WINDOWEND_TIES, and fall through to the code for

         * that case.  If not (nothing was specified, or ONLY was)

         * return NULL without advancing the subplan or the position

         * variable, but change the state machine to record having

         * done so.

         *

         * Once at the end, ideally, we would shut down parallel

         * resources; but that would destroy the parallel context

         * which might be required for rescans.  To do that, we'll

         * need to find a way to pass down more information about

         * whether rescans are possible.

         */

        if (!node->noCount &&

          node->position - node->offset >= node->count)

        {

          if (node->limitOption == LIMIT_OPTION_COUNT)

          {

            node->lstate = LIMIT_WINDOWEND;

            return NULL;

          }

          else

          {

            node->lstate = LIMIT_WINDOWEND_TIES;

            /* we'll fall through to the next case */

          }

        }

        else

        {

          /*

           * Get next tuple from subplan, if any.

           */

          slot = ExecProcNode(outerPlan);

          if (TupIsNull(slot))

          {

            node->lstate = LIMIT_SUBPLANEOF;

            return NULL;

          }

          /*

           * If WITH TIES is active, and this is the last in-window

           * tuple, save it to be used in subsequent WINDOWEND_TIES

           * processing.

           */

          if (node->limitOption == LIMIT_OPTION_WITH_TIES &&

            node->position - node->offset == node->count - 1)

          {

            ExecCopySlot(node->last_slot, slot);

          }

          node->subSlot = slot;

          node->position++;

          break;

        }

      }

      else

      {

        /*

         * Backwards scan, so check for stepping off start of window.

         * As above, only change state-machine status if so.

         */

        if (node->position <= node->offset + 1)

        {

          node->lstate = LIMIT_WINDOWSTART;

          return NULL;

        }

        /*

         * Get previous tuple from subplan; there should be one!

         */

        slot = ExecProcNode(outerPlan);

        if (TupIsNull(slot))

          elog(ERROR, "LIMIT subplan failed to run backwards");

        node->subSlot = slot;

        node->position--;

        break;

      }

      Assert(node->lstate == LIMIT_WINDOWEND_TIES);

      /* FALL THRU */

    case LIMIT_WINDOWEND_TIES:

      if (ScanDirectionIsForward(direction))

      {

        /*

         * Advance the subplan until we find the first row with

         * different ORDER BY pathkeys.

         */

        slot = ExecProcNode(outerPlan);

        if (TupIsNull(slot))

        {

          node->lstate = LIMIT_SUBPLANEOF;

          return NULL;

        }

        /*

         * Test if the new tuple and the last tuple match. If so we

         * return the tuple.

         */

        econtext->ecxt_innertuple = slot;

        econtext->ecxt_outertuple = node->last_slot;

        if (ExecQualAndReset(node->eqfunction, econtext))

        {

          node->subSlot = slot;

          node->position++;

        }

        else

        {

          node->lstate = LIMIT_WINDOWEND;

          return NULL;

        }

      }

      else

      {

        /*

         * Backwards scan, so check for stepping off start of window.

         * Change only state-machine status if so.

         */

        if (node->position <= node->offset + 1)

        {

          node->lstate = LIMIT_WINDOWSTART;

          return NULL;

        }

        /*

         * Get previous tuple from subplan; there should be one! And

         * change state-machine status.

         */

        slot = ExecProcNode(outerPlan);

        if (TupIsNull(slot))

          elog(ERROR, "LIMIT subplan failed to run backwards");

        node->subSlot = slot;

        node->position--;

        node->lstate = LIMIT_INWINDOW;

      }

      break;

    case LIMIT_SUBPLANEOF:

      if (ScanDirectionIsForward(direction))

        return NULL;

      /*

       * Backing up from subplan EOF, so re-fetch previous tuple; there

       * should be one!  Note previous tuple must be in window.

       */

      slot = ExecProcNode(outerPlan);

      if (TupIsNull(slot))

        elog(ERROR, "LIMIT subplan failed to run backwards");

      node->subSlot = slot;

      node->lstate = LIMIT_INWINDOW;

      /* position does not change 'cause we didn't advance it before */

      break;

    case LIMIT_WINDOWEND:

      if (ScanDirectionIsForward(direction))

        return NULL;

      /*

       * We already past one position to detect ties so re-fetch

       * previous tuple; there should be one!  Note previous tuple must

       * be in window.

       */

      if (node->limitOption == LIMIT_OPTION_WITH_TIES)

      {

        slot = ExecProcNode(outerPlan);

        if (TupIsNull(slot))

          elog(ERROR, "LIMIT subplan failed to run backwards");

        node->subSlot = slot;

        node->lstate = LIMIT_INWINDOW;

      }

      else

      {

        /*

         * Backing up from window end: simply re-return the last tuple

         * fetched from the subplan.

         */

        slot = node->subSlot;

        node->lstate = LIMIT_INWINDOW;

        /* position does not change 'cause we didn't advance it before */

      }

      break;

    case LIMIT_WINDOWSTART:

      if (!ScanDirectionIsForward(direction))

        return NULL;

      /*

       * Advancing after having backed off window start: simply

       * re-return the last tuple fetched from the subplan.

       */

      slot = node->subSlot;

      node->lstate = LIMIT_INWINDOW;

      /* position does not change 'cause we didn't change it before */

      break;

    default:

      elog(ERROR, "impossible LIMIT state: %d",

         (int) node->lstate);

      slot = NULL;   /* keep compiler quiet */

      break;

  }

  /* Return the current tuple */

  Assert(!TupIsNull(slot));

  return slot;

}

/*

 * Evaluate the limit/offset expressions --- done at startup or rescan.

 *

 * This is also a handy place to reset the current-position state info.

 */

static void

recompute_limits(LimitState *node)

{

  ExprContext *econtext = node->ps.ps_ExprContext;

  Datum   val;

  bool    isNull;

  if (node->limitOffset)

  {

    val = ExecEvalExprSwitchContext(node->limitOffset,

                    econtext,

                    &isNull);

    /* Interpret NULL offset as no offset */

    if (isNull)

      node->offset = 0;

    else

    {

      node->offset = DatumGetInt64(val);

      if (node->offset < 0)

        ereport(ERROR,

            (errcode(ERRCODE_INVALID_ROW_COUNT_IN_RESULT_OFFSET_CLAUSE),

             errmsg("OFFSET must not be negative")));

    }

  }

  else

  {

    /* No OFFSET supplied */

    node->offset = 0;

  }

  if (node->limitCount)

  {

    val = ExecEvalExprSwitchContext(node->limitCount,

                    econtext,

                    &isNull);

    /* Interpret NULL count as no count (LIMIT ALL) */

    if (isNull)

    {

      node->count = 0;

      node->noCount = true;

    }

    else

    {

      node->count = DatumGetInt64(val);

      if (node->count < 0)

        ereport(ERROR,

            (errcode(ERRCODE_INVALID_ROW_COUNT_IN_LIMIT_CLAUSE),

             errmsg("LIMIT must not be negative")));

      node->noCount = false;

    }

  }

  else

  {

    /* No COUNT supplied */

    node->count = 0;

    node->noCount = true;

  }

  /* Reset position to start-of-scan */

  node->position = 0;

  node->subSlot = NULL;

  /* Set state-machine state */

  node->lstate = LIMIT_RESCAN;

  /*

   * Notify child node about limit.  Note: think not to "optimize" by

   * skipping ExecSetTupleBound if compute_tuples_needed returns < 0.  We

   * must update the child node anyway, in case this is a rescan and the

   * previous time we got a different result.

   */

  ExecSetTupleBound(compute_tuples_needed(node), outerPlanState(node));

}

/*

 * Compute the maximum number of tuples needed to satisfy this Limit node.

 * Return a negative value if there is not a determinable limit.

 */

static int64

compute_tuples_needed(LimitState *node)

{

  if ((node->noCount) || (node->limitOption == LIMIT_OPTION_WITH_TIES))

    return -1;

  /* Note: if this overflows, we'll return a negative value, which is OK */

  return node->count + node->offset;

}

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

 *    ExecInitLimit

 *

 *    This initializes the limit node state structures and

 *    the node's subplan.

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

 */

LimitState *

ExecInitLimit(Limit *node, EState *estate, int eflags)

{

  LimitState *limitstate;

  Plan     *outerPlan;

  /* check for unsupported flags */

  Assert(!(eflags & EXEC_FLAG_MARK));

  /*

   * create state structure

   */

  limitstate = makeNode(LimitState);

  limitstate->ps.plan = (Plan *) node;

  limitstate->ps.state = estate;

  limitstate->ps.ExecProcNode = ExecLimit;

  limitstate->lstate = LIMIT_INITIAL;

  /*

   * Miscellaneous initialization

   *

   * Limit nodes never call ExecQual or ExecProject, but they need an

   * exprcontext anyway to evaluate the limit/offset parameters in.

   */

  ExecAssignExprContext(estate, &limitstate->ps);

  /*

   * initialize outer plan

   */

  outerPlan = outerPlan(node);

  outerPlanState(limitstate) = ExecInitNode(outerPlan, estate, eflags);

  /*

   * initialize child expressions

   */

  limitstate->limitOffset = ExecInitExpr((Expr *) node->limitOffset,

                       (PlanState *) limitstate);

  limitstate->limitCount = ExecInitExpr((Expr *) node->limitCount,

                      (PlanState *) limitstate);

  limitstate->limitOption = node->limitOption;

  /*

   * Initialize result type.

   */

  ExecInitResultTypeTL(&limitstate->ps);

  limitstate->ps.resultopsset = true;

  limitstate->ps.resultops = ExecGetResultSlotOps(outerPlanState(limitstate),

                          &limitstate->ps.resultopsfixed);

  /*

   * limit nodes do no projections, so initialize projection info for this

   * node appropriately

   */

  limitstate->ps.ps_ProjInfo = NULL;

  /*

   * Initialize the equality evaluation, to detect ties.

   */

  if (node->limitOption == LIMIT_OPTION_WITH_TIES)

  {

    TupleDesc desc;

    const TupleTableSlotOps *ops;

    desc = ExecGetResultType(outerPlanState(limitstate));

    ops = ExecGetResultSlotOps(outerPlanState(limitstate), NULL);

    limitstate->last_slot = ExecInitExtraTupleSlot(estate, desc, ops);

    limitstate->eqfunction = execTuplesMatchPrepare(desc,

                            node->uniqNumCols,

                            node->uniqColIdx,

                            node->uniqOperators,

                            node->uniqCollations,

                            &limitstate->ps);

  }

  return limitstate;

}

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

 *    ExecEndLimit

 *

 *    This shuts down the subplan and frees resources allocated

 *    to this node.

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

 */

void

ExecEndLimit(LimitState *node)

{

  ExecEndNode(outerPlanState(node));

}

void

ExecReScanLimit(LimitState *node)

{

  PlanState  *outerPlan = outerPlanState(node);

  /*

   * Recompute limit/offset in case parameters changed, and reset the state

   * machine.  We must do this before rescanning our child node, in case

   * it's a Sort that we are passing the parameters down to.

   */

  recompute_limits(node);

  /*

   * if chgParam of subnode is not null then plan will be re-scanned by

   * first ExecProcNode.

   */

  if (outerPlan->chgParam == NULL)

    ExecReScan(outerPlan);

}