Skip to content

Fix SortedMap's reverse iterator (without start key). #6020

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Feb 24, 2022
Merged
Show file tree
Hide file tree
Changes from 2 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
2 changes: 1 addition & 1 deletion packages/firestore/src/util/sorted_map.ts
Original file line number Diff line number Diff line change
Expand Up @@ -197,7 +197,7 @@ export class SortedMapIterator<K, V> {
while (!node.isEmpty()) {
cmp = startKey ? comparator(node.key, startKey) : 1;
// flip the comparison if we're going in reverse
if (isReverse) {
if (startKey && isReverse) {
Copy link
Contributor

Choose a reason for hiding this comment

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

I can't get my head around this - why do we not need to do this when we iterate without start key? I can debug this to understand more, but I wonder if you have a succinct explanation why this is the correct fix.

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Sure. It boils down to the iteration logic below (line 204-227).

First, description of the bug:

If this constructor is called with isReverse = true and startKey = null, then cmp will always be -1 in all iterations of the while loop. Notice that this.nodeStack.push is never called if cmp < 0. So, reverse traversal with no startKey will always result in iterating zero nodes.

More explanation

The purpose of this while loop is to initialize a stack that will be used for traversal. The while loop always starts at the root node of a red-black tree.

In the absence of a startKey,

  • For forward iteration (iterating from smallest node to largest node): the while loop should start from the root and keep going to the left child repeatedly until we reach the left-most leaf node (pushing these nodes onto the stack along the way).
  • For reverse iteration (iterating from largest node to smallest node): the while loop should start from the root and keep going to the right child repeatedly until we reach the right-most leaf node (pushing these nodes onto the stack along the way).

In the presence of a startKey,

We should start at the root of the tree, and first search for the startKey. We should not push anything onto the stack until we find the startKey (or find a node that's larger than startKey for forward iteration, or find a node that's smaller than startKey for reverse iteration). At that point, we should start pushing nodes onto the stack.

Copy link
Contributor Author

@ehsannas ehsannas Feb 23, 2022

Choose a reason for hiding this comment

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

In other words, the code can also be written as follows:

if(startKey) {
  while (!node.isEmpty()) {
      cmp = comparator(node.key, startKey);
      // flip the comparison if we're going in reverse
      if (isReverse) {
        cmp *= -1;
      }
      if (cmp < 0) {
        // This node is less than our start key. ignore it
        if (this.isReverse) {
          node = node.left;
        } else {
          node = node.right;
        }
      } else if (cmp === 0) {
        // This node is exactly equal to our start key. Push it on the stack,
        // but stop iterating;
        this.nodeStack.push(node);
        break;
      } else {
        // This node is greater than our start key, add it to the stack and move
        // to the next one
        this.nodeStack.push(node);
        if (this.isReverse) {
          node = node.right;
        } else {
          node = node.left;
        }
      }
  }
} else {
  // There is no startKey, go all the way.
  while (!node.isEmpty()) {
        this.nodeStack.push(node);
        if (this.isReverse) {
          node = node.right;
        } else {
          node = node.left;
        }
  }
}

Copy link
Contributor

Choose a reason for hiding this comment

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

Okay, thanks for the explanation. I didn't quite catch that this same code was doing two completely different things (finding the start key and then building the stack). I think we could probably fix this to run two iterations, but that's a P3 and this PR is good to ship.

cmp *= -1;
}

Expand Down
98 changes: 78 additions & 20 deletions packages/firestore/test/unit/util/sorted_map.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,11 @@ import { expect } from 'chai';

import { primitiveComparator } from '../../../src/util/misc';
import { forEach } from '../../../src/util/obj';
import { LLRBNode, SortedMap } from '../../../src/util/sorted_map';
import {
LLRBNode,
SortedMap,
SortedMapIterator
} from '../../../src/util/sorted_map';

function shuffle(arr: number[]): void {
for (let i = arr.length - 1; i > 0; i--) {
Expand All @@ -32,6 +36,41 @@ function shuffle(arr: number[]): void {
// Many of these were adapted from the mugs source code.
// http://mads379.github.com/mugs/
describe('SortedMap Tests', () => {
/** Returns a map with (i, i) entries for i in [0, 9] inclusive */
function getMapOfNumbersZeroToNine(): SortedMap<number, number> {
const items: number[] = [];
for (let i = 0; i < 10; i++) {
items.push(i);
}
shuffle(items);

let map = new SortedMap<number, number>(primitiveComparator);
for (let i = 0; i < 10; i++) {
map = map.insert(items[i], items[i]);
}
return map;
}

/**
* Validates that the given iterator covers the given range by counting from
* the given `from` argument to the given `to` argument (inclusive).
*/
function validateIteratesRangeInclusive(
iterator: SortedMapIterator<number, number>,
from: number,
to: number
): void {
const countUp = from <= to;
let expected = from;
while (iterator.hasNext()) {
Copy link
Contributor

Choose a reason for hiding this comment

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

FWIW From the method contract, it seems like this method should iterate until it reaches to and validate that hasNext remains true.

Copy link
Contributor Author

Choose a reason for hiding this comment

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

will rename to validateIteratorCoversRange

const n = iterator.getNext();
expect(n.key).to.equal(expected);
expect(n.value).to.equal(expected);
expected = countUp ? expected + 1 : expected - 1;
}
expect(expected).to.equal(countUp ? to + 1 : to - 1);
}

it('Create node', () => {
const map = new SortedMap(primitiveComparator).insert('key', 'value');
expect(map.root.left.isEmpty()).to.equal(true);
Expand Down Expand Up @@ -405,27 +444,46 @@ describe('SortedMap Tests', () => {
expect(() => iterator.getNext()).to.throw();
});

it('SortedMapIterator test with 10 items.', () => {
const items: number[] = [];
for (let i = 0; i < 10; i++) {
items.push(i);
}
shuffle(items);
it('forward iterator without start key', () => {
const iterator = getMapOfNumbersZeroToNine().getIterator();
validateIteratesRangeInclusive(iterator, 0, 9);
});

let map = new SortedMap(primitiveComparator);
for (let i = 0; i < 10; i++) {
map = map.insert(items[i], items[i]);
}
it('forward iterator with start key.', () => {
const iterator = getMapOfNumbersZeroToNine().getIteratorFrom(5);
validateIteratesRangeInclusive(iterator, 5, 9);
});

const iterator = map.getIterator();
let expected = 0;
while (iterator.hasNext()) {
const n = iterator.getNext();
expect(n.key).to.equal(expected);
expect(n.value).to.equal(expected);
expected++;
}
expect(expected).to.equal(10);
it('forward iterator with start key larger than max key', () => {
const iterator = getMapOfNumbersZeroToNine().getIteratorFrom(50);
expect(iterator.hasNext()).to.equal(false);
expect(() => iterator.getNext()).to.throw();
});

it('forward iterator with start key smaller than min key', () => {
const iterator = getMapOfNumbersZeroToNine().getIteratorFrom(-50);
validateIteratesRangeInclusive(iterator, 0, 9);
});

it('reverse iterator without start key', () => {
const iterator = getMapOfNumbersZeroToNine().getReverseIterator();
validateIteratesRangeInclusive(iterator, 9, 0);
});

it('reverse iterator with start key.', () => {
const iterator = getMapOfNumbersZeroToNine().getReverseIteratorFrom(7);
validateIteratesRangeInclusive(iterator, 7, 0);
});

it('reverse iterator with start key smaller than min key', () => {
const iterator = getMapOfNumbersZeroToNine().getReverseIteratorFrom(-50);
expect(iterator.hasNext()).to.equal(false);
expect(() => iterator.getNext()).to.throw();
});

it('reverse iterator with start key larger than max key', () => {
const iterator = getMapOfNumbersZeroToNine().getReverseIteratorFrom(50);
validateIteratesRangeInclusive(iterator, 9, 0);
});

it('SortedMap.indexOf returns index.', () => {
Expand Down