fix(overmind-react): improve logic and fix fast refresh

Author: christianalfoni

packages/node_modules/overmind-react/src/index.ts

@@ -17,7 +17,10 @@ import {
 const IS_PRODUCTION = ENVIRONMENT === 'production'
 const IS_TEST = ENVIRONMENT === 'test'
 const isNode =
-  !IS_TEST && typeof process !== "undefined" && process.title && process.title.includes('node')
+  !IS_TEST &&
+  typeof process !== 'undefined' &&
+  process.title &&
+  process.title.includes('node')
 
 export type IReactComponent<P = any> =
   | react.StatelessComponent<P>
@@ -60,15 +63,7 @@ export const Provider: react.ProviderExoticComponent<
 > = context.Provider
 
 function useForceRerender() {
-  const [{ flushId }, setTick] = react.useState({ tick: 0, flushId: 0 })
-
-  const forceRerender = react.useCallback((flushId?) => {
-    setTick((current) => ({
-      ...current,
-      tick: current.tick + 1,
-      flushId: flushId || current.flushId,
-    }))
-  }, [])
+  const [flushId, forceRerender] = react.useState(-1)
 
   return {
     flushId,
@@ -101,53 +96,35 @@ const useState = <Context extends IContext<any>>(
     return overmind.state
   }
 
-  const trackingRef = react.useRef<any>(null)
-
+  const mountedRef = react.useRef<any>(false)
   const { flushId, forceRerender } = useForceRerender()
+  const tree = react.useMemo(
+    () => (overmind as any).proxyStateTreeInstance.getTrackStateTree(),
+    [flushId]
+  )
 
-  if (!trackingRef.current) {
-    const tree = (overmind as any).proxyStateTreeInstance.getTrackStateTree()
-    trackingRef.current = {
-      tree,
-      hasUpdatedBeforeCommit: false,
-      stopTrackingTask: unstable_scheduleCallback(
-        unstable_getCurrentPriorityLevel(),
-        () => {
-          trackingRef.current.tree.stopTracking()
-        }
-      ),
-      state: cb ? cb(tree.state) : tree.state,
-    }
-  }
+  const state = cb ? cb(tree.state) : tree.state
 
   if (IS_PRODUCTION) {
-    react.useLayoutEffect(() => {
-      trackingRef.current.mounted = true
-
-      if (trackingRef.current.hasUpdatedBeforeCommit) {
-        forceRerender()
-      }
+    react.useLayoutEffect(
+      () => {
+        mountedRef.current = true
+        tree.stopTracking()
 
-      return () => {
-        ;(overmind as any).proxyStateTreeInstance.disposeTree(
-          trackingRef.current.tree
-        )
-      }
-    }, [])
+        return () => {
+          tree.dispose()
+        }
+      },
+      [tree]
+    )
 
-    react.useLayoutEffect(() => {
-      if (trackingRef.current.stopTrackingTask) {
-        unstable_cancelCallback(trackingRef.current.stopTrackingTask)
-        trackingRef.current.stopTrackingTask = null
-      }
-      trackingRef.current.tree.stopTracking()
-    })
-    trackingRef.current.tree.track(() => {
-      if (trackingRef.current.mounted) {
-        forceRerender()
-      } else {
-        trackingRef.current.hasUpdatedBeforeCommit = true
+    tree.track((_, __, flushId) => {
+      if (!mountedRef.current) {
+        // This one is not dealt with by the useLayoutEffect
+        tree.dispose()
+        return
       }
+      forceRerender(flushId)
     })
   } else {
     const component = useCurrentComponent()
@@ -162,22 +139,15 @@ const useState = <Context extends IContext<any>>(
     )
 
     react.useLayoutEffect(() => {
-      trackingRef.current.mounted = true
+      mountedRef.current = true
       overmind.eventHub.emitAsync(EventType.COMPONENT_ADD, {
         componentId: component.__componentId,
         componentInstanceId,
         name,
-        paths: Array.from(trackingRef.current.tree.pathDependencies) as any,
+        paths: Array.from(tree.pathDependencies) as any,
       })
 
-      if (trackingRef.current.hasUpdatedBeforeCommit) {
-        forceRerender()
-      }
-
       return () => {
-        ;(overmind as any).proxyStateTreeInstance.disposeTree(
-          trackingRef.current.tree
-        )
         overmind.eventHub.emitAsync(EventType.COMPONENT_REMOVE, {
           componentId: component.__componentId,
           componentInstanceId,
@@ -186,30 +156,35 @@ const useState = <Context extends IContext<any>>(
       }
     }, [])
 
-    react.useLayoutEffect(() => {
-      if (trackingRef.current.stopTrackingTask) {
-        unstable_cancelCallback(trackingRef.current.stopTrackingTask)
-        trackingRef.current.stopTrackingTask = null
-      }
-      trackingRef.current.tree.stopTracking()
-      overmind.eventHub.emitAsync(EventType.COMPONENT_UPDATE, {
-        componentId: component.__componentId,
-        componentInstanceId,
-        name,
-        flushId,
-        paths: Array.from(trackingRef.current.tree.pathDependencies) as any,
-      })
-    })
-    trackingRef.current.tree.track(() => {
-      if (trackingRef.current.mounted) {
-        forceRerender()
-      } else {
-        trackingRef.current.hasUpdatedBeforeCommit = true
+    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()
+        }
+      },
+      [tree]
+    )
+    tree.track((_, __, flushId) => {
+      if (!mountedRef.current) {
+        // This one is not dealt with by the useLayoutEffect
+        tree.dispose()
+        return
       }
+      forceRerender(flushId)
     })
   }
 
-  return trackingRef.current.state
+  return state
 }
 
 const useActions = <Context extends IContext<any>>(): Context['actions'] => {