From 3ee8b97f9b1662aaa9c2eaedc13fb0066f5f0ba0 Mon Sep 17 00:00:00 2001
From: Matteo Collina <hello@matteocollina.com>
Date: Tue, 21 May 2024 16:20:50 +0200
Subject: [PATCH] added store.enterWith()

Signed-off-by: Matteo Collina <hello@matteocollina.com>
---
 README.md            | 33 ++++++++++++++++++++-------------
 asyncforge.d.ts      |  1 +
 asyncforge.js        |  4 ++++
 asyncforge.test-d.ts |  1 +
 example.mjs          |  7 +++++++
 test/memo.test.js    | 24 ++++++++++++++++++++++++
 6 files changed, 57 insertions(+), 13 deletions(-)

diff --git a/README.md b/README.md
index bc82aa0..708deb4 100644
--- a/README.md
+++ b/README.md
@@ -50,6 +50,13 @@ create(() => {
         console.log('b', b())
       })
     })
+
+    setImmediate(() => {
+      store.enterWith()
+      console.log('-- fourth event loop turn --')
+      console.log('a', a())
+      console.log('b', b())
+    })
   })
 })
 ```
@@ -59,25 +66,25 @@ create(() => {
 You can call the `asyncforge` functions in a type-safe way:
 
 ```ts
-import { start, memo, setAll } from "asyncforge";
+import { create, memo } from "asyncforge";
 const memoNum = memo<number>();
+const test = memo<string>();
 
-// This is okay for TypeScript, since you're passing a number
-memoNum.set(123);
-
-// This will not build
-memoNum.set('wrong');
+const store = create()
 
-// The `result` var will be of type `number`
-const result = memoNum()
+store.run(() => {
+  // This is okay for TypeScript, since you're passing a number
+  memoNum.set(123);
 
-const test = memo<string>();
 
-// This is correct, since `test.key` is a `symbol`
-setAll({ [test.key]: 42 });
+  // This will not build
+  memoNum.set('wrong');
+})
 
-// This will fail
-setAll({ 'wrong': 42 });
+store.run(() => {
+  // The `result` var will be of type `number`
+  const result = memoNum()
+})
 ```
 
 ## License
diff --git a/asyncforge.d.ts b/asyncforge.d.ts
index 97885ba..c6118b6 100644
--- a/asyncforge.d.ts
+++ b/asyncforge.d.ts
@@ -1,6 +1,7 @@
 declare namespace asyncforge {
   interface Store {
     run<T>(fn: () => T) : T;
+    enterWith() : void;
   }
 
   export function create () : Store;
diff --git a/asyncforge.js b/asyncforge.js
index 2bfe3e4..22f8ee0 100644
--- a/asyncforge.js
+++ b/asyncforge.js
@@ -14,6 +14,10 @@ class Store {
   run (fn) {
     return asyncLocalStorage.run(this.#internal, fn)
   }
+
+  enterWith () {
+    asyncLocalStorage.enterWith(this.#internal)
+  }
 }
 
 function create (fn) {
diff --git a/asyncforge.test-d.ts b/asyncforge.test-d.ts
index 15791a3..95d15f4 100644
--- a/asyncforge.test-d.ts
+++ b/asyncforge.test-d.ts
@@ -7,6 +7,7 @@ expectType<() => Store>(create);
 const store = create();
 
 expectType<number>(store.run(() => 42));
+expectType<void>(store.enterWith());
 
 // memo
 const memoNum = memo<number>();
diff --git a/example.mjs b/example.mjs
index 07dc9fe..7720771 100644
--- a/example.mjs
+++ b/example.mjs
@@ -34,5 +34,12 @@ create(() => {
         console.log('b', b())
       })
     })
+
+    setImmediate(() => {
+      store.enterWith()
+      console.log('-- fourth event loop turn --')
+      console.log('a', a())
+      console.log('b', b())
+    })
   })
 })
diff --git a/test/memo.test.js b/test/memo.test.js
index 3a77b68..18c9708 100644
--- a/test/memo.test.js
+++ b/test/memo.test.js
@@ -116,3 +116,27 @@ test('run multiple times', async (t) => {
 
   await p.completed
 })
+
+test('enterWith', async (t) => {
+  const p = tspl(t, { plan: 5 })
+  const a = memo()
+
+  p.throws(a, /asyncforge store has not been created/)
+  p.throws(() => a.set('foo'), /asyncforge store has not been created/)
+
+  create().enterWith()
+
+  p.deepStrictEqual(a(), undefined)
+  a.set({ value: 'bar' })
+  p.deepStrictEqual(a(), { value: 'bar' })
+
+  setImmediate(() => {
+    p.deepStrictEqual(a(), { value: 'bar' })
+  })
+
+  queueMicrotask(() => {
+    p.deepStrictEqual(a(), { value: 'bar' })
+  })
+
+  await p.completed
+})