From c8968a9a98f0554f1164b5857a78910cf61f1da4 Mon Sep 17 00:00:00 2001
From: artalar <artalarut@proton.me>
Date: Wed, 25 Oct 2023 12:43:39 +0300
Subject: [PATCH] fix(core): conditional deps duplication

This change improves memory and CPU usage for most cases and has 2-10 times CPU degradation for hundreds and thousands of spy in one atom (such as reducing under atomized array).
---
 packages/core/src/atom.test.ts | 32 ++++++++++++++++++++++++++++++++
 packages/core/src/atom.ts      | 29 +++++++++++------------------
 2 files changed, 43 insertions(+), 18 deletions(-)

diff --git a/packages/core/src/atom.test.ts b/packages/core/src/atom.test.ts
index 232c0dce3..880b0f711 100644
--- a/packages/core/src/atom.test.ts
+++ b/packages/core/src/atom.test.ts
@@ -675,6 +675,38 @@ test('ctx collision', () => {
   ;`👍` //?
 })
 
+test('conditional deps duplication', () => {
+  const listAtom = atom([1, 2, 3])
+
+  const filterAtom = atom<'odd' | 'even'>('odd')
+
+  const filteredListAtom = atom((ctx) => {
+    if (ctx.spy(filterAtom) === 'odd') {
+      return ctx.spy(listAtom).filter((n) => n % 2 === 1)
+    } else if (ctx.spy(filterAtom) === 'even') {
+      return ctx.spy(listAtom).filter((n) => n % 2 === 0)
+    }
+    return ctx.spy(listAtom)
+  })
+
+  const ctx = createCtx()
+
+  const track = mockFn()
+
+  ctx.subscribe(filteredListAtom, track)
+  assert.equal(track.lastInput(), [1, 3])
+
+  filterAtom(ctx, 'even')
+  assert.equal(track.lastInput(), [2])
+
+  filterAtom(ctx, 'odd')
+  assert.equal(track.lastInput(), [1, 3])
+
+  filterAtom(ctx, 'even')
+  assert.equal(track.lastInput(), [2])
+  ;`👍` //?
+})
+
 // test(`maximum call stack`, () => {
 //   const atoms = new Map<AtomProto, Atom>()
 //   let i = 0
diff --git a/packages/core/src/atom.ts b/packages/core/src/atom.ts
index 267b7865c..ad9610155 100644
--- a/packages/core/src/atom.ts
+++ b/packages/core/src/atom.ts
@@ -337,8 +337,7 @@ export const createCtx = ({
 
   let actualizePubs = (patchCtx: Ctx, patch: AtomCache) => {
     let { proto, pubs } = patch
-    let toDisconnect = new Set<AtomProto>()
-    let toConnect = new Set<AtomProto>()
+    let isDepsChanged = false
 
     if (
       pubs.length === 0 ||
@@ -358,11 +357,7 @@ export const createCtx = ({
               ? pubs[newPubs.length - 1]
               : undefined
           let isDepChanged = prevDepPatch?.proto !== depPatch.proto
-
-          if (isDepChanged) {
-            if (prevDepPatch) toDisconnect.add(prevDepPatch.proto)
-            toConnect.add(depProto)
-          }
+          isDepsChanged ||= isDepChanged
 
           let state =
             depProto.isAction && !isDepChanged
@@ -370,7 +365,7 @@ export const createCtx = ({
               : depPatch.state
 
           if (cb && (isDepChanged || !Object.is(state, prevDepPatch!.state))) {
-            if (depProto.isAction) (state as any[]).forEach((call) => cb(call))
+            if (depProto.isAction) for (const call of state) cb(call)
             else cb(state, prevDepPatch?.state)
           } else {
             return state
@@ -383,18 +378,16 @@ export const createCtx = ({
       patch.state = patch.proto.computer!(patchCtx as CtxSpy, patch.state)
       patch.pubs = newPubs
 
-      for (let i = newPubs.length; i < pubs.length; i++) {
-        toDisconnect.add(pubs[i]!.proto)
-      }
-
-      if (toDisconnect.size + toConnect.size && isConnected(patch)) {
-        for (let depProto of toDisconnect) {
-          toConnect.has(depProto) ||
+      if ((isDepsChanged || pubs.length > newPubs.length) && isConnected(patch)) {
+        for (let { proto: depProto } of pubs) {
+          if (newPubs.every((dep) => dep.proto !== depProto)) {
             disconnect(proto, depProto.patch ?? read(depProto)!)
+          }
         }
-
-        for (let depProto of toConnect) {
-          connect(proto, depProto.patch ?? read(depProto)!)
+        for (let { proto: depProto } of newPubs) {
+          if (pubs.every((dep) => dep.proto !== depProto)) {
+            connect(proto, depProto.patch ?? read(depProto)!)
+          }
         }
       }
     }