chore!: remove store.batch, add in temporary scheduler, skip intentio…

…nally broken tests

packages/store/src/derived.ts

@@ -1,4 +1,10 @@
 import { Store } from './store'
+import {
+  __depsThatHaveWrittenThisTick,
+  __derivedToStore,
+  __storeToDerived,
+  __whatStoreIsCurrentlyInUse,
+} from './scheduler'
 import type { Listener } from './types'
 
 export type UnwrapDerivedOrStore<T> =
@@ -65,10 +71,6 @@ export class Derived<
    * @private
    */
   _store: Store<TState>
-  /**
-   * @private
-   */
-  rootStores = new Set<Store<unknown>>()
   options: DerivedOptions<TState, TArr>
 
   /**
@@ -77,30 +79,6 @@ export class Derived<
    */
   _subscriptions: Array<() => void> = []
 
-  /**
-   * What store called the current update, if any
-   * @private
-   */
-  _whatStoreIsCurrentlyInUse: Store<unknown> | null = null
-
-  /**
-   * This is here to solve the pyramid dependency problem where:
-   *       A
-   *      / \
-   *     B   C
-   *      \ /
-   *       D
-   *
-   * Where we deeply traverse this tree, how do we avoid D being recomputed twice; once when B is updated, once when C is.
-   *
-   * To solve this, we create linkedDeps that allows us to sync avoid writes to the state until all of the deps have been
-   * resolved.
-   *
-   * This is a record of stores, because derived stores are not able to write values to, but stores are
-   */
-  storeToDerived = new Map<Store<unknown>, Set<Derived<unknown>>>()
-  derivedToStore = new Map<Derived<unknown>, Set<Store<unknown>>>()
-
   getDepVals = () => {
     const prevDepVals = [] as Array<unknown>
     const currDepVals = [] as Array<unknown>
@@ -130,27 +108,6 @@ export class Derived<
       onSubscribe: options.onSubscribe?.bind(this) as never,
       onUpdate: options.onUpdate,
     })
-
-    const updateStoreToDerived = (
-      store: Store<unknown>,
-      dep: Derived<unknown>,
-    ) => {
-      const prevDerivesForStore = this.storeToDerived.get(store) || new Set()
-      prevDerivesForStore.add(dep)
-      this.storeToDerived.set(store, prevDerivesForStore)
-    }
-    for (const dep of options.deps) {
-      if (dep instanceof Derived) {
-        this.derivedToStore.set(dep, dep.rootStores)
-        for (const store of dep.rootStores) {
-          this.rootStores.add(store)
-          updateStoreToDerived(store, dep)
-        }
-      } else if (dep instanceof Store) {
-        this.rootStores.add(dep)
-        updateStoreToDerived(dep, this as Derived<unknown>)
-      }
-    }
   }
 
   get state() {
@@ -166,20 +123,46 @@ export class Derived<
     return this._store.prevState
   }
 
-  mount = () => {
-    let __depsThatHaveWrittenThisTick = [] as Array<
-      Derived<unknown> | Store<unknown>
-    >
+  registerOnGraph(
+    deps: ReadonlyArray<Derived<any> | Store<any>> = this.options.deps,
+  ) {
+    for (const dep of deps) {
+      if (dep instanceof Derived) {
+        // Go into the deps of the derived and find the root store(s) that it depends on deeply
+        // Then, register this derived as a related derived to the store
+        this.registerOnGraph(dep.options.deps)
+      } else if (dep instanceof Store) {
+        // Register the derived as related derived to the store
+        let relatedLinkedDerivedVals = __storeToDerived.get(dep)
+        if (!relatedLinkedDerivedVals) {
+          relatedLinkedDerivedVals = new Set()
+          __storeToDerived.set(dep, relatedLinkedDerivedVals)
+        }
+        relatedLinkedDerivedVals.add(this as never)
 
+        // Register the store as a related store to this derived
+        let relatedStores = __derivedToStore.get(this as never)
+        if (!relatedStores) {
+          relatedStores = new Set()
+          __derivedToStore.set(this as never, relatedStores)
+        }
+        relatedStores.add(dep)
+      }
+    }
+  }
+
+  unregisterFromGraph() {}
+
+  mount = () => {
+    this.registerOnGraph()
     for (const dep of this.options.deps) {
       const isDepAStore = dep instanceof Store
       let relatedLinkedDerivedVals: null | Set<Derived<unknown>> = null
 
       const unsub = dep.subscribe(() => {
-        const store = isDepAStore ? dep : dep._whatStoreIsCurrentlyInUse
-        this._whatStoreIsCurrentlyInUse = store
+        const store = isDepAStore ? dep : __whatStoreIsCurrentlyInUse.current
         if (store) {
-          relatedLinkedDerivedVals = this.storeToDerived.get(store) ?? null
+          relatedLinkedDerivedVals = __storeToDerived.get(store) ?? null
         }
 
         __depsThatHaveWrittenThisTick.push(dep)
@@ -192,9 +175,9 @@ export class Derived<
             this._store.setState(() => this.options.fn(this.getDepVals()))
           }
 
-          // Cleanup the deps that have written this tick
-          __depsThatHaveWrittenThisTick = []
-          this._whatStoreIsCurrentlyInUse = null
+          // // Cleanup the deps that have written this tick
+          // __depsThatHaveWrittenThisTick = []
+          __whatStoreIsCurrentlyInUse.current = null
           return
         }
       })
@@ -203,6 +186,7 @@ export class Derived<
     }
 
     return () => {
+      this.unregisterFromGraph()
       for (const cleanup of this._subscriptions) {
         cleanup()
       }

packages/store/src/index.ts

@@ -2,3 +2,4 @@ export * from './derived'
 export * from './effect'
 export * from './store'
 export * from './types'
+export * from './scheduler'

packages/store/src/scheduler.ts

@@ -0,0 +1,38 @@
+import type { Store } from './store'
+import type { Derived } from './derived'
+
+/**
+ * What store called the current update, if any
+ * @private
+ */
+export const __whatStoreIsCurrentlyInUse = {
+  current: null as Store<unknown> | null,
+}
+
+/**
+ * This is here to solve the pyramid dependency problem where:
+ *       A
+ *      / \
+ *     B   C
+ *      \ /
+ *       D
+ *
+ * Where we deeply traverse this tree, how do we avoid D being recomputed twice; once when B is updated, once when C is.
+ *
+ * To solve this, we create linkedDeps that allows us to sync avoid writes to the state until all of the deps have been
+ * resolved.
+ *
+ * This is a record of stores, because derived stores are not able to write values to, but stores are
+ */
+export const __storeToDerived = new WeakMap<
+  Store<unknown>,
+  Set<Derived<unknown>>
+>()
+export const __derivedToStore = new WeakMap<
+  Derived<unknown>,
+  Set<Store<unknown>>
+>()
+
+export const __depsThatHaveWrittenThisTick = [] as Array<
+  Derived<unknown> | Store<unknown>
+>

packages/store/src/store.ts

@@ -1,3 +1,4 @@
+import { __whatStoreIsCurrentlyInUse } from './scheduler'
 import type { AnyUpdater, Listener } from './types'
 
 export interface StoreOptions<
@@ -31,14 +32,6 @@ export class Store<
   state: TState
   prevState: TState
   options?: StoreOptions<TState, TUpdater>
-  /**
-   * @private
-   */
-  _batching = false
-  /**
-   * @private
-   */
-  _flushing = 0
 
   constructor(initialState: TState, options?: StoreOptions<TState, TUpdater>) {
     this.prevState = initialState
@@ -56,6 +49,7 @@ export class Store<
   }
 
   setState = (updater: TUpdater) => {
+    __whatStoreIsCurrentlyInUse.current = this as never
     const previous = this.state
     this.state = this.options?.updateFn
       ? this.options.updateFn(previous)(updater)
@@ -67,25 +61,4 @@ export class Store<
     // Attempt to flush
     this._flush()
   }
-
-  /**
-   * @private
-   */
-  _flush = () => {
-    if (this._batching) return
-    const flushId = ++this._flushing
-    for (const listener of this.listeners) {
-      if (this._flushing !== flushId) continue
-      listener({ prevVal: this.prevState, currentVal: this.state })
-      this.prevState = this.state
-    }
-  }
-
-  batch = (cb: () => void) => {
-    if (this._batching) return cb()
-    this._batching = true
-    cb()
-    this._batching = false
-    this._flush()
-  }
 }

packages/store/tests/derived.test.ts

@@ -238,6 +238,33 @@ describe('Derived', () => {
     })
   })
 
+  test.skip('derivedFn should receive old and new dep values for similar derived values', () => {
+    const count = new Store(12)
+    const halfCount = new Derived({
+      deps: [count],
+      fn: () => count.state / 2,
+    })
+    halfCount.mount()
+    const fn = vi.fn()
+    const derived = new Derived({
+      deps: [count, halfCount],
+      fn: ({ prevDepVals, currDepVals }) => {
+        fn({ prevDepVals, currDepVals })
+        return void 0
+      },
+    })
+    derived.mount()
+    expect(fn).toBeCalledWith({
+      prevDepVals: undefined,
+      currDepVals: [12, 6],
+    })
+    count.setState(() => 24)
+    expect(fn).toBeCalledWith({
+      prevDepVals: [12, 6],
+      currDepVals: [24, 12],
+    })
+  })
+
   test('derivedFn should receive the old value', () => {
     const count = new Store(12)
     const date = new Date()

packages/store/tests/scheduler.test.ts

@@ -0,0 +1,31 @@
+import { describe, expect, test } from 'vitest'
+import { Derived, Store, __derivedToStore, __storeToDerived } from '../src'
+
+describe('Scheduler logic', () => {
+  test('Should build a graph properly', () => {
+    const count = new Store<any, any>(10)
+
+    const halfCount = new Derived<any, any>({
+      deps: [count],
+      fn: () => {
+        return count.state / 2
+      },
+    })
+
+    halfCount.registerOnGraph()
+
+    const doubleHalfCount = new Derived<any, any>({
+      deps: [halfCount],
+      fn: () => {
+        return halfCount.state * 2
+      },
+    })
+
+    doubleHalfCount.registerOnGraph()
+
+    expect(__storeToDerived.get(count)).toContain(halfCount)
+    expect(__derivedToStore.get(halfCount)).toContain(count)
+    expect(__storeToDerived.get(count)).toContain(doubleHalfCount)
+    expect(__derivedToStore.get(doubleHalfCount)).toContain(count)
+  })
+})

packages/store/tests/store.test.ts

@@ -53,19 +53,19 @@ describe('store', () => {
     expect(typeof store.state).toEqual('number')
   })
 
-  test('Batch prevents listeners from being called during repeated setStates', () => {
+  test.skip('Batch prevents listeners from being called during repeated setStates', () => {
     const store = new Store(0)
 
     const listener = vi.fn()
 
     const unsub = store.subscribe(listener)
 
-    store.batch(() => {
-      store.setState(() => 1)
-      store.setState(() => 2)
-      store.setState(() => 3)
-      store.setState(() => 4)
-    })
+    // store.batch(() => {
+    //   store.setState(() => 1)
+    //   store.setState(() => 2)
+    //   store.setState(() => 3)
+    //   store.setState(() => 4)
+    // })
 
     expect(store.state).toEqual(4)
     // Listener is only called once because of batching