LibJS: Disable optimization in IteratorNextUnpack if next() is redefined

81b6a11 regressed correctness by always bypassing the `next()` method
resolution for built-in iterators, causing incorrect behavior when
`next()` was redefined on built-in prototypes. This change fixes the
issue by storing a flag on built-in prototypes indicating whether
`next()` has ever been redefined.

Libraries/LibJS/Bytecode/Interpreter.cpp

@@ -1769,7 +1769,7 @@ class PropertyNameIterator final
 public:
     virtual ~PropertyNameIterator() override = default;
 
-    BuiltinIterator* as_builtin_iterator() override { return this; }
+    BuiltinIterator* as_builtin_iterator_if_next_is_not_redefined() override { return this; }
     ThrowCompletionOr<void> next(VM&, bool& done, Value& value) override
     {
         while (true) {
@@ -3175,7 +3175,7 @@ ThrowCompletionOr<void> IteratorNextUnpack::execute_impl(Bytecode::Interpreter&
 
     Value value;
     bool done = false;
-    if (auto* builtin_iterator = iterator_record.iterator->as_builtin_iterator()) {
+    if (auto* builtin_iterator = iterator_record.iterator->as_builtin_iterator_if_next_is_not_redefined()) {
         TRY(builtin_iterator->next(vm, done, value));
     } else {
         auto result = TRY(iterator_next(vm, iterator_record));

Libraries/LibJS/Runtime/ArrayIterator.cpp

@@ -6,6 +6,7 @@
 
 #include <LibJS/Runtime/Array.h>
 #include <LibJS/Runtime/ArrayIterator.h>
+#include <LibJS/Runtime/ArrayIteratorPrototype.h>
 #include <LibJS/Runtime/TypedArray.h>
 
 namespace JS {
@@ -28,6 +29,8 @@ ArrayIterator::ArrayIterator(Value array, Object::PropertyKind iteration_kind, O
     , m_array(array)
     , m_iteration_kind(iteration_kind)
 {
+    auto& array_iterator_prototype = as<ArrayIteratorPrototype>(prototype);
+    m_next_method_was_redefined = array_iterator_prototype.next_method_was_redefined();
 }
 
 void ArrayIterator::visit_edges(Cell::Visitor& visitor)

Libraries/LibJS/Runtime/ArrayIterator.h

@@ -21,7 +21,12 @@ public:
 
     virtual ~ArrayIterator() override = default;
 
-    BuiltinIterator* as_builtin_iterator() override { return this; }
+    BuiltinIterator* as_builtin_iterator_if_next_is_not_redefined() override
+    {
+        if (m_next_method_was_redefined)
+            return nullptr;
+        return this;
+    }
     ThrowCompletionOr<void> next(VM&, bool& done, Value& value) override;
 
 private:

Libraries/LibJS/Runtime/ArrayIteratorPrototype.h

@@ -19,10 +19,20 @@ public:
     virtual void initialize(Realm&) override;
     virtual ~ArrayIteratorPrototype() override = default;
 
+    bool next_method_was_redefined() const { return m_next_method_was_redefined; }
+    void set_next_method_was_redefined() { m_next_method_was_redefined = true; }
+
+    virtual bool is_array_iterator_prototype() const override { return true; }
+
 private:
     explicit ArrayIteratorPrototype(Realm&);
 
     JS_DECLARE_NATIVE_FUNCTION(next);
+
+    bool m_next_method_was_redefined { false };
 };
 
+template<>
+inline bool Object::fast_is<ArrayIteratorPrototype>() const { return is_array_iterator_prototype(); }
+
 }

Libraries/LibJS/Runtime/Iterator.cpp

@@ -272,7 +272,7 @@ static Completion iterator_close_impl(VM& vm, IteratorRecord const& iterator_rec
     auto iterator = iterator_record.iterator;
 
     // OPTIMIZATION: "return" method is not defined on any of iterators we treat as built-in.
-    if (iterator->as_builtin_iterator())
+    if (iterator->as_builtin_iterator_if_next_is_not_redefined())
         return completion;
 
     // 3. Let innerResult be Completion(GetMethod(iterator, "return")).

Libraries/LibJS/Runtime/Iterator.h

@@ -71,6 +71,9 @@ class BuiltinIterator {
 public:
     virtual ~BuiltinIterator() = default;
     virtual ThrowCompletionOr<void> next(VM&, bool& done, Value& value) = 0;
+
+protected:
+    bool m_next_method_was_redefined { false };
 };
 
 // 7.4.12 IfAbruptCloseIterator ( value, iteratorRecord ), https://tc39.es/ecma262/#sec-ifabruptcloseiterator

Libraries/LibJS/Runtime/MapIterator.cpp

@@ -6,6 +6,7 @@
 
 #include <LibJS/Runtime/Array.h>
 #include <LibJS/Runtime/MapIterator.h>
+#include <LibJS/Runtime/MapIteratorPrototype.h>
 
 namespace JS {
 
@@ -22,6 +23,8 @@ MapIterator::MapIterator(Map& map, Object::PropertyKind iteration_kind, Object&
     , m_iteration_kind(iteration_kind)
     , m_iterator(static_cast<Map const&>(map).begin())
 {
+    auto& map_iterator_prototype = as<MapIteratorPrototype>(prototype);
+    m_next_method_was_redefined = map_iterator_prototype.next_method_was_redefined();
 }
 
 void MapIterator::visit_edges(Cell::Visitor& visitor)

Libraries/LibJS/Runtime/MapIterator.h

@@ -22,7 +22,12 @@ public:
 
     virtual ~MapIterator() override = default;
 
-    BuiltinIterator* as_builtin_iterator() override { return this; }
+    BuiltinIterator* as_builtin_iterator_if_next_is_not_redefined() override
+    {
+        if (m_next_method_was_redefined)
+            return nullptr;
+        return this;
+    }
     ThrowCompletionOr<void> next(VM&, bool& done, Value& value) override;
 
 private:

Libraries/LibJS/Runtime/MapIteratorPrototype.h

@@ -19,10 +19,20 @@ public:
     virtual void initialize(Realm&) override;
     virtual ~MapIteratorPrototype() override = default;
 
+    bool next_method_was_redefined() const { return m_next_method_was_redefined; }
+    void set_next_method_was_redefined() { m_next_method_was_redefined = true; }
+
+    virtual bool is_map_iterator_prototype() const override { return true; }
+
 private:
     MapIteratorPrototype(Realm&);
 
     JS_DECLARE_NATIVE_FUNCTION(next);
+
+    bool m_next_method_was_redefined { false };
 };
 
+template<>
+inline bool Object::fast_is<MapIteratorPrototype>() const { return is_map_iterator_prototype(); }
+
 }

Libraries/LibJS/Runtime/Object.cpp

@@ -10,15 +10,19 @@
 #include <LibJS/Runtime/AbstractOperations.h>
 #include <LibJS/Runtime/Accessor.h>
 #include <LibJS/Runtime/Array.h>
+#include <LibJS/Runtime/ArrayIteratorPrototype.h>
 #include <LibJS/Runtime/ClassFieldDefinition.h>
 #include <LibJS/Runtime/ECMAScriptFunctionObject.h>
 #include <LibJS/Runtime/Error.h>
 #include <LibJS/Runtime/GlobalObject.h>
+#include <LibJS/Runtime/MapIteratorPrototype.h>
 #include <LibJS/Runtime/NativeFunction.h>
 #include <LibJS/Runtime/Object.h>
 #include <LibJS/Runtime/PropertyDescriptor.h>
 #include <LibJS/Runtime/ProxyObject.h>
+#include <LibJS/Runtime/SetIteratorPrototype.h>
 #include <LibJS/Runtime/Shape.h>
+#include <LibJS/Runtime/StringIteratorPrototype.h>
 #include <LibJS/Runtime/Value.h>
 
 namespace JS {
@@ -957,6 +961,18 @@ ThrowCompletionOr<bool> Object::internal_set(PropertyKey const& property_key, Va
     VERIFY(!value.is_special_empty_value());
     VERIFY(!receiver.is_special_empty_value());
 
+    if (receiver.is_object() && property_key == vm().names.next) {
+        auto& receiver_object = receiver.as_object();
+        if (auto* array_iterator_prototype = as_if<ArrayIteratorPrototype>(receiver_object))
+            array_iterator_prototype->set_next_method_was_redefined();
+        else if (auto* map_iterator_prototype = as_if<MapIteratorPrototype>(receiver_object))
+            map_iterator_prototype->set_next_method_was_redefined();
+        else if (auto* set_iterator_prototype = as_if<SetIteratorPrototype>(receiver_object))
+            set_iterator_prototype->set_next_method_was_redefined();
+        else if (auto* string_iterator_prototype = as_if<StringIteratorPrototype>(receiver_object))
+            string_iterator_prototype->set_next_method_was_redefined();
+    }
+
     // 2. Let ownDesc be ? O.[[GetOwnProperty]](P).
     auto own_descriptor = TRY(internal_get_own_property(property_key));
 

Libraries/LibJS/Runtime/Object.h

@@ -212,7 +212,12 @@ public:
     virtual bool is_array_iterator() const { return false; }
     virtual bool is_raw_json_object() const { return false; }
 
-    virtual BuiltinIterator* as_builtin_iterator() { return nullptr; }
+    virtual BuiltinIterator* as_builtin_iterator_if_next_is_not_redefined() { return nullptr; }
+
+    virtual bool is_array_iterator_prototype() const { return false; }
+    virtual bool is_map_iterator_prototype() const { return false; }
+    virtual bool is_set_iterator_prototype() const { return false; }
+    virtual bool is_string_iterator_prototype() const { return false; }
 
     // B.3.7 The [[IsHTMLDDA]] Internal Slot, https://tc39.es/ecma262/#sec-IsHTMLDDA-internal-slot
     virtual bool is_htmldda() const { return false; }

Libraries/LibJS/Runtime/SetIterator.cpp

@@ -6,6 +6,7 @@
 
 #include <LibJS/Runtime/Array.h>
 #include <LibJS/Runtime/SetIterator.h>
+#include <LibJS/Runtime/SetIteratorPrototype.h>
 
 namespace JS {
 
@@ -22,6 +23,8 @@ SetIterator::SetIterator(Set& set, Object::PropertyKind iteration_kind, Object&
     , m_iteration_kind(iteration_kind)
     , m_iterator(static_cast<Set const&>(set).begin())
 {
+    auto& set_iterator_prototype = as<SetIteratorPrototype>(prototype);
+    m_next_method_was_redefined = set_iterator_prototype.next_method_was_redefined();
 }
 
 void SetIterator::visit_edges(Cell::Visitor& visitor)

Libraries/LibJS/Runtime/SetIterator.h

@@ -22,7 +22,13 @@ public:
 
     virtual ~SetIterator() override = default;
 
-    BuiltinIterator* as_builtin_iterator() override { return this; }
+    BuiltinIterator* as_builtin_iterator_if_next_is_not_redefined() override
+    {
+        if (m_next_method_was_redefined)
+            return nullptr;
+        return this;
+    }
+
     ThrowCompletionOr<void> next(VM&, bool& done, Value& value) override;
 
 private:

Libraries/LibJS/Runtime/SetIteratorPrototype.h

@@ -19,10 +19,20 @@ public:
     virtual void initialize(Realm&) override;
     virtual ~SetIteratorPrototype() override = default;
 
+    bool next_method_was_redefined() const { return m_next_method_was_redefined; }
+    void set_next_method_was_redefined() { m_next_method_was_redefined = true; }
+
+    virtual bool is_set_iterator_prototype() const override { return true; }
+
 private:
     explicit SetIteratorPrototype(Realm&);
 
     JS_DECLARE_NATIVE_FUNCTION(next);
+
+    bool m_next_method_was_redefined { false };
 };
 
+template<>
+inline bool Object::fast_is<SetIteratorPrototype>() const { return is_set_iterator_prototype(); }
+
 }

Libraries/LibJS/Runtime/StringIterator.cpp

@@ -7,6 +7,7 @@
 #include <AK/Utf8View.h>
 #include <LibJS/Runtime/GlobalObject.h>
 #include <LibJS/Runtime/StringIterator.h>
+#include <LibJS/Runtime/StringIteratorPrototype.h>
 
 namespace JS {
 
@@ -22,6 +23,8 @@ StringIterator::StringIterator(String string, Object& prototype)
     , m_string(move(string))
     , m_iterator(Utf8View(m_string).begin())
 {
+    auto& string_iterator_prototype = as<StringIteratorPrototype>(prototype);
+    m_next_method_was_redefined = string_iterator_prototype.next_method_was_redefined();
 }
 
 ThrowCompletionOr<void> StringIterator::next(VM& vm, bool& done, Value& value)

Libraries/LibJS/Runtime/StringIterator.h

@@ -23,7 +23,12 @@ public:
 
     virtual ~StringIterator() override = default;
 
-    BuiltinIterator* as_builtin_iterator() override { return this; }
+    BuiltinIterator* as_builtin_iterator_if_next_is_not_redefined() override
+    {
+        if (m_next_method_was_redefined)
+            return nullptr;
+        return this;
+    }
     ThrowCompletionOr<void> next(VM&, bool& done, Value& value) override;
 
 private:

Libraries/LibJS/Runtime/StringIteratorPrototype.h

@@ -20,10 +20,20 @@ public:
     virtual void initialize(Realm&) override;
     virtual ~StringIteratorPrototype() override = default;
 
+    bool next_method_was_redefined() const { return m_next_method_was_redefined; }
+    void set_next_method_was_redefined() { m_next_method_was_redefined = true; }
+
+    virtual bool is_string_iterator_prototype() const override { return true; }
+
 private:
     explicit StringIteratorPrototype(Realm&);
 
     JS_DECLARE_NATIVE_FUNCTION(next);
+
+    bool m_next_method_was_redefined { false };
 };
 
+template<>
+inline bool Object::fast_is<StringIteratorPrototype>() const { return is_string_iterator_prototype(); }
+
 }

Libraries/LibJS/Tests/redefine-next-in-builtin-iterators.js

@@ -0,0 +1,46 @@
+describe("redefine next() in built in iterators", () => {
+    test("redefine next() in ArrayIteratorPrototype", () => {
+        let arrayIteratorPrototype = Object.getPrototypeOf([].values());
+        let originalNext = arrayIteratorPrototype.next;
+        let counter = 0;
+        arrayIteratorPrototype.next = function () {
+            counter++;
+            return originalNext.apply(this, arguments);
+        };
+        for (let i of [1, 2, 3]) {
+        }
+        expect(counter).toBe(4);
+    });
+
+    test("redefine next() in MapIteratorPrototype", () => {
+        let m = new Map([
+            [1, 1],
+            [2, 2],
+            [3, 3],
+        ]);
+        let mapIteratorPrototype = Object.getPrototypeOf(m.values());
+        let originalNext = mapIteratorPrototype.next;
+        let counter = 0;
+        mapIteratorPrototype.next = function () {
+            counter++;
+            return originalNext.apply(this, arguments);
+        };
+        for (let v of m.values()) {
+        }
+        expect(counter).toBe(4);
+    });
+
+    test("redefine next() in SetIteratorPrototype", () => {
+        let s = new Set([1, 2, 3]);
+        let setIteratorPrototype = Object.getPrototypeOf(s.values());
+        let originalNext = setIteratorPrototype.next;
+        let counter = 0;
+        setIteratorPrototype.next = function () {
+            counter++;
+            return originalNext.apply(this, arguments);
+        };
+        for (let v of s.values()) {
+        }
+        expect(counter).toBe(4);
+    });
+});