Optimize run to avoid quadratic map cloning

.github/workflows/test.yml

@@ -0,0 +1,44 @@
+name: Node.js CI
+
+on:
+  push:
+    branches: [ master ]
+  pull_request:
+    branches: [ master ]
+
+jobs:
+  test:
+    name: "Test"
+    runs-on: ubuntu-latest
+
+    strategy:
+      matrix:
+        node-version: [14.x, 16.x, 18.x]
+
+    steps:
+      - uses: actions/checkout@v3
+      - name: Use Node.js ${{ matrix.node-version }}
+        uses: actions/setup-node@v3
+        with:
+          node-version: ${{ matrix.node-version }}
+          cache: 'npm'
+      - run: npm ci
+      - run: npm test
+
+  lint:
+    name: "Lint"
+    runs-on: ubuntu-latest
+
+    strategy:
+      matrix:
+        node-version: [18.x]
+
+    steps:
+      - uses: actions/checkout@v3
+      - name: Use Node.js ${{ matrix.node-version }}
+        uses: actions/setup-node@v3
+        with:
+          node-version: ${{ matrix.node-version }}
+          cache: 'npm'
+      - run: npm ci
+      - run: npm run lint

package.json

@@ -5,6 +5,7 @@
   "description": "Async Context proposal for JavaScript",
   "scripts": {
     "build": "mkdir -p build && ecmarkup spec.html build/index.html",
+    "lint": "tsc -p tsconfig.json",
     "test": "mocha"
   },
   "repository": "legendecas/proposal-async-context",
@@ -26,6 +27,7 @@
     "@types/mocha": "10.0.1",
     "@types/node": "18.11.18",
     "ecmarkup": "^3.1.1",
-    "mocha": "10.2.0"
+    "mocha": "10.2.0",
+    "typescript": "4.9.4"
   }
 }

src/fork.ts

@@ -0,0 +1,68 @@
+import type { Mapping } from "./mapping";
+import type { AsyncContext } from "./index";
+
+/**
+ * FrozenRevert holds a frozen Mapping that will be simply restored when the
+ * revert is run.
+ *
+ * This is used when we already know that the mapping is frozen, so that
+ * reverting will not attempt to mutate the Mapping (and allocate a new
+ * mapping) as a Revert would.
+ */
+export class FrozenRevert {
+  #mapping: Mapping;
+
+  constructor(mapping: Mapping) {
+    this.#mapping = mapping;
+  }
+
+  /**
+   * The Storage container will call restore when it wants to revert its
+   * current Mapping to the state at the start of the fork.
+   *
+   * For FrozenRevert, that's as simple as returning the known-frozen Mapping,
+   * because we know it can't have been modified.
+   */
+  restore(_current: Mapping): Mapping {
+    return this.#mapping;
+  }
+}
+
+/**
+ * Revert holds the information on how to undo a modification to our Mappings,
+ * and will attempt to modify the current state when we attempt to restore it
+ * to its prior state.
+ *
+ * This is used when we know that the Mapping is unfrozen at start, because
+ * it's possible that no one will snapshot this Mapping before we restore. In
+ * that case, we can simply modify the Mapping without cloning. If someone did
+ * snapshot it, then modifying will clone the current state and we restore the
+ * clone to the prior state.
+ */
+export class Revert<T> {
+  #key: AsyncContext<T>;
+  #has: boolean;
+  #prev: T | undefined;
+
+  constructor(mapping: Mapping, key: AsyncContext<T>) {
+    this.#key = key;
+    this.#has = mapping.has(key);
+    this.#prev = mapping.get(key);
+  }
+
+  /**
+   * The Storage container will call restore when it wants to revert its
+   * current Mapping to the state at the start of the fork.
+   *
+   * For Revert, we mutate the known-unfrozen-at-start mapping (which may
+   * reallocate if anyone has since taken a snapshot) in the hopes that we
+   * won't need to reallocate.
+   */
+  restore(current: Mapping): Mapping {
+    if (this.#has) {
+      return current.set(this.#key, this.#prev);
+    } else {
+      return current.delete(this.#key);
+    }
+  }
+}

src/index.ts

@@ -1,45 +1,37 @@
-type AnyFunc = (...args: any) => any;
-type Storage = Map<AsyncContext<unknown>, unknown>;
+import { Storage } from "./storage";
 
-let __storage__: Storage = new Map();
+type AnyFunc<T> = (this: T, ...args: any) => any;
 
 export class AsyncContext<T> {
-  static wrap<F extends AnyFunc>(fn: F): F {
-    const current = __storage__;
+  static wrap<F extends AnyFunc<any>>(fn: F): F {
+    const snapshot = Storage.snapshot();
 
-    function wrap(...args: Parameters<F>): ReturnType<F> {
-      return run(fn, current, this, args);
-    };
+    function wrap(this: ThisType<F>, ...args: Parameters<F>): ReturnType<F> {
+      const head = Storage.switch(snapshot);
+      try {
+        return fn.apply(this, args);
+      } finally {
+        Storage.restore(head);
+      }
+    }
 
     return wrap as unknown as F;
   }
 
-  run<F extends AnyFunc>(
+  run<F extends AnyFunc<null>>(
     value: T,
     fn: F,
     ...args: Parameters<F>
   ): ReturnType<F> {
-    const next = new Map(__storage__);
-    next.set(this, value);
-    return run(fn, next, null, args);
+    const revert = Storage.set(this, value);
+    try {
+      return fn.apply(null, args);
+    } finally {
+      Storage.restore(revert);
+    }
   }
 
-  get(): T {
-    return __storage__.get(this) as T;
-  }
-}
-
-function run<F extends AnyFunc>(
-  fn: F,
-  next: Storage,
-  binding: ThisType<F>,
-  args: Parameters<F>
-): ReturnType<F> {
-  const previous = __storage__;
-  try {
-    __storage__ = next;
-    return fn.apply(binding, args);
-  } finally {
-    __storage__ = previous;
+  get(): T | undefined {
+    return Storage.get(this);
   }
 }

src/mapping.ts

@@ -0,0 +1,71 @@
+import type { AsyncContext } from "./index";
+
+/**
+ * Stores all AsyncContext data, and tracks whether any snapshots have been
+ * taken of the current data.
+ */
+export class Mapping {
+  #data: Map<AsyncContext<unknown>, unknown> | null;
+
+  /**
+   * If a snapshot of this data is taken, then further modifications cannot be
+   * made directly. Instead, set/delete will clone this Mapping and modify
+   * _that_ instance.
+   */
+  #frozen: boolean;
+
+  constructor(data: Map<AsyncContext<unknown>, unknown> | null) {
+    this.#data = data;
+    this.#frozen = data === null;
+  }
+
+  has<T>(key: AsyncContext<T>): boolean {
+    return this.#data?.has(key) || false;
+  }
+
+  get<T>(key: AsyncContext<T>): T | undefined {
+    return this.#data?.get(key) as T | undefined;
+  }
+
+  /**
+   * Like the standard Map.p.set, except that we will allocate a new Mapping
+   * instance if this instance is frozen.
+   */
+  set<T>(key: AsyncContext<T>, value: T): Mapping {
+    const mapping = this.#fork();
+    mapping.#data!.set(key, value);
+    return mapping;
+  }
+
+  /**
+   * Like the standard Map.p.delete, except that we will allocate a new Mapping
+   * instance if this instance is frozen.
+   */
+  delete<T>(key: AsyncContext<T>): Mapping {
+    const mapping = this.#fork();
+    mapping.#data!.delete(key);
+    return mapping;
+  }
+
+  /**
+   * Prevents further modifications to this Mapping.
+   */
+  freeze(): void {
+    this.#frozen = true;
+  }
+
+  isFrozen(): boolean {
+    return this.#frozen;
+  }
+
+  /**
+   * We only need to fork if the Mapping is frozen (someone has a snapshot of
+   * the current data), else we can just modify our data directly.
+   */
+  #fork(): Mapping {
+    if (this.#frozen) {
+      return new Mapping(new Map(this.#data));
+    }
+    return this;
+  }
+}

src/storage.ts

@@ -0,0 +1,70 @@
+import { Mapping } from "./mapping";
+import { FrozenRevert, Revert } from "./fork";
+import type { AsyncContext } from "./index";
+
+/**
+ * Storage is the (internal to the language) storage container of all
+ * AsyncContext data.
+ *
+ * None of the methods here are exposed to users, they're only exposed to the AsyncContext class.
+ */
+export class Storage {
+  static #current: Mapping = new Mapping(null);
+
+  /**
+   * Get retrieves the current value assigned to the AsyncContext.
+   */
+  static get<T>(key: AsyncContext<T>): T | undefined {
+    return this.#current.get(key);
+  }
+
+  /**
+   * Set assigns a new value to the AsyncContext, returning a revert that can
+   * undo the modification at a later time.
+   */
+  static set<T>(key: AsyncContext<T>, value: T): FrozenRevert | Revert<T> {
+    // If the Mappings are frozen (someone has snapshot it), then modifying the
+    // mappings will return a clone containing the modification.
+    const current = this.#current;
+    const undo = current.isFrozen()
+      ? new FrozenRevert(current)
+      : new Revert(current, key);
+    this.#current = this.#current.set(key, value);
+    return undo;
+  }
+
+  /**
+   * Restore will, well, restore the global storage state to state at the time
+   * the revert was created.
+   */
+  static restore<T>(revert: FrozenRevert | Revert<T>): void {
+    this.#current = revert.restore(this.#current);
+  }
+
+  /**
+   * Snapshot freezes the current storage state, and returns a new revert which
+   * can restore the global storage state to the state at the time of the
+   * snapshot.
+   */
+  static snapshot(): FrozenRevert {
+    this.#current.freeze();
+    return new FrozenRevert(this.#current);
+  }
+
+  /**
+   * Switch swaps the global storage state to the state at the time of a
+   * snapshot, completely replacing the current state (and making it impossible
+   * for the current state to be modified until the snapshot is reverted).
+   */
+  static switch(snapshot: FrozenRevert): FrozenRevert {
+    const previous = this.#current;
+    this.#current = snapshot.restore(previous);
+
+    // Technically, previous may not be frozen. But we know its state cannot
+    // change, because the only way to modify it is to restore it to the
+    // Storage container, and the only way to do that is to have snapshot it.
+    // So it's either snapshot (and frozen), or it's not and thus cannot be
+    // modified.
+    return new FrozenRevert(previous);
+  }
+}