Make ObjectSet compliant with Set api

nitzanhen · Dec 4, 2023 · ea6c288 · ea6c288
1 parent 98ef688
commit ea6c288
Show file tree

Hide file tree

Showing 3 changed files with 37 additions and 20 deletions.
diff --git a/src/ImmutableMap.ts b/src/ImmutableMap.ts
@@ -61,8 +61,4 @@ export class ImmutableMap<K, V> {
   get [Symbol.toStringTag]() {
     return 'ImmutableMap';
   };
-
-  clone(): ImmutableMap<K, V> {
-    return new ImmutableMap(this.map);
-  }
 }
diff --git a/src/ImmutableSet.ts b/src/ImmutableSet.ts
@@ -13,7 +13,7 @@ export class ImmutableSet<T> {
   }
 
   add(value: T): ImmutableSet<T> {
-    const set = this.set.clone();
+    const set = new ObjectSet(this.set);
     set.add(value);
     return new ImmutableSet(set);
   }
@@ -23,23 +23,31 @@ export class ImmutableSet<T> {
   }
 
   delete(value: T): ImmutableSet<T> {
-    const set = this.set.clone();
+    const set = new ObjectSet(this.set);
     set.delete(value);
     return new ImmutableSet(set);
   }
 
   clear(): ImmutableSet<T> {
-    const set = this.set.clone();
+    const set = new ObjectSet(this.set);
     set.clear();
     return new ImmutableSet(set);
   }
 
+  *entries(): Iterable<[T, T]> {
+    yield* this.set.entries();
+  }
+
+  *keys(): Iterable<T> {
+    yield* this.set.keys();
+  }
+
   *values(): Iterable<T> {
     yield* this.set.values();
   }
   [Symbol.iterator] = this.values;
 
-  clone(): ImmutableSet<T> {
-    return new ImmutableSet(this.set);
+  forEach(callbackfn: (value: T, key: T, set: Set<T>) => void, thisArg?: any): void {
+    this.set.forEach(callbackfn, thisArg);
   }
 }
diff --git a/src/ObjectSet.ts b/src/ObjectSet.ts
@@ -1,7 +1,7 @@
 import { ObjectMap, ObjectMapOptions } from './ObjectMap';
 
-export class ObjectSet<T> {
-  protected map: ObjectMap<T, true>;
+export class ObjectSet<T> implements Set<T> {
+  protected map: ObjectMap<T, T>;
 
   constructor(iterable?: Iterable<T>, options: ObjectMapOptions = {}) {
     this.map = new ObjectMap(
@@ -14,40 +14,53 @@ export class ObjectSet<T> {
     return this.map.size;
   }
 
-  add(value: T) {
-    this.map.set(value, true);
+  add(value: T): this {
+    this.map.set(value, value);
+    return this;
   }
 
   has(value: T) {
     return this.map.has(value);
   }
 
   delete(value: T): boolean {
-    return !!this.map.delete(value)
+    return this.map.delete(value)
   }
 
   clear() {
     this.map.clear();
   }
 
-  *values() {
+  *entries(): IterableIterator<[T, T]> {
+    yield* this.map.entries();
+  }
+
+  *keys(): IterableIterator<T> {
     yield* this.map.keys();
   }
+
+  *values(): IterableIterator<T> {
+    yield* this.map.values();
+  }
   [Symbol.iterator] = this.values;
 
-  clone() {
-    return new ObjectSet(this);
+  forEach(callbackfn: (value: T, key: T, set: ObjectSet<T>) => void, thisArg?: any): void {
+    this.map.forEach((v, k) => callbackfn(v, k, this), thisArg);
+  }
+
+  get [Symbol.toStringTag]() {
+    return 'ObjectSet';
   }
 }
 
 /**
- * Turns a "set" iterator to a "map" iterator by putting each value into a tuple, with the value as `true`.
+ * Turns a "set" iterator to a "map" iterator by doubling each value into a tuple.
  */
-function* toMapIterable<T>(iterable?: Iterable<T>): Iterable<[T, true]> {
+function* toMapIterable<T>(iterable?: Iterable<T>): Iterable<[T, T]> {
   if (!iterable) {
     return undefined;
   }
   for (const value of iterable) {
-    yield [value, true];
+    yield [value, value];
   }
 }