fix(reactivity): fix dep memory leak

Author: skirtles-code

packages/reactivity/__tests__/effect.spec.ts

@@ -12,7 +12,7 @@ import {
   readonly,
   ReactiveEffectRunner
 } from '../src/index'
-import { ITERATE_KEY } from '../src/effect'
+import { getDepFromReactive, ITERATE_KEY } from '../src/effect'
 
 describe('reactivity/effect', () => {
   it('should run the passed function once (wrapped by a effect)', () => {
@@ -999,4 +999,68 @@ describe('reactivity/effect', () => {
       expect(has).toBe(false)
     })
   })
+
+  describe('empty dep cleanup', () => {
+    it('should remove the dep when the effect is stopped', () => {
+      const obj = reactive({ prop: 1 })
+      expect(getDepFromReactive(toRaw(obj), 'prop')).toBeUndefined()
+      const runner = effect(() => obj.prop)
+      const dep = getDepFromReactive(toRaw(obj), 'prop')
+      expect(dep).toHaveLength(1)
+      obj.prop = 2
+      expect(getDepFromReactive(toRaw(obj), 'prop')).toBe(dep)
+      expect(dep).toHaveLength(1)
+      stop(runner)
+      expect(getDepFromReactive(toRaw(obj), 'prop')).toBeUndefined()
+      obj.prop = 3
+      runner()
+      expect(getDepFromReactive(toRaw(obj), 'prop')).toBeUndefined()
+    })
+
+    it('should only remove the dep when the last effect is stopped', () => {
+      const obj = reactive({ prop: 1 })
+      expect(getDepFromReactive(toRaw(obj), 'prop')).toBeUndefined()
+      const runner1 = effect(() => obj.prop)
+      const dep = getDepFromReactive(toRaw(obj), 'prop')
+      expect(dep).toHaveLength(1)
+      const runner2 = effect(() => obj.prop)
+      expect(getDepFromReactive(toRaw(obj), 'prop')).toBe(dep)
+      expect(dep).toHaveLength(2)
+      obj.prop = 2
+      expect(getDepFromReactive(toRaw(obj), 'prop')).toBe(dep)
+      expect(dep).toHaveLength(2)
+      stop(runner1)
+      expect(getDepFromReactive(toRaw(obj), 'prop')).toBe(dep)
+      expect(dep).toHaveLength(1)
+      obj.prop = 3
+      expect(getDepFromReactive(toRaw(obj), 'prop')).toBe(dep)
+      expect(dep).toHaveLength(1)
+      stop(runner2)
+      expect(getDepFromReactive(toRaw(obj), 'prop')).toBeUndefined()
+      obj.prop = 4
+      runner1()
+      runner2()
+      expect(getDepFromReactive(toRaw(obj), 'prop')).toBeUndefined()
+    })
+
+    it('should remove the dep when it is no longer used by the effect', () => {
+      const obj = reactive<{ a: number; b: number; c: 'a' | 'b' }>({
+        a: 1,
+        b: 2,
+        c: 'a'
+      })
+      expect(getDepFromReactive(toRaw(obj), 'prop')).toBeUndefined()
+      effect(() => obj[obj.c])
+      const depC = getDepFromReactive(toRaw(obj), 'c')
+      expect(getDepFromReactive(toRaw(obj), 'a')).toHaveLength(1)
+      expect(getDepFromReactive(toRaw(obj), 'b')).toBeUndefined()
+      expect(depC).toHaveLength(1)
+      obj.c = 'b'
+      obj.a = 4
+      expect(getDepFromReactive(toRaw(obj), 'a')).toBeUndefined()
+      expect(getDepFromReactive(toRaw(obj), 'b')).toHaveLength(1)
+      expect(getDepFromReactive(toRaw(obj), 'c')).toBe(depC)
+      expect(depC).toHaveLength(1)
+    })
+  })
 })

packages/reactivity/src/dep.ts

@@ -1,6 +1,10 @@
-import { ReactiveEffect, trackOpBit } from './effect'
+import { ReactiveEffect, removeEffectFromDep, trackOpBit } from './effect'
 
-export type Dep = Set<ReactiveEffect> & TrackedMarkers
+export type Dep = Set<ReactiveEffect> &
+  TrackedMarkers & {
+    target?: unknown
+    key?: unknown
+  }
 
 /**
  * wasTracked and newTracked maintain the status for several levels of effect
@@ -18,10 +22,16 @@ type TrackedMarkers = {
   n: number
 }
 
-export const createDep = (effects?: ReactiveEffect[]): Dep => {
-  const dep = new Set<ReactiveEffect>(effects) as Dep
+export function createDep(): Dep
+export function createDep(target: unknown, key: unknown): Dep
+export function createDep(target?: unknown, key?: unknown): Dep {
+  const dep = new Set<ReactiveEffect>() as Dep
   dep.w = 0
   dep.n = 0
+  if (target) {
+    dep.target = target
+    dep.key = key
+  }
   return dep
 }
 
@@ -44,7 +54,7 @@ export const finalizeDepMarkers = (effect: ReactiveEffect) => {
     for (let i = 0; i < deps.length; i++) {
       const dep = deps[i]
       if (wasTracked(dep) && !newTracked(dep)) {
-        dep.delete(effect)
+        removeEffectFromDep(dep, effect)
       } else {
         deps[ptr++] = dep
       }

packages/reactivity/src/effect.ts

@@ -143,12 +143,23 @@ function cleanupEffect(effect: ReactiveEffect) {
   const { deps } = effect
   if (deps.length) {
     for (let i = 0; i < deps.length; i++) {
-      deps[i].delete(effect)
+      removeEffectFromDep(deps[i], effect)
     }
     deps.length = 0
   }
 }
 
+export function removeEffectFromDep(dep: Dep, effect: ReactiveEffect) {
+  dep.delete(effect)
+  if (dep.target && dep.size === 0) {
+    const depsMap = targetMap.get(dep.target)
+    if (depsMap) {
+      depsMap.delete(dep.key)
+    }
+    dep.target = dep.key = null
+  }
+}
+
 export interface DebuggerOptions {
   onTrack?: (event: DebuggerEvent) => void
   onTrigger?: (event: DebuggerEvent) => void
@@ -252,7 +263,7 @@ export function track(target: object, type: TrackOpTypes, key: unknown) {
     }
     let dep = depsMap.get(key)
     if (!dep) {
-      depsMap.set(key, (dep = createDep()))
+      depsMap.set(key, (dep = createDep(target, key)))
     }
 
     const eventInfo = __DEV__
@@ -383,19 +394,19 @@ export function trigger(
       }
     }
     if (__DEV__) {
-      triggerEffects(createDep(effects), eventInfo)
+      triggerEffects(new Set(effects), eventInfo)
     } else {
-      triggerEffects(createDep(effects))
+      triggerEffects(new Set(effects))
     }
   }
 }
 
 export function triggerEffects(
-  dep: Dep | ReactiveEffect[],
+  dep: Set<ReactiveEffect>,
   debuggerEventExtraInfo?: DebuggerEventExtraInfo
 ) {
   // spread into array for stabilization
-  const effects = isArray(dep) ? dep : [...dep]
+  const effects = [...dep]
   for (const effect of effects) {
     if (effect.computed) {
       triggerEffect(effect, debuggerEventExtraInfo)