feat(util): 🎸 improve AvlSet iteration strategies

Author: streamich

src/util/trees/avl/AvlMap.ts

@@ -98,9 +98,8 @@ 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));
   }
@@ -115,7 +114,7 @@ export class AvlMap<K, V> implements Printable {
   public iterator0(): (() => undefined | AvlNode<K, V>) {
     let curr = this.first();
     return () => {
-      if (!curr) return undefined;
+      if (!curr) return;
       const value = curr;
       curr = next(curr as HeadlessNode) as AvlNode<K, V> | undefined;
       return value;

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__/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]);
+  });
+});