From 7a4bfa231f336a134ccbd6db9e0dbe835b292dd2 Mon Sep 17 00:00:00 2001
From: Christian Alfoni <christianalfoni@gmail.com>
Date: Tue, 5 Mar 2024 13:49:26 +0100
Subject: [PATCH] proxy state tree green

---
 README.md                                     |   4 +
 packages/overmind-react/src/index.test.tsx    |   2 +
 packages/overmind-react/src/index.ts          | 159 ++-----------
 packages/overmind/src/Overmind.ts             |  17 +-
 packages/overmind/src/operators.ts            |  10 +-
 packages/proxy-state-tree/src/Proxyfier.ts    |   5 +-
 .../proxy-state-tree/src/TrackStateTree.ts    |  29 ++-
 packages/proxy-state-tree/src/index.test.ts   | 216 ++++++++++--------
 packages/proxy-state-tree/src/index.ts        |   9 +-
 packages/proxy-state-tree/src/types.ts        |   7 +-
 10 files changed, 192 insertions(+), 266 deletions(-)

diff --git a/README.md b/README.md
index c36c1e5c..cbc5c7a9 100644
--- a/README.md
+++ b/README.md
@@ -6,6 +6,10 @@ Visit website for more information: [www.overmindjs.org](https://overmindjs.org)
 
 <https://gist.github.com/christianalfoni/f1c4bfe320dcb24c403635d9bca3fa40>
 
+## Contributing
+
+// List all packages and state how to develop and test them
+
 ## Release procedure
 
 ```sh
diff --git a/packages/overmind-react/src/index.test.tsx b/packages/overmind-react/src/index.test.tsx
index ef0e0e89..0f007c83 100644
--- a/packages/overmind-react/src/index.test.tsx
+++ b/packages/overmind-react/src/index.test.tsx
@@ -88,6 +88,7 @@ describe('React', () => {
 
     const FooComponent: React.FunctionComponent = () => {
       const state = useState((state) => state.foo[0])
+      console.log('WTF?')
       renderCount++
 
       return <h1>{state.foo}</h1>
@@ -110,6 +111,7 @@ describe('React', () => {
     act(() => {
       app.actions.doThat()
     })
+
     expect(renderCount).toBe(2)
 
     // This is not showing the expected result, but logging the rendering does, so must be the
diff --git a/packages/overmind-react/src/index.ts b/packages/overmind-react/src/index.ts
index 7caf47d4..251250d4 100644
--- a/packages/overmind-react/src/index.ts
+++ b/packages/overmind-react/src/index.ts
@@ -121,7 +121,7 @@ class ReactTrackerV18 {
   }
 }
 
-const useStateV18 = <Context extends IContext<{ state: {} }>>(
+const useState = <Context extends IContext<{ state: {} }>>(
   cb?: (state: Context['state']) => any
 ): Context['state'] => {
   const overmind = React.useContext(context) as Overmind<any>
@@ -134,30 +134,23 @@ const useStateV18 = <Context extends IContext<{ state: {} }>>(
     return overmind.state
   }
 
-  const ref = React.useRef(null)
-
-  if (!ref.current) {
-    // @ts-ignore
-    ref.current = new ReactTrackerV18(overmind.getTrackStateTree())
-  }
-
-  const tracker = ref.current as any
+  const { flushId, forceRerender } = useForceRerender()
 
-  const snapshot = React.useSyncExternalStore(
-    tracker.subscribe,
-    tracker.getState,
-    tracker.getState
-  )
-  const mountedRef = React.useRef<any>(false)
-  // @ts-ignore
-  const state = cb ? cb(snapshot.state) : snapshot.state
+  const trackStateTree = (
+    overmind as any
+  ).proxyStateTreeInstance.getTrackStateTree()
+  const state = cb ? cb(trackStateTree.state) : trackStateTree.state
 
-  tracker.track()
+  trackStateTree.track()
 
   if (IS_PRODUCTION) {
-    React.useLayoutEffect(() => {
-      tracker.stopTracking()
-    }, [tracker])
+    React.useEffect(
+      () =>
+        trackStateTree.subscribe((_, __, flushId) => {
+          forceRerender(flushId)
+        }),
+      []
+    )
   } else {
     const component = useCurrentComponent()
     const name = getDisplayName(component)
@@ -170,135 +163,29 @@ const useStateV18 = <Context extends IContext<{ state: {} }>>(
       currentComponentInstanceId++
     )
 
-    React.useLayoutEffect(() => {
-      mountedRef.current = true
-      overmind.eventHub.emitAsync(EventType.COMPONENT_ADD, {
-        componentId: component.__componentId,
-        componentInstanceId,
-        name,
-        paths: Array.from(tracker.tree.pathDependencies) as any,
-      })
-
-      return () => {
-        mountedRef.current = false
-        overmind.eventHub.emitAsync(EventType.COMPONENT_REMOVE, {
-          componentId: component.__componentId,
-          componentInstanceId,
-          name,
-        })
-      }
-    }, [])
-
-    React.useLayoutEffect(() => {
-      tracker.stopTracking()
-
+    React.useEffect(() => {
       overmind.eventHub.emitAsync(EventType.COMPONENT_UPDATE, {
         componentId: component.__componentId,
         componentInstanceId,
         name,
-        flushId: 0,
-        paths: Array.from(tracker.tree.pathDependencies) as any,
+        paths: Array.from(trackStateTree.pathDependencies) as any,
       })
-    }, [tracker])
-  }
-
-  return state
-}
-
-const useState = <Context extends IContext<{ state: {} }>>(
-  cb?: (state: Context['state']) => any
-): Context['state'] => {
-  const overmind = React.useContext(context) as Overmind<any>
-
-  if (!(overmind as any).mode) {
-    throwMissingContextError()
-  }
-
-  if (isNode || (overmind as any).mode.mode === MODE_SSR) {
-    return overmind.state
-  }
 
-  const mountedRef = React.useRef<any>(false)
-  const { flushId, forceRerender } = useForceRerender()
-  const tree = React.useMemo(
-    () => (overmind as any).proxyStateTreeInstance.getTrackStateTree(),
-    [flushId]
-  )
-
-  const state = cb ? cb(tree.state) : tree.state
-
-  if (IS_PRODUCTION) {
-    React.useLayoutEffect(() => {
-      mountedRef.current = true
-      tree.stopTracking()
-
-      return () => {
-        tree.dispose()
-      }
-    }, [tree])
-
-    tree.track((_, __, flushId) => {
-      if (!mountedRef.current) {
-        // This one is not dealt with by the useLayoutEffect
-        tree.dispose()
-        return
-      }
-      forceRerender(flushId)
-    })
-  } else {
-    const component = useCurrentComponent()
-    const name = getDisplayName(component)
-    component.__componentId =
-      typeof component.__componentId === 'undefined'
-        ? nextComponentId++
-        : component.__componentId
-
-    const { current: componentInstanceId } = React.useRef<any>(
-      currentComponentInstanceId++
-    )
-
-    React.useLayoutEffect(() => {
-      mountedRef.current = true
-      overmind.eventHub.emitAsync(EventType.COMPONENT_ADD, {
-        componentId: component.__componentId,
-        componentInstanceId,
-        name,
-        paths: Array.from(tree.pathDependencies) as any,
+      const dispose = trackStateTree.subscribe((_, __, flushId) => {
+        forceRerender(flushId)
       })
 
       return () => {
-        mountedRef.current = false
-        overmind.eventHub.emitAsync(EventType.COMPONENT_REMOVE, {
+        overmind.eventHub.emitAsync(EventType.COMPONENT_UPDATE, {
           componentId: component.__componentId,
           componentInstanceId,
           name,
+          paths: [],
         })
-      }
-    }, [])
-
-    React.useLayoutEffect(() => {
-      tree.stopTracking()
 
-      overmind.eventHub.emitAsync(EventType.COMPONENT_UPDATE, {
-        componentId: component.__componentId,
-        componentInstanceId,
-        name,
-        flushId,
-        paths: Array.from(tree.pathDependencies) as any,
-      })
-
-      return () => {
-        tree.dispose()
+        dispose()
       }
-    }, [tree])
-    tree.track((_, __, flushId) => {
-      if (!mountedRef.current) {
-        // This one is not dealt with by the useLayoutEffect
-        tree.dispose()
-        return
-      }
-      forceRerender(flushId)
-    })
+    }, [])
   }
 
   return state
@@ -349,7 +236,7 @@ export const createStateHook: <
   Context extends IContext<{ state: {} }>,
 >() => StateHook<Context> = () =>
   // eslint-disable-next-line dot-notation
-  typeof React['useSyncExternalStore'] === 'function' ? useStateV18 : useState
+  useState
 
 export const createActionsHook: <
   Context extends IContext<{ actions: {} }>,
diff --git a/packages/overmind/src/Overmind.ts b/packages/overmind/src/Overmind.ts
index 240d150e..55d4600a 100644
--- a/packages/overmind/src/Overmind.ts
+++ b/packages/overmind/src/Overmind.ts
@@ -894,21 +894,16 @@ export class Overmind<ThisConfig extends IConfiguration>
     } else {
       const tree = this.proxyStateTreeInstance.getTrackStateTree()
       let returnValue
+      let disposer
       const updateReaction = () => {
-        tree.trackScope(
-          () => (returnValue = stateCallback(tree.state as any)),
-          () => {
-            updateReaction()
-            updateCallback(returnValue)
-          }
-        )
+        tree.trackScope(() => (returnValue = stateCallback(tree.state as any)))
+        disposer = tree.subscribe(() => {
+          updateReaction()
+          updateCallback(returnValue)
+        })
       }
 
       updateReaction()
-
-      disposer = () => {
-        tree.dispose()
-      }
     }
 
     if (options.immediate) {
diff --git a/packages/overmind/src/operators.ts b/packages/overmind/src/operators.ts
index 4a2fd3d8..54fdd9be 100644
--- a/packages/overmind/src/operators.ts
+++ b/packages/overmind/src/operators.ts
@@ -558,6 +558,8 @@ export function throttle<T>(ms: number): IOperator<T, T> {
   )
 }
 
+// waitUntil((state) => state.foo === 'bar')
+
 export function waitUntil<T, C extends IContext<{}>>(
   operation: (state: C['state']) => boolean
 ): IOperator<T, T> {
@@ -568,13 +570,17 @@ export function waitUntil<T, C extends IContext<{}>>(
       if (err) next(err, value)
       else {
         const tree = context.execution.getTrackStateTree()
+        let disposer
         const test = () => {
+          disposer?.()
+
           if (operation(tree.state)) {
-            tree.dispose()
             next(null, value)
+          } else {
+            disposer = tree.subscribe(test)
           }
         }
-        tree.trackScope(test, test)
+        tree.trackScope(test)
       }
     }
   )
diff --git a/packages/proxy-state-tree/src/Proxyfier.ts b/packages/proxy-state-tree/src/Proxyfier.ts
index 724ac93e..920fed60 100644
--- a/packages/proxy-state-tree/src/Proxyfier.ts
+++ b/packages/proxy-state-tree/src/Proxyfier.ts
@@ -104,9 +104,7 @@ export class Proxifier {
     }
 
     if (this.isDefaultProxifier()) {
-      const trackStateTree = this.tree.root
-        .currentTree as ITrackStateTree<any>
-
+      const trackStateTree = this.tree.root.currentTree as ITrackStateTree<any>
       if (!trackStateTree) {
         return
       }
@@ -294,6 +292,7 @@ export class Proxifier {
         }
 
         const trackingTree = proxifier.getTrackingTree()
+
         const targetValue = target[prop]
         const nestedPath = proxifier.concat(path, prop)
         const currentTree = trackingTree || proxifier.tree
diff --git a/packages/proxy-state-tree/src/TrackStateTree.ts b/packages/proxy-state-tree/src/TrackStateTree.ts
index c0ed3474..36dc942f 100644
--- a/packages/proxy-state-tree/src/TrackStateTree.ts
+++ b/packages/proxy-state-tree/src/TrackStateTree.ts
@@ -19,8 +19,6 @@ export class TrackStateTree<T extends object> implements ITrackStateTree<T> {
     this.state = root.state
   }
 
-  // Does not seem to be used
-  /*
   trackPaths() {
     const paths = new Set<string>()
     const listener = (path) => {
@@ -37,10 +35,11 @@ export class TrackStateTree<T extends object> implements ITrackStateTree<T> {
       return paths
     }
   }
-  */
 
   track() {
     this.root.changeTrackStateTree(this)
+
+    return this
   }
 
   canMutate() {
@@ -57,28 +56,26 @@ export class TrackStateTree<T extends object> implements ITrackStateTree<T> {
 
   subscribe(cb: ITrackCallback) {
     this.root.changeTrackStateTree(null)
-    console.log('Adddig', this.pathDependencies)
     for (const path of this.pathDependencies) {
       this.root.addPathDependency(path, cb)
     }
     return () => {
-      console.log('Removing', this.pathDependencies)
       for (const path of this.pathDependencies) {
         this.root.removePathDependency(path, cb)
       }
     }
   }
 
-  /*
-  trackScope(scope: ITrackScopedCallback<T>, cb?: ITrackCallback) {
-    const previousPreviousTree = this.master.previousTree
-    const previousCurrentTree = this.master.currentTree
-    this.master.currentTree = this
-    this.track(cb)
-    const result = scope(this)
-    this.master.currentTree = previousCurrentTree
-    this.master.previousTree = previousPreviousTree
-    return result
+  trackScope(scope: ITrackScopedCallback<T>) {
+    const previousPreviousTree = this.root.previousTree
+    const previousCurrentTree = this.root.currentTree
+
+    this.root.currentTree = this
+    this.track()
+    scope(this)
+    this.root.currentTree = previousCurrentTree
+    this.root.previousTree = previousPreviousTree
+
+    return this
   }
-*/
 }
diff --git a/packages/proxy-state-tree/src/index.test.ts b/packages/proxy-state-tree/src/index.test.ts
index 77a279f2..2855f997 100644
--- a/packages/proxy-state-tree/src/index.test.ts
+++ b/packages/proxy-state-tree/src/index.test.ts
@@ -47,7 +47,6 @@ describe('TrackStateTree', () => {
     mutationTree.flush()
     expect(reactionCount).toBe(1)
   })
-  /*
   test('should allow tracking by mutation', () => {
     let reactionCount = 0
     const tree = new ProxyStateTree({
@@ -71,28 +70,29 @@ describe('TrackStateTree', () => {
 
     expect(reactionCount).toBe(1)
   })
-  test('should stop tracking on command', () => {
+  test('should stop tracking when creating mutation trees', () => {
     let reactionCount = 0
     const tree = new ProxyStateTree({
       foo: 'bar',
       bar: 'baz',
     })
 
-    const accessTree = tree
-      .getTrackStateTree()
-      .track((mutations, paths, flushId) => {
-        reactionCount++
-        expect(flushId).toBe(0)
-      })
+    const accessTree = tree.getTrackStateTree()
 
-    accessTree.state.foo
+    accessTree.track()
 
-    accessTree.stopTracking()
+    accessTree.state.foo
 
-    accessTree.state.bar
+    // This actually happens last in the component via useEffect
+    accessTree.subscribe((mutations, paths, flushId) => {
+      reactionCount++
+      expect(flushId).toBe(0)
+    })
 
     const mutationTree = tree.getMutationTree()
 
+    accessTree.state.bar
+
     mutationTree.state.foo = 'bar2'
     mutationTree.flush()
     mutationTree.state.bar = 'baz2'
@@ -108,8 +108,8 @@ describe('TrackStateTree', () => {
     const accessTreeA = tree.getTrackStateTreeWithProxifier()
     const accessTreeB = tree.getTrackStateTreeWithProxifier()
 
-    accessTreeA.track(() => {})
-    accessTreeB.track(() => {})
+    accessTreeA.track()
+    accessTreeB.track()
 
     accessTreeA.state.foo
     accessTreeB.state.bar
@@ -144,7 +144,8 @@ describe('CLASSES', () => {
         user: new User(),
       }
       const tree = new ProxyStateTree(state)
-      const trackStateTree = tree.getTrackStateTree().track(() => {})
+      const trackStateTree = tree.getTrackStateTree()
+      trackStateTree.track()
 
       expect(trackStateTree.state.user[IS_PROXY]).toBeTruthy()
     })
@@ -157,9 +158,10 @@ describe('CLASSES', () => {
         user: new User(),
       }
       const tree = new ProxyStateTree(state)
-      expect(tree.getTrackStateTree().track(() => {}).state.user.name).toBe(
-        'Bob'
-      )
+      const trackStateTree = tree.getTrackStateTree()
+      trackStateTree.track()
+
+      expect(trackStateTree.state.user.name).toBe('Bob')
     })
 
     test('should track access properties', () => {
@@ -170,10 +172,13 @@ describe('CLASSES', () => {
         user: new User(),
       }
       const tree = new ProxyStateTree(state)
-      const trackStateTree = tree.getTrackStateTree().track(() => {})
+      const trackStateTree = tree.getTrackStateTree()
+      trackStateTree.track()
 
       expect(trackStateTree.state.user.name).toBe('Bob')
 
+      trackStateTree.subscribe(() => {})
+
       expect(Object.keys((tree as any).pathDependencies)).toEqual([
         'user',
         'user.name',
@@ -192,10 +197,12 @@ describe('CLASSES', () => {
         user: new User(),
       }
       const tree = new ProxyStateTree(state)
-      const trackStateTree = tree.getTrackStateTree().track(() => {})
+      const trackStateTree = tree.getTrackStateTree().track()
 
       expect(trackStateTree.state.user.name).toBe('Bob Saget')
 
+      trackStateTree.subscribe(() => {})
+
       expect(Object.keys((tree as any).pathDependencies)).toEqual([
         'user',
         'user.firstName',
@@ -213,7 +220,7 @@ describe('CLASSES', () => {
       }
       const tree = new ProxyStateTree(state)
       expect(() => {
-        tree.getTrackStateTree().track(() => {}).state.user.name = 'bar2'
+        tree.getTrackStateTree().track().state.user.name = 'bar2'
       }).toThrow()
     })
 
@@ -240,7 +247,8 @@ describe('CLASSES', () => {
       }
       const tree = new ProxyStateTree(state)
       const mutationTree = tree.getMutationTree()
-      const accessTree = tree.getTrackStateTree().track(() => {
+      const accessTree = tree.getTrackStateTree().track()
+      accessTree.subscribe(() => {
         renderCount++
       })
       accessTree.addTrackingPath('user.name')
@@ -336,7 +344,7 @@ describe('OBJECTS', () => {
         foo: {},
       }
       const tree = new ProxyStateTree(state)
-      const trackStateTree = tree.getTrackStateTree().track(() => {})
+      const trackStateTree = tree.getTrackStateTree().track()
 
       expect(trackStateTree.state.foo[IS_PROXY]).toBeTruthy()
     })
@@ -348,7 +356,7 @@ describe('OBJECTS', () => {
         },
       }
       const tree = new ProxyStateTree(state)
-      expect(tree.getTrackStateTree().track(() => {}).state.foo.bar).toBe('baz')
+      expect(tree.getTrackStateTree().track().state.foo.bar).toBe('baz')
     })
 
     test('should track access properties', () => {
@@ -358,10 +366,12 @@ describe('OBJECTS', () => {
         },
       }
       const tree = new ProxyStateTree(state)
-      const trackStateTree = tree.getTrackStateTree().track(() => {})
+      const trackStateTree = tree.getTrackStateTree().track()
 
       expect(trackStateTree.state.foo.bar).toBe('baz')
 
+      trackStateTree.subscribe(() => {})
+
       expect(Object.keys((tree as any).pathDependencies)).toEqual([
         'foo',
         'foo.bar',
@@ -376,7 +386,7 @@ describe('OBJECTS', () => {
       }
       const tree = new ProxyStateTree(state)
       expect(() => {
-        tree.getTrackStateTree().track(() => {}).state.foo = 'bar2'
+        tree.getTrackStateTree().track().state.foo = 'bar2'
       }).toThrow()
     })
 
@@ -399,7 +409,8 @@ describe('OBJECTS', () => {
       }
       const tree = new ProxyStateTree(state)
       const mutationTree = tree.getMutationTree()
-      const accessTree = tree.getTrackStateTree().track(() => {
+      const accessTree = tree.getTrackStateTree().track()
+      accessTree.subscribe(() => {
         renderCount++
       })
       accessTree.addTrackingPath('foo')
@@ -463,7 +474,7 @@ describe('ARRAYS', () => {
         foo: [],
       }
       const tree = new ProxyStateTree(state)
-      const trackStateTree = tree.getTrackStateTree().track(() => {})
+      const trackStateTree = tree.getTrackStateTree().track()
       const foo = trackStateTree.state.foo
 
       expect(foo[IS_PROXY]).toBeTruthy()
@@ -474,7 +485,7 @@ describe('ARRAYS', () => {
         foo: 'bar',
       }
       const tree = new ProxyStateTree(state)
-      expect(tree.getTrackStateTree().track(() => {}).state.foo[0]).toBe('b')
+      expect(tree.getTrackStateTree().track().state.foo[0]).toBe('b')
     })
 
     test('should track access properties', () => {
@@ -482,8 +493,9 @@ describe('ARRAYS', () => {
         foo: ['bar'],
       }
       const tree = new ProxyStateTree(state)
-      const accessTree = tree.getTrackStateTree().track(() => {})
+      const accessTree = tree.getTrackStateTree().track()
       expect(accessTree.state.foo[0]).toBe('bar')
+      accessTree.subscribe(() => {})
       expect(Object.keys((tree as any).pathDependencies)).toEqual([
         'foo',
         'foo.0',
@@ -500,24 +512,16 @@ describe('ARRAYS', () => {
         bar: 'baz',
       })
 
-      tree.getTrackStateTree().trackScope(
-        (treeA) => {
-          treeA.state.foo.forEach((item) => {
-            tree.getTrackStateTree().trackScope(
-              (treeB) => {
-                item.title // eslint-disable-line
-                expect(treeB.pathDependencies).toEqual(new Set(['foo.0.title']))
-              },
-              () => {}
-            )
+      tree.getTrackStateTree().trackScope((treeA) => {
+        treeA.state.foo.forEach((item) => {
+          tree.getTrackStateTree().trackScope((treeB) => {
+            item.title // eslint-disable-line
+            expect(treeB.pathDependencies).toEqual(new Set(['foo.0.title']))
           })
-          treeA.state.bar
-          expect(treeA.pathDependencies).toEqual(
-            new Set(['foo', 'foo.0', 'bar'])
-          )
-        },
-        () => {}
-      )
+        })
+        treeA.state.bar
+        expect(treeA.pathDependencies).toEqual(new Set(['foo', 'foo.0', 'bar']))
+      })
     })
     test('should correctly keep track of changing indexes', () => {
       expect.assertions(4)
@@ -529,8 +533,9 @@ describe('ARRAYS', () => {
       const accessTree = tree.getTrackStateTree()
 
       function trackPaths() {
-        accessTree.track(trackPaths)
+        accessTree.track()
         accessTree.state.items.map((item) => item.title)
+        accessTree.subscribe(trackPaths)
 
         if (iterations === 1) {
           expect(Array.from(accessTree.pathDependencies)).toEqual([
@@ -560,6 +565,7 @@ describe('ARRAYS', () => {
       }
 
       trackPaths()
+
       const mutationTree = tree.getMutationTree()
       mutationTree.state.items.unshift({
         title: 'foo',
@@ -584,7 +590,7 @@ describe('ARRAYS', () => {
       foo: [],
     }
     const tree = new ProxyStateTree(state)
-    const trackStateTree = tree.getTrackStateTree().track(() => {})
+    const trackStateTree = tree.getTrackStateTree().track()
     const foo = trackStateTree.state.foo
 
     expect(foo[IS_PROXY]).toBeTruthy()
@@ -593,7 +599,7 @@ describe('ARRAYS', () => {
   describe('MUTATIONS', () => {
     test('should throw when mutating without tracking mutations', () => {
       const state = {
-        foo: [],
+        foo: [] as string[],
       }
       const tree = new ProxyStateTree(state)
       expect(() => {
@@ -737,7 +743,7 @@ describe('FUNCTIONS', () => {
       foo: () => 'bar',
     }
     const tree = new ProxyStateTree(state)
-    const accessTree = tree.getTrackStateTree().track(() => {})
+    const accessTree = tree.getTrackStateTree().track()
     expect(accessTree.state.foo).toBe('bar')
   })
 
@@ -751,7 +757,7 @@ describe('FUNCTIONS', () => {
       },
     }
     const tree = new ProxyStateTree(state)
-    const accessTree = tree.getTrackStateTree().track(() => {})
+    const accessTree = tree.getTrackStateTree().track()
 
     expect(accessTree.state.foo).toBe('bar')
   })
@@ -771,7 +777,7 @@ describe('FUNCTIONS', () => {
         target[prop]('foo', proxyStateTree, path),
     })
 
-    const accessTree = tree.getTrackStateTree().track(() => {})
+    const accessTree = tree.getTrackStateTree().track()
 
     expect(accessTree.state.foo).toBe('bar')
   })
@@ -788,19 +794,20 @@ describe('FUNCTIONS', () => {
         return proxyStateTree.state.items.map((item) => item)
       },
     })
-    const accessTree = tree
-      .getTrackStateTree()
-      .track((mutations, paths, flushId) => {
-        reactionCount++
-        expect(flushId).toBe(0)
-      })
-    const mutationTree = tree.getMutationTree()
+    const accessTree = tree.getTrackStateTree().track()
 
     // @ts-ignore
     accessTree.state.foo.forEach((item) => {
       item.title
     })
 
+    accessTree.subscribe((mutations, paths, flushId) => {
+      reactionCount++
+      expect(flushId).toBe(0)
+    })
+
+    const mutationTree = tree.getMutationTree()
+
     mutationTree.state.items.push({
       title: 'bar',
     })
@@ -816,15 +823,17 @@ describe('REACTIONS', () => {
     const tree = new ProxyStateTree({
       foo: 'bar',
     })
-    const accessTree = tree
-      .getTrackStateTree()
-      .track((mutations, paths, flushId) => {
-        reactionCount++
-        expect(flushId).toBe(0)
-      })
-    const mutationTree = tree.getMutationTree()
+    const accessTree = tree.getTrackStateTree().track()
+
     accessTree.state.foo // eslint-disable-line
 
+    accessTree.subscribe((mutations, paths, flushId) => {
+      reactionCount++
+      expect(flushId).toBe(0)
+    })
+
+    const mutationTree = tree.getMutationTree()
+
     mutationTree.state.foo = 'bar2'
     mutationTree.flush()
     expect(reactionCount).toBe(1)
@@ -862,8 +871,10 @@ describe('REACTIONS', () => {
     }
 
     const mutationTree = tree.getMutationTree()
-    const accessTree = tree.getTrackStateTree().track(render)
+    const accessTree = tree.getTrackStateTree().track()
+
     accessTree.addTrackingPath('foo')
+    accessTree.subscribe(render)
     mutationTree.state.foo = 'bar2'
     mutationTree.flush()
     expect(renderCount).toBe(1)
@@ -875,14 +886,16 @@ describe('REACTIONS', () => {
       foo: 'bar',
       bar: 'baz',
     })
-    const accessTree = tree.getTrackStateTree().track(() => {
-      reactionCount++
-    })
-    const mutationTree = tree.getMutationTree()
+    const accessTree = tree.getTrackStateTree().track()
 
     accessTree.state.foo // eslint-disable-line
     accessTree.state.bar // eslint-disable-line
 
+    accessTree.subscribe(() => {
+      reactionCount++
+    })
+    const mutationTree = tree.getMutationTree()
+
     mutationTree.state.foo = 'bar2'
     mutationTree.state.bar = 'baz2'
 
@@ -895,17 +908,24 @@ describe('REACTIONS', () => {
       foo: 'bar',
       bar: 'baz',
     })
-    const accessTree = tree.getTrackStateTree().track(render)
-    const mutationTree = tree.getMutationTree()
+    const accessTree = tree.getTrackStateTree()
+
     function render() {
+      accessTree.track()
+
       if (accessTree.state.foo === 'bar') {
+        accessTree.subscribe(render)
         return
       }
 
       accessTree.state.bar // eslint-disable-line
+      accessTree.subscribe(render)
     }
 
     render()
+
+    const mutationTree = tree.getMutationTree()
+
     mutationTree.state.foo = 'bar2'
     mutationTree.flush()
     expect((tree.pathDependencies as any).foo.size).toBe(1)
@@ -917,7 +937,8 @@ describe('REACTIONS', () => {
       foo: 'bar',
       bar: 'baz',
     })
-    const accessTree = tree.getTrackStateTree().track(render)
+    const accessTree = tree.getTrackStateTree().track()
+    accessTree.subscribe(render)
     const mutationTree = tree.getMutationTree()
     function render() {
       if (accessTree.state.foo === 'bar') {
@@ -976,7 +997,7 @@ describe('ITERATIONS', () => {
         },
       ],
     })
-    const accessTree = tree.getTrackStateTree().track(() => {})
+    const accessTree = tree.getTrackStateTree().track()
 
     accessTree.state.items.forEach((item) => {
       item.title // eslint-disable-line
@@ -994,7 +1015,7 @@ describe('ITERATIONS', () => {
         bar: 'baz',
       },
     })
-    const accessTree = tree.getTrackStateTree().track(() => {})
+    const accessTree = tree.getTrackStateTree().track()
     Object.keys(accessTree.state.items).forEach((key) => {
       accessTree.state.items[key] // eslint-disable-line
     })
@@ -1010,11 +1031,12 @@ describe('ITERATIONS', () => {
     const tree = new ProxyStateTree({
       items: {} as any,
     })
-    const accessTree = tree.getTrackStateTree().track(update)
-    const mutationTree = tree.getMutationTree()
+    const accessTree = tree.getTrackStateTree()
+
     function update() {
+      accessTree.track()
       accessTree.state.items // eslint-disable-line
-
+      accessTree.subscribe(update)
       if (runCount === 1) {
         expect(accessTree.pathDependencies).toEqual(new Set(['items']))
       }
@@ -1029,6 +1051,8 @@ describe('ITERATIONS', () => {
     }
 
     update()
+
+    const mutationTree = tree.getMutationTree()
     mutationTree.state.items.foo = 'bar'
     mutationTree.flush()
     mutationTree.state.items.bar = 'baz'
@@ -1049,15 +1073,18 @@ describe('ITERATIONS', () => {
         },
       ],
     })
-    const accessTree = tree.getTrackStateTree().track(() => {
+    const accessTree = tree.getTrackStateTree().track()
+
+    accessTree.state.items.map((item) => item.title)
+
+    accessTree.subscribe(() => {
       reactionCount++
     })
-    const mutationTree = tree.getMutationTree()
-    accessTree.state.items.map((item) => item.title)
 
     expect(accessTree.pathDependencies).toEqual(
       new Set(['items', 'items.0', 'items.0.title', 'items.1', 'items.1.title'])
     )
+    const mutationTree = tree.getMutationTree()
 
     mutationTree.state.items.push({
       title: 'mip',
@@ -1078,15 +1105,18 @@ describe('ITERATIONS', () => {
         },
       ],
     })
-    const accessTree = tree.getTrackStateTree().track(() => {
+    const accessTree = tree.getTrackStateTree().track()
+
+    accessTree.state.items.map((item) => item.title)
+
+    accessTree.subscribe(() => {
       reactionCount++
     })
-    const mutationTree = tree.getMutationTree()
 
-    accessTree.state.items.map((item) => item.title)
     expect(accessTree.pathDependencies).toEqual(
       new Set(['items', 'items.0', 'items.0.title', 'items.1', 'items.1.title'])
     )
+    const mutationTree = tree.getMutationTree()
     mutationTree.state.items[0].title = 'baz'
     mutationTree.flush()
     expect(reactionCount).toBe(1)
@@ -1097,15 +1127,18 @@ describe('ITERATIONS', () => {
     const tree = new ProxyStateTree({
       items: [1, 2],
     })
-    const accessTree = tree.getTrackStateTree().track(() => {
+    const accessTree = tree.getTrackStateTree().track()
+
+    accessTree.state.items.map((item) => item)
+
+    accessTree.subscribe(() => {
       reactionCount++
     })
-    const mutationTree = tree.getMutationTree()
 
-    accessTree.state.items.map((item) => item)
     expect(accessTree.pathDependencies).toEqual(
       new Set(['items', 'items.0', 'items.1'])
     )
+    const mutationTree = tree.getMutationTree()
     mutationTree.state.items[0] = 99
     mutationTree.flush()
     expect(reactionCount).toBe(1)
@@ -1118,7 +1151,7 @@ describe('RESCOPING', () => {
       foo: 'bar',
     })
 
-    const accessTree = tree.getTrackStateTree().track(() => {})
+    const accessTree = tree.getTrackStateTree().track()
     const mutationTree = tree.getMutationTree()
 
     const state = accessTree.state
@@ -1147,9 +1180,10 @@ describe('GETTER', () => {
     const mutationTree = tree.getMutationTree()
 
     function render() {
-      accessTree.track(render)
+      accessTree.track()
       accessTree.state.user.fullName
       renderCount++
+      accessTree.subscribe(render)
     }
 
     render()
@@ -1158,6 +1192,4 @@ describe('GETTER', () => {
     mutationTree.flush()
     expect(renderCount).toBe(2)
   })
-
-  */
 })
diff --git a/packages/proxy-state-tree/src/index.ts b/packages/proxy-state-tree/src/index.ts
index ad34f8c7..095037d6 100644
--- a/packages/proxy-state-tree/src/index.ts
+++ b/packages/proxy-state-tree/src/index.ts
@@ -79,7 +79,7 @@ export class ProxyStateTree<T extends object> implements IProxyStateTree<T> {
 
   getMutationTree(): IMutationTree<T> {
     // We never want to do tracking when we want to do mutations
-    this.changeTrackStateTree(null)
+    this.clearTrackStateTree()
 
     if (!this.options.devmode) {
       return (this.mutationTree =
@@ -111,11 +111,14 @@ export class ProxyStateTree<T extends object> implements IProxyStateTree<T> {
     this.currentTree = tree
   }
 
+  clearTrackStateTree() {
+    this.previousTree = null
+    this.currentTree = null
+  }
+
   disposeTree(tree: IMutationTree<T> | ITrackStateTree<T>) {
     if (tree instanceof MutationTree) {
       ;(tree as IMutationTree<T>).dispose()
-    } else if (tree instanceof TrackStateTree) {
-      // (tree as ITrackStateTree<T>).dispose()
     }
   }
 
diff --git a/packages/proxy-state-tree/src/types.ts b/packages/proxy-state-tree/src/types.ts
index 85e083a3..ffe01525 100644
--- a/packages/proxy-state-tree/src/types.ts
+++ b/packages/proxy-state-tree/src/types.ts
@@ -52,8 +52,8 @@ export interface ITrackScopedCallback<T extends object> {
 export interface ITrackStateTree<T extends object> {
   addTrackingPath(path: string): void
   subscribe(cb: ITrackCallback): () => void
-  track(): void
-  // trackScope(scope: ITrackScopedCallback<T>, callback?: ITrackCallback): any
+  track(): ITrackStateTree<T>
+  trackScope(scope: ITrackScopedCallback<T>): any
   canTrack(): boolean
   canMutate(): boolean
   root: IProxyStateTree<T>
@@ -92,7 +92,8 @@ export interface IProxyStateTree<T extends object> {
   removePathDependency(path: string, callback: ITrackCallback): void
   getTrackStateTree(): ITrackStateTree<T>
   getMutationTree(): IMutationTree<T>
-  changeTrackStateTree(tree: ITrackStateTree<T>): void
+  changeTrackStateTree(tree: ITrackStateTree<T> | null): void
+  clearTrackStateTree(): void
   disposeTree(proxy: TTree): void
   onMutation(cb: IMutationCallback): void
   flush(