From 25e2f85840c79059d9d323f2694a2aad4ba1e498 Mon Sep 17 00:00:00 2001
From: Justin Ridgewell <justin@ridgewell.name>
Date: Fri, 13 Jan 2023 21:03:19 -0500
Subject: [PATCH 01/16] Optimize AsyncContext to avoid quadratic cloning

---
 src/fork.ts                 |  69 +++++++++++++++++++
 src/index.ts                |  49 ++++++--------
 src/mapping.ts              |  70 +++++++++++++++++++
 src/storage.ts              |  78 ++++++++++++++++++++++
 tests/async-context.test.ts | 129 ++----------------------------------
 tsconfig.json               |   8 +++
 6 files changed, 250 insertions(+), 153 deletions(-)
 create mode 100644 src/fork.ts
 create mode 100644 src/mapping.ts
 create mode 100644 src/storage.ts
 create mode 100644 tsconfig.json

diff --git a/src/fork.ts b/src/fork.ts
new file mode 100644
index 0000000..1265181
--- /dev/null
+++ b/src/fork.ts
@@ -0,0 +1,69 @@
+import type { Mapping } from './mapping';
+import type { AsyncContext } from './index';
+
+/**
+ * FrozenFork holds a frozen Mapping that will be simply restored when the fork is
+ * rejoined.
+ *
+ * This is used when we already know that the mapping is frozen, so that
+ * rejoining will not attempt to mutate the Mapping (and allocate a new
+ * mapping) as an OwnedFork would.
+ */
+export class FrozenFork {
+  #mapping: Mapping;
+
+  constructor(mapping: Mapping) {
+    this.#mapping = mapping;
+  }
+
+  /**
+   * The Storage container will call join when it wants to restore its current
+   * Mapping to the state at the start of the fork.
+   *
+   * For FrozenFork, that's as simple as returning the known-frozen Mapping,
+   * because we know it can't have been modified.
+   */
+  join(): Mapping {
+    return this.#mapping;
+  }
+}
+
+/**
+ * OwnedFork holds an unfrozen Mapping that we will attempt when rejoining to
+ * 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 rejoin. In
+ * that case, we can simply modify the Mapping (without reallocation) to
+ * restore it to its prior state. If someone does snapshot it, then modifying
+ * will clone the current state and we restore the clone to the prior state.
+ */
+export class OwnedFork<T> {
+  #mapping: Mapping;
+  #key: AsyncContext<T>;
+  #has: boolean;
+  #prev: T | undefined;
+
+  constructor(mapping: Mapping, key: AsyncContext<T>) {
+    this.#mapping = mapping;
+    this.#key = key;
+    this.#has = mapping.has(key);
+    this.#prev = mapping.get(key);
+  }
+
+  /**
+   * The Storage container will call join when it wants to restore its current
+   * Mapping to the state at the start of the fork.
+   *
+   * For OwnedFork, 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.
+   */
+  join(): Mapping {
+    if (this.#has) {
+      return this.#mapping.set(this.#key, this.#prev);
+    } else {
+      return this.#mapping.delete(this.#key);
+    }
+  }
+}
diff --git a/src/index.ts b/src/index.ts
index 0ba6e2a..f143250 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -1,45 +1,38 @@
-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 fork = Storage.restore(snapshot);
+      try {
+        return fn.apply(this, args);
+      } finally {
+        Storage.join(fork);
+      }
+    }
 
     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 fork = Storage.fork(this);
+    Storage.set(this, value);
+    try {
+      return fn.apply(null, args);
+    } finally {
+      Storage.join(fork);
+    }
   }
 
   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;
+    return Storage.get(this);
   }
 }
diff --git a/src/mapping.ts b/src/mapping.ts
new file mode 100644
index 0000000..fc361e5
--- /dev/null
+++ b/src/mapping.ts
@@ -0,0 +1,70 @@
+import type { AsyncContext } from "./index";
+
+/**
+ * Stores all AsyncContext data, and tracks whether any snapshots have been
+ * taken of the current data.
+ */
+export class Mapping {
+  /**
+   * 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 = false;
+
+  #data: Map<AsyncContext<unknown>, unknown>;
+
+  constructor(data: Map<AsyncContext<unknown>, unknown>) {
+    this.#data = data;
+  }
+
+  has<T>(key: AsyncContext<T>): boolean {
+    return this.#data.has(key);
+  }
+
+  get<T>(key: AsyncContext<T>): T {
+    return this.#data.get(key) as T;
+  }
+
+  /**
+   * 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() {
+    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;
+  }
+}
diff --git a/src/storage.ts b/src/storage.ts
new file mode 100644
index 0000000..716a80a
--- /dev/null
+++ b/src/storage.ts
@@ -0,0 +1,78 @@
+import { Mapping } from "./mapping";
+import { FrozenFork, OwnedFork } 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(new Map());
+
+  /**
+   * Get retrieves the current value assigned to the AsyncContext.
+   */
+  static get<T>(key: AsyncContext<T>): T {
+    return this.#current.get(key);
+  }
+
+  /**
+   * Set assigns a new value to the AsyncContext.
+   */
+  static set<T>(key: AsyncContext<T>, value: T) {
+    // If the Mappings are frozen (someone has snapshot it), then modifying the
+    // mappings will return a clone containing the modification.
+    this.#current = this.#current.set(key, value);
+  }
+
+  /**
+   * Fork is called before modifying the global storage state (either by
+   * replacing the current mappings or assigning a new value to an individual
+   * AsyncContext).
+   *
+   * The Fork instance returned will be able to restore the mappings to the
+   * unmodified state.
+   */
+  static fork<T>(key: AsyncContext<T>): OwnedFork<T> | FrozenFork {
+    if (this.#current.isFrozen()) {
+      return new FrozenFork(this.#current);
+    }
+    return new OwnedFork(this.#current, key);
+  }
+
+  /**
+   * Join will restore the global storage state to state at the time of the
+   * fork.
+   */
+  static join<T>(fork: FrozenFork | OwnedFork<T>) {
+    this.#current = fork.join();
+  }
+
+  /**
+   * Snapshot freezes the current storage state, and returns a new fork which
+   * can restore the global storage state to the state at the time of the
+   * snapshot.
+   */
+  static snapshot(): FrozenFork {
+    this.#current.freeze();
+    return new FrozenFork(this.#current);
+  }
+
+  /**
+   * Restore restores the global storage state to the state at the time of the
+   * snapshot.
+   */
+  static restore(snapshot: FrozenFork): FrozenFork {
+    const previous = this.#current;
+    this.#current = snapshot.join();
+
+    // 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 FrozenFork(previous);
+  }
+}
diff --git a/tests/async-context.test.ts b/tests/async-context.test.ts
index 339630a..7b4ca27 100644
--- a/tests/async-context.test.ts
+++ b/tests/async-context.test.ts
@@ -1,11 +1,7 @@
 import { AsyncContext } from "../src/index";
-import { then, nativeThen } from "../src/promise-polyfill";
 import assert from "node:assert/strict";
 
 type Value = { id: number };
-function sleep(ms: number): Promise<void> {
-  return new Promise((r) => setTimeout(r, ms));
-}
 
 describe("sync", () => {
   describe("run and get", () => {
@@ -103,13 +99,14 @@ describe("sync", () => {
       const second = { id: 2 };
 
       const actual = ctx.run(first, () => {
-        const wrapped = ctx.run(second, () => {
+        const firstWrap = AsyncContext.wrap(() => ctx.get());
+        const secondWrap = ctx.run(second, () => {
           return AsyncContext.wrap(() => ctx.get());
         });
-        return [ctx.get(), wrapped(), ctx.get()];
+        return [ctx.get(), firstWrap(), secondWrap(), ctx.get()];
       });
 
-      assert.deepStrictEqual(actual, [first, second, first]);
+      assert.deepStrictEqual(actual, [first, first, second, first]);
     });
 
     it("wrap out of order", () => {
@@ -129,121 +126,3 @@ describe("sync", () => {
     });
   });
 });
-
-describe("async via promises", () => {
-  beforeEach(() => {
-    Promise.prototype.then = then;
-  });
-  afterEach(() => {
-    Promise.prototype.then = nativeThen;
-  });
-
-  describe("run and get", () => {
-    it("get returns current context value", async () => {
-      const ctx = new AsyncContext<Value>();
-      const expected = { id: 1 };
-
-      const actual = await ctx.run(expected, () => {
-        return Promise.resolve().then(() => ctx.get());
-      });
-
-      assert.equal(actual, expected);
-    });
-
-    it("get within nesting contexts", async () => {
-      const ctx = new AsyncContext<Value>();
-      const first = { id: 1 };
-      const second = { id: 2 };
-
-      const actual = await ctx.run(first, () => {
-        return Promise.resolve([])
-          .then((temp) => {
-            temp.push(ctx.get());
-            return temp;
-          })
-          .then((temp) => {
-            return ctx.run(second, () => {
-              return Promise.resolve().then(() => {
-                temp.push(ctx.get());
-                return temp;
-              });
-            });
-          })
-          .then((temp) => {
-            temp.push(ctx.get());
-            return temp;
-          });
-      });
-
-      assert.deepStrictEqual(actual, [first, second, first]);
-    });
-
-    it("get within nesting different contexts", async () => {
-      const a = new AsyncContext<Value>();
-      const b = new AsyncContext<Value>();
-      const first = { id: 1 };
-      const second = { id: 2 };
-
-      const actual = await a.run(first, () => {
-        return Promise.resolve([])
-          .then((temp) => {
-            temp.push(a.get(), b.get());
-            return temp;
-          })
-          .then((temp) => {
-            return b.run(second, () => {
-              return Promise.resolve().then(() => {
-                temp.push(a.get(), b.get());
-                return temp;
-              });
-            });
-          })
-          .then((temp) => {
-            temp.push(a.get(), b.get());
-            return temp;
-          });
-      });
-
-      assert.deepStrictEqual(actual, [
-        first,
-        undefined,
-        first,
-        second,
-        first,
-        undefined,
-      ]);
-    });
-
-    it("get out of order", async () => {
-      const ctx = new AsyncContext<Value>();
-      const first = { id: 1 };
-      const second = { id: 2 };
-
-      const firstRun = ctx.run(first, () => {
-        return [
-          sleep(10).then(() => ctx.get()),
-          sleep(20).then(() => ctx.get()),
-          sleep(30).then(() => ctx.get()),
-        ];
-      });
-      const secondRun = ctx.run(second, () => {
-        return [
-          sleep(25).then(() => ctx.get()),
-          sleep(15).then(() => ctx.get()),
-          sleep(5).then(() => ctx.get()),
-        ];
-      });
-
-      const actual = await Promise.all(firstRun.concat(secondRun));
-
-      assert.deepStrictEqual(actual, [
-        first,
-        first,
-        first,
-        second,
-        second,
-        second,
-      ]);
-    });
-  });
-});
diff --git a/tsconfig.json b/tsconfig.json
new file mode 100644
index 0000000..cb21c27
--- /dev/null
+++ b/tsconfig.json
@@ -0,0 +1,8 @@
+{
+  "compilerOptions": {
+    "strict": true,
+    "target": "esnext",
+    "allowSyntheticDefaultImports": true,
+    "moduleResolution": "nodenext",
+  }
+}

From 1337cbfe6b0255e6cc5767f797cf6bf620b116f9 Mon Sep 17 00:00:00 2001
From: Justin Ridgewell <justin@ridgewell.name>
Date: Fri, 13 Jan 2023 22:57:59 -0500
Subject: [PATCH 02/16] Revamp tests

---
 tests/async-context.test.ts | 148 ++++++++++++++++++++++++++----------
 1 file changed, 107 insertions(+), 41 deletions(-)

diff --git a/tests/async-context.test.ts b/tests/async-context.test.ts
index 7b4ca27..e92c0af 100644
--- a/tests/async-context.test.ts
+++ b/tests/async-context.test.ts
@@ -26,9 +26,9 @@ describe("sync", () => {
       const ctx = new AsyncContext<Value>();
       const expected = { id: 1 };
 
-      const actual = ctx.run(expected, () => ctx.get());
-
-      assert.equal(actual, expected);
+      ctx.run(expected, () => {
+        assert.equal(ctx.get(), expected);
+      });
     });
 
     it("get within nesting contexts", () => {
@@ -36,11 +36,14 @@ describe("sync", () => {
       const first = { id: 1 };
       const second = { id: 2 };
 
-      const actual = ctx.run(first, () => {
-        return [ctx.get(), ctx.run(second, () => ctx.get()), ctx.get()];
+      ctx.run(first, () => {
+        assert.equal(ctx.get(), first);
+        ctx.run(second, () => {
+          assert.equal(ctx.get(), second);
+        });
+        assert.equal(ctx.get(), first);
       });
-
-      assert.deepStrictEqual(actual, [first, second, first]);
+      assert.equal(ctx.get(), undefined);
     });
 
     it("get within nesting different contexts", () => {
@@ -49,24 +52,18 @@ describe("sync", () => {
       const first = { id: 1 };
       const second = { id: 2 };
 
-      const actual = a.run(first, () => {
-        return [
-          a.get(),
-          b.get(),
-          ...b.run(second, () => [a.get(), b.get()]),
-          a.get(),
-          b.get(),
-        ];
+      a.run(first, () => {
+        assert.equal(a.get(), first);
+        assert.equal(b.get(), undefined);
+        b.run(second, () => {
+          assert.equal(a.get(), first);
+          assert.equal(b.get(), second);
+        });
+        assert.equal(a.get(), first);
+        assert.equal(b.get(), undefined);
       });
-
-      assert.deepStrictEqual(actual, [
-        first,
-        undefined,
-        first,
-        second,
-        first,
-        undefined,
-      ]);
+      assert.equal(a.get(), undefined);
+      assert.equal(b.get(), undefined);
     });
   });
 
@@ -75,22 +72,24 @@ describe("sync", () => {
       const ctx = new AsyncContext<Value>();
       const wrapped = AsyncContext.wrap(() => ctx.get());
 
-      const actual = ctx.run({ id: 1 }, () => wrapped());
-
-      assert.equal(actual, undefined);
+      ctx.run({ id: 1 }, () => {
+        assert.equal(wrapped(), undefined);
+      });
     });
 
     it("stores current state", () => {
       const ctx = new AsyncContext<Value>();
       const expected = { id: 1 };
 
-      const wrapped = ctx.run(expected, () => {
-        return AsyncContext.wrap(() => ctx.get());
+      const wrap = ctx.run(expected, () => {
+        const wrap = AsyncContext.wrap(() => ctx.get());
+        assert.equal(wrap(), expected);
+        assert.equal(ctx.get(), expected);
+        return wrap;
       });
 
-      const actual = wrapped();
-
-      assert.equal(actual, expected);
+      assert.equal(wrap(), expected);
+      assert.equal(ctx.get(), undefined);
     });
 
     it("wrap within nesting contexts", () => {
@@ -98,15 +97,76 @@ describe("sync", () => {
       const first = { id: 1 };
       const second = { id: 2 };
 
-      const actual = ctx.run(first, () => {
-        const firstWrap = AsyncContext.wrap(() => ctx.get());
+      const [firstWrap, secondWrap] = ctx.run(first, () => {
+        const firstWrap = AsyncContext.wrap(() => {
+          assert.equal(ctx.get(), first);
+        });
+        firstWrap();
+
         const secondWrap = ctx.run(second, () => {
-          return AsyncContext.wrap(() => ctx.get());
+          const secondWrap = AsyncContext.wrap(() => {
+            firstWrap();
+            assert.equal(ctx.get(), second);
+          });
+          firstWrap();
+          secondWrap();
+          assert.equal(ctx.get(), second);
+
+          return secondWrap;
+        });
+
+        firstWrap();
+        secondWrap();
+        assert.equal(ctx.get(), first);
+
+        return [firstWrap, secondWrap];
+      });
+
+      firstWrap();
+      secondWrap();
+      assert.equal(ctx.get(), undefined);
+    });
+
+    it("wrap within nesting different contexts", () => {
+      const a = new AsyncContext<Value>();
+      const b = new AsyncContext<Value>();
+      const first = { id: 1 };
+      const second = { id: 2 };
+
+      const [firstWrap, secondWrap] = a.run(first, () => {
+        const firstWrap = AsyncContext.wrap(() => {
+          assert.equal(a.get(), first);
+          assert.equal(b.get(), undefined);
         });
-        return [ctx.get(), firstWrap(), secondWrap(), ctx.get()];
+        firstWrap();
+
+        const secondWrap = b.run(second, () => {
+          const secondWrap = AsyncContext.wrap(() => {
+            firstWrap();
+            assert.equal(a.get(), first);
+            assert.equal(b.get(), second);
+          });
+
+          firstWrap();
+          secondWrap();
+          assert.equal(a.get(), first);
+          assert.equal(b.get(), second);
+
+          return secondWrap;
+        });
+
+        firstWrap();
+        secondWrap();
+        assert.equal(a.get(), first);
+        assert.equal(b.get(), undefined);
+
+        return [firstWrap, secondWrap];
       });
 
-      assert.deepStrictEqual(actual, [first, first, second, first]);
+      firstWrap();
+      secondWrap();
+      assert.equal(a.get(), undefined);
+      assert.equal(b.get(), undefined);
     });
 
     it("wrap out of order", () => {
@@ -115,14 +175,20 @@ describe("sync", () => {
       const second = { id: 2 };
 
       const firstWrap = ctx.run(first, () => {
-        return AsyncContext.wrap(() => ctx.get());
+        return AsyncContext.wrap(() => {
+          assert.equal(ctx.get(), first);
+        });
       });
       const secondWrap = ctx.run(second, () => {
-        return AsyncContext.wrap(() => ctx.get());
+        return AsyncContext.wrap(() => {
+          assert.equal(ctx.get(), second);
+        });
       });
-      const actual = [firstWrap(), secondWrap(), firstWrap(), secondWrap()];
 
-      assert.deepStrictEqual(actual, [first, second, first, second]);
+      firstWrap();
+      secondWrap();
+      firstWrap();
+      secondWrap();
     });
   });
 });

From f149be2745efd263ed6c6a5c5fb7492215c2e0b9 Mon Sep 17 00:00:00 2001
From: Justin Ridgewell <justin@ridgewell.name>
Date: Fri, 13 Jan 2023 23:48:30 -0500
Subject: [PATCH 03/16] Fix bug with mutating mappings frozen during deep run

---
 src/fork.ts                 |  10 +-
 src/storage.ts              |   2 +-
 tests/async-context.test.ts | 203 ++++++++++++++++++++++++++++++++++++
 3 files changed, 208 insertions(+), 7 deletions(-)

diff --git a/src/fork.ts b/src/fork.ts
index 1265181..7bae0b7 100644
--- a/src/fork.ts
+++ b/src/fork.ts
@@ -23,7 +23,7 @@ export class FrozenFork {
    * For FrozenFork, that's as simple as returning the known-frozen Mapping,
    * because we know it can't have been modified.
    */
-  join(): Mapping {
+  join(_current: Mapping): Mapping {
     return this.#mapping;
   }
 }
@@ -39,13 +39,11 @@ export class FrozenFork {
  * will clone the current state and we restore the clone to the prior state.
  */
 export class OwnedFork<T> {
-  #mapping: Mapping;
   #key: AsyncContext<T>;
   #has: boolean;
   #prev: T | undefined;
 
   constructor(mapping: Mapping, key: AsyncContext<T>) {
-    this.#mapping = mapping;
     this.#key = key;
     this.#has = mapping.has(key);
     this.#prev = mapping.get(key);
@@ -59,11 +57,11 @@ export class OwnedFork<T> {
    * reallocate if anyone has since taken a snapshot) in the hopes that we
    * won't need to reallocate.
    */
-  join(): Mapping {
+  join(current: Mapping): Mapping {
     if (this.#has) {
-      return this.#mapping.set(this.#key, this.#prev);
+      return current.set(this.#key, this.#prev);
     } else {
-      return this.#mapping.delete(this.#key);
+      return current.delete(this.#key);
     }
   }
 }
diff --git a/src/storage.ts b/src/storage.ts
index 716a80a..c01f3cc 100644
--- a/src/storage.ts
+++ b/src/storage.ts
@@ -47,7 +47,7 @@ export class Storage {
    * fork.
    */
   static join<T>(fork: FrozenFork | OwnedFork<T>) {
-    this.#current = fork.join();
+    this.#current = fork.join(this.#current);
   }
 
   /**
diff --git a/tests/async-context.test.ts b/tests/async-context.test.ts
index e92c0af..8ce3218 100644
--- a/tests/async-context.test.ts
+++ b/tests/async-context.test.ts
@@ -169,6 +169,209 @@ describe("sync", () => {
       assert.equal(b.get(), undefined);
     });
 
+    it("wrap within nesting different contexts, 2", () => {
+      const a = new AsyncContext<Value>();
+      const b = new AsyncContext<Value>();
+      const c = new AsyncContext<Value>();
+      const first = { id: 1 };
+      const second = { id: 2 };
+      const third = { id: 3 };
+
+      const wrap = a.run(first, () => {
+        const wrap = b.run(second, () => {
+          const wrap = c.run(third, () => {
+            debugger;
+            return AsyncContext.wrap(() => {
+              assert.equal(a.get(), first);
+              assert.equal(b.get(), second);
+              assert.equal(c.get(), third);
+            });
+          });
+          assert.equal(a.get(), first);
+          assert.equal(b.get(), second);
+          assert.equal(c.get(), undefined);
+          return wrap;
+        });
+        assert.equal(a.get(), first);
+        assert.equal(b.get(), undefined);
+        assert.equal(c.get(), undefined);
+
+        return wrap;
+      });
+
+      assert.equal(a.get(), undefined);
+      assert.equal(b.get(), undefined);
+      assert.equal(c.get(), undefined);
+      wrap();
+      assert.equal(a.get(), undefined);
+      assert.equal(b.get(), undefined);
+      assert.equal(c.get(), undefined);
+    });
+
+    it("wrap within nesting different contexts, 3", () => {
+      const a = new AsyncContext<Value>();
+      const b = new AsyncContext<Value>();
+      const c = new AsyncContext<Value>();
+      const first = { id: 1 };
+      const second = { id: 2 };
+      const third = { id: 3 };
+
+      const wrap = a.run(first, () => {
+        const wrap = b.run(second, () => {
+          AsyncContext.wrap(() => {});
+
+          const wrap = c.run(third, () => {
+            debugger;
+            return AsyncContext.wrap(() => {
+              assert.equal(a.get(), first);
+              assert.equal(b.get(), second);
+              assert.equal(c.get(), third);
+            });
+          });
+          assert.equal(a.get(), first);
+          assert.equal(b.get(), second);
+          assert.equal(c.get(), undefined);
+          return wrap;
+        });
+        assert.equal(a.get(), first);
+        assert.equal(b.get(), undefined);
+        assert.equal(c.get(), undefined);
+
+        return wrap;
+      });
+
+      assert.equal(a.get(), undefined);
+      assert.equal(b.get(), undefined);
+      assert.equal(c.get(), undefined);
+      wrap();
+      assert.equal(a.get(), undefined);
+      assert.equal(b.get(), undefined);
+      assert.equal(c.get(), undefined);
+    });
+
+    it("wrap within nesting different contexts, 4", () => {
+      const a = new AsyncContext<Value>();
+      const b = new AsyncContext<Value>();
+      const c = new AsyncContext<Value>();
+      const first = { id: 1 };
+      const second = { id: 2 };
+      const third = { id: 3 };
+
+      const wrap = a.run(first, () => {
+        AsyncContext.wrap(() => {});
+
+        const wrap = b.run(second, () => {
+          const wrap = c.run(third, () => {
+            debugger;
+            return AsyncContext.wrap(() => {
+              assert.equal(a.get(), first);
+              assert.equal(b.get(), second);
+              assert.equal(c.get(), third);
+            });
+          });
+          assert.equal(a.get(), first);
+          assert.equal(b.get(), second);
+          assert.equal(c.get(), undefined);
+          return wrap;
+        });
+        assert.equal(a.get(), first);
+        assert.equal(b.get(), undefined);
+        assert.equal(c.get(), undefined);
+
+        return wrap;
+      });
+
+      assert.equal(a.get(), undefined);
+      assert.equal(b.get(), undefined);
+      assert.equal(c.get(), undefined);
+      wrap();
+      assert.equal(a.get(), undefined);
+      assert.equal(b.get(), undefined);
+      assert.equal(c.get(), undefined);
+    });
+
+    it("wrap within nesting different contexts, 5", () => {
+      const a = new AsyncContext<Value>();
+      const b = new AsyncContext<Value>();
+      const c = new AsyncContext<Value>();
+      const first = { id: 1 };
+      const second = { id: 2 };
+      const third = { id: 3 };
+
+      const wrap = a.run(first, () => {
+        const wrap = b.run(second, () => {
+          const wrap = c.run(third, () => {
+            return AsyncContext.wrap(() => {
+              assert.equal(a.get(), first);
+              assert.equal(b.get(), second);
+              assert.equal(c.get(), third);
+            });
+          });
+
+          AsyncContext.wrap(() => {});
+
+          assert.equal(a.get(), first);
+          assert.equal(b.get(), second);
+          assert.equal(c.get(), undefined);
+          return wrap;
+        });
+        assert.equal(a.get(), first);
+        assert.equal(b.get(), undefined);
+        assert.equal(c.get(), undefined);
+
+        return wrap;
+      });
+
+      assert.equal(a.get(), undefined);
+      assert.equal(b.get(), undefined);
+      assert.equal(c.get(), undefined);
+      wrap();
+      assert.equal(a.get(), undefined);
+      assert.equal(b.get(), undefined);
+      assert.equal(c.get(), undefined);
+    });
+
+    it("wrap within nesting different contexts, 6", () => {
+      const a = new AsyncContext<Value>();
+      const b = new AsyncContext<Value>();
+      const c = new AsyncContext<Value>();
+      const first = { id: 1 };
+      const second = { id: 2 };
+      const third = { id: 3 };
+
+      const wrap = a.run(first, () => {
+        const wrap = b.run(second, () => {
+          const wrap = c.run(third, () => {
+            return AsyncContext.wrap(() => {
+              assert.equal(a.get(), first);
+              assert.equal(b.get(), second);
+              assert.equal(c.get(), third);
+            });
+          });
+          assert.equal(a.get(), first);
+          assert.equal(b.get(), second);
+          assert.equal(c.get(), undefined);
+          return wrap;
+        });
+
+        AsyncContext.wrap(() => {});
+
+        assert.equal(a.get(), first);
+        assert.equal(b.get(), undefined);
+        assert.equal(c.get(), undefined);
+
+        return wrap;
+      });
+
+      assert.equal(a.get(), undefined);
+      assert.equal(b.get(), undefined);
+      assert.equal(c.get(), undefined);
+      wrap();
+      assert.equal(a.get(), undefined);
+      assert.equal(b.get(), undefined);
+      assert.equal(c.get(), undefined);
+    });
+
     it("wrap out of order", () => {
       const ctx = new AsyncContext<Value>();
       const first = { id: 1 };

From 57cce61817be6aaf065e54da31e0486235a5bc39 Mon Sep 17 00:00:00 2001
From: Justin Ridgewell <justin@ridgewell.name>
Date: Fri, 13 Jan 2023 23:50:00 -0500
Subject: [PATCH 04/16] Prettier

---
 src/fork.ts | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/fork.ts b/src/fork.ts
index 7bae0b7..1a1272e 100644
--- a/src/fork.ts
+++ b/src/fork.ts
@@ -1,5 +1,5 @@
-import type { Mapping } from './mapping';
-import type { AsyncContext } from './index';
+import type { Mapping } from "./mapping";
+import type { AsyncContext } from "./index";
 
 /**
  * FrozenFork holds a frozen Mapping that will be simply restored when the fork is

From 1c3b3aa6d9b5931aa248503f314eb90f54bb7a51 Mon Sep 17 00:00:00 2001
From: Justin Ridgewell <justin@ridgewell.name>
Date: Sat, 14 Jan 2023 14:37:56 -0500
Subject: [PATCH 05/16] Do now freeze initial state, ensure we restore frozen
 state

---
 src/fork.ts    | 16 ++++++++--------
 src/index.ts   |  2 +-
 src/storage.ts | 29 ++++++++++++++++++-----------
 3 files changed, 27 insertions(+), 20 deletions(-)

diff --git a/src/fork.ts b/src/fork.ts
index 1a1272e..78cf460 100644
--- a/src/fork.ts
+++ b/src/fork.ts
@@ -10,9 +10,9 @@ import type { AsyncContext } from "./index";
  * mapping) as an OwnedFork would.
  */
 export class FrozenFork {
-  #mapping: Mapping;
+  #mapping: Mapping | undefined;
 
-  constructor(mapping: Mapping) {
+  constructor(mapping: Mapping | undefined) {
     this.#mapping = mapping;
   }
 
@@ -23,20 +23,20 @@ export class FrozenFork {
    * For FrozenFork, that's as simple as returning the known-frozen Mapping,
    * because we know it can't have been modified.
    */
-  join(_current: Mapping): Mapping {
+  join(_current: Mapping | undefined): Mapping | undefined {
     return this.#mapping;
   }
 }
 
 /**
- * OwnedFork holds an unfrozen Mapping that we will attempt when rejoining to
- * attempt to restore it to its prior state.
+ * OwnedFork holds an unfrozen Mapping that we will attempt to modify when
+ * rejoining to 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 rejoin. In
- * that case, we can simply modify the Mapping (without reallocation) to
- * restore it to its prior state. If someone does snapshot it, then modifying
- * will clone the current state and we restore the clone to the prior state.
+ * that case, we can simply modify the Mapping (without cloning) to restore it
+ * to its prior state. If someone does snapshot it, then modifying will clone
+ * the current state and we restore the clone to the prior state.
  */
 export class OwnedFork<T> {
   #key: AsyncContext<T>;
diff --git a/src/index.ts b/src/index.ts
index f143250..faadd20 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -32,7 +32,7 @@ export class AsyncContext<T> {
     }
   }
 
-  get(): T {
+  get(): T | undefined {
     return Storage.get(this);
   }
 }
diff --git a/src/storage.ts b/src/storage.ts
index c01f3cc..d62dd03 100644
--- a/src/storage.ts
+++ b/src/storage.ts
@@ -9,22 +9,23 @@ import type { AsyncContext } from "./index";
  * 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(new Map());
+  static #current: Mapping | undefined = undefined;
 
   /**
    * Get retrieves the current value assigned to the AsyncContext.
    */
-  static get<T>(key: AsyncContext<T>): T {
-    return this.#current.get(key);
+  static get<T>(key: AsyncContext<T>): T | undefined {
+    return this.#current?.get(key);
   }
 
   /**
    * Set assigns a new value to the AsyncContext.
    */
   static set<T>(key: AsyncContext<T>, value: T) {
+    const current = this.#current || new Mapping(new Map());
     // If the Mappings are frozen (someone has snapshot it), then modifying the
     // mappings will return a clone containing the modification.
-    this.#current = this.#current.set(key, value);
+    this.#current = current.set(key, value);
   }
 
   /**
@@ -35,11 +36,12 @@ export class Storage {
    * The Fork instance returned will be able to restore the mappings to the
    * unmodified state.
    */
-  static fork<T>(key: AsyncContext<T>): OwnedFork<T> | FrozenFork {
-    if (this.#current.isFrozen()) {
-      return new FrozenFork(this.#current);
+  static fork<T>(key: AsyncContext<T>): FrozenFork | OwnedFork<T> {
+    const current = this.#current;
+    if (current === undefined || current.isFrozen()) {
+      return new FrozenFork(current);
     }
-    return new OwnedFork(this.#current, key);
+    return new OwnedFork(current, key);
   }
 
   /**
@@ -47,7 +49,12 @@ export class Storage {
    * fork.
    */
   static join<T>(fork: FrozenFork | OwnedFork<T>) {
-    this.#current = fork.join(this.#current);
+    // The only way for #current to be undefined at a join is if we're in the
+    // we've snapshot the initial empty state with `wrap` and restored it. In
+    // which case, we're operating on a FrozenFork, and the param doesn't
+    // matter. The only other call to join is in the `run` case, and that
+    // guarantees that we have a mappings.
+    this.#current = fork.join(this.#current!);
   }
 
   /**
@@ -56,7 +63,7 @@ export class Storage {
    * snapshot.
    */
   static snapshot(): FrozenFork {
-    this.#current.freeze();
+    this.#current?.freeze();
     return new FrozenFork(this.#current);
   }
 
@@ -66,7 +73,7 @@ export class Storage {
    */
   static restore(snapshot: FrozenFork): FrozenFork {
     const previous = this.#current;
-    this.#current = snapshot.join();
+    this.#current = snapshot.join(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

From 4df3e56c200edd28434d648d5dda22311f9243b9 Mon Sep 17 00:00:00 2001
From: Justin Ridgewell <justin@ridgewell.name>
Date: Sat, 14 Jan 2023 14:41:20 -0500
Subject: [PATCH 06/16] Test PRs

---
 .github/workflows/test.yml | 25 +++++++++++++++++++++++++
 1 file changed, 25 insertions(+)
 create mode 100644 .github/workflows/test.yml

diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
new file mode 100644
index 0000000..861e5cb
--- /dev/null
+++ b/.github/workflows/test.yml
@@ -0,0 +1,25 @@
+name: Node.js CI
+
+on:
+  push:
+    branches: [ main ]
+  pull_request:
+    branches: [ main ]
+
+jobs:
+  build:
+    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

From 29878f47c64e5623b4ef39b16bb8752dc95c0b41 Mon Sep 17 00:00:00 2001
From: Justin Ridgewell <justin@ridgewell.name>
Date: Sat, 14 Jan 2023 14:44:42 -0500
Subject: [PATCH 07/16] Cleanup

---
 src/mapping.ts | 6 +++---
 src/storage.ts | 4 ++--
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/src/mapping.ts b/src/mapping.ts
index fc361e5..1c7ed4e 100644
--- a/src/mapping.ts
+++ b/src/mapping.ts
@@ -22,8 +22,8 @@ export class Mapping {
     return this.#data.has(key);
   }
 
-  get<T>(key: AsyncContext<T>): T {
-    return this.#data.get(key) as T;
+  get<T>(key: AsyncContext<T>): T | undefined {
+    return this.#data.get(key) as T | undefined;
   }
 
   /**
@@ -53,7 +53,7 @@ export class Mapping {
     this.#frozen = true;
   }
 
-  isFrozen() {
+  isFrozen(): boolean {
     return this.#frozen;
   }
 
diff --git a/src/storage.ts b/src/storage.ts
index d62dd03..2381b10 100644
--- a/src/storage.ts
+++ b/src/storage.ts
@@ -21,7 +21,7 @@ export class Storage {
   /**
    * Set assigns a new value to the AsyncContext.
    */
-  static set<T>(key: AsyncContext<T>, value: T) {
+  static set<T>(key: AsyncContext<T>, value: T): void {
     const current = this.#current || new Mapping(new Map());
     // If the Mappings are frozen (someone has snapshot it), then modifying the
     // mappings will return a clone containing the modification.
@@ -48,7 +48,7 @@ export class Storage {
    * Join will restore the global storage state to state at the time of the
    * fork.
    */
-  static join<T>(fork: FrozenFork | OwnedFork<T>) {
+  static join<T>(fork: FrozenFork | OwnedFork<T>): void {
     // The only way for #current to be undefined at a join is if we're in the
     // we've snapshot the initial empty state with `wrap` and restored it. In
     // which case, we're operating on a FrozenFork, and the param doesn't

From bd5c6ae24fe9b3960adab20dcf2c97da6391c8c9 Mon Sep 17 00:00:00 2001
From: Justin Ridgewell <justin@ridgewell.name>
Date: Sat, 14 Jan 2023 14:51:48 -0500
Subject: [PATCH 08/16] Check types

---
 .github/workflows/test.yml | 21 ++++++++++++++++++++-
 package-lock.json          | 22 +++++++++++++++++++++-
 package.json               |  4 +++-
 tsconfig.json              |  1 +
 4 files changed, 45 insertions(+), 3 deletions(-)

diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 861e5cb..bc16e76 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -7,7 +7,8 @@ on:
     branches: [ main ]
 
 jobs:
-  build:
+  test:
+    name: "Test"
     runs-on: ubuntu-latest
 
     strategy:
@@ -23,3 +24,21 @@ jobs:
           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
diff --git a/package-lock.json b/package-lock.json
index 94f6a94..5faa9f1 100644
--- a/package-lock.json
+++ b/package-lock.json
@@ -14,7 +14,8 @@
         "@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"
       }
     },
     "node_modules/@babel/code-frame": {
@@ -3449,6 +3450,19 @@
         "node": ">=8"
       }
     },
+    "node_modules/typescript": {
+      "version": "4.9.4",
+      "resolved": "https://registry.npmjs.org/typescript/-/typescript-4.9.4.tgz",
+      "integrity": "sha512-Uz+dTXYzxXXbsFpM86Wh3dKCxrQqUcVMxwU54orwlJjOpO3ao8L7j5lH+dWfTwgCwIuM9GQ2kvVotzYJMXTBZg==",
+      "dev": true,
+      "bin": {
+        "tsc": "bin/tsc",
+        "tsserver": "bin/tsserver"
+      },
+      "engines": {
+        "node": ">=4.2.0"
+      }
+    },
     "node_modules/underscore": {
       "version": "1.6.0",
       "resolved": "https://registry.npmjs.org/underscore/-/underscore-1.6.0.tgz",
@@ -6288,6 +6302,12 @@
       "integrity": "sha512-4dbzIzqvjtgiM5rw1k5rEHtBANKmdudhGyBEajN01fEyhaAIhsoKNy6y7+IN93IfpFtwY9iqi7kD+xwKhQsNJA==",
       "dev": true
     },
+    "typescript": {
+      "version": "4.9.4",
+      "resolved": "https://registry.npmjs.org/typescript/-/typescript-4.9.4.tgz",
+      "integrity": "sha512-Uz+dTXYzxXXbsFpM86Wh3dKCxrQqUcVMxwU54orwlJjOpO3ao8L7j5lH+dWfTwgCwIuM9GQ2kvVotzYJMXTBZg==",
+      "dev": true
+    },
     "underscore": {
       "version": "1.6.0",
       "resolved": "https://registry.npmjs.org/underscore/-/underscore-1.6.0.tgz",
diff --git a/package.json b/package.json
index e89f3bd..09b3fe8 100644
--- a/package.json
+++ b/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"
   }
 }
diff --git a/tsconfig.json b/tsconfig.json
index cb21c27..c5b879a 100644
--- a/tsconfig.json
+++ b/tsconfig.json
@@ -4,5 +4,6 @@
     "target": "esnext",
     "allowSyntheticDefaultImports": true,
     "moduleResolution": "nodenext",
+    "noEmit": true
   }
 }

From 9f962309aff64e39d8688bfe684eef9deabe610e Mon Sep 17 00:00:00 2001
From: Justin Ridgewell <justin@ridgewell.name>
Date: Sat, 14 Jan 2023 14:52:22 -0500
Subject: [PATCH 09/16] Fix branch name

---
 .github/workflows/test.yml | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index bc16e76..56b3bfc 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -2,9 +2,9 @@ name: Node.js CI
 
 on:
   push:
-    branches: [ main ]
+    branches: [ master ]
   pull_request:
-    branches: [ main ]
+    branches: [ master ]
 
 jobs:
   test:

From 02a10a46dda9f8551a8bc12c629c79ba2b222dc7 Mon Sep 17 00:00:00 2001
From: Justin Ridgewell <justin@ridgewell.name>
Date: Sat, 14 Jan 2023 14:54:08 -0500
Subject: [PATCH 10/16] Fix node 14 tests

---
 tests/async-context.test.ts | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/async-context.test.ts b/tests/async-context.test.ts
index 8ce3218..f72fdfd 100644
--- a/tests/async-context.test.ts
+++ b/tests/async-context.test.ts
@@ -1,5 +1,5 @@
 import { AsyncContext } from "../src/index";
-import assert from "node:assert/strict";
+import { strict as assert } from "assert";
 
 type Value = { id: number };
 

From 1873d4330b68bb49eb95aa97afbf18b6f03d8c30 Mon Sep 17 00:00:00 2001
From: Justin Ridgewell <justin@ridgewell.name>
Date: Sat, 14 Jan 2023 16:10:08 -0500
Subject: [PATCH 11/16] Isolate initial state to Mapping

---
 src/fork.ts    |  6 +++---
 src/mapping.ts | 17 +++++++++--------
 src/storage.ts | 16 +++++-----------
 3 files changed, 17 insertions(+), 22 deletions(-)

diff --git a/src/fork.ts b/src/fork.ts
index 78cf460..902cd5f 100644
--- a/src/fork.ts
+++ b/src/fork.ts
@@ -10,9 +10,9 @@ import type { AsyncContext } from "./index";
  * mapping) as an OwnedFork would.
  */
 export class FrozenFork {
-  #mapping: Mapping | undefined;
+  #mapping: Mapping;
 
-  constructor(mapping: Mapping | undefined) {
+  constructor(mapping: Mapping) {
     this.#mapping = mapping;
   }
 
@@ -23,7 +23,7 @@ export class FrozenFork {
    * For FrozenFork, that's as simple as returning the known-frozen Mapping,
    * because we know it can't have been modified.
    */
-  join(_current: Mapping | undefined): Mapping | undefined {
+  join(_current: Mapping): Mapping {
     return this.#mapping;
   }
 }
diff --git a/src/mapping.ts b/src/mapping.ts
index 1c7ed4e..bcd2dcf 100644
--- a/src/mapping.ts
+++ b/src/mapping.ts
@@ -5,25 +5,26 @@ import type { AsyncContext } from "./index";
  * 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 = false;
-
-  #data: Map<AsyncContext<unknown>, unknown>;
+  #frozen: boolean;
 
-  constructor(data: Map<AsyncContext<unknown>, unknown>) {
+  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);
+    return this.#data?.has(key) || false;
   }
 
   get<T>(key: AsyncContext<T>): T | undefined {
-    return this.#data.get(key) as T | undefined;
+    return this.#data?.get(key) as T | undefined;
   }
 
   /**
@@ -32,7 +33,7 @@ export class Mapping {
    */
   set<T>(key: AsyncContext<T>, value: T): Mapping {
     const mapping = this.#fork();
-    mapping.#data.set(key, value);
+    mapping.#data!.set(key, value);
     return mapping;
   }
 
@@ -42,7 +43,7 @@ export class Mapping {
    */
   delete<T>(key: AsyncContext<T>): Mapping {
     const mapping = this.#fork();
-    mapping.#data.delete(key);
+    mapping.#data!.delete(key);
     return mapping;
   }
 
diff --git a/src/storage.ts b/src/storage.ts
index 2381b10..fc1dcef 100644
--- a/src/storage.ts
+++ b/src/storage.ts
@@ -9,23 +9,22 @@ import type { AsyncContext } from "./index";
  * None of the methods here are exposed to users, they're only exposed to the AsyncContext class.
  */
 export class Storage {
-  static #current: Mapping | undefined = undefined;
+  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);
+    return this.#current.get(key);
   }
 
   /**
    * Set assigns a new value to the AsyncContext.
    */
   static set<T>(key: AsyncContext<T>, value: T): void {
-    const current = this.#current || new Mapping(new Map());
     // If the Mappings are frozen (someone has snapshot it), then modifying the
     // mappings will return a clone containing the modification.
-    this.#current = current.set(key, value);
+    this.#current = this.#current.set(key, value);
   }
 
   /**
@@ -38,7 +37,7 @@ export class Storage {
    */
   static fork<T>(key: AsyncContext<T>): FrozenFork | OwnedFork<T> {
     const current = this.#current;
-    if (current === undefined || current.isFrozen()) {
+    if (current.isFrozen()) {
       return new FrozenFork(current);
     }
     return new OwnedFork(current, key);
@@ -49,12 +48,7 @@ export class Storage {
    * fork.
    */
   static join<T>(fork: FrozenFork | OwnedFork<T>): void {
-    // The only way for #current to be undefined at a join is if we're in the
-    // we've snapshot the initial empty state with `wrap` and restored it. In
-    // which case, we're operating on a FrozenFork, and the param doesn't
-    // matter. The only other call to join is in the `run` case, and that
-    // guarantees that we have a mappings.
-    this.#current = fork.join(this.#current!);
+    this.#current = fork.join(this.#current);
   }
 
   /**

From 859156984bd697df80452c8a432b8dde68782ddc Mon Sep 17 00:00:00 2001
From: Justin Ridgewell <justin@ridgewell.name>
Date: Sat, 14 Jan 2023 17:21:03 -0500
Subject: [PATCH 12/16] Test run within wrap

---
 tests/async-context.test.ts | 51 +++++++++++++++++++++++++++++++++++++
 1 file changed, 51 insertions(+)

diff --git a/tests/async-context.test.ts b/tests/async-context.test.ts
index f72fdfd..ff3a1ab 100644
--- a/tests/async-context.test.ts
+++ b/tests/async-context.test.ts
@@ -92,6 +92,57 @@ describe("sync", () => {
       assert.equal(ctx.get(), undefined);
     });
 
+    it("runs within wrap", () => {
+      const ctx = new AsyncContext<Value>();
+      const first = { id: 1 };
+      const second = { id: 2 };
+
+      const wrap = ctx.run(first, () => {
+        const wrap = AsyncContext.wrap(() => {
+          assert.equal(ctx.get(), first);
+          ctx.run(second, () => {
+            assert.equal(ctx.get(), second);
+          });
+          assert.equal(ctx.get(), first);
+        });
+        assert.equal(ctx.get(), first);
+        return wrap;
+      });
+
+      wrap();
+      assert.equal(ctx.get(), undefined);
+    });
+
+    it("runs different context within wrap", () => {
+      const a = new AsyncContext<Value>();
+      const b = new AsyncContext<Value>();
+      const first = { id: 1 };
+      const second = { id: 2 };
+
+      const wrap = a.run(first, () => {
+        const wrap = AsyncContext.wrap(() => {
+          assert.equal(a.get(), first);
+          assert.equal(b.get(), undefined);
+
+          b.run(second, () => {
+            assert.equal(a.get(), first);
+            assert.equal(b.get(), second);
+          });
+
+          assert.equal(a.get(), first);
+          assert.equal(b.get(), undefined);
+        });
+
+        assert.equal(a.get(), first);
+        assert.equal(b.get(), undefined);
+        return wrap;
+      });
+
+      wrap();
+      assert.equal(a.get(), undefined);
+      assert.equal(b.get(), undefined);
+    });
+
     it("wrap within nesting contexts", () => {
       const ctx = new AsyncContext<Value>();
       const first = { id: 1 };

From 168d8fb40da0c4b02346cea932187d8d4742c827 Mon Sep 17 00:00:00 2001
From: Justin Ridgewell <justin@ridgewell.name>
Date: Sat, 14 Jan 2023 20:07:13 -0500
Subject: [PATCH 13/16] Remove debuggers

---
 tests/async-context.test.ts | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/tests/async-context.test.ts b/tests/async-context.test.ts
index ff3a1ab..f1f1d12 100644
--- a/tests/async-context.test.ts
+++ b/tests/async-context.test.ts
@@ -231,7 +231,6 @@ describe("sync", () => {
       const wrap = a.run(first, () => {
         const wrap = b.run(second, () => {
           const wrap = c.run(third, () => {
-            debugger;
             return AsyncContext.wrap(() => {
               assert.equal(a.get(), first);
               assert.equal(b.get(), second);
@@ -272,7 +271,6 @@ describe("sync", () => {
           AsyncContext.wrap(() => {});
 
           const wrap = c.run(third, () => {
-            debugger;
             return AsyncContext.wrap(() => {
               assert.equal(a.get(), first);
               assert.equal(b.get(), second);
@@ -313,7 +311,6 @@ describe("sync", () => {
 
         const wrap = b.run(second, () => {
           const wrap = c.run(third, () => {
-            debugger;
             return AsyncContext.wrap(() => {
               assert.equal(a.get(), first);
               assert.equal(b.get(), second);

From b31b1ff23abb04f96c2d0ff2dc2cd80e7dc3b70a Mon Sep 17 00:00:00 2001
From: Justin Ridgewell <justin@ridgewell.name>
Date: Sun, 15 Jan 2023 14:14:38 -0500
Subject: [PATCH 14/16] Ensure we test both initial and run states

---
 tests/async-context.test.ts | 194 +++++++++++++++++++++++++++++++-----
 1 file changed, 169 insertions(+), 25 deletions(-)

diff --git a/tests/async-context.test.ts b/tests/async-context.test.ts
index f1f1d12..f4b5376 100644
--- a/tests/async-context.test.ts
+++ b/tests/async-context.test.ts
@@ -3,9 +3,34 @@ import { strict as assert } from "assert";
 
 type Value = { id: number };
 
+const _it = it;
+it = (() => {
+  throw new Error("use `test` function");
+}) as any;
+
+// Test both from the initial state, and from a run state.
+// This is because the initial state might be "frozen", and
+// that can cause different code paths.
+function test(name: string, fn: () => void) {
+  _it(name, () => {
+    fn();
+
+    // Ensure we're running from a new state, which won't be frozen.
+    const throwaway = new AsyncContext<null>();
+    throwaway.run(null, fn);
+
+    throwaway.run(null, () => {
+      AsyncContext.wrap(() => {});
+
+      // Ensure we're running from a new state, which is frozen.
+      fn();
+    });
+  });
+}
+
 describe("sync", () => {
   describe("run and get", () => {
-    it("has initial undefined state", () => {
+    test("has initial undefined state", () => {
       const ctx = new AsyncContext<Value>();
 
       const actual = ctx.get();
@@ -13,7 +38,7 @@ describe("sync", () => {
       assert.equal(actual, undefined);
     });
 
-    it("return value", () => {
+    test("return value", () => {
       const ctx = new AsyncContext<Value>();
       const expected = { id: 1 };
 
@@ -22,7 +47,7 @@ describe("sync", () => {
       assert.equal(actual, expected);
     });
 
-    it("get returns current context value", () => {
+    test("get returns current context value", () => {
       const ctx = new AsyncContext<Value>();
       const expected = { id: 1 };
 
@@ -31,7 +56,7 @@ describe("sync", () => {
       });
     });
 
-    it("get within nesting contexts", () => {
+    test("get within nesting contexts", () => {
       const ctx = new AsyncContext<Value>();
       const first = { id: 1 };
       const second = { id: 2 };
@@ -46,7 +71,7 @@ describe("sync", () => {
       assert.equal(ctx.get(), undefined);
     });
 
-    it("get within nesting different contexts", () => {
+    test("get within nesting different contexts", () => {
       const a = new AsyncContext<Value>();
       const b = new AsyncContext<Value>();
       const first = { id: 1 };
@@ -68,7 +93,7 @@ describe("sync", () => {
   });
 
   describe("wrap", () => {
-    it("stores initial undefined state", () => {
+    test("stores initial undefined state", () => {
       const ctx = new AsyncContext<Value>();
       const wrapped = AsyncContext.wrap(() => ctx.get());
 
@@ -77,7 +102,7 @@ describe("sync", () => {
       });
     });
 
-    it("stores current state", () => {
+    test("stores current state", () => {
       const ctx = new AsyncContext<Value>();
       const expected = { id: 1 };
 
@@ -92,35 +117,107 @@ describe("sync", () => {
       assert.equal(ctx.get(), undefined);
     });
 
-    it("runs within wrap", () => {
+    test("runs within wrap", () => {
       const ctx = new AsyncContext<Value>();
       const first = { id: 1 };
       const second = { id: 2 };
 
-      const wrap = ctx.run(first, () => {
-        const wrap = AsyncContext.wrap(() => {
+      const [wrap1, wrap2] = ctx.run(first, () => {
+        const wrap1 = AsyncContext.wrap(() => {
           assert.equal(ctx.get(), first);
+
           ctx.run(second, () => {
             assert.equal(ctx.get(), second);
           });
+
           assert.equal(ctx.get(), first);
         });
         assert.equal(ctx.get(), first);
-        return wrap;
+
+        ctx.run(second, () => {
+          assert.equal(ctx.get(), second);
+        });
+
+        const wrap2 = AsyncContext.wrap(() => {
+          assert.equal(ctx.get(), first);
+
+          ctx.run(second, () => {
+            assert.equal(ctx.get(), second);
+          });
+
+          assert.equal(ctx.get(), first);
+        });
+        assert.equal(ctx.get(), first);
+        return [wrap1, wrap2];
       });
 
-      wrap();
+      wrap1();
+      wrap2();
+      assert.equal(ctx.get(), undefined);
+    });
+
+    test("runs within wrap", () => {
+      const ctx = new AsyncContext<Value>();
+      const first = { id: 1 };
+      const second = { id: 2 };
+
+      const [wrap1, wrap2] = ctx.run(first, () => {
+        const wrap1 = AsyncContext.wrap(() => {
+          assert.equal(ctx.get(), first);
+
+          ctx.run(second, () => {
+            assert.equal(ctx.get(), second);
+          });
+
+          assert.equal(ctx.get(), first);
+        });
+        assert.equal(ctx.get(), first);
+
+        ctx.run(second, () => {
+          assert.equal(ctx.get(), second);
+        });
+
+        const wrap2 = AsyncContext.wrap(() => {
+          assert.equal(ctx.get(), first);
+
+          ctx.run(second, () => {
+            assert.equal(ctx.get(), second);
+          });
+
+          assert.equal(ctx.get(), first);
+        });
+        assert.equal(ctx.get(), first);
+        return [wrap1, wrap2];
+      });
+
+      wrap1();
+      wrap2();
       assert.equal(ctx.get(), undefined);
     });
 
-    it("runs different context within wrap", () => {
+    test("runs different context within wrap", () => {
       const a = new AsyncContext<Value>();
       const b = new AsyncContext<Value>();
       const first = { id: 1 };
       const second = { id: 2 };
 
-      const wrap = a.run(first, () => {
-        const wrap = AsyncContext.wrap(() => {
+      const [wrap1, wrap2] = a.run(first, () => {
+        const wrap1 = AsyncContext.wrap(() => {
+          assert.equal(a.get(), first);
+          assert.equal(b.get(), undefined);
+
+          b.run(second, () => {
+            assert.equal(a.get(), first);
+            assert.equal(b.get(), second);
+          });
+
+          assert.equal(a.get(), first);
+          assert.equal(b.get(), undefined);
+        });
+
+        a.run(second, () => {});
+
+        const wrap2 = AsyncContext.wrap(() => {
           assert.equal(a.get(), first);
           assert.equal(b.get(), undefined);
 
@@ -135,15 +232,62 @@ describe("sync", () => {
 
         assert.equal(a.get(), first);
         assert.equal(b.get(), undefined);
-        return wrap;
+        return [wrap1, wrap2];
       });
 
-      wrap();
+      wrap1();
+      wrap2();
+      assert.equal(a.get(), undefined);
+      assert.equal(b.get(), undefined);
+    });
+
+    test("runs different context within wrap, 2", () => {
+      const a = new AsyncContext<Value>();
+      const b = new AsyncContext<Value>();
+      const first = { id: 1 };
+      const second = { id: 2 };
+
+      const [wrap1, wrap2] = a.run(first, () => {
+        const wrap1 = AsyncContext.wrap(() => {
+          assert.equal(a.get(), first);
+          assert.equal(b.get(), undefined);
+
+          b.run(second, () => {
+            assert.equal(a.get(), first);
+            assert.equal(b.get(), second);
+          });
+
+          assert.equal(a.get(), first);
+          assert.equal(b.get(), undefined);
+        });
+
+        b.run(second, () => {});
+
+        const wrap2 = AsyncContext.wrap(() => {
+          assert.equal(a.get(), first);
+          assert.equal(b.get(), undefined);
+
+          b.run(second, () => {
+            assert.equal(a.get(), first);
+            assert.equal(b.get(), second);
+          });
+
+          assert.equal(a.get(), first);
+          assert.equal(b.get(), undefined);
+        });
+
+        assert.equal(a.get(), first);
+        assert.equal(b.get(), undefined);
+        return [wrap1, wrap2];
+      });
+
+      wrap1();
+      wrap2();
       assert.equal(a.get(), undefined);
       assert.equal(b.get(), undefined);
     });
 
-    it("wrap within nesting contexts", () => {
+    test("wrap within nesting contexts", () => {
       const ctx = new AsyncContext<Value>();
       const first = { id: 1 };
       const second = { id: 2 };
@@ -178,7 +322,7 @@ describe("sync", () => {
       assert.equal(ctx.get(), undefined);
     });
 
-    it("wrap within nesting different contexts", () => {
+    test("wrap within nesting different contexts", () => {
       const a = new AsyncContext<Value>();
       const b = new AsyncContext<Value>();
       const first = { id: 1 };
@@ -220,7 +364,7 @@ describe("sync", () => {
       assert.equal(b.get(), undefined);
     });
 
-    it("wrap within nesting different contexts, 2", () => {
+    test("wrap within nesting different contexts, 2", () => {
       const a = new AsyncContext<Value>();
       const b = new AsyncContext<Value>();
       const c = new AsyncContext<Value>();
@@ -258,7 +402,7 @@ describe("sync", () => {
       assert.equal(c.get(), undefined);
     });
 
-    it("wrap within nesting different contexts, 3", () => {
+    test("wrap within nesting different contexts, 3", () => {
       const a = new AsyncContext<Value>();
       const b = new AsyncContext<Value>();
       const c = new AsyncContext<Value>();
@@ -298,7 +442,7 @@ describe("sync", () => {
       assert.equal(c.get(), undefined);
     });
 
-    it("wrap within nesting different contexts, 4", () => {
+    test("wrap within nesting different contexts, 4", () => {
       const a = new AsyncContext<Value>();
       const b = new AsyncContext<Value>();
       const c = new AsyncContext<Value>();
@@ -338,7 +482,7 @@ describe("sync", () => {
       assert.equal(c.get(), undefined);
     });
 
-    it("wrap within nesting different contexts, 5", () => {
+    test("wrap within nesting different contexts, 5", () => {
       const a = new AsyncContext<Value>();
       const b = new AsyncContext<Value>();
       const c = new AsyncContext<Value>();
@@ -379,7 +523,7 @@ describe("sync", () => {
       assert.equal(c.get(), undefined);
     });
 
-    it("wrap within nesting different contexts, 6", () => {
+    test("wrap within nesting different contexts, 6", () => {
       const a = new AsyncContext<Value>();
       const b = new AsyncContext<Value>();
       const c = new AsyncContext<Value>();
@@ -420,7 +564,7 @@ describe("sync", () => {
       assert.equal(c.get(), undefined);
     });
 
-    it("wrap out of order", () => {
+    test("wrap out of order", () => {
       const ctx = new AsyncContext<Value>();
       const first = { id: 1 };
       const second = { id: 2 };

From 235e78044c19a16b0631c7315b4cc1f41155b90c Mon Sep 17 00:00:00 2001
From: Justin Ridgewell <justin@ridgewell.name>
Date: Thu, 19 Jan 2023 00:06:15 -0500
Subject: [PATCH 15/16] Apply suggestions from code review

Co-authored-by: Chengzhong Wu <legendecas@gmail.com>
---
 src/storage.ts | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/storage.ts b/src/storage.ts
index fc1dcef..4dd7ee9 100644
--- a/src/storage.ts
+++ b/src/storage.ts
@@ -57,7 +57,7 @@ export class Storage {
    * snapshot.
    */
   static snapshot(): FrozenFork {
-    this.#current?.freeze();
+    this.#current.freeze();
     return new FrozenFork(this.#current);
   }
 

From 60c3700fa7403cb19bd328c7fdb166490367a0b7 Mon Sep 17 00:00:00 2001
From: Justin Ridgewell <justin@ridgewell.name>
Date: Thu, 19 Jan 2023 00:34:52 -0500
Subject: [PATCH 16/16] Update naming and comments

---
 src/fork.ts    | 41 +++++++++++++++++++-------------------
 src/index.ts   |  9 ++++-----
 src/storage.ts | 53 +++++++++++++++++++++-----------------------------
 3 files changed, 47 insertions(+), 56 deletions(-)

diff --git a/src/fork.ts b/src/fork.ts
index 902cd5f..0b30195 100644
--- a/src/fork.ts
+++ b/src/fork.ts
@@ -2,14 +2,14 @@ import type { Mapping } from "./mapping";
 import type { AsyncContext } from "./index";
 
 /**
- * FrozenFork holds a frozen Mapping that will be simply restored when the fork is
- * rejoined.
+ * 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
- * rejoining will not attempt to mutate the Mapping (and allocate a new
- * mapping) as an OwnedFork would.
+ * reverting will not attempt to mutate the Mapping (and allocate a new
+ * mapping) as a Revert would.
  */
-export class FrozenFork {
+export class FrozenRevert {
   #mapping: Mapping;
 
   constructor(mapping: Mapping) {
@@ -17,28 +17,29 @@ export class FrozenFork {
   }
 
   /**
-   * The Storage container will call join when it wants to restore its current
-   * Mapping to the state at the start of the fork.
+   * The Storage container will call restore when it wants to revert its
+   * current Mapping to the state at the start of the fork.
    *
-   * For FrozenFork, that's as simple as returning the known-frozen Mapping,
+   * For FrozenRevert, that's as simple as returning the known-frozen Mapping,
    * because we know it can't have been modified.
    */
-  join(_current: Mapping): Mapping {
+  restore(_current: Mapping): Mapping {
     return this.#mapping;
   }
 }
 
 /**
- * OwnedFork holds an unfrozen Mapping that we will attempt to modify when
- * rejoining to attempt to restore it to its prior state.
+ * 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 rejoin. In
- * that case, we can simply modify the Mapping (without cloning) to restore it
- * to its prior state. If someone does snapshot it, then modifying will clone
- * the current state and we restore the clone to the prior state.
+ * 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 OwnedFork<T> {
+export class Revert<T> {
   #key: AsyncContext<T>;
   #has: boolean;
   #prev: T | undefined;
@@ -50,14 +51,14 @@ export class OwnedFork<T> {
   }
 
   /**
-   * The Storage container will call join when it wants to restore its current
-   * Mapping to the state at the start of the fork.
+   * The Storage container will call restore when it wants to revert its
+   * current Mapping to the state at the start of the fork.
    *
-   * For OwnedFork, we mutate the known-unfrozen-at-start mapping (which may
+   * 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.
    */
-  join(current: Mapping): Mapping {
+  restore(current: Mapping): Mapping {
     if (this.#has) {
       return current.set(this.#key, this.#prev);
     } else {
diff --git a/src/index.ts b/src/index.ts
index faadd20..ec5977b 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -7,11 +7,11 @@ export class AsyncContext<T> {
     const snapshot = Storage.snapshot();
 
     function wrap(this: ThisType<F>, ...args: Parameters<F>): ReturnType<F> {
-      const fork = Storage.restore(snapshot);
+      const head = Storage.switch(snapshot);
       try {
         return fn.apply(this, args);
       } finally {
-        Storage.join(fork);
+        Storage.restore(head);
       }
     }
 
@@ -23,12 +23,11 @@ export class AsyncContext<T> {
     fn: F,
     ...args: Parameters<F>
   ): ReturnType<F> {
-    const fork = Storage.fork(this);
-    Storage.set(this, value);
+    const revert = Storage.set(this, value);
     try {
       return fn.apply(null, args);
     } finally {
-      Storage.join(fork);
+      Storage.restore(revert);
     }
   }
 
diff --git a/src/storage.ts b/src/storage.ts
index 4dd7ee9..d67fcb1 100644
--- a/src/storage.ts
+++ b/src/storage.ts
@@ -1,5 +1,5 @@
 import { Mapping } from "./mapping";
-import { FrozenFork, OwnedFork } from "./fork";
+import { FrozenRevert, Revert } from "./fork";
 import type { AsyncContext } from "./index";
 
 /**
@@ -19,61 +19,52 @@ export class Storage {
   }
 
   /**
-   * Set assigns a new value to the AsyncContext.
+   * 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): void {
+  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.
-    this.#current = this.#current.set(key, value);
-  }
-
-  /**
-   * Fork is called before modifying the global storage state (either by
-   * replacing the current mappings or assigning a new value to an individual
-   * AsyncContext).
-   *
-   * The Fork instance returned will be able to restore the mappings to the
-   * unmodified state.
-   */
-  static fork<T>(key: AsyncContext<T>): FrozenFork | OwnedFork<T> {
     const current = this.#current;
-    if (current.isFrozen()) {
-      return new FrozenFork(current);
-    }
-    return new OwnedFork(current, key);
+    const undo = current.isFrozen()
+      ? new FrozenRevert(current)
+      : new Revert(current, key);
+    this.#current = this.#current.set(key, value);
+    return undo;
   }
 
   /**
-   * Join will restore the global storage state to state at the time of the
-   * fork.
+   * Restore will, well, restore the global storage state to state at the time
+   * the revert was created.
    */
-  static join<T>(fork: FrozenFork | OwnedFork<T>): void {
-    this.#current = fork.join(this.#current);
+  static restore<T>(revert: FrozenRevert | Revert<T>): void {
+    this.#current = revert.restore(this.#current);
   }
 
   /**
-   * Snapshot freezes the current storage state, and returns a new fork which
+   * 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(): FrozenFork {
+  static snapshot(): FrozenRevert {
     this.#current.freeze();
-    return new FrozenFork(this.#current);
+    return new FrozenRevert(this.#current);
   }
 
   /**
-   * Restore restores the global storage state to the state at the time of the
-   * snapshot.
+   * 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 restore(snapshot: FrozenFork): FrozenFork {
+  static switch(snapshot: FrozenRevert): FrozenRevert {
     const previous = this.#current;
-    this.#current = snapshot.join(previous);
+    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 FrozenFork(previous);
+    return new FrozenRevert(previous);
   }
 }