refactor(reactivity): simplify Effect, EffectScope cleanup behavior

Author: johnsoncodehk

packages/reactivity/__tests__/effectScope.spec.ts

@@ -30,7 +30,8 @@ describe('reactivity/effect/scope', () => {
 
   it('should work w/ active property', () => {
     const scope = effectScope()
-    scope.run(() => 1)
+    const src = computed(() => 1)
+    scope.run(() => src.value)
     expect(scope.active).toBe(true)
     scope.stop()
     expect(scope.active).toBe(false)
@@ -170,17 +171,17 @@ describe('reactivity/effect/scope', () => {
 
     scope.stop()
 
+    expect(getEffectsCount(scope)).toBe(0)
+
     scope.run(() => {
       effect(() => (doubled = counter.num * 2))
     })
 
-    expect('[Vue warn] cannot run an inactive effect scope.').toHaveBeenWarned()
-
-    expect(getEffectsCount(scope)).toBe(0)
+    expect(getEffectsCount(scope)).toBe(1)
 
     counter.num = 7
     expect(dummy).toBe(0)
-    expect(doubled).toBe(undefined)
+    expect(doubled).toBe(14)
   })
 
   it('should fire onScopeDispose hook', () => {
@@ -359,7 +360,7 @@ describe('reactivity/effect/scope', () => {
     expect(cleanupCalls).toBe(1)
 
     expect(getEffectsCount(scope)).toBe(0)
-    expect(scope.cleanups.length).toBe(0)
+    expect(scope.cleanups).toBe(0)
   })
 })
 

packages/reactivity/src/effect.ts

@@ -3,6 +3,7 @@ import type { TrackOpTypes, TriggerOpTypes } from './constants'
 import { setupOnTrigger } from './debug'
 import { activeEffectScope } from './effectScope'
 import {
+  type Dependency,
   type Link,
   type Subscriber,
   SubscriberFlags,
@@ -51,26 +52,24 @@ export enum EffectFlags {
   ALLOW_RECURSE = 1 << 7,
   PAUSED = 1 << 8,
   NOTIFIED = 1 << 9,
-  STOP = 1 << 10,
 }
 
-export class ReactiveEffect<T = any> implements ReactiveEffectOptions {
+export class ReactiveEffect<T = any>
+  implements ReactiveEffectOptions, Dependency, Subscriber
+{
   // Subscriber
   deps: Link | undefined = undefined
   depsTail: Link | undefined = undefined
   flags: number = 0
+  cleanups: number = 0
 
   // Dependency
   subs: Link | undefined = undefined
   subsTail: Link | undefined = undefined
 
-  /**
-   * @internal
-   */
-  cleanup?: () => void = undefined
-
-  onStop?: () => void
+  // dev only
   onTrack?: (event: DebuggerEvent) => void
+  // dev only
   onTrigger?: (event: DebuggerEvent) => void
 
   // @ts-expect-error
@@ -80,13 +79,13 @@ export class ReactiveEffect<T = any> implements ReactiveEffectOptions {
     if (fn !== undefined) {
       this.fn = fn
     }
-    if (activeEffectScope && activeEffectScope.active) {
+    if (activeEffectScope) {
       link(this, activeEffectScope)
     }
   }
 
   get active(): boolean {
-    return !(this.flags & EffectFlags.STOP)
+    return this.deps !== undefined
   }
 
   pause(): void {
@@ -122,12 +121,10 @@ export class ReactiveEffect<T = any> implements ReactiveEffectOptions {
   }
 
   run(): T {
-    let flags = this.flags
-    if (flags & EffectFlags.STOP) {
-      // stopped during cleanup
-      return this.fn()
+    const cleanups = this.cleanups
+    if (cleanups) {
+      cleanup(this, cleanups)
     }
-    cleanupEffect(this)
     const prevSub = activeSub
     setActiveSub(this)
     startTracking(this)
@@ -143,7 +140,7 @@ export class ReactiveEffect<T = any> implements ReactiveEffectOptions {
       }
       setActiveSub(prevSub)
       endTracking(this)
-      flags = this.flags
+      const flags = this.flags
       if (
         (flags & (SubscriberFlags.Recursed | EffectFlags.ALLOW_RECURSE)) ===
         (SubscriberFlags.Recursed | EffectFlags.ALLOW_RECURSE)
@@ -155,17 +152,15 @@ export class ReactiveEffect<T = any> implements ReactiveEffectOptions {
   }
 
   stop(): void {
-    const flags = this.flags
-    if (!(flags & EffectFlags.STOP)) {
-      this.flags = flags | EffectFlags.STOP
-      startTracking(this)
-      endTracking(this)
-      cleanupEffect(this)
-      this.onStop && this.onStop()
-
-      if (this.subs !== undefined) {
-        unlink(this.subs)
-      }
+    const sub = this.subs
+    const cleanups = this.cleanups
+    if (sub !== undefined) {
+      unlink(sub)
+    }
+    startTracking(this)
+    endTracking(this)
+    if (cleanups) {
+      cleanup(this, cleanups)
     }
   }
 
@@ -205,6 +200,15 @@ export function effect<T = any>(
 
   const e = new ReactiveEffect(fn)
   if (options) {
+    const onStop = options.onStop
+    if (onStop !== undefined) {
+      options.onStop = undefined
+      const stop = e.stop.bind(e)
+      e.stop = () => {
+        stop()
+        onStop()
+      }
+    }
     extend(e, options)
   }
   try {
@@ -276,6 +280,32 @@ export function resetTracking(): void {
   }
 }
 
+const cleanupCbs = new WeakMap()
+
+export function onCleanup(
+  sub: Subscriber & { cleanups: number },
+  cb: () => void,
+): void {
+  const cbs = cleanupCbs.get(sub)
+  if (cbs === undefined) {
+    cleanupCbs.set(sub, [cb])
+    sub.cleanups = 1
+  } else {
+    cbs[sub.cleanups!++] = cb
+  }
+}
+
+export function cleanup(
+  sub: Subscriber & { cleanups: number },
+  length: number,
+): void {
+  const cbs = cleanupCbs.get(sub)!
+  for (let i = 0; i < length; ++i) {
+    cbs[i]()
+  }
+  sub.cleanups = 0
+}
+
 /**
  * Registers a cleanup function for the current active effect.
  * The cleanup function is called right before the next effect run, or when the
@@ -289,8 +319,9 @@ export function resetTracking(): void {
  * an active effect.
  */
 export function onEffectCleanup(fn: () => void, failSilently = false): void {
-  if (activeSub instanceof ReactiveEffect) {
-    activeSub.cleanup = fn
+  const e = activeSub
+  if (e instanceof ReactiveEffect) {
+    onCleanup(e, () => cleanupEffect(fn))
   } else if (__DEV__ && !failSilently) {
     warn(
       `onEffectCleanup() was called when there was no active effect` +
@@ -299,18 +330,14 @@ export function onEffectCleanup(fn: () => void, failSilently = false): void {
   }
 }
 
-function cleanupEffect(e: ReactiveEffect) {
-  const cleanup = e.cleanup
-  if (cleanup !== undefined) {
-    e.cleanup = undefined
-    // run cleanup without active effect
-    const prevSub = activeSub
-    activeSub = undefined
-    try {
-      cleanup()
-    } finally {
-      activeSub = prevSub
-    }
+function cleanupEffect(fn: () => void) {
+  // run cleanup without active effect
+  const prevSub = activeSub
+  activeSub = undefined
+  try {
+    fn()
+  } finally {
+    activeSub = prevSub
   }
 }
 

packages/reactivity/src/effectScope.ts

@@ -1,9 +1,11 @@
-import { EffectFlags } from './effect'
+import { EffectFlags, cleanup, onCleanup } from './effect'
 import {
   type Dependency,
   type Link,
   type Subscriber,
+  endTracking,
   link,
+  startTracking,
   unlink,
 } from './system'
 import { warn } from './warning'
@@ -15,28 +17,20 @@ export class EffectScope implements Subscriber, Dependency {
   deps: Link | undefined = undefined
   depsTail: Link | undefined = undefined
   flags: number = 0
+  cleanups: number = 0
 
   // Dependency
   subs: Link | undefined = undefined
   subsTail: Link | undefined = undefined
 
-  /**
-   * @internal
-   */
-  cleanups: (() => void)[] = []
-  /**
-   * @internal
-   */
-  cleanupsLength = 0
-
   constructor(detached = false) {
     if (!detached && activeEffectScope) {
       link(this, activeEffectScope)
     }
   }
 
   get active(): boolean {
-    return !(this.flags & EffectFlags.STOP)
+    return this.deps !== undefined
   }
 
   notify(): void {}
@@ -70,35 +64,25 @@ export class EffectScope implements Subscriber, Dependency {
   }
 
   run<T>(fn: () => T): T | undefined {
-    const flags = this.flags
-    if (!(flags & EffectFlags.STOP)) {
-      const prevEffectScope = activeEffectScope
-      try {
-        activeEffectScope = this
-        return fn()
-      } finally {
-        activeEffectScope = prevEffectScope
-      }
-    } else if (__DEV__) {
-      warn(`cannot run an inactive effect scope.`)
+    const prevEffectScope = activeEffectScope
+    try {
+      activeEffectScope = this
+      return fn()
+    } finally {
+      activeEffectScope = prevEffectScope
     }
   }
 
   stop(): void {
-    const flags = this.flags
-    if (!(flags & EffectFlags.STOP)) {
-      this.flags = flags | EffectFlags.STOP
-      while (this.deps !== undefined) {
-        unlink(this.deps)
-      }
-      const l = this.cleanupsLength
-      for (let i = 0; i < l; i++) {
-        this.cleanups[i]()
-      }
-      this.cleanupsLength = 0
-      if (this.subs !== undefined) {
-        unlink(this.subs)
-      }
+    const sub = this.subs
+    const cleanups = this.cleanups
+    if (sub !== undefined) {
+      unlink(sub)
+    }
+    startTracking(this)
+    endTracking(this)
+    if (cleanups) {
+      cleanup(this, cleanups)
     }
   }
 }
@@ -141,8 +125,8 @@ export function setCurrentScope(
  * @see {@link https://vuejs.org/api/reactivity-advanced.html#onscopedispose}
  */
 export function onScopeDispose(fn: () => void, failSilently = false): void {
-  if (activeEffectScope) {
-    activeEffectScope.cleanups[activeEffectScope.cleanupsLength++] = fn
+  if (activeEffectScope !== undefined) {
+    onCleanup(activeEffectScope, fn)
   } else if (__DEV__ && !failSilently) {
     warn(
       `onScopeDispose() is called when there is no active effect scope` +