Change the RETURNING algorithm so that outputs accumulate in an ephemeral

table until all modifications have been completed, and only then do results
start being returned.  This should help prevent problems with interleaved
sqlite3_step() calls on two separate DML statements.  It also seems to be
closer to how PostgreSQL works, which might prevent compatibility problems.

FossilOrigin-Name: c4615eb28c3dd2d473daf104f32e60d02799f3158d9d275a899c39129cc71401

8 files changed, 156 insertions, 93 deletions

diff --git a/src/build.c b/src/build.c
index d59282779..ce96a71dc 100644
--- a/src/build.c
+++ b/src/build.c
@@ -155,6 +155,22 @@ void sqlite3FinishCoding(Parse *pParse){
   assert( !pParse->isMultiWrite 
        || sqlite3VdbeAssertMayAbort(v, pParse->mayAbort));
   if( v ){
+    if( pParse->bReturning ){
+      Returning *pReturning = pParse->u1.pReturning;
+      int addrRewind;
+      int i;
+      int reg;
+
+      addrRewind =
+         sqlite3VdbeAddOp1(v, OP_Rewind, pReturning->iRetCur);
+      reg = pReturning->iRetReg;
+      for(i=0; i<pReturning->nRetCol; i++){
+        sqlite3VdbeAddOp3(v, OP_Column, pReturning->iRetCur, i, reg+i);
+      }
+      sqlite3VdbeAddOp2(v, OP_ResultRow, reg, i);
+      sqlite3VdbeAddOp2(v, OP_Next, pReturning->iRetCur, addrRewind+1);
+      sqlite3VdbeJumpHere(v, addrRewind);
+    }
     sqlite3VdbeAddOp0(v, OP_Halt);
 
 #if SQLITE_USER_AUTHENTICATION
@@ -232,6 +248,11 @@ void sqlite3FinishCoding(Parse *pParse){
         }
       }
 
+      if( pParse->bReturning ){
+        Returning *pRet = pParse->u1.pReturning;
+        sqlite3VdbeAddOp2(v, OP_OpenEphemeral, pRet->iRetCur, pRet->nRetCol);
+      }
+
       /* Finally, jump back to the beginning of the executable code. */
       sqlite3VdbeGoto(v, 1);
     }
@@ -1213,7 +1234,8 @@ void sqlite3StartTable(
     }else
 #endif
     {
-      pParse->addrCrTab =
+      assert( !pParse->bReturning );
+      pParse->u1.addrCrTab =
          sqlite3VdbeAddOp3(v, OP_CreateBtree, iDb, reg2, BTREE_INTKEY);
     }
     sqlite3OpenSchemaTable(pParse, iDb);
@@ -1291,6 +1313,7 @@ void sqlite3AddReturning(Parse *pParse, ExprList *pList){
     sqlite3ExprListDelete(db, pList);
     return;
   }
+  pParse->u1.pReturning = pRet;
   pRet->pParse = pParse;
   pRet->pReturnEL = pList;
   sqlite3ParserAddCleanup(pParse,
@@ -1304,10 +1327,7 @@ void sqlite3AddReturning(Parse *pParse, ExprList *pList){
   pRet->retTrig.step_list = &pRet->retTStep;
   pRet->retTStep.op = TK_RETURNING;
   pRet->retTStep.pTrig = &pRet->retTrig;
-  pRet->retTStep.pSelect = &pRet->retSel;
-  pRet->retSel.op = TK_ALL;
-  pRet->retSel.pEList = pList;
-  pRet->retSel.pSrc = (SrcList*)&pRet->retSrcList;
+  pRet->retTStep.pExprList = pList;
   pHash = &(db->aDb[1].pSchema->trigHash);
   assert( sqlite3HashFind(pHash, RETURNING_TRIGGER_NAME)==0 );
   if( sqlite3HashInsert(pHash, RETURNING_TRIGGER_NAME, &pRet->retTrig)
@@ -2147,9 +2167,10 @@ static void convertToWithoutRowidTable(Parse *pParse, Table *pTab){
   /* Convert the P3 operand of the OP_CreateBtree opcode from BTREE_INTKEY
   ** into BTREE_BLOBKEY.
   */
-  if( pParse->addrCrTab ){
+  assert( !pParse->bReturning );
+  if( pParse->u1.addrCrTab ){
     assert( v );
-    sqlite3VdbeChangeP3(v, pParse->addrCrTab, BTREE_BLOBKEY);
+    sqlite3VdbeChangeP3(v, pParse->u1.addrCrTab, BTREE_BLOBKEY);
   }
 
   /* Locate the PRIMARY KEY index.  Or, if this table was originally
diff --git a/src/insert.c b/src/insert.c
index 2f22121f7..7522a7269 100644
--- a/src/insert.c
+++ b/src/insert.c
@@ -1143,11 +1143,6 @@ void sqlite3Insert(
       sqlite3VdbeAddOp1(v, OP_MustBeInt, regCols); VdbeCoverage(v);
     }
 
-    /* Cannot have triggers on a virtual table. If it were possible,
-    ** this block would have to account for hidden column.
-    */
-    assert( !IsVirtual(pTab) );
-
     /* Copy the new data already generated. */
     assert( pTab->nNVCol>0 );
     sqlite3VdbeAddOp3(v, OP_Copy, regRowid+1, regCols+1, pTab->nNVCol-1);
diff --git a/src/resolve.c b/src/resolve.c
index 5074a2881..84ba82a11 100644
--- a/src/resolve.c
+++ b/src/resolve.c
@@ -371,7 +371,7 @@ static int lookupName(
     ** it is a new.* or old.* trigger argument reference.  Or
     ** maybe it is an excluded.* from an upsert.
     */
-    if( zDb==0 && cntTab==0 ){
+    if( cntTab==0 && zDb==0 ){
       pTab = 0;
 #ifndef SQLITE_OMIT_TRIGGER
       if( pParse->pTriggerTab!=0 ){
@@ -434,22 +434,27 @@ static int lookupName(
           }else
 #endif /* SQLITE_OMIT_UPSERT */
           {
-#ifndef SQLITE_OMIT_TRIGGER
-            if( iCol<0 ){
-              pExpr->affExpr = SQLITE_AFF_INTEGER;
-            }else if( pExpr->iTable==0 ){
-              testcase( iCol==31 );
-              testcase( iCol==32 );
-              pParse->oldmask |= (iCol>=32 ? 0xffffffff : (((u32)1)<<iCol));
-            }else{
-              testcase( iCol==31 );
-              testcase( iCol==32 );
-              pParse->newmask |= (iCol>=32 ? 0xffffffff : (((u32)1)<<iCol));
-            }
             pExpr->y.pTab = pTab;
-            pExpr->iColumn = (i16)iCol;
-            eNewExprOp = TK_TRIGGER;
+            if( iCol<0 ) pExpr->affExpr = SQLITE_AFF_INTEGER;
+            if( pParse->bReturning ){
+              eNewExprOp = TK_REGISTER;
+              pExpr->iTable = pNC->uNC.iBaseReg + (pTab->nCol+1)*pExpr->iTable
+                                   + iCol + 1;
+            }else{
+              pExpr->iColumn = (i16)iCol;
+              eNewExprOp = TK_TRIGGER;
+#ifndef SQLITE_OMIT_TRIGGER
+              if( pExpr->iTable==0 ){
+                testcase( iCol==31 );
+                testcase( iCol==32 );
+                pParse->oldmask |= (iCol>=32 ? 0xffffffff : (((u32)1)<<iCol));
+              }else{
+                testcase( iCol==31 );
+                testcase( iCol==32 );
+                pParse->newmask |= (iCol>=32 ? 0xffffffff : (((u32)1)<<iCol));
+              }
 #endif /* SQLITE_OMIT_TRIGGER */
+            }
           }
         }
       }
@@ -630,7 +635,10 @@ static int lookupName(
 lookupname_end:
   if( cnt==1 ){
     assert( pNC!=0 );
-    if( !ExprHasProperty(pExpr, EP_Alias) ){
+    if( pParse->db->xAuth
+     && !ExprHasProperty(pExpr, EP_Alias)
+     && pExpr->op!=TK_REGISTER
+    ){
       sqlite3AuthRead(pParse, pExpr, pSchema, pNC->pSrcList);
     }
     /* Increment the nRef value on all name contexts from TopNC up to
diff --git a/src/sqliteInt.h b/src/sqliteInt.h
index 3cadb7f53..7b22b6b3b 100644
--- a/src/sqliteInt.h
+++ b/src/sqliteInt.h
@@ -3037,6 +3037,7 @@ struct NameContext {
     ExprList *pEList;    /* Optional list of result-set columns */
     AggInfo *pAggInfo;   /* Information about aggregates at this level */
     Upsert *pUpsert;     /* ON CONFLICT clause information from an upsert */
+    int iBaseReg;        /* For TK_REGISTER when parsing RETURNING */
   } uNC;
   NameContext *pNext;  /* Next outer name context.  NULL for outermost */
   int nRef;            /* Number of names resolved by this context */
@@ -3065,6 +3066,7 @@ struct NameContext {
 #define NC_UEList    0x00080  /* True if uNC.pEList is used */
 #define NC_UAggInfo  0x00100  /* True if uNC.pAggInfo is used */
 #define NC_UUpsert   0x00200  /* True if uNC.pUpsert is used */
+#define NC_UBaseReg  0x00400  /* True if uNC.iBaseReg is used */
 #define NC_MinMaxAgg 0x01000  /* min/max aggregates seen.  See note above */
 #define NC_Complex   0x02000  /* True if a function or subquery seen */
 #define NC_AllowWin  0x04000  /* Window functions are allowed here */
@@ -3425,7 +3427,10 @@ struct Parse {
   Table *pTriggerTab;  /* Table triggers are being coded for */
   Parse *pParentParse; /* Parent parser if this parser is nested */
   AggInfo *pAggList;   /* List of all AggInfo objects */
-  int addrCrTab;       /* Address of OP_CreateBtree opcode on CREATE TABLE */
+  union {
+    int addrCrTab;         /* Address of OP_CreateBtree on CREATE TABLE */
+    Returning *pReturning; /* The RETURNING clause */
+  } u1;
   u32 nQueryLoop;      /* Est number of iterations of a query (10*log2(N)) */
   u32 oldmask;         /* Mask of old.* columns referenced */
   u32 newmask;         /* Mask of new.* columns referenced */
@@ -3647,7 +3652,7 @@ struct TriggerStep {
   char *zTarget;       /* Target table for DELETE, UPDATE, INSERT */
   SrcList *pFrom;      /* FROM clause for UPDATE statement (if any) */
   Expr *pWhere;        /* The WHERE clause for DELETE or UPDATE steps */
-  ExprList *pExprList; /* SET clause for UPDATE */
+  ExprList *pExprList; /* SET clause for UPDATE, or RETURNING clause */
   IdList *pIdList;     /* Column names for INSERT */
   Upsert *pUpsert;     /* Upsert clauses on an INSERT */
   char *zSpan;         /* Original SQL text of this command */
@@ -3663,8 +3668,9 @@ struct Returning {
   ExprList *pReturnEL;  /* List of expressions to return */
   Trigger retTrig;      /* The transient trigger that implements RETURNING */
   TriggerStep retTStep; /* The trigger step */
-  Select retSel;        /* The SELECT statement that implements RETURNING */
-  u64 retSrcList;       /* The empty FROM clause of the SELECT */
+  int iRetCur;          /* Transient table holding RETURNING results */
+  int nRetCol;          /* Number of in pReturnEL after expansion */
+  int iRetReg;          /* Register array for holding a row of RETURNING */
 };
 
 /*
diff --git a/src/trigger.c b/src/trigger.c
index ef0a90b24..0893f790d 100644
--- a/src/trigger.c
+++ b/src/trigger.c
@@ -70,6 +70,8 @@ Trigger *sqlite3TriggerList(Parse *pParse, Table *pTab){
         pTrig->pNext = pList;
         pList = pTrig;
       }else if( pTrig->op==TK_RETURNING ){
+        assert( pParse->bReturning );
+        assert( &(pParse->u1.pReturning->retTrig) == pTrig );
         pTrig->table = pTab->zName;
         pTrig->pTabSchema = pTab->pSchema;
         pTrig->pNext = pList;
@@ -759,16 +761,21 @@ Trigger *sqlite3TriggersExist(
         ** us what time of trigger it should be. */
         assert( sqlite3IsToplevel(pParse) );
         p->op = op;
-        mask |= TRIGGER_AFTER;
-        if( IsVirtual(pTab) && op!=TK_INSERT ){
-          sqlite3ErrorMsg(pParse,
-             "%s RETURNING is not available on virtual tables",
-             op==TK_DELETE ? "DELETE" : "UPDATE");
+        if( IsVirtual(pTab) ){
+          if( op!=TK_INSERT ){
+            sqlite3ErrorMsg(pParse,
+              "%s RETURNING is not available on virtual tables",
+              op==TK_DELETE ? "DELETE" : "UPDATE");
+          }
+          p->tr_tm = TRIGGER_BEFORE;
+        }else{
+          p->tr_tm = TRIGGER_AFTER;
         }
+        mask |= p->tr_tm;
       }else if( p->bReturning && p->op==TK_INSERT && op==TK_UPDATE
                 && sqlite3IsToplevel(pParse) ){
         /* Also fire a RETURNING trigger for an UPSERT */
-        mask |= TRIGGER_AFTER;
+        mask |= p->tr_tm;
       }
       p = p->pNext;
     }while( p );
@@ -830,6 +837,7 @@ static ExprList *sqlite3ExpandReturning(
   ExprList *pNew = 0;
   sqlite3 *db = pParse->db;
   int i;
+
   for(i=0; i<pList->nExpr; i++){
     Expr *pOldExpr = pList->a[i].pExpr;
     if( ALWAYS(pOldExpr!=0) && pOldExpr->op==TK_ASTERISK ){
@@ -854,10 +862,72 @@ static ExprList *sqlite3ExpandReturning(
       }
     }
   }
+  if( !db->mallocFailed && !pParse->colNamesSet ){
+    Vdbe *v = pParse->pVdbe;
+    assert( v!=0 );
+    sqlite3VdbeSetNumCols(v, pNew->nExpr);
+    for(i=0; i<pNew->nExpr; i++){
+      sqlite3VdbeSetColName(v, i, COLNAME_NAME, pNew->a[i].zEName,
+                            SQLITE_TRANSIENT);
+    }
+  }
   return pNew;
 }
 
 /*
+** Generate code for the RETURNING trigger.  Unlike other triggers
+** that invoke a subprogram in the bytecode, the code for RETURNING
+** is generated in-line.
+*/
+static void codeReturningTrigger(
+  Parse *pParse,       /* Parse context */
+  Trigger *pTrigger,   /* The trigger step that defines the RETURNING */
+  Table *pTab,         /* The table to code triggers from */
+  int regIn            /* The first in an array of registers */
+){
+  Vdbe *v = pParse->pVdbe;
+  ExprList *pNew;
+  Returning *pReturning;
+
+  assert( v!=0 );
+  assert( pParse->bReturning );
+  pReturning = pParse->u1.pReturning;
+  assert( pTrigger == &(pReturning->retTrig) );
+  pNew = sqlite3ExpandReturning(pParse, pReturning->pReturnEL, pTab);
+  if( pNew ){
+    NameContext sNC;
+    memset(&sNC, 0, sizeof(sNC));
+    if( pReturning->nRetCol==0 ){
+      pReturning->nRetCol = pNew->nExpr;
+      pReturning->iRetCur = pParse->nTab++;
+    }
+    sNC.pParse = pParse;
+    sNC.uNC.iBaseReg = regIn;
+    sNC.ncFlags = NC_UBaseReg;
+    pParse->eTriggerOp = pTrigger->op;
+    pParse->pTriggerTab = pTab;
+    if( sqlite3ResolveExprListNames(&sNC, pNew)==SQLITE_OK ){
+      int i;
+      int nCol = pNew->nExpr;
+      int reg = pParse->nMem+1;
+      pParse->nMem += nCol+2;
+      pReturning->iRetReg = reg;
+      for(i=0; i<nCol; i++){
+        sqlite3ExprCodeFactorable(pParse, pNew->a[i].pExpr, reg+i);
+      }
+      sqlite3VdbeAddOp3(v, OP_MakeRecord, reg, i, reg+i);
+      sqlite3VdbeAddOp2(v, OP_NewRowid, pReturning->iRetCur, reg+i+1);
+      sqlite3VdbeAddOp3(v, OP_Insert, pReturning->iRetCur, reg+i, reg+i+1);
+    }
+    sqlite3ExprListDelete(pParse->db, pNew);
+    pParse->eTriggerOp = 0;
+    pParse->pTriggerTab = 0;
+  }
+}
+
+
+
+/*
 ** Generate VDBE code for the statements inside the body of a single 
 ** trigger.
 */
@@ -928,7 +998,7 @@ static int codeTriggerProgram(
         sqlite3VdbeAddOp0(v, OP_ResetCount);
         break;
       }
-      case TK_SELECT: {
+      default: assert( pStep->op==TK_SELECT ); {
         SelectDest sDest;
         Select *pSelect = sqlite3SelectDup(db, pStep->pSelect, 0);
         sqlite3SelectDestInit(&sDest, SRT_Discard, 0);
@@ -936,26 +1006,6 @@ static int codeTriggerProgram(
         sqlite3SelectDelete(db, pSelect);
         break;
       }
-      default: assert( pStep->op==TK_RETURNING ); {
-        Select *pSelect = pStep->pSelect;
-        ExprList *pList = pSelect->pEList;
-        SelectDest sDest;
-        Select *pNew;
-        pSelect->pEList =
-           sqlite3ExpandReturning(pParse, pList, pParse->pTriggerTab);
-        sqlite3SelectDestInit(&sDest, SRT_Output, 0);
-        pNew = sqlite3SelectDup(db, pSelect, 0);
-        if( pNew ){
-          sqlite3Select(pParse, pNew, &sDest);
-          if( pNew->selFlags & (SF_Aggregate|SF_HasAgg|SF_WinRewrite) ){
-            sqlite3ErrorMsg(pParse, "aggregates not allowed in RETURNING");
-          }
-          sqlite3SelectDelete(db, pNew);
-        }
-        sqlite3ExprListDelete(db, pSelect->pEList);
-        pStep->pSelect->pEList = pList;
-        break;
-      }
     } 
   }
 
@@ -1046,7 +1096,6 @@ static TriggerPrg *codeRowTrigger(
   pSubParse->pToplevel = pTop;
   pSubParse->zAuthContext = pTrigger->zName;
   pSubParse->eTriggerOp = pTrigger->op;
-  pSubParse->bReturning = pTrigger->bReturning;
   pSubParse->nQueryLoop = pParse->nQueryLoop;
   pSubParse->disableVtab = pParse->disableVtab;
 
@@ -1096,9 +1145,6 @@ static TriggerPrg *codeRowTrigger(
     if( db->mallocFailed==0 && pParse->nErr==0 ){
       pProgram->aOp = sqlite3VdbeTakeOpArray(v, &pProgram->nOp, &pTop->nMaxArg);
     }
-    if( pTrigger->bReturning ){
-      sqlite3VdbeColumnInfoXfer(sqlite3ParseToplevel(pParse)->pVdbe, v);
-    }
     pProgram->nMem = pSubParse->nMem;
     pProgram->nCsr = pSubParse->nTab;
     pProgram->token = (void *)pTrigger;
@@ -1208,7 +1254,7 @@ void sqlite3CodeRowTriggerDirect(
 **   ...            ...
 **   reg+N          OLD.* value of right-most column of pTab
 **   reg+N+1        NEW.rowid
-**   reg+N+2        OLD.* value of left-most column of pTab
+**   reg+N+2        NEW.* value of left-most column of pTab
 **   ...            ...
 **   reg+N+N+1      NEW.* value of right-most column of pTab
 **
@@ -1261,12 +1307,12 @@ void sqlite3CodeRowTrigger(
     if( (p->op==op || (p->bReturning && p->op==TK_INSERT && op==TK_UPDATE))
      && p->tr_tm==tr_tm 
      && checkColumnOverlap(p->pColumns, pChanges)
-     && (sqlite3IsToplevel(pParse) || !p->bReturning)
     ){
-      u8 origOp = p->op;
-      p->op = op;
-      sqlite3CodeRowTriggerDirect(pParse, p, pTab, reg, orconf, ignoreJump);
-      p->op = origOp;
+      if( !p->bReturning ){
+        sqlite3CodeRowTriggerDirect(pParse, p, pTab, reg, orconf, ignoreJump);
+      }else if( sqlite3IsToplevel(pParse) ){
+        codeReturningTrigger(pParse, p, pTab, reg);
+      }
     }
   }
 }
@@ -1311,13 +1357,18 @@ u32 sqlite3TriggerColmask(
 
   assert( isNew==1 || isNew==0 );
   for(p=pTrigger; p; p=p->pNext){
-    if( p->op==op && (tr_tm&p->tr_tm)
+    if( p->op==op
+     && (tr_tm&p->tr_tm)
      && checkColumnOverlap(p->pColumns,pChanges)
     ){
-      TriggerPrg *pPrg;
-      pPrg = getRowTrigger(pParse, p, pTab, orconf);
-      if( pPrg ){
-        mask |= pPrg->aColmask[isNew];
+      if( p->bReturning ){
+        mask = 0xffffffff;
+      }else{
+        TriggerPrg *pPrg;
+        pPrg = getRowTrigger(pParse, p, pTab, orconf);
+        if( pPrg ){
+          mask |= pPrg->aColmask[isNew];
+        }
       }
     }
   }
diff --git a/src/vdbe.h b/src/vdbe.h
index 48be53df7..17f11fdd7 100644
--- a/src/vdbe.h
+++ b/src/vdbe.h
@@ -259,7 +259,6 @@ void sqlite3VdbeResetStepResult(Vdbe*);
 void sqlite3VdbeRewind(Vdbe*);
 int sqlite3VdbeReset(Vdbe*);
 void sqlite3VdbeSetNumCols(Vdbe*,int);
-void sqlite3VdbeColumnInfoXfer(Vdbe*,Vdbe*);
 int sqlite3VdbeSetColName(Vdbe*, int, int, const char *, void(*)(void*));
 void sqlite3VdbeCountChanges(Vdbe*);
 sqlite3 *sqlite3VdbeDb(Vdbe*);
diff --git a/src/vdbeaux.c b/src/vdbeaux.c
index 1c77b2ce9..7ee4f629e 100644
--- a/src/vdbeaux.c
+++ b/src/vdbeaux.c
@@ -2596,23 +2596,6 @@ void sqlite3VdbeSetNumCols(Vdbe *p, int nResColumn){
 }
 
 /*
-** Transfer the column count and name information from one Vdbe to
-** another.
-*/
-void sqlite3VdbeColumnInfoXfer(Vdbe *pTo, Vdbe *pFrom){
-  sqlite3 *db = pTo->db;
-  assert( db==pFrom->db );
-  if( pTo->nResColumn ){
-    releaseMemArray(pTo->aColName, pTo->nResColumn*COLNAME_N);
-    sqlite3DbFree(db, pTo->aColName);
-  }
-  pTo->aColName = pFrom->aColName;
-  pFrom->aColName = 0;
-  pTo->nResColumn = pFrom->nResColumn;
-  pFrom->nResColumn = 0;
-}
-
-/*
 ** Set the name of the idx'th column to be returned by the SQL statement.
 ** zName must be a pointer to a nul terminated string.
 **
diff --git a/src/vtab.c b/src/vtab.c
index b2c01f2fa..400ae4096 100644
--- a/src/vtab.c
+++ b/src/vtab.c
@@ -828,7 +828,7 @@ int sqlite3_declare_vtab(sqlite3 *db, const char *zCreateTable){
       Table *pNew = sParse.pNewTable;
       Index *pIdx;
       pTab->aCol = pNew->aCol;
-      pTab->nCol = pNew->nCol;
+      pTab->nNVCol = pTab->nCol = pNew->nCol;
       pTab->tabFlags |= pNew->tabFlags & (TF_WithoutRowid|TF_NoVisibleRowid);
       pNew->nCol = 0;
       pNew->aCol = 0;