diff --git a/Userland/Libraries/LibJS/Runtime/IteratorPrototype.cpp b/Userland/Libraries/LibJS/Runtime/IteratorPrototype.cpp index 405b7eebbe..a7a6d79561 100644 --- a/Userland/Libraries/LibJS/Runtime/IteratorPrototype.cpp +++ b/Userland/Libraries/LibJS/Runtime/IteratorPrototype.cpp @@ -34,6 +34,7 @@ ThrowCompletionOr IteratorPrototype::initialize(Realm& realm) define_native_function(realm, vm.names.filter, filter, 1, attr); define_native_function(realm, vm.names.take, take, 1, attr); define_native_function(realm, vm.names.drop, drop, 1, attr); + define_native_function(realm, vm.names.flatMap, flat_map, 1, attr); return {}; } @@ -310,4 +311,138 @@ JS_DEFINE_NATIVE_FUNCTION(IteratorPrototype::drop) return result; } +class FlatMapIterator : public Cell { + JS_CELL(FlatMapIterator, Cell); + +public: + ThrowCompletionOr next(VM& vm, IteratorRecord const& iterated, IteratorHelper& iterator, FunctionObject& mapper) + { + if (m_inner_iterator.has_value()) + return next_inner_iterator(vm, iterated, iterator, mapper); + return next_outer_iterator(vm, iterated, iterator, mapper); + } + +private: + FlatMapIterator() = default; + + virtual void visit_edges(Visitor& visitor) override + { + Base::visit_edges(visitor); + + if (m_inner_iterator.has_value()) + visitor.visit(m_inner_iterator->iterator); + } + + ThrowCompletionOr next_outer_iterator(VM& vm, IteratorRecord const& iterated, IteratorHelper& iterator, FunctionObject& mapper) + { + // i. Let next be ? IteratorStep(iterated). + auto next = TRY(iterator_step(vm, iterated)); + + // ii. If next is false, return undefined. + if (!next) + return iterator.result(js_undefined()); + + // iii. Let value be ? IteratorValue(next). + auto value = TRY(iterator_value(vm, *next)); + + // iv. Let mapped be Completion(Call(mapper, undefined, « value, 𝔽(counter) »)). + auto mapped = call(vm, mapper, js_undefined(), value, Value { iterator.counter() }); + + // v. IfAbruptCloseIterator(mapped, iterated). + if (mapped.is_error()) + return iterator.close_result(mapped.release_error()); + + // vi. Let innerIterator be Completion(GetIteratorFlattenable(mapped)). + auto inner_iterator = get_iterator_flattenable(vm, mapped.release_value()); + + // vii. IfAbruptCloseIterator(innerIterator, iterated). + if (inner_iterator.is_error()) + return iterator.close_result(inner_iterator.release_error()); + + // viii. Let innerAlive be true. + m_inner_iterator = inner_iterator.release_value(); + + // x. Set counter to counter + 1. + // NOTE: We do this step early to ensure it occurs before returning. + iterator.increment_counter(); + + // ix. Repeat, while innerAlive is true, + return next_inner_iterator(vm, iterated, iterator, mapper); + } + + ThrowCompletionOr next_inner_iterator(VM& vm, IteratorRecord const& iterated, IteratorHelper& iterator, FunctionObject& mapper) + { + VERIFY(m_inner_iterator.has_value()); + + // 1. Let innerNext be Completion(IteratorStep(innerIterator)). + auto inner_next = iterator_step(vm, *m_inner_iterator); + + // 2. IfAbruptCloseIterator(innerNext, iterated). + if (inner_next.is_error()) + return iterator.close_result(inner_next.release_error()); + + // 3. If innerNext is false, then + if (!inner_next.value()) { + // a. Set innerAlive to false. + m_inner_iterator.clear(); + + return next_outer_iterator(vm, iterated, iterator, mapper); + } + // 4. Else, + else { + // a. Let innerValue be Completion(IteratorValue(innerNext)). + auto inner_value = iterator_value(vm, *inner_next.release_value()); + + // b. IfAbruptCloseIterator(innerValue, iterated). + if (inner_value.is_error()) + return iterator.close_result(inner_value.release_error()); + + // c. Let completion be Completion(Yield(innerValue)). + // d. If completion is an abrupt completion, then + // i. Let backupCompletion be Completion(IteratorClose(innerIterator, completion)). + // ii. IfAbruptCloseIterator(backupCompletion, iterated). + // iii. Return ? IteratorClose(completion, iterated). + return inner_value.release_value(); + } + } + + Optional m_inner_iterator; +}; + +// 3.1.3.6 Iterator.prototype.flatMap ( mapper ), https://tc39.es/proposal-iterator-helpers/#sec-iteratorprototype.flatmap +JS_DEFINE_NATIVE_FUNCTION(IteratorPrototype::flat_map) +{ + auto& realm = *vm.current_realm(); + + auto mapper = vm.argument(0); + + // 1. Let O be the this value. + // 2. If O is not an Object, throw a TypeError exception. + auto object = TRY(this_object(vm)); + + // 3. If IsCallable(mapper) is false, throw a TypeError exception. + if (!mapper.is_function()) + return vm.throw_completion(ErrorType::NotAFunction, "mapper"sv); + + // 4. Let iterated be ? GetIteratorDirect(O). + auto iterated = TRY(get_iterator_direct(vm, object)); + + auto flat_map_iterator = MUST_OR_THROW_OOM(vm.heap().allocate(realm)); + + // 5. Let closure be a new Abstract Closure with no parameters that captures iterated and mapper and performs the following steps when called: + IteratorHelper::Closure closure = [flat_map_iterator, mapper = NonnullGCPtr { mapper.as_function() }](auto& iterator) mutable -> ThrowCompletionOr { + auto& vm = iterator.vm(); + + auto const& iterated = iterator.underlying_iterator(); + return flat_map_iterator->next(vm, iterated, iterator, *mapper); + }; + + // 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))); + + // 8. Return result. + return result; +} + } diff --git a/Userland/Libraries/LibJS/Runtime/IteratorPrototype.h b/Userland/Libraries/LibJS/Runtime/IteratorPrototype.h index b1aeefa40b..7cf718ba87 100644 --- a/Userland/Libraries/LibJS/Runtime/IteratorPrototype.h +++ b/Userland/Libraries/LibJS/Runtime/IteratorPrototype.h @@ -26,6 +26,7 @@ private: JS_DECLARE_NATIVE_FUNCTION(filter); JS_DECLARE_NATIVE_FUNCTION(take); JS_DECLARE_NATIVE_FUNCTION(drop); + JS_DECLARE_NATIVE_FUNCTION(flat_map); }; } diff --git a/Userland/Libraries/LibJS/Tests/builtins/Iterator/Iterator.prototype.flatMap.js b/Userland/Libraries/LibJS/Tests/builtins/Iterator/Iterator.prototype.flatMap.js new file mode 100644 index 0000000000..c42cb4499a --- /dev/null +++ b/Userland/Libraries/LibJS/Tests/builtins/Iterator/Iterator.prototype.flatMap.js @@ -0,0 +1,141 @@ +describe("errors", () => { + test("called with non-callable object", () => { + expect(() => { + Iterator.prototype.flatMap(Symbol.hasInstance); + }).toThrowWithMessage(TypeError, "mapper is not a function"); + }); + + test("iterator's next method throws", () => { + function TestError() {} + + class TestIterator extends Iterator { + next() { + throw new TestError(); + } + } + + expect(() => { + const iterator = new TestIterator().flatMap(() => 0); + iterator.next(); + }).toThrow(TestError); + }); + + test("value returned by iterator's next method throws", () => { + function TestError() {} + + class TestIterator extends Iterator { + next() { + return { + done: false, + get value() { + throw new TestError(); + }, + }; + } + } + + expect(() => { + const iterator = new TestIterator().flatMap(() => 0); + iterator.next(); + }).toThrow(TestError); + }); + + test("mapper function throws", () => { + function TestError() {} + + class TestIterator extends Iterator { + next() { + return { + done: false, + value: 1, + }; + } + } + + expect(() => { + const iterator = new TestIterator().flatMap(() => { + throw new TestError(); + }); + iterator.next(); + }).toThrow(TestError); + }); + + test("inner mapper value is not an object", () => { + function* generator() { + yield "Well hello"; + yield "friends :^)"; + } + + expect(() => { + const iterator = generator().flatMap(() => Symbol.hasInstance); + iterator.next(); + }).toThrowWithMessage(TypeError, "obj is not an object"); + }); +}); + +describe("normal behavior", () => { + test("length is 1", () => { + expect(Iterator.prototype.flatMap).toHaveLength(1); + }); + + test("mapper function sees every value", () => { + function* generator() { + yield "Well hello"; + yield "friends :^)"; + } + + let count = 0; + + const iterator = generator().flatMap((value, index) => { + ++count; + + switch (index) { + case 0: + expect(value).toBe("Well hello"); + break; + case 1: + expect(value).toBe("friends :^)"); + break; + default: + expect().fail(`Unexpected mapper invocation: value=${value} index=${index}`); + break; + } + + return value.split(" ").values(); + }); + + for (const i of iterator) { + } + + expect(count).toBe(2); + }); + + test("inner values are yielded one at a time", () => { + function* generator() { + yield "Well hello"; + yield "friends :^)"; + } + + const iterator = generator().flatMap(value => value.split(" ").values()); + + let value = iterator.next(); + expect(value.value).toBe("Well"); + expect(value.done).toBeFalse(); + + value = iterator.next(); + expect(value.value).toBe("hello"); + expect(value.done).toBeFalse(); + + value = iterator.next(); + expect(value.value).toBe("friends"); + expect(value.done).toBeFalse(); + + value = iterator.next(); + expect(value.value).toBe(":^)"); + expect(value.done).toBeFalse(); + + value = iterator.next(); + expect(value.value).toBeUndefined(); + expect(value.done).toBeTrue(); + }); +});