Fix Atom dependencies during batch rebuilds

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

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

packages/effect/src/unstable/reactivity/AtomRegistry.ts

@@ -606,6 +606,8 @@ class NodeImpl<A> {
   children = new Set<NodeImpl<any>>()
   listeners = new Set<() => void>()
   skipInvalidation = false
+  building = false
+  invalidatedDuringBuild = false
 
   currentState() {
     switch (this.state) {
@@ -628,7 +630,9 @@ class NodeImpl<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) {
         if (this.preserveInitialValueOnBuild) {
           this.preserveInitialValueOnBuild = false
@@ -735,6 +739,9 @@ class NodeImpl<A> {
   }
 
   invalidate(): void {
+    if (this.building && batchState.phase === BatchPhase.collect) {
+      this.invalidatedDuringBuild = true
+    }
     if (this.state === NodeState.valid) {
       this.state = NodeState.stale
       this.disposeLifetime()
@@ -860,8 +867,9 @@ const LifetimeProto: Omit<Lifetime<any>, "node" | "finalizers" | "disposed" | "i
       return this.node.registry.get(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.AsyncResult<A, E>>, options?: {
@@ -1110,7 +1118,12 @@ export function batch(f: () => void): void {
 
 function batchRebuildNode(node: NodeImpl<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/effect/test/reactivity/Atom.test.ts

@@ -861,6 +861,64 @@ describe.sequential("Atom", () => {
     expect(r.get(derived)).toEqual("2b")
   })
 
+  it.effect("retains method-form dependencies added during a batch rebuild", () =>
+    Effect.gen(function*() {
+      const registry = AtomRegistry.make()
+      const source = Atom.make(Option.none<string>())
+      const gate = yield* Latch.make()
+      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")))
+
+      yield* gate.open
+      yield* Effect.yieldNow
+
+      const result = registry.get(derived) as AsyncResult.AsyncResult<string, "SourceIsNone">
+      assert(AsyncResult.isSuccess(result))
+      assert.strictEqual(result.value, "computed-a")
+    }))
+
+  it("rebuilds an atom invalidated during its own batch rebuild", () => {
+    const registry = AtomRegistry.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 = AtomRegistry.make()
     const state = Atom.make(1).pipe(Atom.keepAlive)