Fix infinite loop in takeInternal with undefined values (#1198)

* test: Add failing test for BTreeIndex infinite loop with undefined values (issue #1186)

This test demonstrates the infinite loop bug that occurs when calling
take() on a BTreeIndex containing items with undefined indexed values.

The bug is in takeInternal() where nextHigherPair(undefined) returns the
minimum pair [undefined, undefined], then key is set to pair[0] (undefined),
causing the same pair to be returned infinitely since the while condition
(pair !== undefined) is always true for arrays.

https://claude.ai/code/session_01RKBKXMoKVe1hSXGy3VVEFo

* ci: apply automated fixes

* Remove broken tests

* Distinguish usage of undefined by introducing a sentinel for starting from the start/end. Also introduce a sentinel such that we never store undefined as a key in the btree.

* Fix infinite loop reproduction test

* Fix some tests

* ci: apply automated fixes

* Fix some tests

* Introduce a normalizeValueForBTree that is separate from normalizeValue because normalizeValue is also used in other places.

* ci: apply automated fixes

* Changeset

* ci: apply automated fixes

---------

Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>

Author: 3 people

.changeset/fix-btree-undefined-infinite-loop.md

@@ -0,0 +1,9 @@
+---
+'@tanstack/db': patch
+---
+
+Fixed infinite loop in `BTreeIndex.takeInternal` when indexed values are `undefined`.
+
+The BTree uses `undefined` as a special parameter meaning "start from beginning/end", which caused an infinite loop when the actual indexed value was `undefined`.
+
+Added `takeFromStart` and `takeReversedFromEnd` methods to explicitly start from the beginning/end, and introduced a sentinel value for storing `undefined` in the BTree.

packages/db/src/collection/change-events.ts

@@ -347,7 +347,7 @@ function getOrderedKeys<T extends object, TKey extends string | number>(
         // Take the keys that match the filter and limit
         // if no limit is provided `index.keyCount` is used,
         // i.e. we will take all keys that match the filter
-        return index.take(limit ?? index.keyCount, undefined, filterFn)
+        return index.takeFromStart(limit ?? index.keyCount, filterFn)
       }
     }
   }

packages/db/src/collection/subscription.ts

@@ -425,6 +425,9 @@ export class CollectionSubscription
       )
     }
 
+    // Check if minValues has a first element (regardless of its value)
+    // This distinguishes between "no min value provided" vs "min value is undefined"
+    const hasMinValue = minValues !== undefined && minValues.length > 0
     // Derive first column value from minValues (used for local index operations)
     const minValue = minValues?.[0]
     // Cast for index operations (index expects string | number)
@@ -436,8 +439,8 @@ export class CollectionSubscription
       ? createFilterFunctionFromExpression(where)
       : undefined
 
-    const filterFn = (key: string | number): boolean => {
-      if (this.sentKeys.has(key)) {
+    const filterFn = (key: string | number | undefined): boolean => {
+      if (key !== undefined && this.sentKeys.has(key)) {
         return false
       }
 
@@ -462,7 +465,7 @@ export class CollectionSubscription
     // For multi-column orderBy, we use the first column value for index operations (wide bounds)
     // This may load some duplicates but ensures we never miss any rows.
     let keys: Array<string | number> = []
-    if (minValueForIndex !== undefined) {
+    if (hasMinValue) {
       // First, get all items with the same FIRST COLUMN value as minValue
       // This provides wide bounds for the local index
       const { expression } = orderBy[0]!
@@ -481,15 +484,16 @@ export class CollectionSubscription
         // Then get items greater than minValue
         const keysGreaterThanMin = index.take(
           limit - keys.length,
-          minValueForIndex,
+          minValueForIndex!,
           filterFn,
         )
         keys.push(...keysGreaterThanMin)
       } else {
-        keys = index.take(limit, minValueForIndex, filterFn)
+        keys = index.take(limit, minValueForIndex!, filterFn)
       }
     } else {
-      keys = index.take(limit, minValueForIndex, filterFn)
+      // No min value provided, start from the beginning
+      keys = index.takeFromStart(limit, filterFn)
     }
 
     const valuesNeeded = () => Math.max(limit - changes.length, 0)
@@ -518,7 +522,7 @@ export class CollectionSubscription
         insertedKeys.add(key) // Track this key
       }
 
-      keys = index.take(valuesNeeded(), biggestObservedValue, filterFn)
+      keys = index.take(valuesNeeded(), biggestObservedValue!, filterFn)
     }
 
     // Track row count for offset-based pagination (before sending to callback)

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

@@ -26,7 +26,7 @@ export interface IndexStats {
 }
 
 export interface IndexInterface<
-  TKey extends string | number = string | number,
+  TKey extends string | number | undefined = string | number | undefined,
 > {
   add: (key: TKey, item: any) => void
   remove: (key: TKey, item: any) => void
@@ -45,12 +45,17 @@ export interface IndexInterface<
 
   take: (
     n: number,
-    from?: TKey,
+    from: TKey,
     filterFn?: (key: TKey) => boolean,
   ) => Array<TKey>
+  takeFromStart: (n: number, filterFn?: (key: TKey) => boolean) => Array<TKey>
   takeReversed: (
     n: number,
-    from?: TKey,
+    from: TKey,
+    filterFn?: (key: TKey) => boolean,
+  ) => Array<TKey>
+  takeReversedFromEnd: (
+    n: number,
     filterFn?: (key: TKey) => boolean,
   ) => Array<TKey>
 
@@ -74,7 +79,7 @@ export interface IndexInterface<
  * Base abstract class that all index types extend
  */
 export abstract class BaseIndex<
-  TKey extends string | number = string | number,
+  TKey extends string | number | undefined = string | number | undefined,
 > implements IndexInterface<TKey> {
   public readonly id: number
   public readonly name?: string
@@ -108,12 +113,20 @@ export abstract class BaseIndex<
   abstract lookup(operation: IndexOperation, value: any): Set<TKey>
   abstract take(
     n: number,
-    from?: TKey,
+    from: TKey,
+    filterFn?: (key: TKey) => boolean,
+  ): Array<TKey>
+  abstract takeFromStart(
+    n: number,
     filterFn?: (key: TKey) => boolean,
   ): Array<TKey>
   abstract takeReversed(
     n: number,
-    from?: TKey,
+    from: TKey,
+    filterFn?: (key: TKey) => boolean,
+  ): Array<TKey>
+  abstract takeReversedFromEnd(
+    n: number,
     filterFn?: (key: TKey) => boolean,
   ): Array<TKey>
   abstract get keyCount(): number

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

@@ -1,6 +1,10 @@
 import { compareKeys } from '@tanstack/db-ivm'
 import { BTree } from '../utils/btree.js'
-import { defaultComparator, normalizeValue } from '../utils/comparison.js'
+import {
+  defaultComparator,
+  denormalizeUndefined,
+  normalizeForBTree,
+} from '../utils/comparison.js'
 import { BaseIndex } from './base-index.js'
 import type { CompareOptions } from '../query/builder/types.js'
 import type { BasicExpression } from '../query/ir.js'
@@ -29,7 +33,7 @@ export interface RangeQueryOptions {
  * This maintains items in sorted order and provides efficient range operations
  */
 export class BTreeIndex<
-  TKey extends string | number = string | number,
+  TKey extends string | number | undefined = string | number | undefined,
 > extends BaseIndex<TKey> {
   public readonly supportedOperations = new Set<IndexOperation>([
     `eq`,
@@ -55,7 +59,16 @@ export class BTreeIndex<
     options?: any,
   ) {
     super(id, expression, name, options)
-    this.compareFn = options?.compareFn ?? defaultComparator
+
+    // Get the base compare function
+    const baseCompareFn = options?.compareFn ?? defaultComparator
+
+    // Wrap it to denormalize sentinels before comparison
+    // This ensures UNDEFINED_SENTINEL is converted back to undefined
+    // before being passed to the baseCompareFn (which can be user-provided and is unaware of the UNDEFINED_SENTINEL)
+    this.compareFn = (a: any, b: any) =>
+      baseCompareFn(denormalizeUndefined(a), denormalizeUndefined(b))
+
     if (options?.compareOptions) {
       this.compareOptions = options!.compareOptions
     }
@@ -78,7 +91,7 @@ export class BTreeIndex<
     }
 
     // Normalize the value for Map key usage
-    const normalizedValue = normalizeValue(indexedValue)
+    const normalizedValue = normalizeForBTree(indexedValue)
 
     // Check if this value already exists
     if (this.valueMap.has(normalizedValue)) {
@@ -111,7 +124,7 @@ export class BTreeIndex<
     }
 
     // Normalize the value for Map key usage
-    const normalizedValue = normalizeValue(indexedValue)
+    const normalizedValue = normalizeForBTree(indexedValue)
 
     if (this.valueMap.has(normalizedValue)) {
       const keySet = this.valueMap.get(normalizedValue)!
@@ -207,7 +220,7 @@ export class BTreeIndex<
    * Performs an equality lookup
    */
   equalityLookup(value: any): Set<TKey> {
-    const normalizedValue = normalizeValue(value)
+    const normalizedValue = normalizeForBTree(value)
     return new Set(this.valueMap.get(normalizedValue) ?? [])
   }
 
@@ -219,10 +232,15 @@ export class BTreeIndex<
     const { from, to, fromInclusive = true, toInclusive = true } = options
     const result = new Set<TKey>()
 
-    const normalizedFrom = normalizeValue(from)
-    const normalizedTo = normalizeValue(to)
-    const fromKey = normalizedFrom ?? this.orderedEntries.minKey()
-    const toKey = normalizedTo ?? this.orderedEntries.maxKey()
+    // Check if from/to were explicitly provided (even if undefined)
+    // vs not provided at all (should use min/max key)
+    const hasFrom = `from` in options
+    const hasTo = `to` in options
+
+    const fromKey = hasFrom
+      ? normalizeForBTree(from)
+      : this.orderedEntries.minKey()
+    const toKey = hasTo ? normalizeForBTree(to) : this.orderedEntries.maxKey()
 
     this.orderedEntries.forRange(
       fromKey,
@@ -250,29 +268,43 @@ export class BTreeIndex<
    */
   rangeQueryReversed(options: RangeQueryOptions = {}): Set<TKey> {
     const { from, to, fromInclusive = true, toInclusive = true } = options
+    const hasFrom = `from` in options
+    const hasTo = `to` in options
+
+    // Swap from/to for reversed query, respecting explicit undefined values
     return this.rangeQuery({
-      from: to ?? this.orderedEntries.maxKey(),
-      to: from ?? this.orderedEntries.minKey(),
+      from: hasTo ? to : this.orderedEntries.maxKey(),
+      to: hasFrom ? from : this.orderedEntries.minKey(),
       fromInclusive: toInclusive,
       toInclusive: fromInclusive,
     })
   }
 
+  /**
+   * Internal method for taking items from the index.
+   * @param n - The number of items to return
+   * @param nextPair - Function to get the next pair from the BTree
+   * @param from - Already normalized! undefined means "start from beginning/end", sentinel means "start from the key undefined"
+   * @param filterFn - Optional filter function
+   * @param reversed - Whether to reverse the order of keys within each value
+   */
   private takeInternal(
     n: number,
     nextPair: (k?: any) => [any, any] | undefined,
-    from?: any,
+    from: any,
     filterFn?: (key: TKey) => boolean,
     reversed: boolean = false,
   ): Array<TKey> {
     const keysInResult: Set<TKey> = new Set()
     const result: Array<TKey> = []
     let pair: [any, any] | undefined
-    let key = normalizeValue(from)
+    let key = from // Use as-is - it's already normalized by the caller
 
     while ((pair = nextPair(key)) !== undefined && result.length < n) {
       key = pair[0]
-      const keys = this.valueMap.get(key)
+      const keys = this.valueMap.get(key) as
+        | Set<Exclude<TKey, undefined>>
+        | undefined
       if (keys && keys.size > 0) {
         // Sort keys for deterministic order, reverse if needed
         const sorted = Array.from(keys).sort(compareKeys)
@@ -291,29 +323,60 @@ export class BTreeIndex<
   }
 
   /**
-   * Returns the next n items after the provided item or the first n items if no from item is provided.
+   * Returns the next n items after the provided item.
    * @param n - The number of items to return
-   * @param from - The item to start from (exclusive). Starts from the smallest item (inclusive) if not provided.
-   * @returns The next n items after the provided key. Returns the first n items if no from item is provided.
+   * @param from - The item to start from (exclusive).
+   * @returns The next n items after the provided key.
    */
-  take(n: number, from?: any, filterFn?: (key: TKey) => boolean): Array<TKey> {
+  take(n: number, from: any, filterFn?: (key: TKey) => boolean): Array<TKey> {
     const nextPair = (k?: any) => this.orderedEntries.nextHigherPair(k)
-    return this.takeInternal(n, nextPair, from, filterFn)
+    // Normalize the from value
+    const normalizedFrom = normalizeForBTree(from)
+    return this.takeInternal(n, nextPair, normalizedFrom, filterFn)
   }
 
   /**
-   * Returns the next n items **before** the provided item (in descending order) or the last n items if no from item is provided.
+   * Returns the first n items from the beginning.
    * @param n - The number of items to return
-   * @param from - The item to start from (exclusive). Starts from the largest item (inclusive) if not provided.
-   * @returns The next n items **before** the provided key. Returns the last n items if no from item is provided.
+   * @param filterFn - Optional filter function
+   * @returns The first n items
+   */
+  takeFromStart(n: number, filterFn?: (key: TKey) => boolean): Array<TKey> {
+    const nextPair = (k?: any) => this.orderedEntries.nextHigherPair(k)
+    // Pass undefined to mean "start from beginning" (BTree's native behavior)
+    return this.takeInternal(n, nextPair, undefined, filterFn)
+  }
+
+  /**
+   * Returns the next n items **before** the provided item (in descending order).
+   * @param n - The number of items to return
+   * @param from - The item to start from (exclusive). Required.
+   * @returns The next n items **before** the provided key.
    */
   takeReversed(
     n: number,
-    from?: any,
+    from: any,
+    filterFn?: (key: TKey) => boolean,
+  ): Array<TKey> {
+    const nextPair = (k?: any) => this.orderedEntries.nextLowerPair(k)
+    // Normalize the from value
+    const normalizedFrom = normalizeForBTree(from)
+    return this.takeInternal(n, nextPair, normalizedFrom, filterFn, true)
+  }
+
+  /**
+   * Returns the last n items from the end.
+   * @param n - The number of items to return
+   * @param filterFn - Optional filter function
+   * @returns The last n items
+   */
+  takeReversedFromEnd(
+    n: number,
     filterFn?: (key: TKey) => boolean,
   ): Array<TKey> {
     const nextPair = (k?: any) => this.orderedEntries.nextLowerPair(k)
-    return this.takeInternal(n, nextPair, from, filterFn, true)
+    // Pass undefined to mean "start from end" (BTree's native behavior)
+    return this.takeInternal(n, nextPair, undefined, filterFn, true)
   }
 
   /**
@@ -323,7 +386,7 @@ export class BTreeIndex<
     const result = new Set<TKey>()
 
     for (const value of values) {
-      const normalizedValue = normalizeValue(value)
+      const normalizedValue = normalizeForBTree(value)
       const keys = this.valueMap.get(normalizedValue)
       if (keys) {
         keys.forEach((key) => result.add(key))
@@ -341,17 +404,25 @@ export class BTreeIndex<
   get orderedEntriesArray(): Array<[any, Set<TKey>]> {
     return this.orderedEntries
       .keysArray()
-      .map((key) => [key, this.valueMap.get(key) ?? new Set()])
+      .map((key) => [
+        denormalizeUndefined(key),
+        this.valueMap.get(key) ?? new Set(),
+      ])
   }
 
   get orderedEntriesArrayReversed(): Array<[any, Set<TKey>]> {
-    return this.takeReversed(this.orderedEntries.size).map((key) => [
-      key,
+    return this.takeReversedFromEnd(this.orderedEntries.size).map((key) => [
+      denormalizeUndefined(key),
       this.valueMap.get(key) ?? new Set(),
     ])
   }
 
   get valueMapData(): Map<any, Set<TKey>> {
-    return this.valueMap
+    // Return a new Map with denormalized keys
+    const result = new Map<any, Set<TKey>>()
+    for (const [key, value] of this.valueMap) {
+      result.set(denormalizeUndefined(key), value)
+    }
+    return result
   }
 }