From 98e79943d2e3043f9880e59218189b8a25d94ccf Mon Sep 17 00:00:00 2001
From: Evan You <yyx990803@gmail.com>
Date: Sun, 28 Oct 2018 12:10:29 -0400
Subject: [PATCH] fix(hooks): fix effect update & cleanup

---
 packages/runtime-core/src/createRenderer.ts | 10 +++-
 packages/runtime-core/src/optional/hooks.ts | 62 ++++++++++++---------
 2 files changed, 44 insertions(+), 28 deletions(-)

diff --git a/packages/runtime-core/src/createRenderer.ts b/packages/runtime-core/src/createRenderer.ts
index 9a46a2733..9e90371b3 100644
--- a/packages/runtime-core/src/createRenderer.ts
+++ b/packages/runtime-core/src/createRenderer.ts
@@ -22,6 +22,7 @@ import {
 } from './componentUtils'
 import { KeepAliveSymbol } from './optional/keepAlive'
 import { pushWarningContext, popWarningContext } from './warning'
+import { handleError, ErrorTypes } from './errorHandling'
 
 interface NodeOps {
   createElement: (tag: string, isSVG?: boolean) => any
@@ -1162,8 +1163,12 @@ export function createRenderer(options: RendererOptions) {
       beforeMount.call($proxy)
     }
 
+    const errorSchedulerHandler = (err: Error) => {
+      handleError(err, instance, ErrorTypes.SCHEDULER)
+    }
+
     const queueUpdate = (instance.$forceUpdate = () => {
-      queueJob(instance._updateHandle, flushHooks)
+      queueJob(instance._updateHandle, flushHooks, errorSchedulerHandler)
     })
 
     instance._updateHandle = autorun(
@@ -1227,7 +1232,7 @@ export function createRenderer(options: RendererOptions) {
       $vnode: prevVNode,
       $parentVNode,
       $proxy,
-      $options: { beforeUpdate, updated }
+      $options: { beforeUpdate }
     } = instance
     if (beforeUpdate) {
       beforeUpdate.call($proxy, prevVNode)
@@ -1256,6 +1261,7 @@ export function createRenderer(options: RendererOptions) {
       }
     }
 
+    const { updated } = instance.$options
     if (updated) {
       // Because the child's update is executed by the scheduler and not
       // synchronously within the parent's update call, the child's updated hook
diff --git a/packages/runtime-core/src/optional/hooks.ts b/packages/runtime-core/src/optional/hooks.ts
index 50d7f4431..9303613df 100644
--- a/packages/runtime-core/src/optional/hooks.ts
+++ b/packages/runtime-core/src/optional/hooks.ts
@@ -1,4 +1,4 @@
-import { ComponentInstance, APIMethods } from '../component'
+import { ComponentInstance, FunctionalComponent } from '../component'
 import { mergeLifecycleHooks, Data } from '../componentOptions'
 import { VNode, Slots } from '../vdom'
 import { observable } from '@vue/observer'
@@ -11,18 +11,31 @@ type Effect = RawEffect & {
 
 type EffectRecord = {
   effect: Effect
+  cleanup: Effect
   deps: any[] | void
 }
 
-type ComponentInstanceWithHook = ComponentInstance & {
-  _state: Record<number, any>
-  _effects: EffectRecord[]
+type HookState = {
+  state: any
+  effects: EffectRecord[]
 }
 
-let currentInstance: ComponentInstanceWithHook | null = null
+let currentInstance: ComponentInstance | null = null
 let isMounting: boolean = false
 let callIndex: number = 0
 
+const hooksState = new WeakMap<ComponentInstance, HookState>()
+
+export function setCurrentInstance(instance: ComponentInstance) {
+  currentInstance = instance
+  isMounting = !currentInstance._mounted
+  callIndex = 0
+}
+
+export function unsetCurrentInstance() {
+  currentInstance = null
+}
+
 export function useState(initial: any) {
   if (!currentInstance) {
     throw new Error(
@@ -30,7 +43,7 @@ export function useState(initial: any) {
     )
   }
   const id = ++callIndex
-  const state = currentInstance._state
+  const { state } = hooksState.get(currentInstance) as HookState
   const set = (newValue: any) => {
     state[id] = newValue
   }
@@ -56,36 +69,35 @@ export function useEffect(rawEffect: Effect, deps?: any[]) {
       }
     }
     const effect: Effect = () => {
-      cleanup()
       const { current } = effect
       if (current) {
-        effect.current = current()
+        cleanup.current = current()
+        effect.current = null
       }
     }
     effect.current = rawEffect
-
-    currentInstance._effects[id] = {
+    ;(hooksState.get(currentInstance) as HookState).effects[id] = {
       effect,
+      cleanup,
       deps
     }
 
     injectEffect(currentInstance, 'mounted', effect)
     injectEffect(currentInstance, 'unmounted', cleanup)
-    if (!deps) {
-      injectEffect(currentInstance, 'updated', effect)
-    }
+    injectEffect(currentInstance, 'updated', effect)
   } else {
-    const { effect, deps: prevDeps = [] } = currentInstance._effects[id]
+    const record = (hooksState.get(currentInstance) as HookState).effects[id]
+    const { effect, cleanup, deps: prevDeps = [] } = record
+    record.deps = deps
     if (!deps || deps.some((d, i) => d !== prevDeps[i])) {
+      cleanup()
       effect.current = rawEffect
-    } else {
-      effect.current = null
     }
   }
 }
 
 function injectEffect(
-  instance: ComponentInstanceWithHook,
+  instance: ComponentInstance,
   key: string,
   effect: Effect
 ) {
@@ -95,21 +107,19 @@ function injectEffect(
     : effect
 }
 
-export function withHooks<T extends APIMethods['render']>(render: T): T {
+export function withHooks<T extends FunctionalComponent>(render: T): T {
   return {
     displayName: render.name,
     created() {
-      const { _self } = this
-      _self._state = observable({})
-      _self._effects = []
+      hooksState.set(this._self, {
+        state: observable({}),
+        effects: []
+      })
     },
     render(props: Data, slots: Slots, attrs: Data, parentVNode: VNode) {
-      const { _self } = this
-      callIndex = 0
-      currentInstance = _self
-      isMounting = !_self._mounted
+      setCurrentInstance(this._self)
       const ret = render(props, slots, attrs, parentVNode)
-      currentInstance = null
+      unsetCurrentInstance()
       return ret
     }
   } as any