#include "sqlite3ext.h"
SQLITE_EXTENSION_INIT1
#include <stdlib.h>
#include <string.h>
#include <assert.h>
#include <stdio.h>
#include <ctype.h>
#ifndef SQLITE_OMIT_VIRTUALTABLE
typedef struct closure_vtab closure_vtab;
typedef struct closure_cursor closure_cursor;
typedef struct closure_queue closure_queue;
typedef struct closure_avl closure_avl;
struct closure_avl {
sqlite3_int64 id;
int iGeneration;
closure_avl *pList;
closure_avl *pBefore;
closure_avl *pAfter;
closure_avl *pUp;
short int height;
short int imbalance;
};
static void closureAvlRecomputeHeight(closure_avl *p){
short int hBefore = p->pBefore ? p->pBefore->height : 0;
short int hAfter = p->pAfter ? p->pAfter->height : 0;
p->imbalance = hBefore - hAfter;
p->height = (hBefore>hAfter ? hBefore : hAfter)+1;
}
static closure_avl *closureAvlRotateBefore(closure_avl *pP){
closure_avl *pB = pP->pBefore;
closure_avl *pY = pB->pAfter;
pB->pUp = pP->pUp;
pB->pAfter = pP;
pP->pUp = pB;
pP->pBefore = pY;
if( pY ) pY->pUp = pP;
closureAvlRecomputeHeight(pP);
closureAvlRecomputeHeight(pB);
return pB;
}
static closure_avl *closureAvlRotateAfter(closure_avl *pP){
closure_avl *pA = pP->pAfter;
closure_avl *pY = pA->pBefore;
pA->pUp = pP->pUp;
pA->pBefore = pP;
pP->pUp = pA;
pP->pAfter = pY;
if( pY ) pY->pUp = pP;
closureAvlRecomputeHeight(pP);
closureAvlRecomputeHeight(pA);
return pA;
}
static closure_avl **closureAvlFromPtr(closure_avl *p, closure_avl **pp){
closure_avl *pUp = p->pUp;
if( pUp==0 ) return pp;
if( pUp->pAfter==p ) return &pUp->pAfter;
return &pUp->pBefore;
}
static closure_avl *closureAvlBalance(closure_avl *p){
closure_avl *pTop = p;
closure_avl **pp;
while( p ){
closureAvlRecomputeHeight(p);
if( p->imbalance>=2 ){
closure_avl *pB = p->pBefore;
if( pB->imbalance<0 ) p->pBefore = closureAvlRotateAfter(pB);
pp = closureAvlFromPtr(p,&p);
p = *pp = closureAvlRotateBefore(p);
}else if( p->imbalance<=(-2) ){
closure_avl *pA = p->pAfter;
if( pA->imbalance>0 ) p->pAfter = closureAvlRotateBefore(pA);
pp = closureAvlFromPtr(p,&p);
p = *pp = closureAvlRotateAfter(p);
}
pTop = p;
p = p->pUp;
}
return pTop;
}
static closure_avl *closureAvlSearch(closure_avl *p, sqlite3_int64 id){
while( p && id!=p->id ){
p = (id<p->id) ? p->pBefore : p->pAfter;
}
return p;
}
static closure_avl *closureAvlFirst(closure_avl *p){
if( p ) while( p->pBefore ) p = p->pBefore;
return p;
}
closure_avl *closureAvlNext(closure_avl *p){
closure_avl *pPrev = 0;
while( p && p->pAfter==pPrev ){
pPrev = p;
p = p->pUp;
}
if( p && pPrev==0 ){
p = closureAvlFirst(p->pAfter);
}
return p;
}
static closure_avl *closureAvlInsert(
closure_avl **ppHead,
closure_avl *pNew
){
closure_avl *p = *ppHead;
if( p==0 ){
p = pNew;
pNew->pUp = 0;
}else{
while( p ){
if( pNew->id<p->id ){
if( p->pBefore ){
p = p->pBefore;
}else{
p->pBefore = pNew;
pNew->pUp = p;
break;
}
}else if( pNew->id>p->id ){
if( p->pAfter ){
p = p->pAfter;
}else{
p->pAfter = pNew;
pNew->pUp = p;
break;
}
}else{
return p;
}
}
}
pNew->pBefore = 0;
pNew->pAfter = 0;
pNew->height = 1;
pNew->imbalance = 0;
*ppHead = closureAvlBalance(p);
return 0;
}
static void closureAvlDestroy(closure_avl *p, void (*xDestroy)(closure_avl*)){
if( p ){
closureAvlDestroy(p->pBefore, xDestroy);
closureAvlDestroy(p->pAfter, xDestroy);
xDestroy(p);
}
}
struct closure_vtab {
sqlite3_vtab base;
char *zDb;
char *zSelf;
char *zTableName;
char *zIdColumn;
char *zParentColumn;
sqlite3 *db;
int nCursor;
};
struct closure_cursor {
sqlite3_vtab_cursor base;
closure_vtab *pVtab;
char *zTableName;
char *zIdColumn;
char *zParentColumn;
closure_avl *pCurrent;
closure_avl *pClosure;
};
struct closure_queue {
closure_avl *pFirst;
closure_avl *pLast;
};
static void queuePush(closure_queue *pQueue, closure_avl *pNode){
pNode->pList = 0;
if( pQueue->pLast ){
pQueue->pLast->pList = pNode;
}else{
pQueue->pFirst = pNode;
}
pQueue->pLast = pNode;
}
static closure_avl *queuePull(closure_queue *pQueue){
closure_avl *p = pQueue->pFirst;
if( p ){
pQueue->pFirst = p->pList;
if( pQueue->pFirst==0 ) pQueue->pLast = 0;
}
return p;
}
static char *closureDequote(const char *zIn){
sqlite3_int64 nIn;
char *zOut;
nIn = strlen(zIn);
zOut = sqlite3_malloc64(nIn+1);
if( zOut ){
char q = zIn[0];
if( q!='[' && q!= '\'' && q!='"' && q!='`' ){
memcpy(zOut, zIn, (size_t)(nIn+1));
}else{
int iOut = 0;
int iIn;
if( q=='[' ) q = ']';
for(iIn=1; iIn<nIn; iIn++){
if( zIn[iIn]==q ) iIn++;
zOut[iOut++] = zIn[iIn];
}
}
assert( (int)strlen(zOut)<=nIn );
}
return zOut;
}
static void closureFree(closure_vtab *p){
if( p ){
sqlite3_free(p->zDb);
sqlite3_free(p->zSelf);
sqlite3_free(p->zTableName);
sqlite3_free(p->zIdColumn);
sqlite3_free(p->zParentColumn);
memset(p, 0, sizeof(*p));
sqlite3_free(p);
}
}
static int closureDisconnect(sqlite3_vtab *pVtab){
closure_vtab *p = (closure_vtab*)pVtab;
assert( p->nCursor==0 );
closureFree(p);
return SQLITE_OK;
}
static const char *closureValueOfKey(const char *zKey, const char *zStr){
int nKey = (int)strlen(zKey);
int nStr = (int)strlen(zStr);
int i;
if( nStr<nKey+1 ) return 0;
if( memcmp(zStr, zKey, nKey)!=0 ) return 0;
for(i=nKey; isspace((unsigned char)zStr[i]); i++){}
if( zStr[i]!='=' ) return 0;
i++;
while( isspace((unsigned char)zStr[i]) ){ i++; }
return zStr+i;
}
static int closureConnect(
sqlite3 *db,
void *pAux,
int argc, const char *const*argv,
sqlite3_vtab **ppVtab,
char **pzErr
){
int rc = SQLITE_OK;
closure_vtab *pNew = 0;
const char *zDb = argv[1];
const char *zVal;
int i;
(void)pAux;
*ppVtab = 0;
pNew = sqlite3_malloc( sizeof(*pNew) );
if( pNew==0 ) return SQLITE_NOMEM;
rc = SQLITE_NOMEM;
memset(pNew, 0, sizeof(*pNew));
pNew->db = db;
pNew->zDb = sqlite3_mprintf("%s", zDb);
if( pNew->zDb==0 ) goto closureConnectError;
pNew->zSelf = sqlite3_mprintf("%s", argv[2]);
if( pNew->zSelf==0 ) goto closureConnectError;
for(i=3; i<argc; i++){
zVal = closureValueOfKey("tablename", argv[i]);
if( zVal ){
sqlite3_free(pNew->zTableName);
pNew->zTableName = closureDequote(zVal);
if( pNew->zTableName==0 ) goto closureConnectError;
continue;
}
zVal = closureValueOfKey("idcolumn", argv[i]);
if( zVal ){
sqlite3_free(pNew->zIdColumn);
pNew->zIdColumn = closureDequote(zVal);
if( pNew->zIdColumn==0 ) goto closureConnectError;
continue;
}
zVal = closureValueOfKey("parentcolumn", argv[i]);
if( zVal ){
sqlite3_free(pNew->zParentColumn);
pNew->zParentColumn = closureDequote(zVal);
if( pNew->zParentColumn==0 ) goto closureConnectError;
continue;
}
*pzErr = sqlite3_mprintf("unrecognized argument: [%s]\n", argv[i]);
closureFree(pNew);
*ppVtab = 0;
return SQLITE_ERROR;
}
rc = sqlite3_declare_vtab(db,
"CREATE TABLE x(id,depth,root HIDDEN,tablename HIDDEN,"
"idcolumn HIDDEN,parentcolumn HIDDEN)"
);
#define CLOSURE_COL_ID 0
#define CLOSURE_COL_DEPTH 1
#define CLOSURE_COL_ROOT 2
#define CLOSURE_COL_TABLENAME 3
#define CLOSURE_COL_IDCOLUMN 4
#define CLOSURE_COL_PARENTCOLUMN 5
if( rc!=SQLITE_OK ){
closureFree(pNew);
}
*ppVtab = &pNew->base;
return rc;
closureConnectError:
closureFree(pNew);
return rc;
}
static int closureOpen(sqlite3_vtab *pVTab, sqlite3_vtab_cursor **ppCursor){
closure_vtab *p = (closure_vtab*)pVTab;
closure_cursor *pCur;
pCur = sqlite3_malloc( sizeof(*pCur) );
if( pCur==0 ) return SQLITE_NOMEM;
memset(pCur, 0, sizeof(*pCur));
pCur->pVtab = p;
*ppCursor = &pCur->base;
p->nCursor++;
return SQLITE_OK;
}
static void closureClearCursor(closure_cursor *pCur){
closureAvlDestroy(pCur->pClosure, (void(*)(closure_avl*))sqlite3_free);
sqlite3_free(pCur->zTableName);
sqlite3_free(pCur->zIdColumn);
sqlite3_free(pCur->zParentColumn);
pCur->zTableName = 0;
pCur->zIdColumn = 0;
pCur->zParentColumn = 0;
pCur->pCurrent = 0;
pCur->pClosure = 0;
}
static int closureClose(sqlite3_vtab_cursor *cur){
closure_cursor *pCur = (closure_cursor *)cur;
closureClearCursor(pCur);
pCur->pVtab->nCursor--;
sqlite3_free(pCur);
return SQLITE_OK;
}
static int closureNext(sqlite3_vtab_cursor *cur){
closure_cursor *pCur = (closure_cursor*)cur;
pCur->pCurrent = closureAvlNext(pCur->pCurrent);
return SQLITE_OK;
}
static int closureInsertNode(
closure_queue *pQueue,
closure_cursor *pCur,
sqlite3_int64 id,
int iGeneration
){
closure_avl *pNew = sqlite3_malloc( sizeof(*pNew) );
if( pNew==0 ) return SQLITE_NOMEM;
memset(pNew, 0, sizeof(*pNew));
pNew->id = id;
pNew->iGeneration = iGeneration;
closureAvlInsert(&pCur->pClosure, pNew);
queuePush(pQueue, pNew);
return SQLITE_OK;
}
static int closureFilter(
sqlite3_vtab_cursor *pVtabCursor,
int idxNum, const char *idxStr,
int argc, sqlite3_value **argv
){
closure_cursor *pCur = (closure_cursor *)pVtabCursor;
closure_vtab *pVtab = pCur->pVtab;
sqlite3_int64 iRoot;
int mxGen = 999999999;
char *zSql;
sqlite3_stmt *pStmt;
closure_avl *pAvl;
int rc = SQLITE_OK;
const char *zTableName = pVtab->zTableName;
const char *zIdColumn = pVtab->zIdColumn;
const char *zParentColumn = pVtab->zParentColumn;
closure_queue sQueue;
(void)idxStr;
(void)argc;
closureClearCursor(pCur);
memset(&sQueue, 0, sizeof(sQueue));
if( (idxNum & 1)==0 ){
return SQLITE_OK;
}
iRoot = sqlite3_value_int64(argv[0]);
if( (idxNum & 0x000f0)!=0 ){
mxGen = sqlite3_value_int(argv[(idxNum>>4)&0x0f]);
if( (idxNum & 0x00002)!=0 ) mxGen--;
}
if( (idxNum & 0x00f00)!=0 ){
zTableName = (const char*)sqlite3_value_text(argv[(idxNum>>8)&0x0f]);
pCur->zTableName = sqlite3_mprintf("%s", zTableName);
}
if( (idxNum & 0x0f000)!=0 ){
zIdColumn = (const char*)sqlite3_value_text(argv[(idxNum>>12)&0x0f]);
pCur->zIdColumn = sqlite3_mprintf("%s", zIdColumn);
}
if( (idxNum & 0x0f0000)!=0 ){
zParentColumn = (const char*)sqlite3_value_text(argv[(idxNum>>16)&0x0f]);
pCur->zParentColumn = sqlite3_mprintf("%s", zParentColumn);
}
zSql = sqlite3_mprintf(
"SELECT \"%w\".\"%w\" FROM \"%w\" WHERE \"%w\".\"%w\"=?1",
zTableName, zIdColumn, zTableName, zTableName, zParentColumn);
if( zSql==0 ){
return SQLITE_NOMEM;
}else{
rc = sqlite3_prepare_v2(pVtab->db, zSql, -1, &pStmt, 0);
sqlite3_free(zSql);
if( rc ){
sqlite3_free(pVtab->base.zErrMsg);
pVtab->base.zErrMsg = sqlite3_mprintf("%s", sqlite3_errmsg(pVtab->db));
return rc;
}
}
if( rc==SQLITE_OK ){
rc = closureInsertNode(&sQueue, pCur, iRoot, 0);
}
while( (pAvl = queuePull(&sQueue))!=0 ){
if( pAvl->iGeneration>=mxGen ) continue;
sqlite3_bind_int64(pStmt, 1, pAvl->id);
while( rc==SQLITE_OK && sqlite3_step(pStmt)==SQLITE_ROW ){
if( sqlite3_column_type(pStmt,0)==SQLITE_INTEGER ){
sqlite3_int64 iNew = sqlite3_column_int64(pStmt, 0);
if( closureAvlSearch(pCur->pClosure, iNew)==0 ){
rc = closureInsertNode(&sQueue, pCur, iNew, pAvl->iGeneration+1);
}
}
}
sqlite3_reset(pStmt);
}
sqlite3_finalize(pStmt);
if( rc==SQLITE_OK ){
pCur->pCurrent = closureAvlFirst(pCur->pClosure);
}
return rc;
}
static int closureColumn(sqlite3_vtab_cursor *cur, sqlite3_context *ctx, int i){
closure_cursor *pCur = (closure_cursor*)cur;
switch( i ){
case CLOSURE_COL_ID: {
sqlite3_result_int64(ctx, pCur->pCurrent->id);
break;
}
case CLOSURE_COL_DEPTH: {
sqlite3_result_int(ctx, pCur->pCurrent->iGeneration);
break;
}
case CLOSURE_COL_ROOT: {
sqlite3_result_null(ctx);
break;
}
case CLOSURE_COL_TABLENAME: {
sqlite3_result_text(ctx,
pCur->zTableName ? pCur->zTableName : pCur->pVtab->zTableName,
-1, SQLITE_TRANSIENT);
break;
}
case CLOSURE_COL_IDCOLUMN: {
sqlite3_result_text(ctx,
pCur->zIdColumn ? pCur->zIdColumn : pCur->pVtab->zIdColumn,
-1, SQLITE_TRANSIENT);
break;
}
case CLOSURE_COL_PARENTCOLUMN: {
sqlite3_result_text(ctx,
pCur->zParentColumn ? pCur->zParentColumn : pCur->pVtab->zParentColumn,
-1, SQLITE_TRANSIENT);
break;
}
}
return SQLITE_OK;
}
static int closureRowid(sqlite3_vtab_cursor *cur, sqlite_int64 *pRowid){
closure_cursor *pCur = (closure_cursor*)cur;
*pRowid = pCur->pCurrent->id;
return SQLITE_OK;
}
static int closureEof(sqlite3_vtab_cursor *cur){
closure_cursor *pCur = (closure_cursor*)cur;
return pCur->pCurrent==0;
}
static int closureBestIndex(
sqlite3_vtab *pTab,
sqlite3_index_info *pIdxInfo
){
int iPlan = 0;
int i;
int idx = 1;
const struct sqlite3_index_constraint *pConstraint;
closure_vtab *pVtab = (closure_vtab*)pTab;
double rCost = 10000000.0;
pConstraint = pIdxInfo->aConstraint;
for(i=0; i<pIdxInfo->nConstraint; i++, pConstraint++){
if( pConstraint->usable==0 ) continue;
if( (iPlan & 1)==0
&& pConstraint->iColumn==CLOSURE_COL_ROOT
&& pConstraint->op==SQLITE_INDEX_CONSTRAINT_EQ
){
iPlan |= 1;
pIdxInfo->aConstraintUsage[i].argvIndex = 1;
pIdxInfo->aConstraintUsage[i].omit = 1;
rCost /= 100.0;
}
if( (iPlan & 0x0000f0)==0
&& pConstraint->iColumn==CLOSURE_COL_DEPTH
&& (pConstraint->op==SQLITE_INDEX_CONSTRAINT_LT
|| pConstraint->op==SQLITE_INDEX_CONSTRAINT_LE
|| pConstraint->op==SQLITE_INDEX_CONSTRAINT_EQ)
){
iPlan |= idx<<4;
pIdxInfo->aConstraintUsage[i].argvIndex = ++idx;
if( pConstraint->op==SQLITE_INDEX_CONSTRAINT_LT ) iPlan |= 0x000002;
rCost /= 5.0;
}
if( (iPlan & 0x000f00)==0
&& pConstraint->iColumn==CLOSURE_COL_TABLENAME
&& pConstraint->op==SQLITE_INDEX_CONSTRAINT_EQ
){
iPlan |= idx<<8;
pIdxInfo->aConstraintUsage[i].argvIndex = ++idx;
pIdxInfo->aConstraintUsage[i].omit = 1;
rCost /= 5.0;
}
if( (iPlan & 0x00f000)==0
&& pConstraint->iColumn==CLOSURE_COL_IDCOLUMN
&& pConstraint->op==SQLITE_INDEX_CONSTRAINT_EQ
){
iPlan |= idx<<12;
pIdxInfo->aConstraintUsage[i].argvIndex = ++idx;
pIdxInfo->aConstraintUsage[i].omit = 1;
}
if( (iPlan & 0x0f0000)==0
&& pConstraint->iColumn==CLOSURE_COL_PARENTCOLUMN
&& pConstraint->op==SQLITE_INDEX_CONSTRAINT_EQ
){
iPlan |= idx<<16;
pIdxInfo->aConstraintUsage[i].argvIndex = ++idx;
pIdxInfo->aConstraintUsage[i].omit = 1;
}
}
if( (pVtab->zTableName==0 && (iPlan & 0x000f00)==0)
|| (pVtab->zIdColumn==0 && (iPlan & 0x00f000)==0)
|| (pVtab->zParentColumn==0 && (iPlan & 0x0f0000)==0)
){
iPlan = 0;
}
if( (iPlan&1)==0 ){
rCost *= 1e30;
for(i=0; i<pIdxInfo->nConstraint; i++, pConstraint++){
pIdxInfo->aConstraintUsage[i].argvIndex = 0;
}
iPlan = 0;
}
pIdxInfo->idxNum = iPlan;
if( pIdxInfo->nOrderBy==1
&& pIdxInfo->aOrderBy[0].iColumn==CLOSURE_COL_ID
&& pIdxInfo->aOrderBy[0].desc==0
){
pIdxInfo->orderByConsumed = 1;
}
pIdxInfo->estimatedCost = rCost;
return SQLITE_OK;
}
static sqlite3_module closureModule = {
0,
closureConnect,
closureConnect,
closureBestIndex,
closureDisconnect,
closureDisconnect,
closureOpen,
closureClose,
closureFilter,
closureNext,
closureEof,
closureColumn,
closureRowid,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0
};
#endif
#ifdef _WIN32
__declspec(dllexport)
#endif
int sqlite3_closure_init(
sqlite3 *db,
char **pzErrMsg,
const sqlite3_api_routines *pApi
){
int rc = SQLITE_OK;
SQLITE_EXTENSION_INIT2(pApi);
(void)pzErrMsg;
#ifndef SQLITE_OMIT_VIRTUALTABLE
rc = sqlite3_create_module(db, "transitive_closure", &closureModule, 0);
#endif
return rc;
}