Fix ordered set with map

__tests__/OrderedSet.ts

@@ -7,7 +7,7 @@
 
 ///<reference path='../resources/jest.d.ts'/>
 
-import { OrderedSet } from '../';
+import { OrderedSet, Map } from '../';
 
 describe('OrderedSet', () => {
   it('provides initial values in a mixed order', () => {
@@ -116,4 +116,30 @@ describe('OrderedSet', () => {
     expect(aNotC.size).toBe(22);
     expect(aNotD.size).toBe(23);
   });
+
+  it('updating a value with ".map()" should keep the set ordered', () => {
+    const first = Map({ id: 1, valid: true });
+    const second = Map({ id: 2, valid: true });
+    const third = Map({ id: 3, valid: true });
+    const initial = OrderedSet([first, second, third]);
+
+    const out = initial.map((t) => {
+      if (2 === t.get('id')) {
+        return t.set('valid', false);
+      }
+      return t;
+    });
+
+    const expected = OrderedSet([
+      Map({ id: 1, valid: true }),
+      Map({ id: 2, valid: false }),
+      Map({ id: 3, valid: true }),
+    ]);
+
+    expect(out).toEqual(expected);
+
+    expect(out.has(first)).toBe(true);
+    expect(out.has(second)).toBe(false);
+    expect(out.has(third)).toBe(true);
+  });
 });

src/Set.js

@@ -82,19 +82,23 @@ export class Set extends SetCollection {
   // @pragma Composition
 
   map(mapper, context) {
-    const removes = [];
-    const adds = [];
-    this.forEach((value) => {
-      const mapped = mapper.call(context, value, value, this);
-      if (mapped !== value) {
-        removes.push(value);
-        adds.push(mapped);
-      }
-    });
-    return this.withMutations((set) => {
-      removes.forEach((value) => set.remove(value));
-      adds.forEach((value) => set.add(value));
-    });
+    // keep track if the set is altered by the map function
+    let didChanges = false;
+
+    const newMap = updateSet(
+      this,
+      this._map.mapEntries(([, v]) => {
+        const mapped = mapper.call(context, v, v, this);
+
+        if (mapped !== v) {
+          didChanges = true;
+        }
+
+        return [mapped, mapped];
+      }, context)
+    );
+
+    return didChanges ? newMap : this;
   }
 
   union(...iters) {