perf(db): harden unchanged index updates

.changeset/quick-trees-rest.md

@@ -2,4 +2,5 @@
 '@tanstack/db': patch
 ---
 
-Skip index writes when an update leaves the indexed value unchanged.
+Skip unchanged index writes, cache index evaluators, and preserve key counts
+after failed removals.

packages/db/src/indexes/base-index.ts

@@ -1,6 +1,7 @@
 import { compileSingleRowExpression } from '../query/compiler/evaluators.js'
 import { comparisonFunctions } from '../query/builder/functions.js'
 import { DEFAULT_COMPARE_OPTIONS, deepEquals } from '../utils.js'
+import type { CompiledSingleRowExpression } from '../query/compiler/evaluators.js'
 import type { RangeQueryOptions } from './btree-index.js'
 import type { CompareOptions } from '../query/builder/types.js'
 import type { BasicExpression, OrderByDirection } from '../query/ir.js'
@@ -99,6 +100,7 @@ export abstract class BaseIndex<
   protected totalLookupTime = 0
   protected lastUpdated = new Date()
   protected compareOptions: CompareOptions
+  private compiledIndexEvaluator: CompiledSingleRowExpression | undefined
   /**
    * Set by subclasses when constructed with a user-supplied comparator, whose
    * ordering may not match the WHERE evaluator's relational operators.
@@ -210,7 +212,8 @@ export abstract class BaseIndex<
   protected abstract initialize(options?: any): void
 
   protected evaluateIndexExpression(item: any): any {
-    const evaluator = compileSingleRowExpression(this.expression)
+    const evaluator = (this.compiledIndexEvaluator ??=
+      compileSingleRowExpression(this.expression))
     return evaluator(item as Record<string, unknown>)
   }
 

packages/db/src/indexes/basic-index.ts

@@ -173,7 +173,8 @@ export class BasicIndex<
 
     if (
       areSameValueZeroEqual(oldValue, newValue) &&
-      this.valueMap.get(newValue)?.has(key)
+      this.valueMap.get(newValue)?.has(key) &&
+      this.indexedKeys.has(key)
     ) {
       return
     }

packages/db/src/utils/comparison.ts

@@ -185,6 +185,10 @@ export const UNDEFINED_SENTINEL = `__TS_DB_BTREE_UNDEFINED_VALUE__`
  * for BTree index operations that need to distinguish undefined values.
  */
 export function normalizeValue(value: any): any {
+  if (typeof value !== `object` || value === null) {
+    return value
+  }
+
   if (value instanceof Date) {
     return value.getTime()
   }

packages/db/tests/index-update-short-circuit.test.ts

@@ -1,4 +1,4 @@
-import { describe, expect, it } from 'vitest'
+import { describe, expect, it, vi } from 'vitest'
 import { BasicIndex } from '../src/indexes/basic-index.js'
 import { BTreeIndex } from '../src/indexes/btree-index.js'
 import { PropRef } from '../src/query/ir.js'
@@ -47,6 +47,7 @@ describe.each(indexTypes)(`%s update`, (_indexName, IndexType) => {
     const index = createIndex()
     index.add(`a`, { value: oldValue })
     const bucket = index.valueMapData.get(normalizeValue(oldValue))
+    expect(bucket).toBeDefined()
 
     index.update(`a`, { value: oldValue }, { value: newValue })
 
@@ -85,4 +86,48 @@ describe.each(indexTypes)(`%s update`, (_indexName, IndexType) => {
     expect(index.lookup(`eq`, `A`)).toEqual(new Set())
     expect(index.lookup(`eq`, `a`)).toEqual(new Set([`a`]))
   })
+
+  it(`preserves the previous error behavior when evaluation fails`, () => {
+    const index = createIndex()
+    index.add(`a`, { value: 1 })
+    const newItem = Object.defineProperty({}, `value`, {
+      get() {
+        throw new Error(`evaluation failed`)
+      },
+    })
+
+    expect(() => index.update(`a`, { value: 1 }, newItem)).toThrow(
+      `evaluation failed`,
+    )
+
+    expect(index.lookup(`eq`, 1)).toEqual(new Set())
+    expect(index.keyCount).toBe(0)
+  })
+})
+
+describe(`BasicIndex update bookkeeping`, () => {
+  it(`repairs indexed key membership after a failed removal`, () => {
+    const index = new BasicIndex<string>(1, new PropRef([`value`]))
+    index.add(`a`, { value: 1 })
+    const itemWithThrowingValue = Object.defineProperty({}, `value`, {
+      get() {
+        throw new Error(`evaluation failed`)
+      },
+    })
+    const warn = vi.spyOn(console, `warn`).mockImplementation(() => {})
+
+    try {
+      index.remove(`a`, itemWithThrowingValue)
+    } finally {
+      warn.mockRestore()
+    }
+
+    expect(index.lookup(`eq`, 1)).toEqual(new Set([`a`]))
+    expect(index.keyCount).toBe(0)
+
+    index.update(`a`, { value: 1 }, { value: 1 })
+
+    expect(index.lookup(`eq`, 1)).toEqual(new Set([`a`]))
+    expect(index.keyCount).toBe(1)
+  })
 })

packages/db/tests/index-update.property.test.ts

@@ -0,0 +1,126 @@
+import { describe, expect } from 'vitest'
+import { fc, test as fcTest } from '@fast-check/vitest'
+import { BasicIndex } from '../src/indexes/basic-index.js'
+import { BTreeIndex } from '../src/indexes/btree-index.js'
+import { PropRef } from '../src/query/ir.js'
+import type { BaseIndex } from '../src/indexes/base-index.js'
+
+type IndexValue = number
+
+type IndexConstructor = new (
+  id: number,
+  expression: PropRef,
+) => BaseIndex<string>
+
+type IndexAction =
+  | { type: `put`; key: string; value: IndexValue }
+  | { type: `delete`; key: string }
+
+const indexTypes: Array<[string, IndexConstructor]> = [
+  [`BasicIndex`, BasicIndex as IndexConstructor],
+  [`BTreeIndex`, BTreeIndex as IndexConstructor],
+]
+
+const arbitraryValue: fc.Arbitrary<IndexValue> = fc.integer({
+  min: -3,
+  max: 3,
+})
+
+const arbitraryAction: fc.Arbitrary<IndexAction> = fc.oneof(
+  fc.record({
+    type: fc.constant(`put` as const),
+    key: fc.integer({ min: 0, max: 7 }).map(String),
+    value: arbitraryValue,
+  }),
+  fc.record({
+    type: fc.constant(`delete` as const),
+    key: fc.integer({ min: 0, max: 7 }).map(String),
+  }),
+)
+
+const probeValues: Array<IndexValue> = [-3, -2, -1, -0, 0, 1, 2, 3, 99]
+const rangeBoundaries: Array<IndexValue> = [-2, 0, 2]
+
+function groupKeysByValue(
+  rows: Map<string, IndexValue>,
+): Map<IndexValue, Set<string>> {
+  const groups = new Map<IndexValue, Set<string>>()
+  for (const [key, value] of rows) {
+    const keys = groups.get(value)
+    if (keys) {
+      keys.add(key)
+    } else {
+      groups.set(value, new Set([key]))
+    }
+  }
+  return groups
+}
+
+function expectIndexMatchesModel(
+  index: BaseIndex<string>,
+  rows: Map<string, IndexValue>,
+): void {
+  const groups = groupKeysByValue(rows)
+
+  expect(index.keyCount).toBe(rows.size)
+  expect(index.indexedKeysSet).toEqual(new Set(rows.keys()))
+  expect(index.valueMapData).toEqual(groups)
+  expect(index.orderedEntriesArray).toEqual(
+    [...groups].sort(([left], [right]) => left - right),
+  )
+
+  for (const value of probeValues) {
+    expect(index.lookup(`eq`, value)).toEqual(groups.get(value) ?? new Set())
+  }
+
+  for (const boundary of rangeBoundaries) {
+    const keysAtOrAbove = new Set(
+      [...rows].filter(([, value]) => value >= boundary).map(([key]) => key),
+    )
+    const keysAtOrBelow = new Set(
+      [...rows].filter(([, value]) => value <= boundary).map(([key]) => key),
+    )
+
+    expect(index.rangeQuery({ from: boundary })).toEqual(keysAtOrAbove)
+    expect(index.rangeQuery({ to: boundary })).toEqual(keysAtOrBelow)
+  }
+}
+
+describe.each(indexTypes)(`%s update properties`, (_indexName, IndexType) => {
+  fcTest.prop([
+    fc.array(arbitraryAction, {
+      minLength: 1,
+      maxLength: 100,
+    }),
+  ])(
+    `matches a reference model across valid operation sequences`,
+    (actions) => {
+      const index = new IndexType(1, new PropRef([`value`]))
+      const rows = new Map<string, IndexValue>()
+
+      for (const action of actions) {
+        if (action.type === `put`) {
+          if (rows.has(action.key)) {
+            index.update(
+              action.key,
+              { value: rows.get(action.key) },
+              { value: action.value },
+            )
+          } else {
+            index.add(action.key, { value: action.value })
+          }
+          rows.set(action.key, action.value)
+        } else if (rows.has(action.key)) {
+          index.remove(action.key, { value: rows.get(action.key) })
+          rows.delete(action.key)
+        }
+
+        expectIndexMatchesModel(index, rows)
+      }
+
+      const rebuilt = new IndexType(2, new PropRef([`value`]))
+      rebuilt.build([...rows].map(([key, value]) => [key, { value }] as const))
+      expectIndexMatchesModel(rebuilt, rows)
+    },
+  )
+})