diff --git a/src/msgpack.c b/src/msgpack.c index 54e1a43..898925c 100644 --- a/src/msgpack.c +++ b/src/msgpack.c @@ -1992,6 +1992,307 @@ static void msgpackObjectFunc( ** ------------------------------------------------------- */ +/* +** --------------------------------------------------------------------------- +** msgpack_extract() multi-path acceleration. +** +** A call such as msgpack_extract(mp, '$.p0', '$.p1', ..., '$.pK') originally +** resolved every path independently with mpLookup(), each of which re-scans the +** container from the start — O(K*N) for a K-path request over an N-element +** container. When the root is a container and every path is a single component +** ($.key or $[idx]), the helpers below gather all requested elements in ONE +** linear walk of the container (O(N)), indexing the requests so each element is +** matched in O(1). The returned bytes are byte-for-byte identical to the +** original per-path resolution. +** +** Paths are parsed locally per call (borrowing the argument text, which is valid +** for the duration of the call). Per-argument sqlite3_set_auxdata() is +** deliberately NOT used here: SQLite stores auxdata in a per-statement linked +** list, so sqlite3_get_auxdata() is O(#auxdata entries); caching one entry per +** path would make a K-path call O(K^2) to look up — worse than simply parsing +** each short path string, which is O(1) per path. +** +** Anything not covered by the fast path (few paths, non-simple paths, +** non-container root, malformed input, OOM) transparently falls back to the +** original per-path mpLookup() loop, so behaviour is unchanged in every case. +** --------------------------------------------------------------------------- +*/ + +/* Parsed msgpack_extract path. All pointers borrow the argument text and are +** valid only for the duration of the function call. */ +typedef struct MpPath { + const char *zPath; /* the raw path text (borrowed) */ + int simple; /* 1 => single component off the root ($.key or $[idx])*/ + int kind; /* 'k' (map key) or 'i' (array index) when simple */ + const char *zKey; /* key bytes (into zPath) when kind=='k' */ + int nKey; /* key length when kind=='k' */ + i64 idx; /* index when kind=='i' */ +} MpPath; + +/* Parse a path string into *p, borrowing zPath (no allocation). */ +static void mpPathParse(MpPath *p, const char *zPath){ + p->zPath = zPath; p->simple = 0; p->kind = 0; p->zKey = 0; p->nKey = 0; p->idx = 0; + if( !zPath || zPath[0]!='$' ) return; + { + int pi = 1; + const char *zKey = 0; int nKey = 0; i64 idx = 0; + int step = mpPathStep(zPath, &pi, &zKey, &nKey, &idx); + if( step=='k' || step=='i' ){ + const char *zk2 = 0; int nk2 = 0; i64 id2 = 0; + if( mpPathStep(zPath, &pi, &zk2, &nk2, &id2)==0 ){ /* nothing after it */ + p->simple = 1; p->kind = step; + p->zKey = zKey; p->nKey = nKey; p->idx = idx; + } + } + } +} + +/* Below this many paths, a linear compare per element is cheaper than building +** an index; at or above it the single pass indexes requests for O(1) match. */ +#define MP_MULTI_HASH_MIN 16 + +/* mpKeyHash() (FNV-1a over key bytes) is defined earlier and reused here. */ + +/* Single linear pass over a map root at offset 0, resolving every "simple" key +** path. Fills aStart/aEnd/aFound (indexed by path). Returns 1 on success, or +** 0 to request the caller fall back (malformed structure or OOM). Byte ranges +** match exactly what mpLookup() would have produced for each path. +** +** The requested keys are indexed in an open-addressing hash (with per-key chains +** for duplicate paths), so matching each of the N map entries is O(1) and the +** whole pass is O(N + nPath) rather than O(N*nPath). Only called for +** nPath >= MP_MULTI_HASH_MIN, where the index build pays off. */ +static int mpScanMap(const u8 *a, u32 n, + const MpPath *paths, int nPath, + u32 *aStart, u32 *aEnd, u8 *aFound){ + u8 b = a[0]; + u32 count, elemOff, j, iCur; + int r, h, remaining = 0; + int *htab, *chain, cap; + if( b>=0x80 && b<=0x8f ){ count = b & 0x0f; elemOff = 1; } + else if( b==MP_MAP16 ){ if(3>n) return 0; count = mpRead16(a+1); elemOff = 3; } + else if( b==MP_MAP32 ){ if(5>n) return 0; count = mpRead32(a+1); elemOff = 5; } + else return 0; + for( r=0; r0; j++ ){ + u8 kb; + const char *kStr = 0; + u32 kLen = 0, kStart = 0, valOff = 0, valEnd; + if( iCur >= n ){ sqlite3_free(htab); sqlite3_free(chain); return 0; } + kb = a[iCur]; + /* Decode the key header (same fast path as mpLookup). For a string key, + ** kStart is the offset of the key bytes; the header bounds checks guarantee + ** kStart<=n, so the value offset is validated below WITHOUT the u32 wrap that + ** a large (malformed) STR32 length could otherwise use to slip past valOff>n + ** and drive an out-of-bounds read in mpKeyHash()/memcmp(). Non-string keys + ** are skipped structurally. */ + if( kb>=0xa0 && kb<=0xbf ){ + kLen = kb & 0x1f; kStart = iCur+1; kStr = (const char*)(a+kStart); + } else if( kb==MP_STR8 && iCur+2<=n ){ + kLen = a[iCur+1]; kStart = iCur+2; kStr = (const char*)(a+kStart); + } else if( kb==MP_STR16 && iCur+3<=n ){ + kLen = mpRead16(a+iCur+1); kStart = iCur+3; kStr = (const char*)(a+kStart); + } else if( kb==MP_STR32 && iCur+5<=n ){ + kLen = mpRead32(a+iCur+1); kStart = iCur+5; kStr = (const char*)(a+kStart); + } + if( kStr ){ + if( kLen > n - kStart ){ sqlite3_free(htab); sqlite3_free(chain); return 0; } + valOff = kStart + kLen; /* guaranteed <= n, no wrap */ + } else { + valOff = mpSkipOne(a, n, iCur); /* non-string key */ + } + if( !valOff || valOff>n ){ sqlite3_free(htab); sqlite3_free(chain); return 0; } + valEnd = mpSkipOne(a, n, valOff); + if( !valEnd ){ sqlite3_free(htab); sqlite3_free(chain); return 0; } + if( kStr ){ + h = (int)(mpKeyHash(kStr,kLen) & (u32)(cap-1)); + while( htab[h]!=-1 ){ + int e = htab[h]; + if( paths[e].nKey==(int)kLen + && memcmp(paths[e].zKey, kStr, (size_t)kLen)==0 ){ + int x; + for( x=e; x!=-1; x=chain[x] ) + if( !aFound[x] ){ aStart[x]=valOff; aEnd[x]=valEnd; aFound[x]=1; remaining--; } + break; + } + h = (h+1) & (cap-1); + } + } + iCur = valEnd; + } + sqlite3_free(htab); sqlite3_free(chain); + return 1; +} + +/* Single linear pass over an array root at offset 0, resolving every "simple" +** index path. Uses a direct index table (idx -> chain of paths) which is much +** cheaper than hashing for the dense small indices typical of a property bag. +** Falls back (returns 0) if the largest requested index is too large to index +** compactly, or on OOM/malformed input. Semantics mirror mpScanMap(). */ +static int mpScanArray(const u8 *a, u32 n, + const MpPath *paths, int nPath, + u32 *aStart, u32 *aEnd, u8 *aFound){ + u8 b = a[0]; + u32 count, elemOff, cur, iCur; + int r, remaining = 0; + i64 maxIdx = -1; + int *slot, *chain, i; + if( b>=0x90 && b<=0x9f ){ count = b & 0x0f; elemOff = 1; } + else if( b==MP_ARRAY16 ){ if(3>n) return 0; count = mpRead16(a+1); elemOff = 3; } + else if( b==MP_ARRAY32 ){ if(5>n) return 0; count = mpRead32(a+1); elemOff = 5; } + else return 0; + for( r=0; r=0 && (u64)paths[r].idx<(u64)count ){ + if( paths[r].idx>maxIdx ) maxIdx = paths[r].idx; + remaining++; + } + } + if( remaining==0 ) return 1; /* nothing in range → all NIL */ + /* Bound the direct-index table; if indices are too sparse, let the caller use + ** the original per-path loop instead of allocating a huge table. */ + if( maxIdx+1 > (i64)(8*nPath + 1024) ) return 0; + + slot = (int*)sqlite3_malloc((int)(sizeof(int)*(maxIdx+1))); + chain = (int*)sqlite3_malloc((int)(sizeof(int)*nPath)); + if( !slot || !chain ){ sqlite3_free(slot); sqlite3_free(chain); return 0; } + for( i=0; i<=(int)maxIdx; i++ ) slot[i] = -1; + for( r=0; r=0 && paths[r].idx<=maxIdx ){ + chain[r] = slot[paths[r].idx]; /* chain duplicate requests for same idx */ + slot[paths[r].idx] = r; + } + } + + iCur = elemOff; + for( cur=0; cur0; cur++ ){ + u32 elemStart = iCur, elemEnd; + if( iCur >= n ){ sqlite3_free(slot); sqlite3_free(chain); return 0; } + elemEnd = mpSkipOne(a, n, iCur); + if( !elemEnd ){ sqlite3_free(slot); sqlite3_free(chain); return 0; } + if( (i64)cur<=maxIdx && slot[cur]>=0 ){ + int x; + for( x=slot[cur]; x!=-1; x=chain[x] ){ + aStart[x]=elemStart; aEnd[x]=elemEnd; aFound[x]=1; remaining--; + } + } + iCur = elemEnd; + } + sqlite3_free(slot); sqlite3_free(chain); + return 1; +} + +/* Implements msgpack_extract() for the multi-path (argc>2) case. +** +** For a container root with enough single-component paths the requested elements +** are gathered in one linear pass (mpScanMap/mpScanArray). Every other case — +** few paths, non-container root, non-simple paths, or OOM — runs the original +** per-path mpLookup() loop, byte-for-byte unchanged and with no extra +** allocation, so it is never a regression. */ +static void mpExtractMulti(sqlite3_context *ctx, int argc, sqlite3_value **argv, + const u8 *a, u32 n){ + int nPath = argc - 1; + MpBuf buf; + u8 *result; + u32 nOut; + int r; + + /* Accelerated single pass: only worthwhile with many paths on a container. */ + if( n>0 && nPath>=MP_MULTI_HASH_MIN ){ + u8 b = a[0]; + int isMap = (b>=0x80 && b<=0x8f) || b==MP_MAP16 || b==MP_MAP32; + int isArray = (b>=0x90 && b<=0x9f) || b==MP_ARRAY16 || b==MP_ARRAY32; + if( isMap || isArray ){ + MpPath *paths = (MpPath*)sqlite3_malloc((int)(sizeof(MpPath) * nPath)); + if( paths ){ + int handled = 0, allSimple = 1; + for( r=0; r= 16 single-component paths on a container root gathers all requested +** elements in one linear walk; the result must equal per-path resolution (including +** duplicate paths, missing keys/out-of-range indices -> NIL), a non-simple path must fall +** back correctly, and a malformed STR32 key length must not read out of bounds. +**-------------------------------------------------------------------------------------*/ +static void test_extract_multipath(sqlite3 *db){ + char obj[4096], me[4096], sql[4400]; + int off, k, ok; + + /* Wide map p0..p39 = i*7+3; extract a 16-path subset (dup + missing key). */ + off = snprintf(obj, sizeof(obj), "msgpack_object("); + for(k=0;k<40;k++) off += snprintf(obj+off, sizeof(obj)-off, "%s'p%d',%d", k?",":"", k, k*7+3); + snprintf(obj+off, sizeof(obj)-off, ")"); + + struct { const char *path; int slot; } reqs[] = { /* slot < 0 => missing (NIL) */ + {"$.p0",0},{"$.p39",39},{"$.p5",5},{"$.p39",39},{"$.nope",-1},{"$.p20",20}, + {"$.p1",1},{"$.p38",38},{"$.p2",2},{"$.p37",37},{"$.p3",3},{"$.p36",36}, + {"$.p4",4},{"$.p35",35},{"$.p6",6},{"$.p34",34}}; + int K = (int)(sizeof(reqs)/sizeof(reqs[0])); + off = snprintf(me, sizeof(me), "msgpack_extract(%s", obj); + for(k=0;k= 40){ + snprintf(sql, sizeof(sql), "SELECT msgpack_extract(%s,'$[%d]') IS NULL", me, k); + if(exec1i(db, sql) != 1) ok = 0; + } else { + snprintf(sql, sizeof(sql), "SELECT msgpack_extract(%s,'$[%d]')", me, k); + if(exec1i(db, sql) != (sqlite3_int64)(idxs[k]*3+1)) ok = 0; + } + } + CHECK("11.4 multi-path array: every element == per-path result", ok); + + /* A nested path among >= 16 paths forces the per-path fallback; still correct. */ + off = snprintf(obj, sizeof(obj), "msgpack_object("); + for(k=0;k<20;k++){ + if(k==3) off += snprintf(obj+off, sizeof(obj)-off, "%s'p3',msgpack_object('x',77)", k?",":""); + else off += snprintf(obj+off, sizeof(obj)-off, "%s'p%d',%d", k?",":"", k, k+100); + } + snprintf(obj+off, sizeof(obj)-off, ")"); + off = snprintf(me, sizeof(me), "msgpack_extract(%s,'$.p0','$.p3.x'", obj); + for(k=1;k<20;k++) if(k!=3) off += snprintf(me+off, sizeof(me)-off, ",'$.p%d'", k); + snprintf(me+off, sizeof(me)-off, ")"); + snprintf(sql, sizeof(sql), "SELECT msgpack_extract(%s,'$[0]')", me); + CHECK("11.5 fallback (nested path): simple key correct", exec1i(db, sql) == 100); + snprintf(sql, sizeof(sql), "SELECT msgpack_extract(%s,'$[1]')", me); + CHECK("11.6 fallback (nested path): nested value correct", exec1i(db, sql) == 77); + + /* Malformed STR32 map key (declared length 0xFFFFFFFF) must not read OOB. */ + off = snprintf(me, sizeof(me), "msgpack_extract(x'81dbffffffff'"); + for(k=0;k<16;k++) off += snprintf(me+off, sizeof(me)-off, ",'$.k%d'", k); + snprintf(me+off, sizeof(me)-off, ")"); + snprintf(sql, sizeof(sql), "SELECT msgpack_array_length(%s)", me); + CHECK("11.7 malformed STR32 key: safe, length == 16", exec1i(db, sql) == 16); + ok = 1; + for(k=0;k<16;k++){ + snprintf(sql, sizeof(sql), "SELECT msgpack_extract(%s,'$[%d]') IS NULL", me, k); + if(exec1i(db, sql) != 1) ok = 0; + } + CHECK("11.8 malformed STR32 key: all elements NIL", ok); +} + int main(void){ sqlite3 *db = NULL; if(sqlite3_open(":memory:", &db) != SQLITE_OK){ fprintf(stderr,"open failed\n"); return 1; } @@ -536,6 +627,7 @@ int main(void){ test_array_length(db); test_error_position(db); test_extract_column_types(db); + test_extract_multipath(db); sqlite3_close(db); printf("\n%d passed, %d failed\n", g_pass, g_fail);