Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
124 changes: 55 additions & 69 deletions packages/db/src/SortedMap.ts
Original file line number Diff line number Diff line change
@@ -1,4 +1,10 @@
import { compareKeys } from '@tanstack/db-ivm'
import { BTree } from './utils/btree.js'

type SortKey<TKey extends string | number, TValue> = {
key: TKey
value: TValue
}

/**
* A Map implementation that keeps its entries sorted based on a comparator function
Expand All @@ -7,7 +13,7 @@ import { compareKeys } from '@tanstack/db-ivm'
*/
export class SortedMap<TKey extends string | number, TValue> {
private map: Map<TKey, TValue>
private sortedKeys: Array<TKey>
private sortedKeys: BTree<SortKey<TKey, TValue>, undefined>
private comparator: ((a: TValue, b: TValue) => number) | undefined

/**
Expand All @@ -18,68 +24,50 @@ export class SortedMap<TKey extends string | number, TValue> {
*/
constructor(comparator?: (a: TValue, b: TValue) => number) {
this.map = new Map<TKey, TValue>()
this.sortedKeys = []
this.comparator = comparator
this.sortedKeys = new BTree<SortKey<TKey, TValue>, undefined>(
(left, right) => this.compareSortKeys(left, right),
)
}

/**
* Finds the index where a key-value pair should be inserted to maintain sort order.
* Uses binary search to find the correct position based on the value (if comparator provided),
* with key-based tie-breaking for deterministic ordering when values compare as equal.
* If no comparator is provided, sorts by key only.
* Runs in O(log n) time.
* Compares sort keys based on value first when a comparator is provided,
* falling back to the collection key for deterministic tie-breaking.
*
* @param key - The key to find position for (used as tie-breaker or primary sort when no comparator)
* @param value - The value to compare against (only used if comparator is provided)
* @returns The index where the key should be inserted
* If no comparator is provided, entries are ordered by key only.
*/
private indexOf(key: TKey, value: TValue): number {
let left = 0
let right = this.sortedKeys.length

// Fast path: no comparator means sort by key only
private compareSortKeys(
left: SortKey<TKey, TValue>,
right: SortKey<TKey, TValue>,
): number {
if (!this.comparator) {
while (left < right) {
const mid = Math.floor((left + right) / 2)
const midKey = this.sortedKeys[mid]!
const keyComparison = compareKeys(key, midKey)
if (keyComparison < 0) {
right = mid
} else if (keyComparison > 0) {
left = mid + 1
} else {
return mid
}
}
return left
return compareKeys(left.key, right.key)
}

// With comparator: sort by value first, then key as tie-breaker
while (left < right) {
const mid = Math.floor((left + right) / 2)
const midKey = this.sortedKeys[mid]!
const midValue = this.map.get(midKey)!
const valueComparison = this.comparator(value, midValue)

if (valueComparison < 0) {
right = mid
} else if (valueComparison > 0) {
left = mid + 1
} else {
// Values are equal, use key as tie-breaker for deterministic ordering
const keyComparison = compareKeys(key, midKey)
if (keyComparison < 0) {
right = mid
} else if (keyComparison > 0) {
left = mid + 1
} else {
// Same key (shouldn't happen during insert, but handle for lookups)
return mid
}
}
const valueComparison = this.comparator(left.value, right.value)
if (valueComparison !== 0) {
return valueComparison
}

return left
return compareKeys(left.key, right.key)
}

private createSortKey(key: TKey, value: TValue): SortKey<TKey, TValue> {
return { key, value }
}
Comment on lines +39 to +57
Copy link
Copy Markdown
Contributor

@coderabbitai coderabbitai Bot May 21, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major | 🏗️ Heavy lift

Avoid using live mutable values as BTree ordering keys.

Line 55 stores the original TValue inside the tree key, and Line 47 compares on that live object. If a caller mutates a comparator-observed field after set(), the ordering of an already-inserted BTree key changes without a delete/reinsert, which can invalidate traversal/search/delete behavior. Please switch this to an immutable sort snapshot, or explicitly enforce/document that comparator-relevant value state must not mutate while the entry is in the map, and add a regression test for that contract.

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@packages/db/src/SortedMap.ts` around lines 39 - 57, The code currently stores
the live TValue in SortKey via createSortKey and uses it in compareSortKeys,
which means mutating that TValue after insertion can silently change tree
ordering; modify createSortKey/SortKey so the tree stores an immutable snapshot
used for ordering (e.g., derive and store a frozen/primitive sort token or a
shallow/deep-cloned comparator-relevant snapshot instead of the original
TValue), update compareSortKeys to compare that immutable snapshot (or derived
token) rather than the live value, and add a regression test that inserts an
entry, mutates the original TValue, then verifies traversal/search/delete
behavior is unchanged unless the entry is explicitly reinserted/deleted;
reference createSortKey, SortKey and compareSortKeys when making the change.

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Took a pass on this and verified the concern against current SortedMap.

I tried the suggested direction by storing an immutable comparator snapshot in SortKey and adding a regression test for post-insert mutation. That change did address the specific mutation hazard, but it also caused broad ordering regressions across existing query/order-by behavior in the current codebase, so it isn’t safe to land as-is here.

Because the review request was to fix only still-valid issues and skip the rest, I reverted that attempt and kept the branch on the validated O(N²) insert fix only. The branch is back to green with the full @tanstack/db test suite passing.

If we want to pursue this further, I think it needs to be handled as a separate follow-up with a narrower design that preserves current ordering semantics.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Seems like the humans are having a chat. I'll hop back into my burrow for now. If you need me again, just tag @coderabbitai in a new comment, and I'll come hopping out!


private *iterateSortKeys(): IterableIterator<SortKey<TKey, TValue>> {
let previous: SortKey<TKey, TValue> | undefined

for (;;) {
const nextPair = this.sortedKeys.nextHigherPair(previous)
if (!nextPair) {
return
}

previous = nextPair[0]
yield previous
}
}

/**
Expand All @@ -91,16 +79,11 @@ export class SortedMap<TKey extends string | number, TValue> {
*/
set(key: TKey, value: TValue): this {
if (this.map.has(key)) {
// Need to remove the old key from the sorted keys array
const oldValue = this.map.get(key)!
const oldIndex = this.indexOf(key, oldValue)
this.sortedKeys.splice(oldIndex, 1)
this.sortedKeys.delete(this.createSortKey(key, oldValue))
}

// Insert the new key at the correct position
const index = this.indexOf(key, value)
this.sortedKeys.splice(index, 0, key)

this.sortedKeys.set(this.createSortKey(key, value), undefined)
this.map.set(key, value)

return this
Expand All @@ -125,8 +108,7 @@ export class SortedMap<TKey extends string | number, TValue> {
delete(key: TKey): boolean {
if (this.map.has(key)) {
const oldValue = this.map.get(key)
const index = this.indexOf(key, oldValue!)
this.sortedKeys.splice(index, 1)
this.sortedKeys.delete(this.createSortKey(key, oldValue!))
return this.map.delete(key)
}

Expand All @@ -148,7 +130,7 @@ export class SortedMap<TKey extends string | number, TValue> {
*/
clear(): void {
this.map.clear()
this.sortedKeys = []
this.sortedKeys.clear()
}

/**
Expand All @@ -164,8 +146,8 @@ export class SortedMap<TKey extends string | number, TValue> {
* @returns An iterator for the map's entries
*/
*[Symbol.iterator](): IterableIterator<[TKey, TValue]> {
for (const key of this.sortedKeys) {
yield [key, this.map.get(key)!] as [TKey, TValue]
for (const sortKey of this.iterateSortKeys()) {
yield [sortKey.key, this.map.get(sortKey.key)!] as [TKey, TValue]
}
}

Expand All @@ -184,7 +166,11 @@ export class SortedMap<TKey extends string | number, TValue> {
* @returns An iterator for the map's keys
*/
keys(): IterableIterator<TKey> {
return this.sortedKeys[Symbol.iterator]()
return function* (this: SortedMap<TKey, TValue>) {
for (const sortKey of this.iterateSortKeys()) {
yield sortKey.key
}
}.call(this)
}

/**
Expand All @@ -194,8 +180,8 @@ export class SortedMap<TKey extends string | number, TValue> {
*/
values(): IterableIterator<TValue> {
return function* (this: SortedMap<TKey, TValue>) {
for (const key of this.sortedKeys) {
yield this.map.get(key)!
for (const sortKey of this.iterateSortKeys()) {
yield this.map.get(sortKey.key)!
}
}.call(this)
}
Expand All @@ -208,8 +194,8 @@ export class SortedMap<TKey extends string | number, TValue> {
forEach(
callbackfn: (value: TValue, key: TKey, map: Map<TKey, TValue>) => void,
): void {
for (const key of this.sortedKeys) {
callbackfn(this.map.get(key)!, key, this.map)
for (const sortKey of this.iterateSortKeys()) {
callbackfn(this.map.get(sortKey.key)!, sortKey.key, this.map)
}
}
}
31 changes: 31 additions & 0 deletions packages/db/tests/SortedMap.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,24 @@ describe(`SortedMap`, () => {
expect(values).toEqual([2, 3])
})

it(`reorders entries when a custom comparator observes an updated value`, () => {
const map = new SortedMap<string, { priority: number }>(
(a, b) => a.priority - b.priority,
)

map.set(`a`, { priority: 1 })
map.set(`b`, { priority: 3 })
map.set(`c`, { priority: 2 })
map.set(`b`, { priority: 0 })

expect(Array.from(map.keys())).toEqual([`b`, `a`, `c`])
expect(Array.from(map.values())).toEqual([
{ priority: 0 },
{ priority: 1 },
{ priority: 2 },
])
})

it(`correctly handles deletions`, () => {
const map = new SortedMap<string, number>()
map.set(`a`, 1)
Expand Down Expand Up @@ -102,6 +120,19 @@ describe(`SortedMap`, () => {
expect(keys).toEqual([`a`, `b`, `c`])
})

it(`keeps deterministic key order for equal comparator values after updates`, () => {
const map = new SortedMap<string, { priority: number }>(
(a, b) => a.priority - b.priority,
)

map.set(`c`, { priority: 1 })
map.set(`a`, { priority: 1 })
map.set(`b`, { priority: 2 })
map.set(`b`, { priority: 1 })

expect(Array.from(map.keys())).toEqual([`a`, `b`, `c`])
})

// Test for Symbol.iterator implementation
it(`supports direct iteration with for...of`, () => {
const map = new SortedMap<string, number>()
Expand Down