Implement preview hooks

lib/client-api/src/client_api.js

@@ -10,6 +10,7 @@ import mergeWith from 'lodash/mergeWith';
 import isEqual from 'lodash/isEqual';
 
 import subscriptionsStore from './subscriptions_store';
+import { applyHooks } from './hooks';
 
 // merge with concatenating arrays, but no duplicates
 const merge = (a, b) =>
@@ -204,7 +205,7 @@ export default class ClientApi {
           parameters: allParam,
         },
         {
-          applyDecorators: this._decorateStory,
+          applyDecorators: applyHooks(this._decorateStory),
           getDecorators: () => [
             ...(allParam.decorators || []),
             ...localDecorators,

lib/client-api/src/hooks.test.js

@@ -0,0 +1,317 @@
+import { FORCE_RE_RENDER } from '@storybook/addon-contexts/src/shared/constants';
+import { defaultDecorateStory } from './client_api';
+import { applyHooks, useEffect, useMemo, useCallback, useRef, useState, useReducer } from './hooks';
+
+jest.mock('@storybook/client-logger', () => ({
+  logger: { warn: jest.fn(), log: jest.fn() },
+}));
+
+const mockChannel = {
+  emit: jest.fn(),
+};
+
+jest.mock('@storybook/addons', () => ({
+  getChannel: () => mockChannel,
+}));
+
+const decorateStory = applyHooks(defaultDecorateStory);
+
+const run = (storyFn, decorators = []) => decorateStory(storyFn, decorators)();
+
+describe('Preview hooks', () => {
+  describe('useEffect', () => {
+    it('triggers the effect from story function', () => {
+      const effect = jest.fn();
+      run(() => {
+        useEffect(effect);
+      });
+      expect(effect).toHaveBeenCalledTimes(1);
+    });
+    it('triggers the effect from decorator', () => {
+      const effect = jest.fn();
+      run(() => {}, [
+        storyFn => {
+          useEffect(effect);
+          return storyFn();
+        },
+      ]);
+      expect(effect).toHaveBeenCalledTimes(1);
+    });
+    it('retriggers the effect if no deps array is provided', () => {
+      const effect = jest.fn();
+      const storyFn = () => {
+        useEffect(effect);
+      };
+      run(storyFn);
+      run(storyFn);
+      expect(effect).toHaveBeenCalledTimes(2);
+    });
+    it("doesn't retrigger the effect if empty deps array is provided", () => {
+      const effect = jest.fn();
+      const storyFn = () => {
+        useEffect(effect, []);
+      };
+      run(storyFn);
+      run(storyFn);
+      expect(effect).toHaveBeenCalledTimes(1);
+    });
+    it('retriggers the effect if some of the deps are changed', () => {
+      const effect = jest.fn();
+      let counter = 0;
+      const storyFn = () => {
+        useEffect(effect, [counter]);
+        counter += 1;
+      };
+      run(storyFn);
+      run(storyFn);
+      expect(effect).toHaveBeenCalledTimes(2);
+    });
+    it("doesn't retrigger the effect if none of the deps are changed", () => {
+      const effect = jest.fn();
+      const storyFn = () => {
+        useEffect(effect, [0]);
+      };
+      run(storyFn);
+      run(storyFn);
+      expect(effect).toHaveBeenCalledTimes(1);
+    });
+    it('performs cleanup when retriggering', () => {
+      const destroy = jest.fn();
+      const storyFn = () => {
+        useEffect(() => destroy);
+      };
+      run(storyFn);
+      run(storyFn);
+      expect(destroy).toHaveBeenCalledTimes(1);
+    });
+    it("doesn't perform cleanup when keeping the current effect", () => {
+      const destroy = jest.fn();
+      const storyFn = () => {
+        useEffect(() => destroy, []);
+      };
+      run(storyFn);
+      run(storyFn);
+      expect(destroy).not.toHaveBeenCalled();
+    });
+    it('performs cleanup when removing the decorator', () => {
+      const destroy = jest.fn();
+      run(() => {}, [
+        storyFn => {
+          useEffect(() => destroy);
+          return storyFn();
+        },
+      ]);
+      run(() => {});
+      expect(destroy).toHaveBeenCalledTimes(1);
+    });
+  });
+  describe('useMemo', () => {
+    it('performs the calculation', () => {
+      let result;
+      const nextCreate = jest.fn(() => 'foo');
+      const storyFn = () => {
+        result = useMemo(nextCreate, []);
+      };
+      run(storyFn);
+      expect(nextCreate).toHaveBeenCalledTimes(1);
+      expect(result).toBe('foo');
+    });
+    it('performs the calculation once if deps are unchanged', () => {
+      let result;
+      const nextCreate = jest.fn(() => 'foo');
+      const storyFn = () => {
+        result = useMemo(nextCreate, []);
+      };
+      run(storyFn);
+      run(storyFn);
+      expect(nextCreate).toHaveBeenCalledTimes(1);
+      expect(result).toBe('foo');
+    });
+    it('performs the calculation again if deps are changed', () => {
+      let result;
+      let counter = 0;
+      const nextCreate = jest.fn(() => counter);
+      const storyFn = () => {
+        counter += 1;
+        result = useMemo(nextCreate, [counter]);
+      };
+      run(storyFn);
+      run(storyFn);
+      expect(nextCreate).toHaveBeenCalledTimes(2);
+      expect(result).toBe(counter);
+    });
+  });
+  describe('useCallback', () => {
+    it('returns the callback', () => {
+      let result;
+      const callback = () => {};
+      const storyFn = () => {
+        result = useCallback(callback, []);
+      };
+      run(storyFn);
+      expect(result).toBe(callback);
+    });
+    it('returns the previous callback reference if deps are unchanged', () => {
+      const callbacks = [];
+      const storyFn = () => {
+        const callback = useCallback(() => {}, []);
+        callbacks.push(callback);
+      };
+      run(storyFn);
+      run(storyFn);
+      expect(callbacks[0]).toBe(callbacks[1]);
+    });
+    it('creates new callback reference if deps are changed', () => {
+      const callbacks = [];
+      let counter = 0;
+      const storyFn = () => {
+        counter += 1;
+        const callback = useCallback(() => {}, [counter]);
+        callbacks.push(callback);
+      };
+      run(storyFn);
+      run(storyFn);
+      expect(callbacks[0]).not.toBe(callbacks[1]);
+    });
+  });
+  describe('useRef', () => {
+    it('attaches initial value', () => {
+      let ref;
+      const storyFn = () => {
+        ref = useRef('foo');
+      };
+      run(storyFn);
+      expect(ref.current).toBe('foo');
+    });
+    it('stores mutations', () => {
+      let refValueFromSecondRender;
+      let counter = 0;
+      const storyFn = () => {
+        counter += 1;
+        const ref = useRef('foo');
+
+        if (counter === 1) {
+          ref.current = 'bar';
+        } else {
+          refValueFromSecondRender = ref.current;
+        }
+      };
+      run(storyFn);
+      run(storyFn);
+      expect(refValueFromSecondRender).toBe('bar');
+    });
+  });
+  describe('useState', () => {
+    it('sets initial state', () => {
+      let state;
+      const storyFn = () => {
+        [state] = useState('foo');
+      };
+      run(storyFn);
+      expect(state).toBe('foo');
+    });
+    it('calculates initial state', () => {
+      let state;
+      const storyFn = () => {
+        [state] = useState(() => 'foo');
+      };
+      run(storyFn);
+      expect(state).toBe('foo');
+    });
+    it('performs synchronous state updates', () => {
+      let state;
+      let setState;
+      const storyFn = jest.fn(() => {
+        [state, setState] = useState('foo');
+        if (state === 'foo') {
+          setState('bar');
+        }
+      });
+      run(storyFn);
+      expect(storyFn).toHaveBeenCalledTimes(2);
+      expect(state).toBe('bar');
+    });
+    it('performs synchronous state updates with updater function', () => {
+      let state;
+      let setState;
+      const storyFn = jest.fn(() => {
+        [state, setState] = useState(0);
+        if (state < 3) {
+          setState(prevState => prevState + 1);
+        }
+      });
+      run(storyFn);
+      expect(storyFn).toHaveBeenCalledTimes(4);
+      expect(state).toBe(3);
+    });
+    it('performs asynchronous state updates', () => {
+      let state;
+      let setState;
+      const storyFn = jest.fn(() => {
+        [state, setState] = useState('foo');
+      });
+      run(storyFn);
+      setState('bar');
+      expect(mockChannel.emit).toHaveBeenCalledWith(FORCE_RE_RENDER);
+      run(storyFn);
+      expect(state).toBe('bar');
+    });
+  });
+  describe('useReducer', () => {
+    it('sets initial state', () => {
+      let state;
+      const storyFn = () => {
+        [state] = useReducer(() => 'bar', 'foo');
+      };
+      run(storyFn);
+      expect(state).toBe('foo');
+    });
+    it('calculates initial state', () => {
+      let state;
+      const storyFn = () => {
+        [state] = useReducer(() => 'bar', 'foo', arg => arg);
+      };
+      run(storyFn);
+      expect(state).toBe('foo');
+    });
+    it('performs synchronous state updates', () => {
+      let state;
+      let dispatch;
+      const storyFn = jest.fn(() => {
+        [state, dispatch] = useReducer((prevState, action) => {
+          switch (action) {
+            case 'INCREMENT':
+              return prevState + 1;
+            default:
+              return prevState;
+          }
+        }, 0);
+        if (state < 3) {
+          dispatch('INCREMENT');
+        }
+      });
+      run(storyFn);
+      expect(storyFn).toHaveBeenCalledTimes(4);
+      expect(state).toBe(3);
+    });
+    it('performs asynchronous state updates', () => {
+      let state;
+      let dispatch;
+      const storyFn = jest.fn(() => {
+        [state, dispatch] = useReducer((prevState, action) => {
+          switch (action) {
+            case 'INCREMENT':
+              return prevState + 1;
+            default:
+              return prevState;
+          }
+        }, 0);
+      });
+      run(storyFn);
+      dispatch('INCREMENT');
+      expect(mockChannel.emit).toHaveBeenCalledWith(FORCE_RE_RENDER);
+      run(storyFn);
+      expect(state).toBe(1);
+    });
+  });
+});

lib/client-api/src/hooks.ts

@@ -0,0 +1,246 @@
+/* eslint-disable import/export */
+import { logger } from '@storybook/client-logger';
+import addons from '@storybook/addons';
+import { FORCE_RE_RENDER } from '@storybook/core-events';
+
+interface Hook {
+  name: string;
+  memoizedState?: any;
+  deps?: any[] | undefined;
+}
+
+interface Effect {
+  create: () => (() => void) | void;
+  destroy?: (() => void) | void;
+}
+
+type AbstractFunction = (...args: any[]) => any;
+
+const hookListsMap = new WeakMap<AbstractFunction, Hook[]>();
+let mountedDecorators = new Set<AbstractFunction>();
+
+let currentHooks: Hook[] = [];
+let nextHookIndex = 0;
+const getNextHook = () => {
+  const hook = currentHooks[nextHookIndex];
+  nextHookIndex += 1;
+  return hook;
+};
+let currentPhase: 'MOUNT' | 'UPDATE' | 'NONE' = 'NONE';
+let currentEffects: Effect[] = [];
+let prevEffects: Effect[] = [];
+let currentDecoratorName: string | null = null;
+let hasUpdates = false;
+
+const triggerEffects = () => {
+  // destroy removed effects
+  prevEffects.forEach(effect => {
+    if (!currentEffects.includes(effect) && effect.destroy) {
+      effect.destroy();
+    }
+  });
+  // trigger added effects
+  currentEffects.forEach(effect => {
+    if (!prevEffects.includes(effect)) {
+      // eslint-disable-next-line no-param-reassign
+      effect.destroy = effect.create();
+    }
+  });
+  prevEffects = currentEffects;
+  currentEffects = [];
+};
+
+const hookify = (fn: AbstractFunction) => (...args: any[]) => {
+  currentDecoratorName = fn.name;
+  if (mountedDecorators.has(fn)) {
+    currentPhase = 'UPDATE';
+    currentHooks = hookListsMap.get(fn) || [];
+  } else {
+    currentPhase = 'MOUNT';
+    currentHooks = [];
+    hookListsMap.set(fn, currentHooks);
+  }
+  nextHookIndex = 0;
+
+  const result = fn(...args);
+
+  if (currentPhase === 'UPDATE' && getNextHook() != null) {
+    throw new Error(
+      'Rendered fewer hooks than expected. This may be caused by an accidental early return statement.'
+    );
+  }
+
+  currentDecoratorName = null;
+  return result;
+};
+
+// Counter to prevent infinite loops.
+let numberOfRenders = 0;
+const RENDER_LIMIT = 25;
+export const applyHooks = (
+  applyDecorators: (storyFn: AbstractFunction, decorators: AbstractFunction[]) => any
+) => (storyFn: AbstractFunction, decorators: AbstractFunction[]) => {
+  const decorated = applyDecorators(hookify(storyFn), decorators.map(hookify));
+  return (...args: any[]) => {
+    hasUpdates = false;
+    let result = decorated(...args);
+    mountedDecorators = new Set([storyFn, ...decorators]);
+    numberOfRenders = 1;
+    while (hasUpdates) {
+      hasUpdates = false;
+      result = decorated(...args);
+      numberOfRenders += 1;
+      if (numberOfRenders > RENDER_LIMIT) {
+        throw new Error(
+          'Too many re-renders. Storybook limits the number of renders to prevent an infinite loop.'
+        );
+      }
+    }
+    currentPhase = 'NONE';
+    triggerEffects();
+    return result;
+  };
+};
+
+const areDepsEqual = (deps: any[], nextDeps: any[]) =>
+  deps.length === nextDeps.length && deps.every((dep, i) => dep === nextDeps[i]);
+
+function useHook(name: string, callback: (hook: Hook) => void, deps?: any[] | undefined): Hook {
+  if (currentPhase === 'MOUNT') {
+    if (deps != null && !Array.isArray(deps)) {
+      logger.warn(
+        `${name} received a final argument that is not an array (instead, received ${deps}). When specified, the final argument must be an array.`
+      );
+    }
+    const hook: Hook = { name, deps };
+    currentHooks.push(hook);
+    callback(hook);
+    return hook;
+  }
+
+  if (currentPhase === 'UPDATE') {
+    const hook = getNextHook();
+    if (hook == null) {
+      throw new Error('Rendered more hooks than during the previous render.');
+    }
+
+    if (hook.name !== name) {
+      logger.warn(
+        `Storybook has detected a change in the order of Hooks${
+          currentDecoratorName ? ` called by ${currentDecoratorName}` : ''
+        }. This will lead to bugs and errors if not fixed.`
+      );
+    }
+
+    if (deps != null && hook.deps == null) {
+      logger.warn(
+        `${name} received a final argument during this render, but not during the previous render. Even though the final argument is optional, its type cannot change between renders.`
+      );
+    }
+
+    if (deps != null && hook.deps != null && deps.length !== hook.deps.length) {
+      logger.warn(`The final argument passed to ${name} changed size between renders. The order and size of this array must remain constant.
+Previous: ${hook.deps}
+Incoming: ${deps}`);
+    }
+
+    if (deps == null || hook.deps == null || !areDepsEqual(deps, hook.deps)) {
+      callback(hook);
+      hook.deps = deps;
+    }
+    return hook;
+  }
+
+  throw new Error(
+    'Storybook preview hooks can only be called inside decorators and story functions.'
+  );
+}
+
+function useMemoLike<T>(name: string, nextCreate: () => T, deps: any[] | undefined): T {
+  const { memoizedState } = useHook(
+    name,
+    hook => {
+      // eslint-disable-next-line no-param-reassign
+      hook.memoizedState = nextCreate();
+    },
+    deps
+  );
+  return memoizedState;
+}
+
+export function useMemo<T>(nextCreate: () => T, deps?: any[]): T {
+  return useMemoLike('useMemo', nextCreate, deps);
+}
+
+export function useCallback<T>(callback: T, deps?: any[]): T {
+  return useMemoLike('useCallback', () => callback, deps);
+}
+
+function useRefLike<T>(name: string, initialValue: T): { current: T } {
+  return useMemoLike(name, () => ({ current: initialValue }), []);
+}
+
+export function useRef<T>(initialValue: T): { current: T } {
+  return useRefLike('useRef', initialValue);
+}
+
+function triggerUpdate() {
+  // Rerun storyFn if updates were triggered synchronously, force rerender otherwise
+  if (currentPhase !== 'NONE') {
+    hasUpdates = true;
+  } else {
+    try {
+      addons.getChannel().emit(FORCE_RE_RENDER);
+    } catch (e) {
+      logger.warn('State updates of Storybook preview hooks work only in browser');
+    }
+  }
+}
+
+function useStateLike<S>(
+  name: string,
+  initialState: (() => S) | S
+): [S, (update: ((prevState: S) => S) | S) => void] {
+  const stateRef = useRefLike(
+    name,
+    // @ts-ignore S type should never be function, but there's no way to tell that to TypeScript
+    typeof initialState === 'function' ? initialState() : initialState
+  );
+  const setState = (update: ((prevState: S) => S) | S) => {
+    // @ts-ignore S type should never be function, but there's no way to tell that to TypeScript
+    stateRef.current = typeof update === 'function' ? update(stateRef.current) : update;
+    triggerUpdate();
+  };
+  return [stateRef.current, setState];
+}
+
+export function useState<S>(
+  initialState: (() => S) | S
+): [S, (update: ((prevState: S) => S) | S) => void] {
+  return useStateLike('useState', initialState);
+}
+
+export function useReducer<S, A>(
+  reducer: (state: S, action: A) => S,
+  initialState: S
+): [S, (action: A) => void];
+export function useReducer<S, I, A>(
+  reducer: (state: S, action: A) => S,
+  initialArg: I,
+  init: (initialArg: I) => S
+): [S, (action: A) => void];
+export function useReducer<S, A>(
+  reducer: (state: S, action: A) => S,
+  initialArg: any,
+  init?: any
+): [S, (action: A) => void] {
+  const initialState: (() => S) | S = init != null ? () => init(initialArg) : initialArg;
+  const [state, setState] = useStateLike('useReducer', initialState);
+  const dispatch = (action: A) => setState(prevState => reducer(prevState, action));
+  return [state, dispatch];
+}
+
+export function useEffect(create: () => (() => void) | void, deps?: any[]): void {
+  const effect = useMemoLike('useEffect', () => ({ create }), deps);
+  currentEffects.push(effect);
+}

lib/client-api/tsconfig.json

@@ -0,0 +1,8 @@
+{
+  "extends": "../../tsconfig.json",
+  "compilerOptions": {
+    "rootDir": "./src"
+  },
+  "include": ["src/**/*"],
+  "exclude": ["src/**/*.test.ts"]
+}