/*-------------------------------------------------------------------------
 *
 * jsonbsubs.c
 *   Subscripting support functions for jsonb.
 *
 * Portions Copyright (c) 1996-2025, PostgreSQL Global Development Group
 * Portions Copyright (c) 1994, Regents of the University of California
 *
 *
 * IDENTIFICATION
 *   src/backend/utils/adt/jsonbsubs.c
 *
 *-------------------------------------------------------------------------
 */
#include "postgres.h"

#include "executor/execExpr.h"
#include "nodes/nodeFuncs.h"
#include "nodes/subscripting.h"
#include "parser/parse_coerce.h"
#include "parser/parse_expr.h"
#include "utils/builtins.h"
#include "utils/jsonb.h"


/* SubscriptingRefState.workspace for jsonb subscripting execution */
typedef struct JsonbSubWorkspace
{
 bool  expectArray; /* jsonb root is expected to be an array */
 Oid     *indexOid;  /* OID of coerced subscript expression, could
 * be only integer or text */
 Datum    *index;   /* Subscript values in Datum format */
} JsonbSubWorkspace;


/*
 * Finish parse analysis of a SubscriptingRef expression for a jsonb.
 *
 * Transform the subscript expressions, coerce them to text,
 * and determine the result type of the SubscriptingRef node.
 */
static void
jsonb_subscript_transform(SubscriptingRef *sbsref,
        List *indirection,
        ParseState *pstate,
        bool isSlice,
        bool isAssignment)
{
 List    *upperIndexpr = NIL;
 ListCell   *idx;

 /*
  * Transform and convert the subscript expressions. Jsonb subscripting
  * does not support slices, look only and the upper index.
 */
 foreach(idx, indirection)
 {
  A_Indices  *ai = lfirst_node(A_Indices, idx);
  Node    *subExpr;

  if (isSlice)
  {
   Node    *expr = ai->uidx ? ai->uidx : ai->lidx;

   ereport(ERROR,
     (errcode(ERRCODE_DATATYPE_MISMATCH),
      errmsg("jsonb subscript does not support slices"),
      parser_errposition(pstate, exprLocation(expr))));
  }

  if (ai->uidx)
  {
   Oid   subExprType = InvalidOid,
      targetType = UNKNOWNOID;

   subExpr = transformExpr(pstate, ai->uidx, pstate->p_expr_kind);
   subExprType = exprType(subExpr);

   if (subExprType != UNKNOWNOID)
   {
    Oid   targets[2] = {INT4OID, TEXTOID};

    /*
     * Jsonb can handle multiple subscript types, but cases when a
     * subscript could be coerced to multiple target types must be
     * avoided, similar to overloaded functions. It could be
     * possibly extend with jsonpath in the future.
 */
    for (int i = 0; i < 2; i++)
    {
     if (can_coerce_type(1, &subExprType, &targets[i], COERCION_IMPLICIT))
     {
      /*
       * One type has already succeeded, it means there are
       * two coercion targets possible, failure.
 */
      if (targetType != UNKNOWNOID)
       ereport(ERROR,
         (errcode(ERRCODE_DATATYPE_MISMATCH),
          errmsg("subscript type %s is not supported", format_type_be(subExprType)),
          errhint("jsonb subscript must be coercible to only one type, integer or text."),
          parser_errposition(pstate, exprLocation(subExpr))));

      targetType = targets[i];
     }
    }

    /*
     * No suitable types were found, failure.
 */
    if (targetType == UNKNOWNOID)
     ereport(ERROR,
       (errcode(ERRCODE_DATATYPE_MISMATCH),
        errmsg("subscript type %s is not supported", format_type_be(subExprType)),
        errhint("jsonb subscript must be coercible to either integer or text."),
        parser_errposition(pstate, exprLocation(subExpr))));
   }
   else
    targetType = TEXTOID;

   /*
    * We known from can_coerce_type that coercion will succeed, so
    * coerce_type could be used. Note the implicit coercion context,
    * which is required to handle subscripts of different types,
    * similar to overloaded functions.
 */
   subExpr = coerce_type(pstate,
          subExpr, subExprType,
          targetType, -1,
          COERCION_IMPLICIT,
          COERCE_IMPLICIT_CAST,
          -1);
   if (subExpr == NULL)
    ereport(ERROR,
      (errcode(ERRCODE_DATATYPE_MISMATCH),
       errmsg("jsonb subscript must have text type"),
       parser_errposition(pstate, exprLocation(subExpr))));
  }
  else
  {
   /*
    * Slice with omitted upper bound. Should not happen as we already
    * errored out on slice earlier, but handle this just in case.
 */
   Assert(isSlice && ai->is_slice);
   ereport(ERROR,
     (errcode(ERRCODE_DATATYPE_MISMATCH),
      errmsg("jsonb subscript does not support slices"),
      parser_errposition(pstate, exprLocation(ai->uidx))));
  }

  upperIndexpr = lappend(upperIndexpr, subExpr);
 }

 /* store the transformed lists into the SubscriptRef node */
 sbsref->refupperindexpr = upperIndexpr;
 sbsref->reflowerindexpr = NIL;

 /* Determine the result type of the subscripting operation; always jsonb */
 sbsref->refrestype = JSONBOID;
 sbsref->reftypmod = -1;
}

/*
 * During execution, process the subscripts in a SubscriptingRef expression.
 *
 * The subscript expressions are already evaluated in Datum form in the
 * SubscriptingRefState's arrays.  Check and convert them as necessary.
 *
 * If any subscript is NULL, we throw error in assignment cases, or in fetch
 * cases set result to NULL and return false (instructing caller to skip the
 * rest of the SubscriptingRef sequence).
 */
static bool
jsonb_subscript_check_subscripts(ExprState *state,
         ExprEvalStep *op,
         ExprContext *econtext)
{
 SubscriptingRefState *sbsrefstate = op->d.sbsref_subscript.state;
 JsonbSubWorkspace *workspace = (JsonbSubWorkspace *) sbsrefstate->workspace;

 /*
  * In case if the first subscript is an integer, the source jsonb is
  * expected to be an array. This information is not used directly, all
  * such cases are handled within corresponding jsonb assign functions. But
  * if the source jsonb is NULL the expected type will be used to construct
  * an empty source.
 */
 if (sbsrefstate->numupper > 0 && sbsrefstate->upperprovided[0] &&
  !sbsrefstate->upperindexnull[0] && workspace->indexOid[0] == INT4OID)
  workspace->expectArray = true;

 /* Process upper subscripts */
 for (int i = 0; i < sbsrefstate->numupper; i++)
 {
  if (sbsrefstate->upperprovided[i])
  {
   /* If any index expr yields NULL, result is NULL or error */
   if (sbsrefstate->upperindexnull[i])
   {
    if (sbsrefstate->isassignment)
     ereport(ERROR,
       (errcode(ERRCODE_NULL_VALUE_NOT_ALLOWED),
        errmsg("jsonb subscript in assignment must not be null")));
    *op->resnull = true;
    return false;
   }

   /*
    * For jsonb fetch and assign functions we need to provide path in
    * text format. Convert if it's not already text.
 */
   if (workspace->indexOid[i] == INT4OID)
   {
    Datum  datum = sbsrefstate->upperindex[i];
    char    *cs = DatumGetCString(DirectFunctionCall1(int4out, datum));

    workspace->index[i] = CStringGetTextDatum(cs);
   }
   else
    workspace->index[i] = sbsrefstate->upperindex[i];
  }
 }

 return true;
}

/*
 * Evaluate SubscriptingRef fetch for a jsonb element.
 *
 * Source container is in step's result variable (it's known not NULL, since
 * we set fetch_strict to true).
 */
static void
jsonb_subscript_fetch(ExprState *state,
       ExprEvalStep *op,
       ExprContext *econtext)
{
 SubscriptingRefState *sbsrefstate = op->d.sbsref.state;
 JsonbSubWorkspace *workspace = (JsonbSubWorkspace *) sbsrefstate->workspace;
 Jsonb    *jsonbSource;

 /* Should not get here if source jsonb (or any subscript) is null */
 Assert(!(*op->resnull));

 jsonbSource = DatumGetJsonbP(*op->resvalue);
 *op->resvalue = jsonb_get_element(jsonbSource,
           workspace->index,
           sbsrefstate->numupper,
           op->resnull,
           false);
}

/*
 * Evaluate SubscriptingRef assignment for a jsonb element assignment.
 *
 * Input container (possibly null) is in result area, replacement value is in
 * SubscriptingRefState's replacevalue/replacenull.
 */
static void
jsonb_subscript_assign(ExprState *state,
        ExprEvalStep *op,
        ExprContext *econtext)
{
 SubscriptingRefState *sbsrefstate = op->d.sbsref.state;
 JsonbSubWorkspace *workspace = (JsonbSubWorkspace *) sbsrefstate->workspace;
 Jsonb    *jsonbSource;
 JsonbValue replacevalue;

 if (sbsrefstate->replacenull)
  replacevalue.type = jbvNull;
 else
  JsonbToJsonbValue(DatumGetJsonbP(sbsrefstate->replacevalue),
        &replacevalue);

 /*
  * In case if the input container is null, set up an empty jsonb and
  * proceed with the assignment.
 */
 if (*op->resnull)
 {
  JsonbValue newSource;

  /*
   * To avoid any surprising results, set up an empty jsonb array in
   * case of an array is expected (i.e. the first subscript is integer),
   * otherwise jsonb object.
 */
  if (workspace->expectArray)
  {
   newSource.type = jbvArray;
   newSource.val.array.nElems = 0;
   newSource.val.array.rawScalar = false;
  }
  else
  {
   newSource.type = jbvObject;
   newSource.val.object.nPairs = 0;
  }

  jsonbSource = JsonbValueToJsonb(&newSource);
  *op->resnull = false;
 }
 else
  jsonbSource = DatumGetJsonbP(*op->resvalue);

 *op->resvalue = jsonb_set_element(jsonbSource,
           workspace->index,
           sbsrefstate->numupper,
           &replacevalue);
 /* The result is never NULL, so no need to change *op->resnull */
}

/*
 * Compute old jsonb element value for a SubscriptingRef assignment
 * expression.  Will only be called if the new-value subexpression
 * contains SubscriptingRef or FieldStore.  This is the same as the
 * regular fetch case, except that we have to handle a null jsonb,
 * and the value should be stored into the SubscriptingRefState's
 * prevvalue/prevnull fields.
 */
static void
jsonb_subscript_fetch_old(ExprState *state,
        ExprEvalStep *op,
        ExprContext *econtext)
{
 SubscriptingRefState *sbsrefstate = op->d.sbsref.state;

 if (*op->resnull)
 {
  /* whole jsonb is null, so any element is too */
  sbsrefstate->prevvalue = (Datum) 0;
  sbsrefstate->prevnull = true;
 }
 else
 {
  Jsonb    *jsonbSource = DatumGetJsonbP(*op->resvalue);

  sbsrefstate->prevvalue = jsonb_get_element(jsonbSource,
               sbsrefstate->upperindex,
               sbsrefstate->numupper,
               &sbsrefstate->prevnull,
               false);
 }
}

/*
 * Set up execution state for a jsonb subscript operation. Opposite to the
 * arrays subscription, there is no limit for number of subscripts as jsonb
 * type itself doesn't have nesting limits.
 */
static void
jsonb_exec_setup(const SubscriptingRef *sbsref,
     SubscriptingRefState *sbsrefstate,
     SubscriptExecSteps *methods)
{
 JsonbSubWorkspace *workspace;
 ListCell   *lc;
 int   nupper = sbsref->refupperindexpr->length;
 char    *ptr;

 /* Allocate type-specific workspace with space for per-subscript data */
 workspace = palloc0(MAXALIGN(sizeof(JsonbSubWorkspace)) +
      nupper * (sizeof(Datum) + sizeof(Oid)));
 workspace->expectArray = false;
 ptr = ((char *) workspace) + MAXALIGN(sizeof(JsonbSubWorkspace));

 /*
  * This coding assumes sizeof(Datum) >= sizeof(Oid), else we might
  * misalign the indexOid pointer
 */
 workspace->index = (Datum *) ptr;
 ptr += nupper * sizeof(Datum);
 workspace->indexOid = (Oid *) ptr;

 sbsrefstate->workspace = workspace;

 /* Collect subscript data types necessary at execution time */
 foreach(lc, sbsref->refupperindexpr)
 {
  Node    *expr = lfirst(lc);
  int   i = foreach_current_index(lc);

  workspace->indexOid[i] = exprType(expr);
 }

 /*
  * Pass back pointers to appropriate step execution functions.
 */
 methods->sbs_check_subscripts = jsonb_subscript_check_subscripts;
 methods->sbs_fetch = jsonb_subscript_fetch;
 methods->sbs_assign = jsonb_subscript_assign;
 methods->sbs_fetch_old = jsonb_subscript_fetch_old;
}

/*
 * jsonb_subscript_handler
 *  Subscripting handler for jsonb.
 *
 */
Datum
jsonb_subscript_handler(PG_FUNCTION_ARGS)
{
 static const SubscriptRoutines sbsroutines = {
  .transform = jsonb_subscript_transform,
  .exec_setup = jsonb_exec_setup,
  .fetch_strict = true, /* fetch returns NULL for NULL inputs */
  .fetch_leakproof = true, /* fetch returns NULL for bad subscript */
  .store_leakproof = false /* ... but assignment throws error */
 };

 PG_RETURN_POINTER(&sbsroutines);
}