Allow iterables in PersistencePromise.map/forEach/waitFor (#1269)

Author: schmidt-sebastian

packages/firestore/src/core/sync_engine.ts

@@ -684,7 +684,7 @@ export class SyncEngine implements RemoteSyncer, SharedClientStateSyncer {
         queryView.targetId
       );
       this.limboDocumentRefs.removeReferencesForId(queryView.targetId);
-      await PersistencePromise.forEach(limboKeys.toArray(), limboKey => {
+      await PersistencePromise.forEach(limboKeys, limboKey => {
         return this.limboDocumentRefs
           .containsKey(null, limboKey)
           .next(isReferenced => {

packages/firestore/src/local/persistence_promise.ts

@@ -170,49 +170,54 @@ export class PersistencePromise<T> {
 
   static waitFor(
     // tslint:disable-next-line:no-any Accept all Promise types in waitFor().
-    all: Array<PersistencePromise<any>>
+    all: Iterable<PersistencePromise<any>>
   ): PersistencePromise<void> {
-    const expectedCount = all.length;
-    if (expectedCount === 0) {
-      return PersistencePromise.resolve();
-    }
-
-    let resolvedCount = 0;
+    const it = all[Symbol.iterator]();
     return new PersistencePromise<void>((resolve, reject) => {
-      for (const promise of all) {
-        promise.next(
+      let expectedCount = 0;
+      let resolvedCount = 0;
+
+      let result = it.next();
+      while (!result.done) {
+        ++expectedCount;
+        result.value.next(
           () => {
             ++resolvedCount;
-            if (resolvedCount === expectedCount) {
+            if (result.done && resolvedCount === expectedCount) {
               resolve();
             }
           },
           err => reject(err)
         );
+        result = it.next();
+      }
+
+      if (resolvedCount === expectedCount) {
+        resolve();
       }
     });
   }
 
-  static map<R>(all: Array<PersistencePromise<R>>): PersistencePromise<R[]> {
+  static map<R>(all: Iterable<PersistencePromise<R>>): PersistencePromise<R[]> {
     const results: R[] = [];
-    const promises: Array<PersistencePromise<void>> = [];
-    for (let i = 0; i < all.length; ++i) {
-      promises[i] = all[i].next(result => {
-        results[i] = result;
-      });
-    }
-    return PersistencePromise.waitFor(promises).next(() => {
-      return results;
-    });
+    return PersistencePromise.forEach(all, result => {
+      results.push(result);
+      return PersistencePromise.resolve();
+    }).next(() => results);
   }
 
   static forEach<T>(
-    elements: T[],
+    all: Iterable<T>,
     callback: (T) => PersistencePromise<void>
   ): PersistencePromise<void> {
+    const it = all[Symbol.iterator]();
+
     let p = PersistencePromise.resolve();
-    for (const element of elements) {
-      p = p.next(() => callback(element));
+    let result = it.next();
+    while (!result.done) {
+      const value = result.value;
+      p = p.next(() => callback(value));
+      result = it.next();
     }
     return p;
   }

packages/firestore/src/util/sorted_map.ts

@@ -168,6 +168,21 @@ export class SortedMap<K, V> {
   getReverseIteratorFrom(key: K): SortedMapIterator<K, V> {
     return new SortedMapIterator<K, V>(this.root, key, this.comparator, true);
   }
+
+  [Symbol.iterator](): Iterator<Entry<K, V>> {
+    const it = this.getIterator();
+    return {
+      next(): IteratorResult<Entry<K, V>> {
+        if (it.hasNext()) {
+          return { done: false, value: it.getNext() };
+        } else {
+          // The TypeScript typings don't allow `undefined` for Iterator<T>,
+          // so we return an empty object instead.
+          return { done: true, value: {} as Entry<K, V> };
+        }
+      }
+    };
+  }
 } // end SortedMap
 
 // An iterator over an LLRBNode.

packages/firestore/src/util/sorted_set.ts

@@ -154,6 +154,21 @@ export class SortedSet<T> {
     return 'SortedSet(' + result.toString() + ')';
   }
 
+  [Symbol.iterator](): Iterator<T> {
+    const it = this.data.getIterator();
+    return {
+      next(): IteratorResult<T> {
+        if (it.hasNext()) {
+          return { done: false, value: it.getNext().key };
+        } else {
+          // The TypeScript typings don't allow `undefined` for Iterator<T>,
+          // so we return an empty object instead.
+          return { done: true, value: {} as T };
+        }
+      }
+    };
+  }
+
   private copy(data: SortedMap<T, boolean>): SortedSet<T> {
     const result = new SortedSet(this.comparator);
     result.data = data;