@dxos/util 0.8.4-main.c85a9c8dae → 0.8.4-main.d05539e30a

package/src/weak.ts

@@ -7,13 +7,20 @@
  */
 export class WeakDictionary<K, V extends object> implements Map<K, V> {
   private readonly _internal = new Map<K, WeakRef<V>>();
-  private readonly _finalization = new FinalizationRegistry((cleanUpCallback: () => void) => {
-    cleanUpCallback();
+  private readonly _finalization = new FinalizationRegistry<{ key: K; ref: WeakRef<V> }>(({ key, ref }) => {
+    if (this._internal.get(key) === ref) {
+      this._internal.delete(key);
+    }
   });
 
   constructor(entries?: [K, V][]) {
-    this._internal = new Map(entries?.map(([key, value]) => [key, new WeakRef(value)]));
-    entries?.forEach(([key, value]) => this._register(key, value));
+    if (entries) {
+      for (const [key, value] of entries) {
+        const ref = new WeakRef(value);
+        this._internal.set(key, ref);
+        this._register(key, value, ref);
+      }
+    }
   }
 
   *entries(): SetIterator<[K, V]> {
@@ -55,11 +62,48 @@ export class WeakDictionary<K, V extends object> implements Map<K, V> {
   }
 
   set(key: K, value: V): this {
-    this._internal.set(key, new WeakRef(value));
-    this._register(key, value);
+    const previous = this._internal.get(key)?.deref();
+    if (previous) {
+      this._unregister(previous);
+    }
+    const ref = new WeakRef(value);
+    this._internal.set(key, ref);
+    this._register(key, value, ref);
     return this;
   }
 
+  /**
+   * Returns the value for the given key if present, otherwise inserts and returns the default value.
+   * @param key - The key to look up or insert.
+   * @param defaultValue - The value to insert if the key is not present.
+   * @returns The existing or newly inserted value.
+   */
+  getOrInsert(key: K, defaultValue: V): V {
+    const existing = this.get(key);
+    if (existing !== undefined) {
+      return existing;
+    }
+    this.set(key, defaultValue);
+    return defaultValue;
+  }
+
+  /**
+   * Returns the value for the given key if present, otherwise computes, inserts, and returns a new value.
+   * The callback is only invoked when the key is missing.
+   * @param key - The key to look up or insert.
+   * @param callbackfn - Function to compute the value if the key is not present.
+   * @returns The existing or newly computed value.
+   */
+  getOrInsertComputed(key: K, callbackfn: (key: K) => V): V {
+    const existing = this.get(key);
+    if (existing !== undefined) {
+      return existing;
+    }
+    const value = callbackfn(key);
+    this.set(key, value);
+    return value;
+  }
+
   has(key: K): boolean {
     return this._internal.has(key) && this._internal.get(key)!.deref() !== undefined;
   }
@@ -96,14 +140,8 @@ export class WeakDictionary<K, V extends object> implements Map<K, V> {
     });
   }
 
-  private _register(key: K, value: V): void {
-    this._finalization.register(
-      value,
-      () => {
-        this._internal.delete(key);
-      },
-      value,
-    );
+  private _register(key: K, value: V, ref: WeakRef<V>): void {
+    this._finalization.register(value, { key, ref }, value);
   }
 
   private _unregister(value: V): void {