Update the polyfill to the latest API

src/fork.ts

@@ -1,5 +1,5 @@
 import type { Mapping } from "./mapping";
-import type { AsyncContext } from "./index";
+import type { Variable } from "./variable";
 
 /**
  * FrozenRevert holds a frozen Mapping that will be simply restored when the
@@ -40,11 +40,11 @@ export class FrozenRevert {
  * clone to the prior state.
  */
 export class Revert<T> {
-  #key: AsyncContext<T>;
+  #key: Variable<T>;
   #has: boolean;
   #prev: T | undefined;
 
-  constructor(mapping: Mapping, key: AsyncContext<T>) {
+  constructor(mapping: Mapping, key: Variable<T>) {
     this.#key = key;
     this.#has = mapping.has(key);
     this.#prev = mapping.get(key);

src/index.ts

@@ -1,37 +1,8 @@
-import { Storage } from "./storage";
-
-type AnyFunc<T> = (this: T, ...args: any) => any;
-
-export class AsyncContext<T> {
-  static wrap<F extends AnyFunc<any>>(fn: F): F {
-    const snapshot = Storage.snapshot();
-
-    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<null>>(
-    value: T,
-    fn: F,
-    ...args: Parameters<F>
-  ): ReturnType<F> {
-    const revert = Storage.set(this, value);
-    try {
-      return fn.apply(null, args);
-    } finally {
-      Storage.restore(revert);
-    }
-  }
-
-  get(): T | undefined {
-    return Storage.get(this);
-  }
-}
+import { Snapshot } from "./snapshot";
+import { Variable } from "./variable";
+
+export const AsyncContext = {
+  Snapshot,
+  Variable,
+};
+export { Snapshot, Variable };

src/mapping.ts

@@ -1,49 +1,48 @@
-import type { AsyncContext } from "./index";
+import type { Variable } from "./variable";
 
 /**
- * Stores all AsyncContext data, and tracks whether any snapshots have been
+ * Stores all Variable data, and tracks whether any snapshots have been
  * taken of the current data.
  */
 export class Mapping {
-  #data: Map<AsyncContext<unknown>, unknown> | null;
+  #data: Map<Variable<unknown>, unknown>;
 
   /**
    * 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;
+  #frozen = false;
 
-  constructor(data: Map<AsyncContext<unknown>, unknown> | null) {
+  constructor(data: Map<Variable<unknown>, unknown>) {
     this.#data = data;
-    this.#frozen = data === null;
   }
 
-  has<T>(key: AsyncContext<T>): boolean {
-    return this.#data?.has(key) || false;
+  has<T>(key: Variable<T>): boolean {
+    return this.#data.has(key) || false;
   }
 
-  get<T>(key: AsyncContext<T>): T | undefined {
-    return this.#data?.get(key) as T | undefined;
+  get<T>(key: Variable<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 {
+  set<T>(key: Variable<T>, value: T): Mapping {
     const mapping = this.#fork();
-    mapping.#data!.set(key, value);
+    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 {
+  delete<T>(key: Variable<T>): Mapping {
     const mapping = this.#fork();
-    mapping.#data!.delete(key);
+    mapping.#data.delete(key);
     return mapping;
   }
 

src/promise-polyfill.ts

@@ -1,12 +1,13 @@
 import { AsyncContext } from "./index";
 
-type AnyFunc = (...args: any) => any;
+import type { AnyFunc } from "./types";
 
 export const nativeThen = Promise.prototype.then;
+const { wrap } = AsyncContext.Snapshot;
 
-function wrapFn<F extends AnyFunc>(fn: F | null | undefined) {
+function wrapFn<F extends AnyFunc<any>>(fn: F | null | undefined) {
   if (typeof fn !== "function") return undefined;
-  return AsyncContext.wrap(fn);
+  return wrap(fn);
 }
 
 export function then<T>(

src/snapshot.ts

@@ -0,0 +1,36 @@
+import { Storage } from "./storage";
+
+import type { FrozenRevert } from "./fork";
+import type { AnyFunc } from "./types";
+
+export class Snapshot {
+  #snapshot = Storage.snapshot();
+
+  static wrap<F extends AnyFunc<any>>(fn: F): F {
+    const snapshot = Storage.snapshot();
+
+    function wrap(this: ThisType<F>, ...args: Parameters<F>): ReturnType<F> {
+      return run(fn, this, args, snapshot);
+    }
+
+    return wrap as unknown as F;
+  }
+
+  run<F extends AnyFunc<null>>(fn: F, ...args: Parameters<F>) {
+    return run(fn, null as any, args, this.#snapshot);
+  }
+}
+
+function run<F extends AnyFunc<any>>(
+  fn: F,
+  context: ThisType<F>,
+  args: any[],
+  snapshot: FrozenRevert
+): ReturnType<F> {
+  const revert = Storage.switch(snapshot);
+  try {
+    return fn.apply(context, args);
+  } finally {
+    Storage.restore(revert);
+  }
+}

src/storage.ts

@@ -1,36 +1,44 @@
 import { Mapping } from "./mapping";
 import { FrozenRevert, Revert } from "./fork";
-import type { AsyncContext } from "./index";
+
+import type { Variable } from "./variable";
 
 /**
  * Storage is the (internal to the language) storage container of all
- * AsyncContext data.
+ * Variable data.
  *
- * None of the methods here are exposed to users, they're only exposed to the AsyncContext class.
+ * None of the methods here are exposed to users, they're only exposed internally.
  */
 export class Storage {
-  static #current: Mapping = new Mapping(null);
+  static #current: Mapping = new Mapping(new Map());
+
+  /**
+   * Has checks if the Variable has a value.
+   */
+  static has<T>(key: Variable<T>): boolean {
+    return this.#current.has(key);
+  }
 
   /**
-   * Get retrieves the current value assigned to the AsyncContext.
+   * Get retrieves the current value assigned to the Variable.
    */
-  static get<T>(key: AsyncContext<T>): T | undefined {
+  static get<T>(key: Variable<T>): T | undefined {
     return this.#current.get(key);
   }
 
   /**
-   * Set assigns a new value to the AsyncContext, returning a revert that can
+   * Set assigns a new value to the Variable, returning a revert that can
    * undo the modification at a later time.
    */
-  static set<T>(key: AsyncContext<T>, value: T): FrozenRevert | Revert<T> {
+  static set<T>(key: Variable<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()
+    const revert = current.isFrozen()
       ? new FrozenRevert(current)
-      : new Revert(current, key);
+      : new Revert<T>(current, key);
     this.#current = this.#current.set(key, value);
-    return undo;
+    return revert;
   }
 
   /**

src/types.ts

@@ -0,0 +1 @@
+export type AnyFunc<T> = (this: T, ...args: any) => any;

src/variable.ts

@@ -0,0 +1,43 @@
+import { Storage } from "./storage";
+
+import type { AnyFunc } from "./types";
+
+export interface VariableOptions<T> {
+  name?: string;
+  defaultValue?: T;
+}
+
+export class Variable<T> {
+  #name = "";
+  #defaultValue: T | undefined;
+
+  constructor(options?: VariableOptions<T>) {
+    if (options) {
+      if ("name" in options) {
+        this.#name = String(options.name);
+      }
+      this.#defaultValue = options.defaultValue;
+    }
+  }
+
+  get name() {
+    return this.#name;
+  }
+
+  run<F extends AnyFunc<null>>(
+    value: T,
+    fn: F,
+    ...args: Parameters<F>
+  ): ReturnType<F> {
+    const revert = Storage.set(this, value);
+    try {
+      return fn.apply(null, args);
+    } finally {
+      Storage.restore(revert);
+    }
+  }
+
+  get(): T | undefined {
+    return Storage.has(this) ? Storage.get(this) : this.#defaultValue;
+  }
+}