{ Iterator, AsyncIterator }.prototype.flatMap supports returning bo…

…th - iterables and iterators tc39/proposal-iterator-helpers#233
zloirock · Oct 21, 2022 · f5f7df5 · f5f7df5
1 parent 43cfb4e
commit f5f7df5
Show file tree

Hide file tree

Showing 9 changed files with 68 additions and 70 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -9,6 +9,7 @@
 - Recent updates of the [iterator helpers proposal](https://github.com/tc39/proposal-iterator-helpers):
   - Added a counter parameter to helpers, [proposal-iterator-helpers/211](https://github.com/tc39/proposal-iterator-helpers/pull/211)
   - Don't await non-objects returned from functions passed to `AsyncIterator` helpers, [proposal-iterator-helpers/239](https://github.com/tc39/proposal-iterator-helpers/pull/239)
+  - `{ Iterator, AsyncIterator }.prototype.flatMap` supports returning both - iterables and iterators, [proposal-iterator-helpers/233](https://github.com/tc39/proposal-iterator-helpers/pull/233)
   - Early exit on broken `.next` in missed cases of `{ Iterator, AsyncIterator }.from`, [proposal-iterator-helpers/232](https://github.com/tc39/proposal-iterator-helpers/pull/232)
 - Added `inverse` option to `core-js-compat`, [#1119](https://github.com/zloirock/core-js/issues/1119)
 - Added `format` option to `core-js-builder`, [#1120](https://github.com/zloirock/core-js/issues/1120)

diff --git a/packages/core-js/internals/async-iterator-create-proxy.js b/packages/core-js/internals/async-iterator-create-proxy.js
@@ -79,11 +79,10 @@ var createAsyncIteratorProxyPrototype = function (IS_ITERATOR) {
       return exit ? state : enqueue(state, function () {
         state.done = true;
         var iterator = state.iterator;
-        var innerIterator = state.innerIterator;
         var returnMethod, result;
         var completion = perform(function () {
-          if (innerIterator) try {
-            iteratorClose(innerIterator, 'return');
+          if (state.inner) try {
+            iteratorClose(state.inner.iterator, 'return');
           } catch (error) {
             return iteratorClose(iterator, 'throw', error);
           }

diff --git a/packages/core-js/internals/get-async-iterator-flattenable.js b/packages/core-js/internals/get-async-iterator-flattenable.js
@@ -0,0 +1,28 @@
+var call = require('../internals/function-call');
+var isCallable = require('../internals/is-callable');
+var toObject = require('../internals/to-object');
+var getIteratorDirect = require('../internals/get-iterator-direct');
+var getIteratorMethod = require('../internals/get-iterator-method');
+var getMethod = require('../internals/get-method');
+var wellKnownSymbol = require('../internals/well-known-symbol');
+var AsyncFromSyncIterator = require('../internals/async-from-sync-iterator');
+
+var ASYNC_ITERATOR = wellKnownSymbol('asyncIterator');
+
+module.exports = function from(obj) {
+  var object = toObject(obj);
+  var alreadyAsync = true;
+  var method = getMethod(object, ASYNC_ITERATOR);
+  var iterator;
+  if (!isCallable(method)) {
+    method = getIteratorMethod(object);
+    alreadyAsync = false;
+  }
+  if (isCallable(method)) {
+    iterator = call(method, object);
+  } else {
+    iterator = object;
+    alreadyAsync = true;
+  }
+  return getIteratorDirect(alreadyAsync ? iterator : new AsyncFromSyncIterator(getIteratorDirect(iterator)));
+};
diff --git a/packages/core-js/internals/get-iterator-flattenable.js b/packages/core-js/internals/get-iterator-flattenable.js
@@ -0,0 +1,11 @@
+var call = require('../internals/function-call');
+var isCallable = require('../internals/is-callable');
+var toObject = require('../internals/to-object');
+var getIteratorDirect = require('../internals/get-iterator-direct');
+var getIteratorMethod = require('../internals/get-iterator-method');
+
+module.exports = function (obj) {
+  var object = toObject(obj);
+  var method = getIteratorMethod(object);
+  return getIteratorDirect(isCallable(method) ? call(method, object) : object);
+};
diff --git a/packages/core-js/internals/iterator-create-proxy.js b/packages/core-js/internals/iterator-create-proxy.js
@@ -44,9 +44,8 @@ var createIteratorProxyPrototype = function (IS_ITERATOR) {
         var returnMethod = getMethod(iterator, 'return');
         return returnMethod ? call(returnMethod, iterator) : createIterResultObject(undefined, true);
       }
-      var innerIterator = state.innerIterator;
-      if (innerIterator) try {
-        iteratorClose(innerIterator, 'return');
+      if (state.inner) try {
+        iteratorClose(state.inner.iterator, 'return');
       } catch (error) {
         return iteratorClose(iterator, 'throw', error);
       }

diff --git a/packages/core-js/modules/esnext.async-iterator.flat-map.js b/packages/core-js/modules/esnext.async-iterator.flat-map.js
@@ -8,14 +8,13 @@ var isObject = require('../internals/is-object');
 var getIteratorDirect = require('../internals/get-iterator-direct');
 var createAsyncIteratorProxy = require('../internals/async-iterator-create-proxy');
 var createIterResultObject = require('../internals/create-iter-result-object');
-var getAsyncIterator = require('../internals/get-async-iterator');
+var getAsyncIteratorFlattenable = require('../internals/get-async-iterator-flattenable');
 var closeAsyncIteration = require('../internals/async-iterator-close');
 
 var AsyncIteratorProxy = createAsyncIteratorProxy(function (Promise) {
   var state = this;
   var iterator = state.iterator;
   var mapper = state.mapper;
-  var innerIterator;
 
   return new Promise(function (resolve, reject) {
     var doneAndReject = function (error) {
@@ -41,8 +40,7 @@ var AsyncIteratorProxy = createAsyncIteratorProxy(function (Promise) {
 
                 var handler = function (mapped) {
                   try {
-                    state.innerIterator = innerIterator = getAsyncIterator(mapped);
-                    state.innerNext = aCallable(innerIterator.next);
+                    state.inner = getAsyncIteratorFlattenable(mapped);
                     innerLoop();
                   } catch (error4) { ifAbruptCloseAsyncIterator(error4); }
                 };
@@ -57,12 +55,13 @@ var AsyncIteratorProxy = createAsyncIteratorProxy(function (Promise) {
     };
 
     var innerLoop = function () {
-      if (innerIterator = state.innerIterator) {
+      var inner = state.inner;
+      if (inner) {
         try {
-          Promise.resolve(anObject(call(state.innerNext, innerIterator))).then(function (result) {
+          Promise.resolve(anObject(call(inner.next, inner.iterator))).then(function (result) {
             try {
               if (anObject(result).done) {
-                state.innerIterator = state.innerNext = null;
+                state.inner = null;
                 outerLoop();
               } else resolve(createIterResultObject(result.value, false));
             } catch (error1) { ifAbruptCloseAsyncIterator(error1); }
@@ -79,8 +78,7 @@ $({ target: 'AsyncIterator', proto: true, real: true, forced: true }, {
   flatMap: function flatMap(mapper) {
     return new AsyncIteratorProxy(getIteratorDirect(this), {
       mapper: aCallable(mapper),
-      innerIterator: null,
-      innerNext: null
+      inner: null
     });
   }
 });
diff --git a/packages/core-js/modules/esnext.async-iterator.from.js b/packages/core-js/modules/esnext.async-iterator.from.js
@@ -1,34 +1,15 @@
 // https://github.com/tc39/proposal-iterator-helpers
 var $ = require('../internals/export');
-var toObject = require('../internals/to-object');
 var isPrototypeOf = require('../internals/object-is-prototype-of');
+var getAsyncIteratorFlattenable = require('../internals/get-async-iterator-flattenable');
 var AsyncIteratorPrototype = require('../internals/async-iterator-prototype');
-var getAsyncIterator = require('../internals/get-async-iterator');
-var getIterator = require('../internals/get-iterator');
-var getIteratorDirect = require('../internals/get-iterator-direct');
-var getIteratorMethod = require('../internals/get-iterator-method');
-var getMethod = require('../internals/get-method');
-var wellKnownSymbol = require('../internals/well-known-symbol');
-var AsyncFromSyncIterator = require('../internals/async-from-sync-iterator');
 var WrapAsyncIterator = require('../internals/async-iterator-wrap');
 
-var ASYNC_ITERATOR = wellKnownSymbol('asyncIterator');
-
 $({ target: 'AsyncIterator', stat: true, forced: true }, {
   from: function from(O) {
-    var object = toObject(O);
-    var usingIterator = getMethod(object, ASYNC_ITERATOR);
-    var iteratorRecord;
-    if (usingIterator) {
-      iteratorRecord = getIteratorDirect(getAsyncIterator(object, usingIterator));
-      if (isPrototypeOf(AsyncIteratorPrototype, iteratorRecord.iterator)) return iteratorRecord.iterator;
-    }
-    if (iteratorRecord === undefined) {
-      usingIterator = getIteratorMethod(object);
-      if (usingIterator) iteratorRecord = getIteratorDirect(new AsyncFromSyncIterator(
-        getIteratorDirect(getIterator(object, usingIterator))
-      ));
-    }
-    return new WrapAsyncIterator(iteratorRecord !== undefined ? iteratorRecord : getIteratorDirect(object));
+    var iteratorRecord = getAsyncIteratorFlattenable(O);
+    return isPrototypeOf(AsyncIteratorPrototype, iteratorRecord.iterator)
+      ? iteratorRecord.iterator
+      : new WrapAsyncIterator(iteratorRecord);
   }
 });
diff --git a/packages/core-js/modules/esnext.iterator.flat-map.js b/packages/core-js/modules/esnext.iterator.flat-map.js
@@ -5,38 +5,28 @@ var call = require('../internals/function-call');
 var aCallable = require('../internals/a-callable');
 var anObject = require('../internals/an-object');
 var getIteratorDirect = require('../internals/get-iterator-direct');
-var getIteratorMethod = require('../internals/get-iterator-method');
+var getIteratorFlattenable = require('../internals/get-iterator-flattenable');
 var createIteratorProxy = require('../internals/iterator-create-proxy');
 var iteratorClose = require('../internals/iterator-close');
 
-var $TypeError = TypeError;
-
 var IteratorProxy = createIteratorProxy(function () {
   var iterator = this.iterator;
   var mapper = this.mapper;
-  var result, mapped, iteratorMethod, innerIterator;
+  var result, inner;
 
   while (true) {
-    if (innerIterator = this.innerIterator) try {
-      result = anObject(call(this.innerNext, innerIterator));
+    if (inner = this.inner) try {
+      result = anObject(call(inner.next, inner.iterator));
       if (!result.done) return result.value;
-      this.innerIterator = this.innerNext = null;
+      this.inner = null;
     } catch (error) { iteratorClose(iterator, 'throw', error); }
 
     result = anObject(call(this.next, iterator));
 
     if (this.done = !!result.done) return;
 
     try {
-      mapped = mapper(result.value, this.counter++);
-      iteratorMethod = getIteratorMethod(mapped);
-
-      if (!iteratorMethod) {
-        throw $TypeError('.flatMap callback should return an iterable object');
-      }
-
-      this.innerIterator = innerIterator = anObject(call(iteratorMethod, mapped));
-      this.innerNext = aCallable(innerIterator.next);
+      this.inner = getIteratorFlattenable(mapper(result.value, this.counter++));
     } catch (error) { iteratorClose(iterator, 'throw', error); }
   }
 });
@@ -45,8 +35,7 @@ $({ target: 'Iterator', proto: true, real: true, forced: true }, {
   flatMap: function flatMap(mapper) {
     return new IteratorProxy(getIteratorDirect(this), {
       mapper: aCallable(mapper),
-      innerIterator: null,
-      innerNext: null
+      inner: null
     });
   }
 });
diff --git a/packages/core-js/modules/esnext.iterator.from.js b/packages/core-js/modules/esnext.iterator.from.js
@@ -1,28 +1,20 @@
 // https://github.com/tc39/proposal-iterator-helpers
 var $ = require('../internals/export');
 var call = require('../internals/function-call');
-var toObject = require('../internals/to-object');
 var isPrototypeOf = require('../internals/object-is-prototype-of');
 var IteratorPrototype = require('../internals/iterators-core').IteratorPrototype;
 var createIteratorProxy = require('../internals/iterator-create-proxy');
-var getIterator = require('../internals/get-iterator');
-var getIteratorDirect = require('../internals/get-iterator-direct');
-var getIteratorMethod = require('../internals/get-iterator-method');
+var getIteratorFlattenable = require('../internals/get-iterator-flattenable');
 
 var IteratorProxy = createIteratorProxy(function () {
   return call(this.next, this.iterator);
 }, true);
 
 $({ target: 'Iterator', stat: true, forced: true }, {
   from: function from(O) {
-    var object = toObject(O);
-    var usingIterator = getIteratorMethod(object);
-    var iteratorRecord;
-    if (usingIterator) {
-      iteratorRecord = getIteratorDirect(getIterator(object, usingIterator));
-      if (isPrototypeOf(IteratorPrototype, iteratorRecord.iterator)) return iteratorRecord.iterator;
-    } else {
-      iteratorRecord = getIteratorDirect(object);
-    } return new IteratorProxy(iteratorRecord);
+    var iteratorRecord = getIteratorFlattenable(O);
+    return isPrototypeOf(IteratorPrototype, iteratorRecord.iterator)
+      ? iteratorRecord.iterator
+      : new IteratorProxy(iteratorRecord);
   }
 });