LibJS: Handle abrupt closures from Iterator.prototype.flatMap

This is in preparation of implementing %IteratorHelperPrototype%.return.
That will invoke GeneratorResumeAbrupt, which will execute the generator
with an abrupt completion. At that time, we must take care to close the
current inner iterator.

Userland/Libraries/LibJS/Runtime/IteratorHelper.cpp

@@ -11,15 +11,16 @@
 
 namespace JS {
 
-ThrowCompletionOr<NonnullGCPtr<IteratorHelper>> IteratorHelper::create(Realm& realm, IteratorRecord underlying_iterator, Closure closure)
+ThrowCompletionOr<NonnullGCPtr<IteratorHelper>> IteratorHelper::create(Realm& realm, IteratorRecord underlying_iterator, Closure closure, Optional<AbruptClosure> abrupt_closure)
 {
-    return TRY(realm.heap().allocate<IteratorHelper>(realm, realm, realm.intrinsics().iterator_helper_prototype(), move(underlying_iterator), move(closure)));
+    return TRY(realm.heap().allocate<IteratorHelper>(realm, realm, realm.intrinsics().iterator_helper_prototype(), move(underlying_iterator), move(closure), move(abrupt_closure)));
 }
 
-IteratorHelper::IteratorHelper(Realm& realm, Object& prototype, IteratorRecord underlying_iterator, Closure closure)
+IteratorHelper::IteratorHelper(Realm& realm, Object& prototype, IteratorRecord underlying_iterator, Closure closure, Optional<AbruptClosure> abrupt_closure)
     : GeneratorObject(realm, prototype, realm.vm().running_execution_context().copy(), "Iterator Helper"sv)
     , m_underlying_iterator(move(underlying_iterator))
     , m_closure(move(closure))
+    , m_abrupt_closure(move(abrupt_closure))
 {
 }
 
@@ -41,9 +42,16 @@ ThrowCompletionOr<Value> IteratorHelper::close_result(VM& vm, Completion complet
     return *TRY(iterator_close(vm, underlying_iterator(), move(completion)));
 }
 
-ThrowCompletionOr<Value> IteratorHelper::execute(VM& vm, JS::Completion const&)
+ThrowCompletionOr<Value> IteratorHelper::execute(VM& vm, JS::Completion const& completion)
 {
     ScopeGuard guard { [&] { vm.pop_execution_context(); } };
+
+    if (completion.is_abrupt()) {
+        if (m_abrupt_closure.has_value())
+            return (*m_abrupt_closure)(vm, *this, completion);
+        return close_result(vm, completion);
+    }
+
     auto result_value = m_closure(vm, *this);
 
     if (result_value.is_throw_completion()) {

Userland/Libraries/LibJS/Runtime/IteratorHelper.h

@@ -19,8 +19,9 @@ class IteratorHelper final : public GeneratorObject {
 
 public:
     using Closure = JS::SafeFunction<ThrowCompletionOr<Value>(VM&, IteratorHelper&)>;
+    using AbruptClosure = JS::SafeFunction<ThrowCompletionOr<Value>(VM&, IteratorHelper&, Completion const&)>;
 
-    static ThrowCompletionOr<NonnullGCPtr<IteratorHelper>> create(Realm&, IteratorRecord, Closure);
+    static ThrowCompletionOr<NonnullGCPtr<IteratorHelper>> create(Realm&, IteratorRecord, Closure, Optional<AbruptClosure> = {});
 
     IteratorRecord const& underlying_iterator() const { return m_underlying_iterator; }
 
@@ -31,13 +32,14 @@ public:
     ThrowCompletionOr<Value> close_result(VM&, Completion);
 
 private:
-    IteratorHelper(Realm&, Object& prototype, IteratorRecord, Closure);
+    IteratorHelper(Realm&, Object& prototype, IteratorRecord, Closure, Optional<AbruptClosure>);
 
     virtual void visit_edges(Visitor&) override;
     virtual ThrowCompletionOr<Value> execute(VM&, JS::Completion const& completion) override;
 
     IteratorRecord m_underlying_iterator; // [[UnderlyingIterator]]
     Closure m_closure;
+    Optional<AbruptClosure> m_abrupt_closure;
 
     size_t m_counter { 0 };
     bool m_done { false };

Userland/Libraries/LibJS/Runtime/IteratorPrototype.cpp

@@ -321,6 +321,24 @@ public:
         return next_outer_iterator(vm, iterated, iterator, mapper);
     }
 
+    // NOTE: This implements step 5.b.ix.4.d of Iterator.prototype.flatMap.
+    ThrowCompletionOr<Value> on_abrupt_completion(VM& vm, IteratorHelper& iterator, Completion const& completion)
+    {
+        VERIFY(m_inner_iterator.has_value());
+
+        // d. If completion is an abrupt completion, then
+
+        // i. Let backupCompletion be Completion(IteratorClose(innerIterator, completion)).
+        auto backup_completion = iterator_close(vm, *m_inner_iterator, completion);
+
+        // ii. IfAbruptCloseIterator(backupCompletion, iterated).
+        if (backup_completion.is_error())
+            return iterator.close_result(vm, backup_completion.release_error());
+
+        // iii. Return ? IteratorClose(completion, iterated).
+        return iterator.close_result(vm, completion);
+    }
+
 private:
     FlatMapIterator() = default;
 
@@ -397,10 +415,7 @@ private:
                 return iterator.close_result(vm, inner_value.release_error());
 
             // c. Let completion be Completion(Yield(innerValue)).
-            // d. If completion is an abrupt completion, then
+            // NOTE: Step d is implemented via on_abrupt_completion.
-            //     i. Let backupCompletion be Completion(IteratorClose(innerIterator, completion)).
-            //     ii. IfAbruptCloseIterator(backupCompletion, iterated).
-            //     iii. Return ? IteratorClose(completion, iterated).
             return inner_value.release_value();
         }
     }
@@ -434,9 +449,13 @@ JS_DEFINE_NATIVE_FUNCTION(IteratorPrototype::flat_map)
         return flat_map_iterator->next(vm, iterated, iterator, *mapper);
     };
 
+    IteratorHelper::AbruptClosure abrupt_closure = [flat_map_iterator](auto& vm, auto& iterator, auto const& completion) -> ThrowCompletionOr<Value> {
+        return flat_map_iterator->on_abrupt_completion(vm, iterator, completion);
+    };
+
     // 6. Let result be CreateIteratorFromClosure(closure, "Iterator Helper", %IteratorHelperPrototype%, « [[UnderlyingIterator]] »).
     // 7. Set result.[[UnderlyingIterator]] to iterated.
-    auto result = TRY(IteratorHelper::create(realm, move(iterated), move(closure)));
+    auto result = TRY(IteratorHelper::create(realm, move(iterated), move(closure), move(abrupt_closure)));
 
     // 8. Return result.
     return result;