Fix Atom dependencies during batch rebuilds (#423)

Author: tim-smart

.changeset/fix-atom-batch-dependencies.md

@@ -0,0 +1,5 @@
+---
+"@effect-atom/atom": patch
+---
+
+Fix Atom dependency tracking and re-entrant invalidation during batch rebuilds.

packages/atom/src/internal/registry.ts

@@ -289,6 +289,8 @@ class Node<A> {
   children = new Set<Node<any>>()
   listeners = new Set<() => void>()
   skipInvalidation = false
+  building = false
+  invalidatedDuringBuild = false
 
   get canBeRemoved(): boolean {
     return !this.atom.keepAlive && this.listeners.size === 0 && this.children.size === 0 &&
@@ -299,7 +301,9 @@ class Node<A> {
   value(): A {
     if ((this.state & NodeFlags.waitingForValue) !== 0) {
       this.lifetime = makeLifetime(this)
+      this.building = true
       const value = this.atom.read(this.lifetime)
+      this.building = false
       if ((this.state & NodeFlags.waitingForValue) !== 0) {
         this.setValue(value)
       }
@@ -383,6 +387,9 @@ class Node<A> {
   }
 
   invalidate(): void {
+    if (this.building && batchState.phase === BatchPhase.collect) {
+      this.invalidatedDuringBuild = true
+    }
     if (this.state === NodeState.valid) {
       this.state = NodeState.stale
       this.disposeLifetime()
@@ -512,8 +519,9 @@ const LifetimeProto: Omit<Lifetime<any>, "node" | "finalizers" | "disposed" | "i
       throw disposedError(this.node.atom)
     }
     const parent = this.node.registry.ensureNode(atom)
+    const value = parent.value()
     this.node.addParent(parent)
-    return parent.value()
+    return value
   },
 
   result<A, E>(this: Lifetime<any>, atom: Atom.Atom<Result.Result<A, E>>, options?: {
@@ -795,7 +803,12 @@ export function batch(f: () => void): void {
 
 function batchRebuildNode(node: Node<any>) {
   if (node.state === NodeState.valid) {
-    return
+    if (!node.invalidatedDuringBuild) {
+      return
+    }
+    node.invalidatedDuringBuild = false
+    node.state = NodeState.stale
+    node.disposeLifetime()
   }
 
   for (const parent of node.parents) {

packages/atom/test/Atom.test.ts

@@ -697,6 +697,63 @@ describe("Atom", () => {
     expect(r.get(derived)).toEqual("2b")
   })
 
+  it("retains method-form dependencies added during a batch rebuild", async () => {
+    const registry = Registry.make()
+    const source = Atom.make(Option.none<string>())
+    const gate = Effect.unsafeMakeLatch()
+    const asyncAtom = Atom.make((get) =>
+      Effect.gen(function*() {
+        const value = get(source)
+        if (Option.isNone(value)) {
+          return yield* Effect.fail("SourceIsNone" as const)
+        }
+        yield* gate.await
+        return `computed-${value.value}`
+      })
+    )
+    const derived = Atom.make((get): unknown => {
+      const value = get.get(source)
+      if (Option.isNone(value)) {
+        return "empty"
+      }
+      return get.get(asyncAtom)
+    })
+
+    registry.subscribe(derived, () => {}, { immediate: true })
+    registry.subscribe(asyncAtom, () => {}, { immediate: true })
+
+    Atom.batch(() => registry.set(source, Option.some("a")))
+
+    gate.unsafeOpen()
+    await Effect.runPromise(Effect.yieldNow())
+
+    const result = registry.get(derived) as Result.Result<string, "SourceIsNone">
+    assert(Result.isSuccess(result))
+    assert.strictEqual(result.value, "computed-a")
+  })
+
+  it("rebuilds an atom invalidated during its own batch rebuild", () => {
+    const registry = Registry.make()
+    const source = Atom.make(0)
+    const enabled = Atom.make(false)
+    const updateSource = Atom.make((get) => {
+      get.set(source, 1)
+    })
+    const derived = Atom.make((get) => {
+      const value = get(source)
+      if (get(enabled)) {
+        get(updateSource)
+      }
+      return value
+    })
+
+    registry.subscribe(derived, () => {}, { immediate: true })
+
+    Atom.batch(() => registry.set(enabled, true))
+
+    assert.strictEqual(registry.get(derived), 1)
+  })
+
   it("nested batch", async () => {
     const r = Registry.make()
     const state = Atom.make(1).pipe(Atom.keepAlive)