[Fix] flatMap: only increment the count when iterating the outer it…

…erator
es-shims · May 1, 2023 · 955d0b0 · 955d0b0
1 parent 3a78767
commit 955d0b0
Show file tree

Hide file tree

Showing 2 changed files with 57 additions and 38 deletions.
diff --git a/Iterator.prototype.flatMap/implementation.js b/Iterator.prototype.flatMap/implementation.js
@@ -59,46 +59,44 @@ module.exports = function flatMap(mapper) {
 			var value = IteratorValue(next); // step 6.b.iii
 		}
 
-		try {
-			if (innerIterator === sentinel) {
-				innerAlive = true; // step 6.b.viii
-				try {
-					var mapped = Call(mapper, void undefined, [value, counter]); // step 6.b.iv
-					// yield mapped // step 6.b.vi
-					innerIterator = GetIteratorFlattenable(mapped); // step 6.b.vi
-				} catch (e) {
-					innerAlive = false;
-					innerIterator = sentinel;
-					closeIfAbrupt(ThrowCompletion(e)); // steps 6.b.v, 6.b.vii
-				}
+		if (innerIterator === sentinel) {
+			innerAlive = true; // step 6.b.viii
+			try {
+				var mapped = Call(mapper, void undefined, [value, counter]); // step 6.b.iv
+				// yield mapped // step 6.b.vi
+				innerIterator = GetIteratorFlattenable(mapped); // step 6.b.vi
+			} catch (e) {
+				innerAlive = false;
+				innerIterator = sentinel;
+				closeIfAbrupt(ThrowCompletion(e)); // steps 6.b.v, 6.b.vii
+			} finally {
+				counter += 1; // step 6.b.x
+			}
+		}
+		// while (innerAlive) { // step 6.b.ix
+		if (innerAlive) {
+			var innerNext;
+			try {
+				innerNext = IteratorStep(innerIterator); // step 6.b.ix.1
+			} catch (e) {
+				innerIterator = sentinel;
+				closeIfAbrupt(ThrowCompletion(e)); // step 6.b.ix.2
+			}
+			if (!innerNext) {
+				innerAlive = false; // step 6.b.ix.3.a
+				innerIterator = sentinel;
+				return closure();
 			}
-			// while (innerAlive) { // step 6.b.ix
-			if (innerAlive) {
-				var innerNext;
-				try {
-					innerNext = IteratorStep(innerIterator); // step 6.b.ix.1
-				} catch (e) {
-					innerIterator = sentinel;
-					closeIfAbrupt(ThrowCompletion(e)); // step 6.b.ix.2
-				}
-				if (!innerNext) {
-					innerAlive = false; // step 6.b.ix.3.a
-					innerIterator = sentinel;
-					return closure();
-				}
-				// step 6.b.ix.4
-				var innerValue;
-				try {
-					innerValue = IteratorValue(innerNext); // step 6.b.ix.4.a
-				} catch (e) {
-					innerAlive = false;
-					innerIterator = sentinel;
-					closeIfAbrupt(ThrowCompletion(e)); // step 6.b.ix.4.b
-				}
-				return innerValue; // step 6.b.ix.4.c
+			// step 6.b.ix.4
+			var innerValue;
+			try {
+				innerValue = IteratorValue(innerNext); // step 6.b.ix.4.a
+			} catch (e) {
+				innerAlive = false;
+				innerIterator = sentinel;
+				closeIfAbrupt(ThrowCompletion(e)); // step 6.b.ix.4.b
 			}
-		} finally {
-			counter += 1; // step 6.b.x
+			return innerValue; // step 6.b.ix.4.c
 		}
 		// }
 		// return void undefined;

diff --git a/test/Iterator.prototype.flatMap.js b/test/Iterator.prototype.flatMap.js
@@ -119,6 +119,27 @@ module.exports = {
 				return ret;
 			}), [0, 1, 2], st, 'test262: test/built-ins/Iterator/prototype/flatMap/iterable-to-iterator-fallback');
 
+			var counts = [];
+			testIterator(flatMap(['a', 'b', 'c', 'd', 'e'][Symbol.iterator](), function (value, count) {
+				counts.push(count);
+
+				if (value === 'a' || value === 'b') {
+					return [0];
+				}
+				if (value === 'c') {
+					return [1, 2];
+				}
+				if (value === 'd') {
+					return [3, 4, 5];
+				}
+				if (value === 'e') {
+					return [6, 7, 8, 9];
+				}
+
+				return st.fail('got unexpected value: ' + debug(v));
+			}), [0, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9], st, 'test262: test/built-ins/Iterator/prototype/flatMap/mapper-args');
+			st.deepEqual(counts, [0, 1, 2, 3, 4], 'count values are as expected');
+
 			st.end();
 		});
 	},