AVL map and set iterator improvements

src/util/buffers/cmpUint8Array2.ts

@@ -1,3 +1,11 @@
+/**
+ * Compares two `Uint8Arrays` byte-by-byte. Returns a negative number if `a` is
+ * less than `b`, a positive number if `a` is greater than `b`, or 0 if `a` is
+ * equal to `b`.
+ *
+ * @returns A negative number if a is less than b, a positive number if a is
+ *         greater than b, or 0 if a is equal to b.
+ */
 export const cmpUint8Array2 = (a: Uint8Array, b: Uint8Array): number => {
   const len1 = a.length;
   const len2 = b.length;

src/util/buffers/cmpUint8Array3.ts

@@ -0,0 +1,19 @@
+/**
+ * Compares two `Uint8Arrays`, first by length, then by each byte. Returns a
+ * negative number if `a` is less than `b`, a positive number if `a` is greater
+ * than `b`, or 0 if `a` is equal to `b`.
+ *
+ * @returns A negative number if a is less than b, a positive number if a is
+ *          greater than b, or 0 if a is equal to b.
+ */
+export const cmpUint8Array3 = (a: Uint8Array, b: Uint8Array): number => {
+  const len1 = a.length;
+  const len2 = b.length;
+  const diff = len1 - len2;
+  if (diff !== 0) return diff;
+  for (let i = 0; i < len1; i++) {
+    const diff = a[i] - b[i];
+    if (diff !== 0) return diff;
+  }
+  return 0;
+};

src/util/trees/avl/AvlMap.ts

@@ -98,13 +98,44 @@ export class AvlMap<K, V> implements Printable {
   }
 
   public forEach(fn: (node: AvlNode<K, V>) => void): void {
-    const root = this.root;
-    if (!root) return;
-    let curr = first(root);
+    let curr = this.first();
+    if (!curr) return;
     do fn(curr!);
     while ((curr = next(curr as HeadlessNode) as AvlNode<K, V> | undefined));
   }
 
+  public first(): AvlNode<K, V> | undefined {
+    const root = this.root;
+    return root ? first(root) : undefined;
+  }
+
+  public readonly next = next;
+
+  public iterator0(): () => undefined | AvlNode<K, V> {
+    let curr = this.first();
+    return () => {
+      if (!curr) return;
+      const value = curr;
+      curr = next(curr as HeadlessNode) as AvlNode<K, V> | undefined;
+      return value;
+    };
+  }
+
+  public iterator(): Iterator<AvlNode<K, V>> {
+    const iterator = this.iterator0();
+    return {
+      next: () => {
+        const value = iterator();
+        const res = <IteratorResult<AvlNode<K, V>>>{value, done: !value};
+        return res;
+      },
+    };
+  }
+
+  public entries(): IterableIterator<AvlNode<K, V>> {
+    return <any>{[Symbol.iterator]: () => this.iterator()};
+  }
+
   public toString(tab: string): string {
     return this.constructor.name + printTree(tab, [(tab) => print(this.root, tab)]);
   }

src/util/trees/avl/AvlSet.ts

@@ -91,13 +91,44 @@ export class AvlSet<V> implements Printable {
   }
 
   public forEach(fn: (node: AvlSetNode<V>) => void): void {
-    const root = this.root;
-    if (!root) return;
-    let curr = first(root);
+    let curr = this.first();
+    if (!curr) return;
     do fn(curr!);
     while ((curr = next(curr as HeadlessNode) as AvlSetNode<V> | undefined));
   }
 
+  public first(): AvlSetNode<V> | undefined {
+    const root = this.root;
+    return root ? first(root) : undefined;
+  }
+
+  public readonly next = next;
+
+  public iterator0(): () => undefined | AvlSetNode<V> {
+    let curr = this.first();
+    return () => {
+      if (!curr) return undefined;
+      const value = curr;
+      curr = next(curr as HeadlessNode) as AvlSetNode<V> | undefined;
+      return value;
+    };
+  }
+
+  public iterator(): Iterator<AvlSetNode<V>> {
+    const iterator = this.iterator0();
+    return {
+      next: () => {
+        const value = iterator();
+        const res = <IteratorResult<AvlSetNode<V>>>{value, done: !value};
+        return res;
+      },
+    };
+  }
+
+  public entries(): IterableIterator<AvlSetNode<V>> {
+    return <any>{[Symbol.iterator]: () => this.iterator()};
+  }
+
   public toString(tab: string): string {
     return this.constructor.name + printTree(tab, [(tab) => print(this.root, tab)]);
   }

src/util/trees/avl/__tests__/AvlMap.spec.ts

@@ -14,3 +14,97 @@ test('smoke test', () => {
   tree.forEach((node) => keys.push(node.k));
   expect(keys).toEqual([1, 3, 4, 4.1, 44]);
 });
+
+describe('.first()/next() iteration', () => {
+  test('for empty map, returns finished iterator', () => {
+    const tree = new AvlMap<string, number>();
+    const entry = tree.first();
+    expect(entry).toEqual(undefined);
+  });
+
+  test('can iterate through map entries', () => {
+    const tree = new AvlMap<string, number>();
+    tree.set('a', 1);
+    tree.set('b', 2);
+    tree.set('c', 3);
+    const list: [string, number][] = [];
+    for (let entry = tree.first(); entry; entry = tree.next(entry)) {
+      list.push([entry.k, entry.v]);
+    }
+    expect(list).toEqual([
+      ['a', 1],
+      ['b', 2],
+      ['c', 3],
+    ]);
+  });
+});
+
+describe('.iterator0()', () => {
+  test('for empty map, returns finished iterator', () => {
+    const tree = new AvlMap<string, number>();
+    const iterator = tree.iterator0();
+    const entry = iterator();
+    expect(entry).toEqual(undefined);
+  });
+
+  test('can iterate through map entries', () => {
+    const tree = new AvlMap<string, number>();
+    tree.set('a', 1);
+    tree.set('b', 2);
+    tree.set('c', 3);
+    const list: [string, number][] = [];
+    const iterator = tree.iterator0();
+    for (let entry = iterator(); entry; entry = iterator()) {
+      list.push([entry.k, entry.v]);
+    }
+    expect(list).toEqual([
+      ['a', 1],
+      ['b', 2],
+      ['c', 3],
+    ]);
+  });
+});
+
+describe('.iterator()', () => {
+  test('for empty map, returns finished iterator', () => {
+    const tree = new AvlMap<string, number>();
+    const iterator = tree.iterator();
+    const entry = iterator.next();
+    expect(entry).toEqual({done: true, value: undefined});
+  });
+
+  test('can iterate through map entries', () => {
+    const tree = new AvlMap<string, number>();
+    tree.set('a', 1);
+    tree.set('b', 2);
+    tree.set('c', 3);
+    const iterator = tree.iterator();
+    const list: [string, number][] = [];
+    for (let entry = iterator.next(); !entry.done; entry = iterator.next()) {
+      list.push([entry.value!.k, entry.value!.v]);
+    }
+    expect(list).toEqual([
+      ['a', 1],
+      ['b', 2],
+      ['c', 3],
+    ]);
+  });
+});
+
+describe('for...of iteration', () => {
+  test('can iterate through map entries', () => {
+    const tree = new AvlMap<string, number>();
+    tree.set('a', 1);
+    tree.set('b', 2);
+    tree.set('c', 3);
+    const list: [string, number][] = [];
+    for (const entry of tree.entries()) {
+      list.push([entry.k, entry.v]);
+    }
+    expect(list).toEqual([
+      ['a', 1],
+      ['b', 2],
+      ['c', 3],
+    ]);
+  });
+});

src/util/trees/avl/__tests__/AvlSet.spec.ts

@@ -64,3 +64,81 @@ test('can store structs', () => {
   expect(set.has(new Struct(3, 3))).toBe(true);
   expect(set.has(new Struct(2, 3))).toBe(true);
 });
+
+describe('.first()/next() iteration', () => {
+  test('for empty map, returns finished iterator', () => {
+    const tree = new AvlSet<string>();
+    const entry = tree.first();
+    expect(entry).toEqual(undefined);
+  });
+
+  test('can iterate through map entries', () => {
+    const tree = new AvlSet<number>();
+    tree.add(1);
+    tree.add(2);
+    tree.add(3);
+    const list: number[] = [];
+    for (let entry = tree.first(); entry; entry = tree.next(entry)) {
+      list.push(entry.k);
+    }
+    expect(list).toEqual([1, 2, 3]);
+  });
+});
+
+describe('.iterator0()', () => {
+  test('for empty map, returns finished iterator', () => {
+    const tree = new AvlSet<number>();
+    const iterator = tree.iterator0();
+    const entry = iterator();
+    expect(entry).toEqual(undefined);
+  });
+
+  test('can iterate through map entries', () => {
+    const tree = new AvlSet<number>();
+    tree.add(1);
+    tree.add(2);
+    tree.add(3);
+    const list: number[] = [];
+    const iterator = tree.iterator0();
+    for (let entry = iterator(); entry; entry = iterator()) {
+      list.push(entry.k);
+    }
+    expect(list).toEqual([1, 2, 3]);
+  });
+});
+
+describe('.iterator()', () => {
+  test('for empty map, returns finished iterator', () => {
+    const tree = new AvlSet<number>();
+    const iterator = tree.iterator();
+    const entry = iterator.next();
+    expect(entry).toEqual({done: true, value: undefined});
+  });
+
+  test('can iterate through map entries', () => {
+    const tree = new AvlSet<number>();
+    tree.add(1);
+    tree.add(2);
+    tree.add(3);
+    const list: number[] = [];
+    const iterator = tree.iterator();
+    for (let entry = iterator.next(); !entry.done; entry = iterator.next()) {
+      list.push(entry.value!.k);
+    }
+    expect(list).toEqual([1, 2, 3]);
+  });
+});
+
+describe('for...of iteration', () => {
+  test('can iterate through map entries', () => {
+    const tree = new AvlSet<number>();
+    tree.add(1);
+    tree.add(2);
+    tree.add(3);
+    const list: number[] = [];
+    for (const entry of tree.entries()) {
+      list.push(entry.k);
+    }
+    expect(list).toEqual([1, 2, 3]);
+  });
+});