fix(watch): fix watcher callback not flushing at right timing

Resolves: #120

Author: liximomo

src/apis/watch.ts

@@ -22,11 +22,17 @@ type MapSources<T> = {
 type FlushMode = 'pre' | 'post' | 'sync';
 
 interface WatcherOption {
-  lazy: boolean;
+  lazy: boolean; // whether or not to delay callcack invoking
   deep: boolean;
   flush: FlushMode;
 }
 
+export interface VueWatcher {
+  lazy: boolean;
+  get(): any;
+  teardown(): void;
+}
+
 let fallbackVM: ComponentInstance;
 
 function flushPreQueue(this: any) {
@@ -84,13 +90,39 @@ function queueFlushJob(vm: any, fn: () => void, mode: Exclude<FlushMode, 'sync'>
   }
 }
 
+function createVueWatcher(
+  vm: ComponentInstance,
+  getter: () => any,
+  callback: (n: any, o: any) => any,
+  options: {
+    deep: boolean;
+    sync: boolean;
+    immediateInvokeCallback?: boolean;
+    noRun?: boolean;
+    before?: () => void;
+  }
+): VueWatcher {
+  const index = vm._watchers.length;
+  // @ts-ignore: use undocumented options
+  vm.$watch(getter, callback, {
+    immediate: options.immediateInvokeCallback,
+    deep: options.deep,
+    lazy: options.noRun,
+    sync: options.sync,
+    before: options.before,
+  });
+
+  return vm._watchers[index];
+}
+
 function createWatcher(
   vm: ComponentInstance,
   source: WatcherSource<unknown> | WatcherSource<unknown>[] | SimpleEffect,
   cb: WatcherCallBack<any> | null,
   options: WatcherOption
 ): () => void {
   const flushMode = options.flush;
+  const isSync = flushMode === 'sync';
   let cleanup: (() => void) | null;
   const registerCleanup: CleanupRegistrator = (fn: () => void) => {
     cleanup = () => {
@@ -101,54 +133,51 @@ function createWatcher(
       }
     };
   };
+  // cleanup before running getter again
+  const runCleanup = () => {
+    if (cleanup) {
+      cleanup();
+      cleanup = null;
+    }
+  };
+  const createScheduler = <T extends Function>(fn: T): T => {
+    if (isSync || /* without a current active instance, ignore pre|post mode */ vm === fallbackVM) {
+      return fn;
+    }
+    return (((...args: any[]) =>
+      queueFlushJob(
+        vm,
+        () => {
+          fn(...args);
+        },
+        flushMode as 'pre' | 'post'
+      )) as any) as T;
+  };
 
   // effect watch
   if (cb === null) {
     const getter = () => (source as SimpleEffect)(registerCleanup);
-    // cleanup before running getter again
-    const runCleanup = () => {
-      if (cleanup) {
-        cleanup();
-      }
-    };
-
-    if (flushMode === 'sync') {
-      // @ts-ignore: use undocumented option "sync" ahd "before"
-      return vm.$watch(getter, noopFn, {
-        immediate: true,
-        deep: options.deep,
-        sync: true,
-        before: runCleanup,
-      });
-    }
-
-    let stopRef: Function;
-    let hasEnded: boolean = false;
-    const doWatch = () => {
-      if (hasEnded) return;
+    const watcher = createVueWatcher(vm, getter, noopFn, {
+      noRun: true, // take control the initial gettet invoking
+      deep: options.deep,
+      sync: isSync,
+      before: runCleanup,
+    });
 
-      // @ts-ignore: use undocumented option "before"
-      stopRef = vm.$watch(getter, noopFn, {
-        immediate: false,
-        deep: options.deep,
-        before: runCleanup,
-      });
-    };
+    // enable the watcher update
+    watcher.lazy = false;
 
-    /* without a current active instance, ignore pre|post mode */
-    if (vm === fallbackVM) {
-      vm.$nextTick(doWatch);
+    const originGet = watcher.get.bind(watcher);
+    if (isSync) {
+      watcher.get();
     } else {
-      queueFlushJob(vm, doWatch, flushMode);
+      vm.$nextTick(originGet);
     }
+    watcher.get = createScheduler(originGet);
 
     return () => {
-      hasEnded = true;
-      if (stopRef) {
-        stopRef();
-        runCleanup();
-        cleanup = null;
-      }
+      watcher.teardown();
+      runCleanup();
     };
   }
 
@@ -163,48 +192,33 @@ function createWatcher(
 
   const applyCb = (n: any, o: any) => {
     // cleanup before running cb again
-    if (cleanup) {
-      cleanup();
-    }
-
+    runCleanup();
     cb(n, o, registerCleanup);
   };
-
-  const callback =
-    flushMode === 'sync' ||
-    /* without a current active instance, ignore pre|post mode */
-    vm === fallbackVM
-      ? applyCb
-      : (n: any, o: any) =>
-          queueFlushJob(
-            vm,
-            () => {
-              applyCb(n, o);
-            },
-            flushMode
-          );
-
-  // `shiftCallback` is used to handle dirty sync effect.
-  // The subsequent callbacks will redirect to `callback`.
-  let shiftCallback = (n: any, o: any) => {
-    shiftCallback = callback;
-    applyCb(n, o);
-  };
+  let callback = createScheduler(applyCb);
+  if (!options.lazy) {
+    const originalCallbck = callback;
+    // `shiftCallback` is used to handle the first sync effect run.
+    // The subsequent callbacks will redirect to `callback`.
+    let shiftCallback = (n: any, o: any) => {
+      shiftCallback = originalCallbck;
+      applyCb(n, o);
+    };
+    callback = (n: any, o: any) => {
+      shiftCallback(n, o);
+    };
+  }
 
   // @ts-ignore: use undocumented option "sync"
-  const stop = vm.$watch(getter, options.lazy ? callback : shiftCallback, {
+  const stop = vm.$watch(getter, callback, {
     immediate: !options.lazy,
     deep: options.deep,
-    // @ts-ignore
-    sync: flushMode === 'sync',
+    sync: isSync,
   });
 
   return () => {
     stop();
-    if (cleanup) {
-      cleanup();
-      cleanup = null;
-    }
+    runCleanup();
   };
 }
 

src/env.d.ts

@@ -1,5 +1,6 @@
 import { VueConstructor } from 'vue';
 import { VfaState } from './vmStateManager';
+import { VueWatcher } from './apis/watch';
 
 declare global {
   interface Window {
@@ -11,6 +12,7 @@ declare module 'vue/types/vue' {
   interface Vue {
     readonly _uid: number;
     readonly _data: Record<string, any>;
+    _watchers: VueWatcher[];
     __secret_vfa_state__?: VfaState;
   }
 

src/setup.ts

@@ -46,8 +46,9 @@ function updateTemplateRef(vm: ComponentInstance) {
   const oldRefKeys = vmStateManager.get(vm, 'refs') || [];
   for (let index = 0; index < oldRefKeys.length; index++) {
     const key = oldRefKeys[index];
-    if (!refs[key]) {
-      (rawBindings[key] as Ref<any>).value = null;
+    const setupValue = rawBindings[key];
+    if (!refs[key] && setupValue && isRef(setupValue)) {
+      setupValue.value = null;
     }
   }
 

test/apis/watch.spec.js

@@ -147,7 +147,7 @@ describe('api/watch', () => {
       .then(done);
   });
 
-  it('should flush after render', done => {
+  it('should flush after render (lazy=true)', done => {
     let rerenderedText;
     const vm = new Vue({
       setup() {
@@ -177,6 +177,35 @@ describe('api/watch', () => {
     }).then(done);
   });
 
+  it('should flush after render (lazy=false)', done => {
+    let rerenderedText;
+    var vm = new Vue({
+      setup() {
+        const a = ref(1);
+        watch(a, (newVal, oldVal) => {
+          spy(newVal, oldVal);
+          if (vm) {
+            rerenderedText = vm.$el.textContent;
+          }
+        });
+        return {
+          a,
+        };
+      },
+      render(h) {
+        return h('div', this.a);
+      },
+    }).$mount();
+    expect(spy).toBeCalledTimes(1);
+    expect(spy).toHaveBeenLastCalledWith(1, undefined);
+    vm.a = 2;
+    waitForUpdate(() => {
+      expect(rerenderedText).toBe('2');
+      expect(spy).toBeCalledTimes(2);
+      expect(spy).toHaveBeenLastCalledWith(2, 1);
+    }).then(done);
+  });
+
   it('should flush before render', done => {
     const vm = new Vue({
       setup() {
@@ -262,6 +291,60 @@ describe('api/watch', () => {
     expect(spy).toHaveBeenNthCalledWith(2, 1, 0);
   });
 
+  it('should run in a expected order', done => {
+    const result = [];
+    var vm = new Vue({
+      setup() {
+        const x = ref(0);
+
+        // prettier-ignore
+        watch(() => { void x.value; result.push('sync getter'); }, { flush: 'sync' });
+        // prettier-ignore
+        watch(() => { void x.value; result.push('pre getter'); }, { flush: 'pre' });
+        // prettier-ignore
+        watch(() => { void x.value; result.push('post getter'); }, { flush: 'post' });
+
+        // prettier-ignore
+        watch(x, () => { result.push('sync callback') }, { flush: 'sync' })
+        // prettier-ignore
+        watch(x, () => { result.push('pre callback') }, { flush: 'pre' })
+        // prettier-ignore
+        watch(x, () => { result.push('post callback') }, { flush: 'post' })
+
+        const inc = () => {
+          result.push('before inc');
+          x.value++;
+          result.push('after inc');
+        };
+
+        return { x, inc };
+      },
+      template: `<div>{{x}}</div>`,
+    }).$mount();
+    expect(result).toEqual(['sync getter', 'sync callback', 'pre callback', 'post callback']);
+    result.length = 0;
+
+    waitForUpdate(() => {
+      expect(result).toEqual(['pre getter', 'post getter']);
+      result.length = 0;
+
+      vm.inc();
+    })
+      .then(() => {
+        expect(result).toEqual([
+          'before inc',
+          'sync getter',
+          'sync callback',
+          'after inc',
+          'pre getter',
+          'pre callback',
+          'post getter',
+          'post callback',
+        ]);
+      })
+      .then(done);
+  });
+
   describe('simple effect', () => {
     let renderedText;
     it('should work', done => {

test/templateRefs.spec.js

@@ -54,7 +54,10 @@ describe('ref', () => {
       vm.value = 'foo';
     })
       .then(() => {
-        // watcher callback gets invoked after this;
+        // vm updated. ref update occures after updated;
+      })
+      .then(() => {
+        // no render cycle, empty tick
       })
       .then(() => {
         expect(dummy1).toBe(null);