Add new interfaces to enable virtual table to process IN operator constraints

all at once, rather than one element at a time.

FossilOrigin-Name: eb84b80e1f6d8c32bf0c9e1731f0233de0160a13f714f766779ae01fdf504e7b

10 files changed, 258 insertions, 12 deletions

diff --git a/src/loadext.c b/src/loadext.c
index a4aad7e74..603516e18 100644
--- a/src/loadext.c
+++ b/src/loadext.c
@@ -489,6 +489,9 @@ static const sqlite3_api_routines sqlite3Apis = {
   sqlite3_error_offset,
   sqlite3_vtab_rhs_value,
   sqlite3_vtab_distinct,
+  sqlite3_vtab_in,
+  sqlite3_vtab_in_first,
+  sqlite3_vtab_in_next
 };
 
 /* True if x is the directory separator character
diff --git a/src/sqlite.h.in b/src/sqlite.h.in
index 82f844d6a..40aa91775 100644
--- a/src/sqlite.h.in
+++ b/src/sqlite.h.in
@@ -9605,6 +9605,90 @@ SQLITE_EXPERIMENTAL const char *sqlite3_vtab_collation(sqlite3_index_info*,int);
 int sqlite3_vtab_distinct(sqlite3_index_info*);
 
 /*
+** CAPI3REF: Identify and handle IN(...) constraints in xBestIndex
+**
+** This API may only be used from within an xBestIndex() callback. The
+** results of calling it from outside of an xBestIndex() callback are
+** undefined.
+**
+** When a column of a virtual table is subject to a "col IN (...)" 
+** constraint, this is passed through to the xBestIndex method of the
+** virtual table as an SQLITE_INDEX_CONSTRAINT_EQ constraint. Then, if
+** the virtual table indicates that it can handle the constraint, SQLite
+** uses the xFilter and xNext methods of the virtual table to query
+** separately for each distinct element in RHS of the IN(...) expression.
+** This API allows a virtual table implementation to determine which 
+** SQLITE_INDEX_CONSTRAINT_EQ constraints are actually IN(...) constraints,
+** and to handle them using a single xFilter scan.
+** 
+** If the second argument passed to this function is not the index of an
+** SQLITE_INDEX_CONSTRAINT_EQ constraint in the aConstraint[] array of the
+** sqlite3_index_info object, or if the SQLITE_INDEX_CONSTRAINT_EQ is not
+** really an IN(...) expression, then this function is a no-op. Zero is
+** returned in this case.
+**
+** Otherwise, if parameter iCons is the index of an SQLITE_INDEX_CONSTRAINT_EQ
+** constraint that is really an IN(...) expression, then this function
+** returns non-zero. In this case, the call also specifies whether SQLite 
+** should invoke xFilter() once for each element on the RHS of the IN(...)
+** expression (the default, if bHandle is zero), or just once for the entire
+** list (if bHandle is non-zero), should the associated 
+** aConstraintUsage[].argvIndex variable be set by xBestIndex.
+**
+** In cases where the list on the RHS of an IN(...) constraint is passed to a
+** single xFilter() call, the (sqlite3_value*) passed appears in most
+** respects to be a NULL value. Except, it may be used with the
+** sqlite3_vtab_in_first() and sqlite3_vtab_in_next() APIs to interate through
+** the list of values.
+*/
+int sqlite3_vtab_in(sqlite3_index_info*, int iCons, int bHandle);
+
+/*
+** CAPI3REF: Visit the first element of an IN(...) list passed to xFilter
+**
+** This API may only be used from within an xFilter() callback. The
+** results of calling it from outside of an xFilter() callback are
+** undefined.
+**
+** If the (sqlite3_value*) passed as the first argument to this function
+** is not a value representing the RHS of an IN(...) operator (see
+** API function sqlite3_vtab_in()), of if the RHS of the IN(...) operator
+** is an empty set, this function sets (*ppOut) to NULL and returns 
+** SQLITE_OK. 
+**
+** Otherwise, if no error occurs, it sets (*ppOut) to point to an object
+** containing the first value in the list and returns SQLITE_OK. The
+** pointer is valid until either the xFilter() call returns or until the 
+** next call to sqlite3_vtab_in_first() or sqlite3_vtab_in_next() with
+** the same first argument.
+**
+** If an error occurs, (*ppOut) is set to NULL and an SQLite error code
+** returned.
+*/
+int sqlite3_vtab_in_first(sqlite3_value *pVal, sqlite3_value **ppOut);
+
+/*
+** CAPI3REF: Visit the next element of an IN(...) list passed to xFilter
+**
+** This function is called after a successful call to sqlite3_vtab_in_first()
+** to visit the next and subsequent elements of an IN(...) list passed
+** to an xFilter callback. Example usage:
+**
+** <pre>
+**   for(rc=sqlite3_vtab_in_first(pList, &pVal);
+**       rc==SQLITE_OK && pVal
+**       rc=sqlite3_vtab_in_next(pList, &pVal)
+**   ){
+**     // do something with pVal
+**   }
+**   if( rc!=SQLITE_OK ){
+**     // an error has occurred
+**   }
+** </pre>
+*/
+int sqlite3_vtab_in_next(sqlite3_value *pVal, sqlite3_value **ppOut);
+
+/*
 ** CAPI3REF: Constraint values in xBestIndex()
 ** METHOD: sqlite3_index_info
 **
diff --git a/src/sqlite3ext.h b/src/sqlite3ext.h
index 88010b9d8..2eac4f3f0 100644
--- a/src/sqlite3ext.h
+++ b/src/sqlite3ext.h
@@ -348,6 +348,9 @@ struct sqlite3_api_routines {
   int (*error_offset)(sqlite3*);
   int (*vtab_rhs_value)(sqlite3_index_info*,int,sqlite3_value**);
   int (*vtab_distinct)(sqlite3_index_info*);
+  int (*vtab_in)(sqlite3_index_info*,int,int);
+  int (*vtab_in_first)(sqlite3_value*,sqlite3_value**);
+  int (*vtab_in_next)(sqlite3_value*,sqlite3_value**);
 };
 
 /*
@@ -663,6 +666,9 @@ typedef int (*sqlite3_loadext_entry)(
 #define sqlite3_error_offset           sqlite3_api->error_offset
 #define sqlite3_vtab_rhs_value         sqlite3_api->vtab_rhs_value
 #define sqlite3_vtab_distinct          sqlite3_api->vtab_distinct
+#define sqlite3_vtab_in                sqlite3_api->vtab_in
+#define sqlite3_vtab_in_first          sqlite3_api->vtab_in_first
+#define sqlite3_vtab_in_next           sqlite3_api->vtab_in_next
 #endif /* !defined(SQLITE_CORE) && !defined(SQLITE_OMIT_LOAD_EXTENSION) */
 
 #if !defined(SQLITE_CORE) && !defined(SQLITE_OMIT_LOAD_EXTENSION)
diff --git a/src/sqliteInt.h b/src/sqliteInt.h
index 6d9365071..374fc9554 100644
--- a/src/sqliteInt.h
+++ b/src/sqliteInt.h
@@ -1231,10 +1231,11 @@ typedef struct With With;
 /*
 ** A bit in a Bitmask
 */
-#define MASKBIT(n)   (((Bitmask)1)<<(n))
-#define MASKBIT64(n) (((u64)1)<<(n))
-#define MASKBIT32(n) (((unsigned int)1)<<(n))
-#define ALLBITS      ((Bitmask)-1)
+#define MASKBIT(n)    (((Bitmask)1)<<(n))
+#define MASKBIT64(n)  (((u64)1)<<(n))
+#define MASKBIT32(n)  (((unsigned int)1)<<(n))
+#define SMASKBIT32(n) ((n)<=31?((unsigned int)1)<<(n):0)
+#define ALLBITS       ((Bitmask)-1)
 
 /* A VList object records a mapping between parameters/variables/wildcards
 ** in the SQL statement (such as $abc, @pqr, or :xyz) and the integer
diff --git a/src/vdbe.c b/src/vdbe.c
index d4ff33141..5643d43ce 100644
--- a/src/vdbe.c
+++ b/src/vdbe.c
@@ -7735,6 +7735,27 @@ case OP_VOpen: {
 #endif /* SQLITE_OMIT_VIRTUALTABLE */
 
 #ifndef SQLITE_OMIT_VIRTUALTABLE
+/* Opcode: VInitIn P1 P2 P3 * *
+** Synopsis: r[P2]=cursor over eph table P1.
+**
+** Initialize register P2 as a value that can be used as an iterator over
+** the contents of ephemeral table P1 by an xFilter callback implementation.
+** Register P3 is used as a cache by the iterator.
+*/
+case OP_VInitIn: {        /* out2 */
+  VdbeCursor *pC;
+  pC = p->apCsr[pOp->p1];
+  pOut = out2Prerelease(p, pOp);
+  pOut->z = (char*)(pC->uc.pCursor);
+  pOut->u.pVal = &aMem[pOp->p3];
+  pOut->uTemp = SQLITE_VTAB_IN_MAGIC;
+  pOut->flags = MEM_Null;
+  break;
+}
+#endif /* SQLITE_OMIT_VIRTUALTABLE */
+
+
+#ifndef SQLITE_OMIT_VIRTUALTABLE
 /* Opcode: VFilter P1 P2 P3 P4 *
 ** Synopsis: iplan=r[P3] zplan='P4'
 **
diff --git a/src/vdbeInt.h b/src/vdbeInt.h
index 376c9edac..9ddb742a4 100644
--- a/src/vdbeInt.h
+++ b/src/vdbeInt.h
@@ -195,6 +195,12 @@ struct VdbeFrame {
 */
 #define VdbeFrameMem(p) ((Mem *)&((u8 *)p)[ROUND8(sizeof(VdbeFrame))])
 
+/* Magic number for Mem.uTemp when it is acting as as the cache for the
+** IN(...) iterator for sqlite3_vtab_in_next()
+*/
+#define SQLITE_VTAB_IN_MAGIC 0xd3ab12ec
+
+
 /*
 ** Internally, the vdbe manipulates nearly all SQL values as Mem
 ** structures. Each Mem struct may cache multiple representations (string,
@@ -207,6 +213,7 @@ struct sqlite3_value {
     int nZero;          /* Extra zero bytes when MEM_Zero and MEM_Blob set */
     const char *zPType; /* Pointer type when MEM_Term|MEM_Subtype|MEM_Null */
     FuncDef *pDef;      /* Used only when flags==MEM_Agg */
+    sqlite3_value *pVal;/* Current value for xFilter IN(...) iterator */
   } u;
   u16 flags;          /* Some combination of MEM_Null, MEM_Str, MEM_Dyn, etc. */
   u8  enc;            /* SQLITE_UTF8, SQLITE_UTF16BE, SQLITE_UTF16LE */
diff --git a/src/vdbeapi.c b/src/vdbeapi.c
index 17df807de..1521eee79 100644
--- a/src/vdbeapi.c
+++ b/src/vdbeapi.c
@@ -847,6 +847,83 @@ int sqlite3_vtab_nochange(sqlite3_context *p){
 }
 
 /*
+** The first argument is an iterator value created by VDBE instruction
+** OP_VInitIn. The iterator is guaranteed to point to a valid entry. This 
+** function attempts to load the current value from the iterator into 
+** object pVal->u.pVal. If successful, (*ppOut) is set to point to 
+** pVal->u.pVal and SQLITE_OK is returned. Otherwise, if an error
+** occurs, an SQLite error code is returned and (*ppOut) is left unchanged.
+*/
+static int vtabInLoadValue(sqlite3_value *pVal, sqlite3_value **ppOut){
+  BtCursor *pCsr = (BtCursor*)pVal->z;
+  sqlite3_value *pOut = pVal->u.pVal;
+  int sz;
+  int rc;
+
+  sz = (int)sqlite3BtreePayloadSize(pCsr);
+  if( sz>pVal->szMalloc ){
+    if( pVal->szMalloc==0 ) pVal->zMalloc = 0;
+    pVal->zMalloc = sqlite3DbReallocOrFree(pVal->db, pVal->zMalloc, sz*2);
+    if( pVal->zMalloc ){
+      pVal->szMalloc = sqlite3DbMallocSize(pVal->db, pVal->zMalloc);
+    }else{
+      pVal->szMalloc = 0;
+      return SQLITE_NOMEM_BKPT;
+    }
+  }
+
+  rc = sqlite3BtreePayload(pCsr, 0, sz, pVal->zMalloc);
+  if( rc==SQLITE_OK ){
+    u32 iSerial;
+    int iOff = 1 + getVarint32((const u8*)&pVal->zMalloc[1], iSerial);
+    sqlite3VdbeSerialGet((const u8*)&pVal->zMalloc[iOff], iSerial, pOut);
+    pOut->enc = ENC(pVal->db);
+    *ppOut = pOut;
+  }
+  return rc;
+}
+
+/*
+** Implementation of sqlite3_vtab_in_first() (if bNext==0) and
+** sqlite3_vtab_in_next() (if bNext!=0).
+*/
+static int vtabInOp(sqlite3_value *pVal, sqlite3_value **ppOut, int bNext){
+  int rc = SQLITE_OK;
+  *ppOut = 0;
+  if( pVal && pVal->uTemp==SQLITE_VTAB_IN_MAGIC ){
+    BtCursor *pCsr = (BtCursor*)pVal->z;
+
+    if( bNext ){
+      rc = sqlite3BtreeNext(pCsr, 0);
+    }else{
+      int dummy = 0;
+      rc = sqlite3BtreeFirst(pCsr, &dummy);
+    }
+
+    if( rc==SQLITE_OK && sqlite3BtreeEof(pCsr)==0 ){
+      rc = vtabInLoadValue(pVal, ppOut);
+    }
+  }
+  return rc;
+}
+
+/*
+** Set the iterator value pVal to point to the first value in the set.
+** Set (*ppOut) to point to this value before returning.
+*/
+int sqlite3_vtab_in_first(sqlite3_value *pVal, sqlite3_value **ppOut){
+  return vtabInOp(pVal, ppOut, 0);
+}
+
+/*
+** Set the iterator value pVal to point to the next value in the set.
+** Set (*ppOut) to point to this value before returning.
+*/
+int sqlite3_vtab_in_next(sqlite3_value *pVal, sqlite3_value **ppOut){
+  return vtabInOp(pVal, ppOut, 1);
+}
+
+/*
 ** Return the current time for a statement.  If the current time
 ** is requested more than once within the same run of a single prepared
 ** statement, the exact same time is returned for each invocation regardless
diff --git a/src/where.c b/src/where.c
index 4efb7084d..0db38ed7c 100644
--- a/src/where.c
+++ b/src/where.c
@@ -33,6 +33,8 @@ struct HiddenIndexInfo {
   WhereClause *pWC;        /* The Where clause being analyzed */
   Parse *pParse;           /* The parsing context */
   int eDistinct;           /* Value to return from sqlite3_vtab_distinct() */
+  u32 mIn;                 /* Mask of terms that are <col> IN (...) */
+  u32 mHandleIn;           /* Terms that vtab will handle as <col> IN (...) */
   sqlite3_value *aRhs[1];  /* RHS values for constraints. MUST BE LAST
                            ** because extra space is allocated to hold up
                            ** to nTerm such values */
@@ -1119,6 +1121,7 @@ static sqlite3_index_info *allocateIndexInfo(
   int nOrderBy;
   sqlite3_index_info *pIdxInfo;
   u16 mNoOmit = 0;
+  u32 mIn = 0;
   const Table *pTab;
   int eDistinct = 0;
   ExprList *pOrderBy = pWInfo->pOrderBy;
@@ -1143,6 +1146,12 @@ static sqlite3_index_info *allocateIndexInfo(
     testcase( pTerm->eOperator & WO_ALL );
     if( (pTerm->eOperator & ~(WO_EQUIV))==0 ) continue;
     if( pTerm->wtFlags & TERM_VNULL ) continue;
+    if( (pTerm->eOperator & WO_IN)!=0
+     && ExprHasProperty(pTerm->pExpr, EP_xIsSelect)==0
+    ){
+      mIn |= SMASKBIT32(i);
+    }
+
     assert( (pTerm->eOperator & (WO_OR|WO_AND))==0 );
     assert( pTerm->u.x.leftColumn>=XN_ROWID );
     assert( pTerm->u.x.leftColumn<pTab->nCol );
@@ -1204,7 +1213,7 @@ static sqlite3_index_info *allocateIndexInfo(
       /* No matches cause a break out of the loop */
       break;
     }
-    if( i==n){
+    if( i==n ){
       nOrderBy = n;
       if( (pWInfo->wctrlFlags & (WHERE_GROUPBY|WHERE_DISTINCTBY)) ){
         eDistinct = 1 + ((pWInfo->wctrlFlags & WHERE_DISTINCTBY)!=0);
@@ -1232,6 +1241,7 @@ static sqlite3_index_info *allocateIndexInfo(
   pHidden->pWC = pWC;
   pHidden->pParse = pParse;
   pHidden->eDistinct = eDistinct;
+  pHidden->mIn = mIn;
   for(i=j=0, pTerm=pWC->a; i<pWC->nTerm; i++, pTerm++){
     u16 op;
     if( (pTerm->wtFlags & TERM_OK)==0 ) continue;
@@ -3499,6 +3509,7 @@ static int whereLoopAddVirtualOne(
   int *pbRetryLimit               /* OUT: Retry without LIMIT/OFFSET */
 ){
   WhereClause *pWC = pBuilder->pWC;
+  HiddenIndexInfo *pHidden = (HiddenIndexInfo*)&pIdxInfo[1];
   struct sqlite3_index_constraint *pIdxCons;
   struct sqlite3_index_constraint_usage *pUsage = pIdxInfo->aConstraintUsage;
   int i;
@@ -3537,6 +3548,7 @@ static int whereLoopAddVirtualOne(
   pIdxInfo->estimatedRows = 25;
   pIdxInfo->idxFlags = 0;
   pIdxInfo->colUsed = (sqlite3_int64)pSrc->colUsed;
+  pHidden->mHandleIn = 0;
 
   /* Invoke the virtual table xBestIndex() method */
   rc = vtabBestIndex(pParse, pSrc->pTab, pIdxInfo);
@@ -3593,7 +3605,9 @@ static int whereLoopAddVirtualOne(
           pNew->u.vtab.bOmitOffset = 1;
         }
       }
-      if( (pTerm->eOperator & WO_IN)!=0 ){
+      if( SMASKBIT32(i) & pHidden->mHandleIn ){ 
+        pNew->u.vtab.mHandleIn |= SMASKBIT32(iTerm);
+      }else if( (pTerm->eOperator & WO_IN)!=0 ){
         /* A virtual table that is constrained by an IN clause may not
         ** consume the ORDER BY clause because (1) the order of IN terms
         ** is not necessarily related to the order of output terms and
@@ -3692,6 +3706,25 @@ const char *sqlite3_vtab_collation(sqlite3_index_info *pIdxInfo, int iCons){
 }
 
 /*
+** Return true if constraint iCons is really an IN(...) constraint, or
+** false otherwise. If iCons is an IN(...) constraint, set (if bHandle!=0)
+** or clear (if bHandle==0) the flag to handle it using an iterator.
+*/
+int sqlite3_vtab_in(sqlite3_index_info *pIdxInfo, int iCons, int bHandle){
+  HiddenIndexInfo *pHidden = (HiddenIndexInfo*)&pIdxInfo[1];
+  u32 m = SMASKBIT32(iCons);
+  if( m & pHidden->mIn ){
+    if( bHandle==0 ){ 
+      pHidden->mHandleIn &= ~m;
+    }else{
+      pHidden->mHandleIn |= m;
+    }
+    return 1;
+  }
+  return 0;
+}
+
+/*
 ** This interface is callable from within the xBestIndex callback only.
 **
 ** If possible, set (*ppVal) to point to an object containing the value 
diff --git a/src/whereInt.h b/src/whereInt.h
index 14c40fcb4..fc3740f51 100644
--- a/src/whereInt.h
+++ b/src/whereInt.h
@@ -128,6 +128,7 @@ struct WhereLoop {
       i8 isOrdered;          /* True if satisfies ORDER BY */
       u16 omitMask;          /* Terms that may be omitted */
       char *idxStr;          /* Index identifier string */
+      u32 mHandleIn;         /* Terms to handle as IN(...) instead of == */
     } vtab;
   } u;
   u32 wsFlags;          /* WHERE_* flags describing the plan */
diff --git a/src/wherecode.c b/src/wherecode.c
index 2beb596e6..8d7163ce0 100644
--- a/src/wherecode.c
+++ b/src/wherecode.c
@@ -1532,8 +1532,15 @@ Bitmask sqlite3WhereCodeOneLoopStart(
       pTerm = pLoop->aLTerm[j];
       if( NEVER(pTerm==0) ) continue;
       if( pTerm->eOperator & WO_IN ){
-        codeEqualityTerm(pParse, pTerm, pLevel, j, bRev, iTarget);
-        addrNotFound = pLevel->addrNxt;
+        if( SMASKBIT32(j) & pLoop->u.vtab.mHandleIn ){
+          int iTab = pParse->nTab++;
+          int iCache = ++pParse->nMem;
+          sqlite3CodeRhsOfIN(pParse, pTerm->pExpr, iTab);
+          sqlite3VdbeAddOp3(v, OP_VInitIn, iTab, iTarget, iCache);
+        }else{
+          codeEqualityTerm(pParse, pTerm, pLevel, j, bRev, iTarget);
+          addrNotFound = pLevel->addrNxt;
+        }
       }else{
         Expr *pRight = pTerm->pExpr->pRight;
         codeExprOrVector(pParse, pRight, iTarget, 1);
@@ -1568,13 +1575,19 @@ Bitmask sqlite3WhereCodeOneLoopStart(
       iIn = 0;
     }
     for(j=nConstraint-1; j>=0; j--){
+      int bIn;     /* True to generate byte code to loop over RHS IN values */
       pTerm = pLoop->aLTerm[j];
-      if( (pTerm->eOperator & WO_IN)!=0 ) iIn--;
+      if( (pTerm->eOperator & WO_IN)!=0
+       && (SMASKBIT32(j) & pLoop->u.vtab.mHandleIn)==0
+      ){
+        bIn = 1;
+      }else{
+        bIn = 0;
+      }
+      if( bIn ) iIn--;
       if( j<16 && (pLoop->u.vtab.omitMask>>j)&1 ){
         disableTerm(pLevel, pTerm);
-      }else if( (pTerm->eOperator & WO_IN)!=0
-        && sqlite3ExprVectorSize(pTerm->pExpr->pLeft)==1
-      ){
+      }else if( bIn && sqlite3ExprVectorSize(pTerm->pExpr->pLeft)==1 ){
         Expr *pCompare;  /* The comparison operator */
         Expr *pRight;    /* RHS of the comparison */
         VdbeOp *pOp;     /* Opcode to access the value of the IN constraint */