feat(util): 🎸 improve AvlMap iteration strategies

Author: streamich

src/util/trees/avl/AvlMap.ts

@@ -105,20 +105,38 @@ export class AvlMap<K, V> implements Printable {
     while ((curr = next(curr as HeadlessNode) as AvlNode<K, V> | undefined));
   }
 
-  public iterator(): Iterator<AvlNode<K, V>> {
+  public first(): AvlNode<K, V> | undefined {
     const root = this.root;
-    if (!root) return {next: () => ({done: true, value: undefined})};
-    let curr = first(root);
+    return root ? first(root) : undefined;
+  }
+
+  public readonly next = next;
+
+  public iterator0(): (() => undefined | AvlNode<K, V>) {
+    let curr = this.first();
+    return () => {
+      if (!curr) return undefined;
+      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: () => {
-        if (!curr) return {done: true, value: undefined};
-        const value = curr;
-        curr = next(curr as HeadlessNode) as AvlNode<K, V> | undefined;
-        return {done: false, value};
+        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/__tests__/AvlMap.spec.ts

@@ -15,6 +15,47 @@ test('smoke test', () => {
   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', () => {
@@ -37,3 +78,17 @@ describe('.iterator()', () => {
     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]]);
+  });
+});