TwoWayWeakMap for https://es.discourse.group/t/how-to-implement-twowayweakmap-using-weakref/890

two-way-weak-map.d.ts

interface TwoWayWeakMap<K extends object, V> extends WeakMap<K, V> {
  keysFor(value: V): K[];
}

interface TwoWayWeakMapConstructor {
  new <K extends object = object, V = any>(
    entries?: readonly [K, V][] | null
  ): TwoWayWeakMap<K, V>;
  readonly prototype: TwoWayWeakMap<object, any>;
}

declare const TwoWayWeakMap: TwoWayWeakMapConstructor;

export default TwoWayWeakMap;

two-way-weak-map.js

const testSet = new WeakSet();

const hasObject = (value) => {
  const type = typeof value;
  switch (type) {
    case "object":
      if (!value) return false;
    // fallthrough
    case "box":
    case "function":
      return true;
    case "boolean":
    case "number":
    case "string":
    case "symbol":
    case "bigint":
    case "undefined":
      return false;
    case "record":
    case "tuple":
    // Use Box.containsBox or other predicate
    // fallthrough for now
    default:
      try {
        testSet.add(value);
        testSet.delete(value);
        return true;
      } catch (err) {}
      return false;
  }
};

const minusZero = Symbol("-0");

const cleanup = ({ keyRefs, value, wr }) => {
  const wrSet = keyRefs.get(value);
  if (wrSet) {
    wrSet.delete(wr);
    console.log(`removed wr for ${String(value)}, left ${wrSet.size}`);
    if (!wrSet.size) {
      keyRefs.delete(value);
    }
  } else {
    console.log("Cleaning up missing set for value", String(value));
  }
};

const fr = new FinalizationRegistry(cleanup);

export default class TwoWayWeakMap extends WeakMap {
  #keyRefsForPrimitiveValues;
  #keyRefsForObjectValues;

  constructor(entries = []) {
    super();
    this.#keyRefsForPrimitiveValues = new Map();
    this.#keyRefsForObjectValues = new WeakMap();
    for (const [key, value] of entries) {
      this.set(key, value);
    }
  }

  /** @override */
  get(key) {
    const { value } = super.get(key) || {};
    return value;
  }

  /** @override */
  set(key, value) {
    const current = super.get(key);
    if (current) {
      if (Object.is(current.value, value)) {
        return this;
      }
      this.delete(key);
    }
    const wr = new WeakRef(key);
    super.set(key, { value, wr });
    if (Object.is(-0, value)) value = minusZero;
    let keyRefs;
    if (hasObject(value)) {
      keyRefs = this.#keyRefsForObjectValues;
    } else {
      keyRefs = this.#keyRefsForPrimitiveValues;
      fr.register(key, { keyRefs, value, wr }, wr);
    }
    let wrSet = keyRefs.get(value);
    if (!wrSet) {
      wrSet = new Set();
      keyRefs.set(value, wrSet);
    }
    wrSet.add(wr);
    return this;
  }

  /** @override */
  delete(key) {
    const current = super.get(key);
    if (current) {
      const { wr, value } = current;
      if (Object.is(-0, value)) value = minusZero;
      let keyRefs;
      if (hasObject(value)) {
        keyRefs = this.#keyRefsForObjectValues;
      } else {
        keyRefs = this.#keyRefsForPrimitiveValues;
        fr.unregister(wr);
      }
      cleanup({ keyRefs, value, wr });
    }
    return super.delete(key);
  }

  keysFor(value) {
    if (Object.is(-0, value)) value = minusZero;
    const valueHasObject = hasObject(value);
    const keyRefs = valueHasObject
      ? this.#keyRefsForObjectValues
      : this.#keyRefsForPrimitiveValues;
    const wrSet = keyRefs.get(value);

    const res = [];
    if (wrSet) {
      for (const wr of wrSet) {
        const key = wr.deref();
        if (key) {
          res.push(key);
        } else {
          if (!valueHasObject) {
            fr.unregister(wr);
          }
          cleanup({ keyRefs, value, wr });
        }
      }
    }
    return res;
  }
}

two-way-weak-map.test.js

// @ts-check

import "./zzz-harness.js";
import assert from "./zzz-assert.js";
import TwoWayWeakMap from "./two-way-weak-map.js";

const allCollections = [];
const fr = new FinalizationRegistry((resolve) => {
  resolve();
});

const addMessage = (promise, msg) => promise.then(() => msg);

const addCollection = (obj, msg) =>
  allCollections.push(
    addMessage(new Promise((resolve) => fr.register(obj, resolve)), msg)
  );

/** @type {TwoWayWeakMap} */
let wm;

{
  const oi = { v: "init" };
  wm = new TwoWayWeakMap([[{}, "init"]]);
  addCollection(oi, "oi");

  const om0 = { v: "m0" };
  wm.set(om0, -0);
  addCollection(om0, "om0");

  const op0 = { v: "p0" };
  wm.set(op0, 0);
  addCollection(op0, "op0");

  const o1 = { v: 1 };
  wm.set(o1, 1);
  wm.set(o1, 1);
  addCollection(o1, "o1");

  const o420 = { v: 420 };
  const o421 = { v: 421 };
  wm.set(o420, 42);
  wm.set(o1, 42);
  wm.set(o421, 42);
  addCollection(o420, "o420");
  addCollection(o421, "o421");

  const oo0 = { om0, op0 };
  om0.o = oo0;
  op0.o = oo0;
  wm.set(oo0, { om0, op0 });
  addCollection(oo0, "oo0");

  wm.set(o1, 1);

  const keysFor1 = wm.keysFor(1);
  assert(keysFor1.length === 1 && keysFor1[0] === o1);

  const keysForM0 = wm.keysFor(-0);
  assert(keysForM0.length === 1 && keysForM0[0] === om0);

  const keysForP0 = wm.keysFor(0);
  assert(keysForP0.length === 1 && keysForP0[0] === op0);

  const keysFor42 = wm.keysFor(42);
  assert(
    keysFor42.length === 2 && keysFor42[0] === o420 && keysFor42[1] === o421
  );
}

const readyGC = () => {
  new Promise((resolve) => fr.register({}, resolve)).then(() =>
    console.log("sentinel collected")
  );
  return queueGCJob();
};

const allCollected = Promise.all(allCollections).then(() => {
  console.log("all keys collected");
  return true;
});

const cleanup = async (tries = 0) => {
  const result = await Promise.race([
    allCollected,
    readyGC().then(() => false),
  ]);
  if (result) {
    return;
  }
  console.log("GC attempt", ++tries);
  gc();
  console.log("Keys for 42:", wm.keysFor(42).length);
  if (tries > 10) throw new Error();

  return cleanup(tries);
};

export const result = cleanup().then(async () => {
  // Make sure the map's internal pending finalization callbacks are called
  await queueGCJob();
  console.log("clean");
});

zzz-assert.js

// Copyright (C) 2017 Ecma International.  All rights reserved.
// This code is governed by the BSD license found in the LICENSE file.
/*---
description: |
    Collection of assertion functions used throughout test262
---*/

export default function assert(mustBeTrue, message) {
  if (mustBeTrue === true) {
    return;
  }

  if (message === undefined) {
    message = "Expected true but got " + String(mustBeTrue);
  }
  $ERROR(message);
}

assert._isSameValue = function (a, b) {
  if (a === b) {
    // Handle +/-0 vs. -/+0
    return a !== 0 || 1 / a === 1 / b;
  }

  // Handle NaN vs. NaN
  return a !== a && b !== b;
};

assert.sameValue = function (actual, expected, message) {
  if (assert._isSameValue(actual, expected)) {
    return;
  }

  if (message === undefined) {
    message = "";
  } else {
    message += " ";
  }

  message +=
    "Expected SameValue(«" +
    String(actual) +
    "», «" +
    String(expected) +
    "») to be true";

  $ERROR(message);
};

assert.notSameValue = function (actual, unexpected, message) {
  if (!assert._isSameValue(actual, unexpected)) {
    return;
  }

  if (message === undefined) {
    message = "";
  } else {
    message += " ";
  }

  message +=
    "Expected SameValue(«" +
    String(actual) +
    "», «" +
    String(unexpected) +
    "») to be false";

  $ERROR(message);
};

assert.throws = function (expectedErrorConstructor, func, message) {
  if (typeof func !== "function") {
    $ERROR(
      "assert.throws requires two arguments: the error constructor " +
        "and a function to run"
    );
    return;
  }
  if (message === undefined) {
    message = "";
  } else {
    message += " ";
  }

  try {
    func();
  } catch (thrown) {
    if (typeof thrown !== "object" || thrown === null) {
      message += "Thrown value was not an object!";
      $ERROR(message);
    } else if (thrown.constructor !== expectedErrorConstructor) {
      message +=
        "Expected a " +
        expectedErrorConstructor.name +
        " but got a " +
        thrown.constructor.name;
      $ERROR(message);
    }
    return;
  }

  message +=
    "Expected a " +
    expectedErrorConstructor.name +
    " to be thrown but no exception was thrown at all";
  $ERROR(message);
};

assert.throws.early = function (err, code) {
  var wrappedCode = "function wrapperFn() { " + code + " }";
  var ieval = eval;

  assert.throws(
    err,
    function () {
      Function(wrappedCode);
    },
    "Function: " + code
  );
};

zzz-harness.js

if (typeof console === "undefined") {
  globalThis.console = {
    log() {
      print(...arguments);
    },
  };
}

globalThis.$ERROR =
  globalThis.$ERROR ||
  ((message) => {
    throw new Error(message);
  });

let defaultEnqueueGCJob;
if (typeof setTimeout === "function") {
  defaultEnqueueGCJob = function () {
    return new Promise(function (resolve) {
      setTimeout(resolve, 0);
    });
  };
} else if (typeof drainJobQueue === "function") {
  defaultEnqueueGCJob = function () {
    return new Promise(function (resolve) {
      drainJobQueue();
      resolve();
    });
  };
} else {
  defaultEnqueueGCJob = function () {
    return Promise.resolve();
  };
}

globalThis.queueGCJob = defaultEnqueueGCJob;

let gc = globalThis.gc || (typeof $262 !== "undefined" ? $262.gc : null);
if (!gc) {
  try {
    const [{ ["default"]: v8 }, { ["default"]: vm }] = await Promise.all([
      import("v8"),
      import("vm"),
    ]);
    v8.setFlagsFromString("--expose_gc");
    gc = vm.runInNewContext("gc");
    v8.setFlagsFromString("--no-expose_gc");
  } catch (err) {
    gc = () => void Array.from({ length: 2 ** 24 }, () => Math.random());
  }
}
globalThis.gc = gc;

export {};