dgb-wallet/packages/p2p/src/gcs.ts

import { sipKey, siphash24Pair, type SipKey } from './siphash.js';
import { decodeVarint } from './varint.js';
import { loadPowWasm } from './pow.js';

// BIP158 Golomb-Coded Set filter operations.
// Parameters for the "basic" filter type (0x00):
export const P = 19;
export const M = 784931n;
const M_NUM = 784931;

// Two implementations of the same thing:
//
//  • native — gcs.c inside ../native/dgbpow.wasm. The watched scripts
//    are uploaded into the module once per script set (the scan driver
//    passes the same array for a whole scan) and each filter is then
//    one call. ~30 ns per script; a few thousand scripts × a year of
//    DigiByte blocks is minutes, not hours.
//  • JS — Number arithmetic throughout (a filter's range N·M is far
//    below 2^53, and the one overflowing step, (h·F) >> 64, is done
//    with 32×32-bit partial products). ~3 µs per script. Used when the
//    module is missing or the script set exceeds its static buffers.
//
// Both are checked against a BigInt reference by scripts/gcs-check.ts.

// Derive the SipHash key from a block hash: first 16 bytes of the
// block hash's LITTLE-endian byte order.
// (Our stored block hashes are hex big-endian display form, so we
// have to reverse them before slicing.)
export function keyFromBlockHash(blockHashHex: string): Buffer {
  const be = Buffer.from(blockHashHex, 'hex');
  const le = Buffer.from(be).reverse();
  return Buffer.from(le.subarray(0, 16));
}

// ---------------------------------------------------------------- native
interface Native {
  mem: Uint8Array;
  scriptsPtr: number; offsetsPtr: number; keyPtr: number; filterPtr: number; matchesPtr: number;
  maxScripts: number; scriptBytes: number; filterBytes: number;
  setCount: (n: number) => void;
  match: (nSet: number, bodyLen: number, wantAll: number) => number;
}
let native: Native | null | undefined; // undefined = not tried yet
let forceJs = false;
/** Test hook: `gcsUseJs(true)` disables the native path for this process. */
export function gcsUseJs(on: boolean): void { forceJs = on; uploaded = null; }
export function gcsBackend(): 'wasm' | 'js' { return !forceJs && getNative() ? 'wasm' : 'js'; }

function getNative(): Native | null {
  if (native !== undefined) return native;
  try {
    const w = loadPowWasm();
    const x = w.exports;
    if (typeof x.gcs_match !== 'function') { native = null; return null; }
    const limitsPtr = (x.gcs_limits_ptr as () => number)();
    const limits = new Uint32Array(w.memory.buffer, limitsPtr, 3);
    native = {
      mem: w.mem,
      scriptsPtr: (x.gcs_scripts_ptr as () => number)(),
      offsetsPtr: (x.gcs_offsets_ptr as () => number)(),
      keyPtr: (x.gcs_key_ptr as () => number)(),
      filterPtr: (x.gcs_filter_ptr as () => number)(),
      matchesPtr: (x.gcs_matches_ptr as () => number)(),
      maxScripts: limits[0], scriptBytes: limits[1], filterBytes: limits[2],
      setCount: x.gcs_set_count as (n: number) => void,
      match: x.gcs_match as (nSet: number, bodyLen: number, wantAll: number) => number,
    };
  } catch {
    native = null;
  }
  return native;
}

// The script set currently uploaded into the module, by array identity
// and length (the scan driver never mutates its array mid-scan, but a
// grown array is re-uploaded anyway).
let uploaded: { arr: Buffer[]; len: number } | null = null;
function ensureUploaded(n: Native, elements: Buffer[]): boolean {
  if (uploaded && uploaded.arr === elements && uploaded.len === elements.length) return true;
  if (elements.length > n.maxScripts) return false;
  let total = 0;
  for (const e of elements) total += e.length;
  if (total > n.scriptBytes) return false;
  const offsets = new Uint32Array(n.mem.buffer, n.offsetsPtr, elements.length + 1);
  let off = 0;
  for (let i = 0; i < elements.length; i++) {
    offsets[i] = off;
    n.mem.set(elements[i], n.scriptsPtr + off);
    off += elements[i].length;
  }
  offsets[elements.length] = off;
  n.setCount(elements.length);
  uploaded = { arr: elements, len: elements.length };
  return true;
}

function nativeMatch(n: Native, filterBody: Buffer, blockHash: string, elements: Buffer[], wantAll: boolean): number | null {
  if (!ensureUploaded(n, elements)) return null;
  const { value: nBig, bytesRead } = decodeVarint(filterBody, 0);
  const nSet = Number(nBig);
  if (nSet === 0) return 0;
  const body = filterBody.subarray(bytesRead);
  if (body.length > n.filterBytes) return null;
  n.mem.set(keyFromBlockHash(blockHash), n.keyPtr);
  n.mem.set(body, n.filterPtr);
  const r = n.match(nSet, body.length, wantAll ? 1 : 0);
  return r < 0 ? null : r; // -1: filter too big for the module's buffers → JS
}

// -------------------------------------------------------------------- JS
// Exact 32×32 → 64-bit product as [hi, lo] uint32 (via 16-bit halves).
const mulOut = new Float64Array(2);
function mul32(a: number, b: number): void {
  const al = a & 0xffff, ah = a >>> 16, bl = b & 0xffff, bh = b >>> 16;
  const ll = al * bl, mid = al * bh + ah * bl, hh = ah * bh;
  const lo = ll + (mid % 65536) * 65536;
  const carry = Math.floor(lo / 4294967296);
  mulOut[1] = lo - carry * 4294967296;
  mulOut[0] = hh + Math.floor(mid / 65536) + carry;
}

// Map an element to the filter's [0, N·M) range: fastReduce of the
// SipHash output over N·M, i.e. (h · F) >> 64. Matches BIP158's
// `hashToRange`. `F` is split once per filter into fh·2^32 + fl.
const hashOut = new Uint32Array(2);
function hashToRange(key: SipKey, element: Buffer, fh: number, fl: number): number {
  siphash24Pair(key, element, hashOut);
  const hi = hashOut[0], lo = hashOut[1];
  // h·F = hi·fh·2^64 + (hi·fl + lo·fh)·2^32 + lo·fl
  mul32(lo, fl); const zHi = mulOut[0];
  mul32(hi, fl); const xHi = mulOut[0], xLo = mulOut[1];
  mul32(lo, fh); const yHi = mulOut[0], yLo = mulOut[1];
  const l = xLo + yLo + zHi;
  const c = Math.floor(l / 4294967296);
  return hi * fh + xHi + yHi + c;
}

function splitF(n: number): { fh: number; fl: number } {
  const F = n * M_NUM;
  const fh = Math.floor(F / 4294967296);
  return { fh, fl: F - fh * 4294967296 };
}

// Decode all values of a Golomb-Rice-coded set (unary quotient, P-bit
// remainder, cumulative) into a sorted Float64Array. The body is what
// `cfilter.filter` returns AFTER the varint N prefix has been consumed.
function decodeGcsValues(body: Buffer, n: number): Float64Array {
  const values = new Float64Array(n);
  const bitLen = body.length * 8;
  let pos = 0;
  let last = 0;
  const bit = (p: number): number => (body[p >>> 3] >> (7 - (p & 7))) & 1;
  for (let i = 0; i < n; i++) {
    let q = 0;
    while (pos < bitLen && bit(pos) === 1) { q++; pos++; }
    if (pos + 1 + P > bitLen) throw new Error('GCS: past end of buffer');
    pos++; // the terminating 0 of the unary prefix
    let r = 0;
    for (let k = 0; k < P; k++) r = r * 2 + bit(pos++);
    last += q * 524288 + r; // q·2^P + r
    values[i] = last;
  }
  return values;
}

function targetsFor(key: SipKey, elements: Buffer[], n: number): Float64Array {
  const { fh, fl } = splitF(n);
  const t = new Float64Array(elements.length);
  for (let i = 0; i < elements.length; i++) t[i] = hashToRange(key, elements[i], fh, fl);
  return t.sort();
}

// ------------------------------------------------------------------ API
// Test whether ANY of `elements` are in the filter. Returns true on
// the first match; false if none matched. This is the primary
// wallet-side operation: pass every scriptPubKey (and any input
// script data for spend detection) we care about.
export function filterMatchesAny(
  filterBody: Buffer,
  blockHash: string,
  elements: Buffer[],
): boolean {
  if (elements.length === 0) return false;
  if (!forceJs) {
    const n = getNative();
    if (n) {
      const r = nativeMatch(n, filterBody, blockHash, elements, false);
      if (r !== null) return r > 0;
    }
  }
  const { value: nBig, bytesRead } = decodeVarint(filterBody, 0);
  const n = Number(nBig);
  if (n === 0) return false;
  const key = sipKey(keyFromBlockHash(blockHash));
  const targets = targetsFor(key, elements, n);
  const setValues = decodeGcsValues(filterBody.subarray(bytesRead), n);
  let ti = 0;
  for (let i = 0; i < n; i++) {
    const v = setValues[i];
    while (ti < targets.length && targets[ti] < v) ti++;
    if (ti === targets.length) return false;
    if (targets[ti] === v) return true;
  }
  return false;
}

// Return the subset of elements that matched the filter — useful when
// probing one filter against many address candidates and needing to
// know WHICH matched. Slower than matchesAny (can't short-circuit).
export function filterMatches(
  filterBody: Buffer,
  blockHash: string,
  elements: Buffer[],
): Buffer[] {
  if (elements.length === 0) return [];
  if (!forceJs) {
    const nat = getNative();
    if (nat) {
      const r = nativeMatch(nat, filterBody, blockHash, elements, true);
      if (r !== null) {
        const idx = new Uint32Array(nat.mem.buffer, nat.matchesPtr, r);
        return Array.from(idx, (i) => elements[i]);
      }
    }
  }
  const { value: nBig, bytesRead } = decodeVarint(filterBody, 0);
  const n = Number(nBig);
  if (n === 0) return [];
  const key = sipKey(keyFromBlockHash(blockHash));
  const { fh, fl } = splitF(n);
  const withHash = elements.map((el) => ({ el, target: hashToRange(key, el, fh, fl) }));
  withHash.sort((a, b) => a.target - b.target);
  const setValues = decodeGcsValues(filterBody.subarray(bytesRead), n);
  const matched: Buffer[] = [];
  let ti = 0;
  for (let i = 0; i < n; i++) {
    const v = setValues[i];
    while (ti < withHash.length && withHash[ti].target < v) ti++;
    if (ti === withHash.length) break;
    while (ti < withHash.length && withHash[ti].target === v) matched.push(withHash[ti++].el);
  }
  return matched;
}