@tanstack/db-ivm 0.1.2 → 0.1.3

package/src/hashing/murmur.ts

@@ -0,0 +1,128 @@
+/*
+ * Implementation of murmur hash based on the Composites polyfill implementation:
+ * https://github.com/tc39/proposal-composites
+ */
+
+const RANDOM_SEED = randomHash()
+const STRING_MARKER = randomHash()
+const BIG_INT_MARKER = randomHash()
+const NEG_BIG_INT_MARKER = randomHash()
+const SYMBOL_MARKER = randomHash()
+
+export type Hash = number
+
+export function randomHash() {
+  return (Math.random() * (2 ** 31 - 1)) >>> 0
+}
+
+export interface Hasher {
+  update: (val: symbol | string | number | bigint) => void
+  digest: () => number
+}
+
+/**
+ * This implementation of Murmur hash uses a random initial seed and random markers.
+ * This means that hashes aren't deterministic across app restarts.
+ * This is intentional in the composites polyfill to be resistent to hash-flooding attacks
+ * where malicious users would precompute lots of different objects whose hashes collide with each other.
+ *
+ * Currently, for ts/db-ivm this is fine because we don't persist client state.
+ * However, when we will introduce persistence we will either need to store the seeds or remove the randomness
+ * to ensure deterministic hashes across app restarts.
+ */
+export class MurmurHashStream implements Hasher {
+  private hash: number = RANDOM_SEED
+  private length = 0
+  private carry = 0
+  private carryBytes = 0
+
+  private _mix(k1: number): void {
+    k1 = Math.imul(k1, 0xcc9e2d51)
+    k1 = (k1 << 15) | (k1 >>> 17)
+    k1 = Math.imul(k1, 0x1b873593)
+    this.hash ^= k1
+    this.hash = (this.hash << 13) | (this.hash >>> 19)
+    this.hash = Math.imul(this.hash, 5) + 0xe6546b64
+  }
+
+  private _writeByte(byte: number): void {
+    this.carry |= (byte & 0xff) << (8 * this.carryBytes)
+    this.carryBytes++
+    this.length++
+
+    if (this.carryBytes === 4) {
+      this._mix(this.carry >>> 0)
+      this.carry = 0
+      this.carryBytes = 0
+    }
+  }
+
+  update(chunk: symbol | string | number | bigint): void {
+    switch (typeof chunk) {
+      case `symbol`: {
+        this.update(SYMBOL_MARKER)
+        const description = chunk.description
+        if (!description) {
+          return
+        }
+
+        for (let i = 0; i < description.length; i++) {
+          const code = description.charCodeAt(i)
+          this._writeByte(code & 0xff)
+          this._writeByte((code >>> 8) & 0xff)
+        }
+        return
+      }
+      case `string`:
+        this.update(STRING_MARKER)
+        for (let i = 0; i < chunk.length; i++) {
+          const code = chunk.charCodeAt(i)
+          this._writeByte(code & 0xff)
+          this._writeByte((code >>> 8) & 0xff)
+        }
+        return
+      case `number`:
+        this._writeByte(chunk & 0xff)
+        this._writeByte((chunk >>> 8) & 0xff)
+        this._writeByte((chunk >>> 16) & 0xff)
+        this._writeByte((chunk >>> 24) & 0xff)
+        return
+      case `bigint`: {
+        let value = chunk
+        if (value < 0n) {
+          value = -value
+          this.update(NEG_BIG_INT_MARKER)
+        } else {
+          this.update(BIG_INT_MARKER)
+        }
+        while (value > 0n) {
+          this._writeByte(Number(value & 0xffn))
+          value >>= 8n
+        }
+        if (chunk === 0n) this._writeByte(0)
+        return
+      }
+      default:
+        throw new TypeError(`Unsupported input type: ${typeof chunk}`)
+    }
+  }
+
+  digest(): number {
+    if (this.carryBytes > 0) {
+      let k1 = this.carry >>> 0
+      k1 = Math.imul(k1, 0xcc9e2d51)
+      k1 = (k1 << 15) | (k1 >>> 17)
+      k1 = Math.imul(k1, 0x1b873593)
+      this.hash ^= k1
+    }
+
+    this.hash ^= this.length
+    this.hash ^= this.hash >>> 16
+    this.hash = Math.imul(this.hash, 0x85ebca6b)
+    this.hash ^= this.hash >>> 13
+    this.hash = Math.imul(this.hash, 0xc2b2ae35)
+    this.hash ^= this.hash >>> 16
+
+    return this.hash >>> 0
+  }
+}

package/src/indexes.ts

@@ -1,5 +1,7 @@
 import { MultiSet } from "./multiset.js"
-import { DefaultMap } from "./utils.js"
+import { HashIndex } from "./hashIndex.js"
+import { ValueIndex } from "./valueIndex.js"
+import { concatIterable, mapIterable } from "./utils.js"
 
 /**
  * A map from a difference collection trace's keys -> (value, multiplicities) that changed.
@@ -7,77 +9,134 @@ import { DefaultMap } from "./utils.js"
  * exploit the key-value structure of the data to run efficiently.
  */
 export class Index<K, V> {
-  #inner: DefaultMap<K, Map<V, number>>
+  /*
+   * This is a hybrid Index that composes a ValueIndex and a HashIndex.
+   * Keys that have only one value are stored in the ValueIndex.
+   * Keys that have multiple values are stored in the HashIndex, the hash distinguishes between the values.
+   * This reduces the amount of hashes we need to compute since often times only a small portion of the keys are updated
+   * so we don't have to hash the keys that are never updated.
+   *
+   * Note: The `valueIndex` and `hashIndex` have disjoint keys.
+   *       When a key that has only one value gets a new distinct value,
+   *       it is added to the `hashIndex` and removed from the `valueIndex` and vice versa.
+   */
+  #valueIndex: ValueIndex<K, V>
+  #hashIndex: HashIndex<K, V>
 
   constructor() {
-    this.#inner = new DefaultMap<K, Map<V, number>>(() => new Map<V, number>())
-    // #inner is a map of:
-    // {
-    //   [key]: Map<V, number>  // Direct value-to-multiplicity mapping
-    // }
+    this.#valueIndex = new ValueIndex<K, V>()
+    this.#hashIndex = new HashIndex<K, V>()
   }
 
   toString(indent = false): string {
-    return `Index(${JSON.stringify(
-      [...this.#inner].map(([k, valueMap]) => [k, [...valueMap]]),
-      undefined,
-      indent ? `  ` : undefined
-    )})`
+    return `Index(\n  ${this.#valueIndex.toString(indent)},\n  ${this.#hashIndex.toString(indent)}\n)`
   }
 
   get(key: K): Array<[V, number]> {
-    const valueMap = this.#inner.get(key)
-    return [...valueMap.entries()]
+    if (this.#valueIndex.has(key)) {
+      return [this.#valueIndex.get(key)!]
+    }
+    return this.#hashIndex.get(key)
   }
 
   getMultiplicity(key: K, value: V): number {
-    const valueMap = this.#inner.get(key)
-    return valueMap.get(value) ?? 0
+    if (this.#valueIndex.has(key)) {
+      return this.#valueIndex.getMultiplicity(key)
+    }
+    return this.#hashIndex.getMultiplicity(key, value)
   }
 
-  entries() {
-    return this.#inner.entries()
+  /**
+   * This returns an iterator that iterates over all key-value pairs.
+   * @returns An iterable of all key-value pairs (and their multiplicities) in the index.
+   */
+  #entries(): Iterable<[K, [V, number]]> {
+    return concatIterable(
+      this.#valueIndex.entries(),
+      this.#hashIndex.entriesIterator()
+    )
   }
 
-  keys() {
-    return this.#inner.keys()
+  /**
+   * This method only iterates over the keys and not over the values.
+   * Hence, it is more efficient than the `#entries` method.
+   * It returns an iterator that you can use if you need to iterate over the values for a given key.
+   * @returns An iterator of all *keys* in the index and their corresponding value iterator.
+   */
+  *#entriesIterators(): Iterable<[K, Iterable<[V, number]>]> {
+    for (const [key, [value, multiplicity]] of this.#valueIndex.entries()) {
+      yield [key, new Map<V, number>([[value, multiplicity]])]
+    }
+    for (const [key, valueMap] of this.#hashIndex.entries()) {
+      yield [
+        key,
+        mapIterable(valueMap, ([_hash, [value, multiplicity]]) => [
+          value,
+          multiplicity,
+        ]),
+      ]
+    }
   }
 
   has(key: K): boolean {
-    return this.#inner.has(key)
+    return this.#valueIndex.has(key) || this.#hashIndex.has(key)
   }
 
   get size(): number {
-    return this.#inner.size
+    return this.#valueIndex.size + this.#hashIndex.size
   }
 
   addValue(key: K, value: [V, number]): void {
-    const [val, multiplicity] = value
-    const valueMap = this.#inner.get(key)
-    const existingMultiplicity = valueMap.get(val) ?? 0
-    const newMultiplicity = existingMultiplicity + multiplicity
-
-    if (multiplicity !== 0) {
-      if (newMultiplicity === 0) {
-        valueMap.delete(val)
-      } else {
-        valueMap.set(val, newMultiplicity)
+    const containedInValueIndex = this.#valueIndex.has(key)
+    const containedInHashIndex = this.#hashIndex.has(key)
+
+    if (containedInHashIndex && containedInValueIndex) {
+      throw new Error(
+        `Key ${key} is contained in both the value index and the hash index. This should never happen because they should have disjoint keysets.`
+      )
+    }
+
+    if (!containedInValueIndex && !containedInHashIndex) {
+      // This is the first time we see the key
+      // Add it to the value index
+      this.#valueIndex.addValue(key, value)
+      return
+    }
+
+    if (containedInValueIndex) {
+      // This key is already in the value index
+      // It could be that it's the same value or a different one
+      // If it's a different value we will need to remove the key from the value index
+      // and add the key and its two values to the hash index
+      try {
+        this.#valueIndex.addValue(key, value)
+      } catch {
+        // This is a different value, need to move the key to the hash index
+        const existingValue = this.#valueIndex.get(key)!
+        this.#valueIndex.delete(key)
+        this.#hashIndex.addValue(key, existingValue)
+        this.#hashIndex.addValue(key, value)
       }
+      return
+    }
+
+    if (containedInHashIndex) {
+      // This key is already in the hash index so it already has two or more values.
+      // However, this new value and multiplicity could cause an existing value to be removed
+      // and lead to the key having only a single value in which case we need to move it back to the value index
+      const singleRemainingValue = this.#hashIndex.addValue(key, value)
+      if (singleRemainingValue) {
+        // The key only has a single remaining value so we need to move it back to the value index
+        this.#hashIndex.delete(key)
+        this.#valueIndex.addValue(key, singleRemainingValue)
+      }
+      return
     }
   }
 
   append(other: Index<K, V>): void {
-    for (const [key, otherValueMap] of other.entries()) {
-      const thisValueMap = this.#inner.get(key)
-      for (const [value, multiplicity] of otherValueMap.entries()) {
-        const existingMultiplicity = thisValueMap.get(value) ?? 0
-        const newMultiplicity = existingMultiplicity + multiplicity
-        if (newMultiplicity === 0) {
-          thisValueMap.delete(value)
-        } else {
-          thisValueMap.set(value, newMultiplicity)
-        }
-      }
+    for (const [key, value] of other.#entries()) {
+      this.addValue(key, value)
     }
   }
 
@@ -87,10 +146,10 @@ export class Index<K, V> {
     // We want to iterate over the smaller of the two indexes to reduce the
     // number of operations we need to do.
     if (this.size <= other.size) {
-      for (const [key, valueMap] of this.entries()) {
+      for (const [key, valueIt] of this.#entriesIterators()) {
         if (!other.has(key)) continue
         const otherValues = other.get(key)
-        for (const [val1, mul1] of valueMap.entries()) {
+        for (const [val1, mul1] of valueIt) {
           for (const [val2, mul2] of otherValues) {
             if (mul1 !== 0 && mul2 !== 0) {
               result.push([[key, [val1, val2]], mul1 * mul2])
@@ -99,10 +158,10 @@ export class Index<K, V> {
         }
       }
     } else {
-      for (const [key, otherValueMap] of other.entries()) {
+      for (const [key, otherValueIt] of other.#entriesIterators()) {
         if (!this.has(key)) continue
         const values = this.get(key)
-        for (const [val2, mul2] of otherValueMap.entries()) {
+        for (const [val2, mul2] of otherValueIt) {
           for (const [val1, mul1] of values) {
             if (mul1 !== 0 && mul2 !== 0) {
               result.push([[key, [val1, val2]], mul1 * mul2])

package/src/multiset.ts

@@ -2,8 +2,8 @@ import {
   DefaultMap,
   chunkedArrayPush,
   globalObjectIdGenerator,
-  hash,
 } from "./utils.js"
+import { hash } from "./hashing/index.js"
 
 export type MultiSetArray<T> = Array<[T, number]>
 export type KeyedData<T> = [key: string, value: T]

package/src/operators/distinct.ts

@@ -1,11 +1,11 @@
 import { DifferenceStreamWriter, UnaryOperator } from "../graph.js"
 import { StreamBuilder } from "../d2.js"
-import { hash } from "../utils.js"
+import { hash } from "../hashing/index.js"
 import { MultiSet } from "../multiset.js"
+import type { Hash } from "../hashing/index.js"
 import type { DifferenceStreamReader } from "../graph.js"
 import type { IStreamBuilder } from "../types.js"
 
-type HashedValue = string
 type Multiplicity = number
 
 /**
@@ -13,7 +13,7 @@ type Multiplicity = number
  */
 export class DistinctOperator<T> extends UnaryOperator<T> {
   #by: (value: T) => any
-  #values: Map<HashedValue, Multiplicity> // keeps track of the number of times each value has been seen
+  #values: Map<Hash, Multiplicity> // keeps track of the number of times each value has been seen
 
   constructor(
     id: number,
@@ -27,7 +27,7 @@ export class DistinctOperator<T> extends UnaryOperator<T> {
   }
 
   run(): void {
-    const updatedValues = new Map<HashedValue, [Multiplicity, T]>()
+    const updatedValues = new Map<Hash, [Multiplicity, T]>()
 
     // Compute the new multiplicity for each value
     for (const message of this.inputMessages()) {

package/src/operators/groupBy.ts

@@ -165,13 +165,16 @@ export function sum<T>(
 /**
  * Creates a count aggregate function
  */
-export function count<T>(): AggregateFunction<T, number, number> {
+export function count<T>(
+  valueExtractor: (value: T) => any = (v) => v
+): AggregateFunction<T, number, number> {
   return {
-    preMap: () => 1,
+    // Count only not-null values (the `== null` comparison gives true for both null and undefined)
+    preMap: (data: T) => (valueExtractor(data) == null ? 0 : 1),
     reduce: (values: Array<[number, number]>) => {
       let totalCount = 0
-      for (const [_, multiplicity] of values) {
-        totalCount += multiplicity
+      for (const [nullMultiplier, multiplicity] of values) {
+        totalCount += nullMultiplier * multiplicity
       }
       return totalCount
     },

package/src/utils.ts

@@ -1,4 +1,15 @@
-import murmurhash from "murmurhash-js"
+/**
+ * Simple assertion function for runtime checks.
+ * Throws an error if the condition is false.
+ */
+export function assert(
+  condition: unknown,
+  message?: string
+): asserts condition {
+  if (!condition) {
+    throw new Error(message || `Assertion failed`)
+  }
+}
 
 /**
  * A map that returns a default value for keys that are not present.
@@ -13,7 +24,8 @@ export class DefaultMap<K, V> extends Map<K, V> {
 
   get(key: K): V {
     if (!this.has(key)) {
-      this.set(key, this.defaultValue())
+      // this.set(key, this.defaultValue())
+      return this.defaultValue()
     }
     return super.get(key)!
   }
@@ -45,52 +57,6 @@ export function chunkedArrayPush(array: Array<unknown>, other: Array<unknown>) {
   }
 }
 
-const hashCache = new WeakMap()
-
-/**
- * Replacer function for JSON.stringify that converts unsupported types to strings
- */
-function hashReplacer(_key: string, value: any): any {
-  if (typeof value === `bigint`) {
-    return String(value)
-  } else if (typeof value === `symbol`) {
-    return String(value)
-  } else if (typeof value === `function`) {
-    return String(value)
-  } else if (value === undefined) {
-    return `undefined`
-  } else if (value instanceof Map) {
-    return `Map(${JSON.stringify(Array.from(value.entries()), hashReplacer)})`
-  } else if (value instanceof Set) {
-    return `Set(${JSON.stringify(Array.from(value.values()), hashReplacer)})`
-  }
-  return value
-}
-
-/**
- * A hash method that caches the hash of a value in a week map
- */
-export function hash(data: any): string {
-  if (
-    data === null ||
-    data === undefined ||
-    (typeof data !== `object` && typeof data !== `function`)
-  ) {
-    // Can't be cached in the weak map because it's not an object
-    const serialized = JSON.stringify(data, hashReplacer)
-    return murmurhash.murmur3(serialized).toString(16)
-  }
-
-  if (hashCache.has(data)) {
-    return hashCache.get(data)
-  }
-
-  const serialized = JSON.stringify(data, hashReplacer)
-  const hashValue = murmurhash.murmur3(serialized).toString(16)
-  hashCache.set(data, hashValue)
-  return hashValue
-}
-
 export function binarySearch<T>(
   array: Array<T>,
   value: T,
@@ -161,3 +127,20 @@ export class ObjectIdGenerator {
  * Global instance for cases where a shared object ID space is needed.
  */
 export const globalObjectIdGenerator = new ObjectIdGenerator()
+
+export function* concatIterable<T>(
+  ...iterables: Array<Iterable<T>>
+): Iterable<T> {
+  for (const iterable of iterables) {
+    yield* iterable
+  }
+}
+
+export function* mapIterable<T, U>(
+  it: Iterable<T>,
+  fn: (t: T) => U
+): Iterable<U> {
+  for (const t of it) {
+    yield fn(t)
+  }
+}

package/src/valueIndex.ts

@@ -0,0 +1,78 @@
+import { hash } from "./hashing/index.js"
+
+/**
+ * A map from a difference collection trace's keys -> (value, multiplicities) that changed.
+ * Used in operations like join and reduce where the operation needs to
+ * exploit the key-value structure of the data to run efficiently.
+ */
+export class ValueIndex<K, V> {
+  #inner: Map<K, [V, number]> // Maps key to the value and its multiplicity
+
+  constructor() {
+    this.#inner = new Map()
+  }
+
+  toString(indent = false): string {
+    return `ValueIndex(${JSON.stringify(
+      [...this.#inner.entries()],
+      undefined,
+      indent ? 2 : undefined
+    )})`
+  }
+
+  get(key: K): [V, number] | undefined {
+    return this.#inner.get(key)
+  }
+
+  getMultiplicity(key: K): number {
+    return this.get(key)?.[1] ?? 0
+  }
+
+  entries() {
+    return this.#inner.entries()
+  }
+
+  has(key: K): boolean {
+    return this.#inner.has(key)
+  }
+
+  delete(key: K): void {
+    this.#inner.delete(key)
+  }
+
+  get size(): number {
+    return this.#inner.size
+  }
+
+  addValue(key: K, v: [V, number]): void {
+    const [value, multiplicity] = v
+
+    if (multiplicity === 0) {
+      return
+    }
+
+    if (this.has(key)) {
+      const [currValue, currMultiplicity] = this.get(key)!
+      if (hash(value) === hash(currValue)) {
+        // Update the multiplicity
+        this.#setMultiplicity(key, value, currMultiplicity + multiplicity)
+        return
+      }
+      // Different value, not allowed.
+      // ValueIndex only supports one value per key.
+      throw new Error(
+        `Cannot add value for key ${key} because it already exists in ValueIndex with a different value`
+      )
+    }
+
+    this.#inner.set(key, [value, multiplicity])
+  }
+
+  #setMultiplicity(key: K, value: V, multiplicity: number): void {
+    if (multiplicity === 0) {
+      this.#inner.delete(key)
+    } else {
+      this.#inner.set(key, [value, multiplicity])
+    }
+  }
+}