From 60a1b974044a7b3cf326808104a6a05bf6122f69 Mon Sep 17 00:00:00 2001
From: Rizumu Ayaka <rizumu@ayaka.moe>
Date: Tue, 9 Jan 2024 20:45:31 +0800
Subject: [PATCH] feat: middleware in baseWatch

---
 .../reactivity/__tests__/baseWatch.spec.ts    | 78 +++++++++++++++++++
 packages/reactivity/src/baseWatch.ts          | 62 +++++++++------
 packages/reactivity/src/index.ts              |  1 +
 3 files changed, 117 insertions(+), 24 deletions(-)

diff --git a/packages/reactivity/__tests__/baseWatch.spec.ts b/packages/reactivity/__tests__/baseWatch.spec.ts
index 02d9e64e0fa..0aab0aee6e8 100644
--- a/packages/reactivity/__tests__/baseWatch.spec.ts
+++ b/packages/reactivity/__tests__/baseWatch.spec.ts
@@ -175,4 +175,82 @@ describe('baseWatch', () => {
     scope.stop()
     expect(calls).toEqual(['sync 2', 'post 2'])
   })
+  test('baseWatch with middleware', async () => {
+    let effectCalls: string[] = []
+    let watchCalls: string[] = []
+    const source = ref(0)
+
+    // effect
+    baseWatch(
+      () => {
+        source.value
+        effectCalls.push('effect')
+        onEffectCleanup(() => effectCalls.push('effect cleanup'))
+      },
+      null,
+      {
+        scheduler,
+        middleware: next => {
+          effectCalls.push('before effect running')
+          next()
+          effectCalls.push('effect ran')
+        },
+      },
+    )
+    // watch
+    baseWatch(
+      () => source.value,
+      () => {
+        watchCalls.push('watch')
+        onEffectCleanup(() => watchCalls.push('watch cleanup'))
+      },
+      {
+        scheduler,
+        middleware: next => {
+          watchCalls.push('before watch running')
+          next()
+          watchCalls.push('watch ran')
+        },
+      },
+    )
+
+    expect(effectCalls).toEqual([])
+    expect(watchCalls).toEqual([])
+    await nextTick()
+    expect(effectCalls).toEqual([
+      'before effect running',
+      'effect',
+      'effect ran',
+    ])
+    expect(watchCalls).toEqual([])
+    effectCalls.length = 0
+    watchCalls.length = 0
+
+    source.value++
+    await nextTick()
+    expect(effectCalls).toEqual([
+      'before effect running',
+      'effect cleanup',
+      'effect',
+      'effect ran',
+    ])
+    expect(watchCalls).toEqual(['before watch running', 'watch', 'watch ran'])
+    effectCalls.length = 0
+    watchCalls.length = 0
+
+    source.value++
+    await nextTick()
+    expect(effectCalls).toEqual([
+      'before effect running',
+      'effect cleanup',
+      'effect',
+      'effect ran',
+    ])
+    expect(watchCalls).toEqual([
+      'before watch running',
+      'watch cleanup',
+      'watch',
+      'watch ran',
+    ])
+  })
 })
diff --git a/packages/reactivity/src/baseWatch.ts b/packages/reactivity/src/baseWatch.ts
index a97f43366b9..f6e20cfa52c 100644
--- a/packages/reactivity/src/baseWatch.ts
+++ b/packages/reactivity/src/baseWatch.ts
@@ -71,6 +71,7 @@ export interface BaseWatchOptions<Immediate = boolean> extends DebuggerOptions {
   deep?: boolean
   once?: boolean
   scheduler?: Scheduler
+  middleware?: BaseWatchMiddleware
   onError?: HandleError
   onWarn?: HandleWarn
 }
@@ -83,6 +84,7 @@ export type Scheduler = (
   effect: ReactiveEffect,
   isInit: boolean,
 ) => void
+export type BaseWatchMiddleware = (next: () => unknown) => any
 export type HandleError = (err: unknown, type: BaseWatchErrorCodes) => void
 export type HandleWarn = (msg: string, ...args: any[]) => void
 
@@ -132,6 +134,7 @@ export function baseWatch(
     scheduler = DEFAULT_SCHEDULER,
     onWarn = __DEV__ ? warn : NOOP,
     onError = DEFAULT_HANDLE_ERROR,
+    middleware,
     onTrack,
     onTrigger,
   }: BaseWatchOptions = EMPTY_OBJ,
@@ -211,6 +214,10 @@ export function baseWatch(
           activeEffect = currentEffect
         }
       }
+      if (middleware) {
+        const baseGetter = getter
+        getter = () => middleware(baseGetter)
+      }
     }
   } else {
     getter = NOOP
@@ -264,31 +271,38 @@ export function baseWatch(
           ? (newValue as any[]).some((v, i) => hasChanged(v, oldValue[i]))
           : hasChanged(newValue, oldValue))
       ) {
-        // cleanup before running cb again
-        if (cleanup) {
-          cleanup()
+        const next = () => {
+          // cleanup before running cb again
+          if (cleanup) {
+            cleanup()
+          }
+          const currentEffect = activeEffect
+          activeEffect = effect
+          try {
+            callWithAsyncErrorHandling(
+              cb!,
+              onError,
+              BaseWatchErrorCodes.WATCH_CALLBACK,
+              [
+                newValue,
+                // pass undefined as the old value when it's changed for the first time
+                oldValue === INITIAL_WATCHER_VALUE
+                  ? undefined
+                  : isMultiSource && oldValue[0] === INITIAL_WATCHER_VALUE
+                    ? []
+                    : oldValue,
+                onEffectCleanup,
+              ],
+            )
+            oldValue = newValue
+          } finally {
+            activeEffect = currentEffect
+          }
         }
-        const currentEffect = activeEffect
-        activeEffect = effect
-        try {
-          callWithAsyncErrorHandling(
-            cb,
-            onError,
-            BaseWatchErrorCodes.WATCH_CALLBACK,
-            [
-              newValue,
-              // pass undefined as the old value when it's changed for the first time
-              oldValue === INITIAL_WATCHER_VALUE
-                ? undefined
-                : isMultiSource && oldValue[0] === INITIAL_WATCHER_VALUE
-                  ? []
-                  : oldValue,
-              onEffectCleanup,
-            ],
-          )
-          oldValue = newValue
-        } finally {
-          activeEffect = currentEffect
+        if (middleware) {
+          middleware(next)
+        } else {
+          next()
         }
       }
     } else {
diff --git a/packages/reactivity/src/index.ts b/packages/reactivity/src/index.ts
index 8fe9c18c63a..88ef249e441 100644
--- a/packages/reactivity/src/index.ts
+++ b/packages/reactivity/src/index.ts
@@ -76,5 +76,6 @@ export {
   traverse,
   BaseWatchErrorCodes,
   type BaseWatchOptions,
+  type BaseWatchMiddleware,
   type Scheduler,
 } from './baseWatch'