Distinguish Set and Map iterators.

paulmillr#387
Xotic750 · Mar 19, 2017 · 7c6322a · 7c6322a
1 parent 61b9685
commit 7c6322a
Show file tree

Hide file tree

Showing 3 changed files with 56 additions and 2 deletions.
diff --git a/es6-shim.js b/es6-shim.js
@@ -2842,7 +2842,11 @@
         };
 
         MapIterator.prototype = {
+          isMapIterator: true,
           next: function next() {
+            if (!this.isMapIterator) {
+              throw new TypeError('Not a MapIterator');
+            }
             var i = this.i;
             var kind = this.kind;
             var head = this.head;
@@ -3225,13 +3229,13 @@
           values: function values() {
             requireSetSlot(this, 'values');
             ensureMap(this);
-            return this['[[SetData]]'].values();
+            return new SetIterator(this['[[SetData]]'].values());
           },
 
           entries: function entries() {
             requireSetSlot(this, 'entries');
             ensureMap(this);
-            return this['[[SetData]]'].entries();
+            return new SetIterator(this['[[SetData]]'].entries());
           },
 
           forEach: function forEach(callback) {
@@ -3251,6 +3255,20 @@
         defineProperty(SetShim.prototype, 'keys', SetShim.prototype.values, true);
         addIterator(SetShim.prototype, SetShim.prototype.values);
 
+        var SetIterator = function SetIterator(it) {
+          this.it = it;
+        };
+        SetIterator.prototype = {
+          isSetIterator: true,
+          next: function next() {
+            if (!this.isSetIterator) {
+              throw new TypeError('Not a SetIterator');
+            }
+            return this.it.next();
+          }
+        };
+        addIterator(SetIterator.prototype);
+
         return SetShim;
       }())
     };

diff --git a/test/map.js b/test/map.js
@@ -600,4 +600,22 @@ describe('Map', function () {
     expect(iterator.next().done).to.equal(true);
     expect(keys).to.eql(['a', 'd', 'e']);
   });
+
+  it('MapIterator identification', function () {
+    var mapEntriesProto = Object.getPrototypeOf(new Map().entries());
+    var setEntriesProto = Object.getPrototypeOf(new Set().entries());
+    expect(mapEntriesProto).to.not.equal(setEntriesProto);
+
+    var fnMapValues = Map.prototype.values;
+    var mapSentinel = new Map([[1, 'MapSentinel']]);
+    var testMap = new Map();
+    var testMapValues = testMap.values();
+    expect(testMapValues.next.call(fnMapValues.call(mapSentinel)).value).to.equal('MapSentinel');
+
+    var testSet = new Set();
+    var testSetValues = testSet.values();
+    expect(function () {
+      return testSetValues.next.call(fnMapValues.call(mapSentinel)).value;
+    }).to['throw'](TypeError);
+  });
 });
diff --git a/test/set.js b/test/set.js
@@ -634,4 +634,22 @@ describe('Set', function () {
 
   it.skip('should throw proper errors when user invokes methods with wrong types of receiver', function () {
   });
+
+  it('SetIterator identification', function () {
+    var mapEntriesProto = Object.getPrototypeOf(new Map().entries());
+    var setEntriesProto = Object.getPrototypeOf(new Set().entries());
+    expect(mapEntriesProto).to.not.equal(setEntriesProto);
+
+    var fnSetValues = Set.prototype.values;
+    var setSentinel = new Set(['SetSentinel']);
+    var testSet1 = new Set();
+    var testSetValues = testSet1.values();
+    expect(testSetValues.next.call(fnSetValues.call(setSentinel)).value).to.equal('SetSentinel');
+
+    var testMap = new Map();
+    var testMapValues = testMap.values();
+    expect(function () {
+      return testMapValues.next.call(fnSetValues.call(setSentinel)).value;
+    }).to['throw'](TypeError);
+  });
 });