Merge pull request #14692 from storybookjs/12255-fix-jsx-decorators-4

Core: Don't recreate a bound story function each time we call a decorated story

examples/official-storybook/stories/core/rendering.stories.js

@@ -1,20 +1,55 @@
-import React, { useRef } from 'react';
+import React, { useEffect, useRef } from 'react';
+import { useArgs } from '@storybook/client-api';
 
 export default {
   title: 'Core/Rendering',
 };
 
 // NOTE: in our example apps each component is mounted twice as we render in strict mode
-let timesMounted = 0;
+let timesCounterMounted = 0;
 export const Counter = () => {
   const countRef = useRef();
 
-  if (!countRef.current) timesMounted += 1;
+  if (!countRef.current) timesCounterMounted += 1;
   countRef.current = (countRef.current || 0) + 1;
 
   return (
     <div>
-      Mounted: {timesMounted}, rendered (this mount): {countRef.current}
+      Mounted: {timesCounterMounted}, rendered (this mount): {countRef.current}
     </div>
   );
 };
+
+// An example to test what happens when the story is remounted due to argChanges
+let timesArgsChangeMounted = 0;
+export const ArgsChange = () => {
+  const countRef = useRef();
+
+  if (!countRef.current) timesArgsChangeMounted += 1;
+  countRef.current = true;
+
+  return (
+    <div>
+      Mounted: {timesArgsChangeMounted} (NOTE: we use strict mode so this number is 2x what you'd
+      expect -- it should be 2, not 4 though!)
+    </div>
+  );
+};
+
+ArgsChange.args = {
+  first: 0,
+};
+
+ArgsChange.decorators = [
+  (StoryFn) => {
+    const [args, updateArgs] = useArgs();
+
+    useEffect(() => {
+      if (args.first === 0) {
+        updateArgs({ first: 1 });
+      }
+    }, []);
+
+    return <StoryFn />;
+  },
+];

lib/addons/src/types.ts

@@ -53,6 +53,7 @@ export interface StoryIdentifier {
   name: StoryName;
 }
 
+export type StoryContextUpdate = Partial<StoryContext>;
 export type StoryContext = StoryIdentifier & {
   [key: string]: any;
   parameters: Parameters;
@@ -93,8 +94,13 @@ export interface OptionsParameter extends Object {
 
 export type StoryGetter = (context: StoryContext) => any;
 
+// This is the type of story function passed to a decorator -- does not rely on being passed any context
+export type PartialStoryFn<ReturnType = unknown> = (p?: StoryContextUpdate) => ReturnType;
+// This is a passArgsFirst: false user story function
 export type LegacyStoryFn<ReturnType = unknown> = (p?: StoryContext) => ReturnType;
+// This is a passArgsFirst: true user story function
 export type ArgsStoryFn<ReturnType = unknown> = (a?: Args, p?: StoryContext) => ReturnType;
+// This is either type of user story function
 export type StoryFn<ReturnType = unknown> = LegacyStoryFn<ReturnType> | ArgsStoryFn<ReturnType>;
 
 export type StoryWrapper = (
@@ -136,16 +142,16 @@ export interface StoryApi<StoryFnReturnType = unknown> {
 }
 
 export type DecoratorFunction<StoryFnReturnType = unknown> = (
-  fn: StoryFn<StoryFnReturnType>,
+  fn: PartialStoryFn<StoryFnReturnType>,
   c: StoryContext
-) => ReturnType<StoryFn<StoryFnReturnType>>;
+) => ReturnType<LegacyStoryFn<StoryFnReturnType>>;
 
 export type LoaderFunction = (c: StoryContext) => Promise<Record<string, any>>;
 
 export type DecorateStoryFunction<StoryFnReturnType = unknown> = (
-  storyFn: StoryFn<StoryFnReturnType>,
+  storyFn: LegacyStoryFn<StoryFnReturnType>,
   decorators: DecoratorFunction<StoryFnReturnType>[]
-) => StoryFn<StoryFnReturnType>;
+) => LegacyStoryFn<StoryFnReturnType>;
 
 export interface ClientStoryApi<StoryFnReturnType = unknown> {
   storiesOf(kind: StoryKind, module: NodeModule): StoryApi<StoryFnReturnType>;

lib/client-api/src/decorators.test.ts

@@ -2,7 +2,7 @@ import { StoryContext } from '@storybook/addons';
 
 import { defaultDecorateStory } from './decorators';
 
-function makeContext(input: Record<string, any>): StoryContext {
+function makeContext(input: Record<string, any> = {}): StoryContext {
   return {
     id: 'id',
     kind: 'kind',
@@ -24,7 +24,7 @@ describe('client-api.decorators', () => {
     const decorated = defaultDecorateStory(() => order.push(4), decorators);
 
     expect(order).toEqual([]);
-    decorated();
+    decorated(makeContext());
     expect(order).toEqual([3, 2, 1, 4]);
   });
 
@@ -42,6 +42,50 @@ describe('client-api.decorators', () => {
     expect(contexts.map((c) => c.k)).toEqual([0, 3, 2, 1]);
   });
 
+  it('does not recreate decorated story functions each time', () => {
+    const decoratedStories = [];
+    const decorators = [
+      (s, c) => {
+        decoratedStories.push = s;
+        return s();
+      },
+    ];
+    const decorated = defaultDecorateStory(() => 0, decorators);
+
+    decorated(makeContext());
+    decorated(makeContext());
+    expect(decoratedStories[0]).toBe(decoratedStories[1]);
+  });
+
+  // NOTE: important point--this test would not work if we called `decoratedOne` twice simultaneously
+  // both story functions would receive {story: 2}. The assumption here is that we'll never render
+  // the same story twice at the same time.
+  it('does not interleave contexts if two decorated stories are call simultaneously', async () => {
+    const contexts = [];
+    let resolve;
+    const fence = new Promise((r) => {
+      resolve = r;
+    });
+    const decorators = [
+      async (s, c) => {
+        // The fence here simulates async-ness in react rendering an element (`<S />` doesn't run `S()` straight away)
+        await fence;
+        s();
+      },
+    ];
+    const decoratedOne = defaultDecorateStory((c) => contexts.push(c), decorators);
+    const decoratedTwo = defaultDecorateStory((c) => contexts.push(c), decorators);
+
+    decoratedOne(makeContext({ value: 1 }));
+    decoratedTwo(makeContext({ value: 2 }));
+
+    resolve();
+    await fence;
+
+    expect(contexts[0].value).toBe(1);
+    expect(contexts[1].value).toBe(2);
+  });
+
   it('merges contexts', () => {
     const contexts = [];
     const decorators = [(s, c) => contexts.push(c) && s({ c: 'd' })];

lib/client-api/src/decorators.ts

@@ -1,10 +1,6 @@
-import { StoryContext, StoryFn } from '@storybook/addons';
+import { StoryContext, StoryContextUpdate, PartialStoryFn, LegacyStoryFn } from '@storybook/addons';
 import { DecoratorFunction } from './types';
 
-interface StoryContextUpdate {
-  [key: string]: any;
-}
-
 const defaultContext: StoryContext = {
   id: 'unspecified',
   name: 'unspecified',
@@ -25,15 +21,44 @@ const defaultContext: StoryContext = {
  * This will override the `foo` property on the `innerContext`, which gets
  * merged in with the default context
  */
-export const decorateStory = (storyFn: StoryFn, decorator: DecoratorFunction) => {
+const bindWithContext = (
-  return (context: StoryContext = defaultContext) =>
+  storyFn: LegacyStoryFn,
-    decorator(
+  getStoryContext: () => StoryContext
-      // You cannot override the parameters key, it is fixed
+): PartialStoryFn =>
-      ({ parameters, ...innerContext }: StoryContextUpdate = {}) =>
+  // (NOTE: You cannot override the parameters key, it is fixed)
-        storyFn({ ...context, ...innerContext }),
+  ({ parameters, ...contextUpdate }: StoryContextUpdate = {}) =>
-      context
+    storyFn({ ...getStoryContext(), ...contextUpdate });
-    );
+
+export const decorateStory = (
+  storyFn: LegacyStoryFn,
+  decorator: DecoratorFunction,
+  getStoryContext: () => StoryContext
+): LegacyStoryFn => {
+  // Bind the partially decorated storyFn so that when it is called it always knows about the story context,
+  // no matter what it is passed directly. This is because we cannot guarantee a decorator will
+  // pass the context down to the next decorated story in the chain.
+  const boundStoryFunction = bindWithContext(storyFn, getStoryContext);
+
+  return (context: StoryContext) => decorator(boundStoryFunction, context);
 };
 
-export const defaultDecorateStory = (storyFn: StoryFn, decorators: DecoratorFunction[]) =>
+export const defaultDecorateStory = (
-  decorators.reduce(decorateStory, storyFn);
+  storyFn: LegacyStoryFn,
+  decorators: DecoratorFunction[]
+): LegacyStoryFn => {
+  // We use a trick to avoid recreating the bound story function inside `decorateStory`.
+  // Instead we pass it a context "getter", which is defined once (at "decoration time")
+  // The getter reads a variable which is scoped to this call of `decorateStory`
+  // (ie to this story), so there is no possibility of overlap.
+  // This will break if you call the same story twice interleaved
+  // (React might do it if you rendered the same story twice in the one ReactDom.render call, for instance)
+  let contextStore: StoryContext;
+  const decoratedWithContextStore = decorators.reduce(
+    (story, decorator) => decorateStory(story, decorator, () => contextStore),
+    storyFn
+  );
+  return (context = defaultContext) => {
+    contextStore = context;
+    return decoratedWithContextStore(context); // Pass the context directly into the first decorator
+  };
+};

lib/client-api/src/story_store.ts

@@ -394,7 +394,9 @@ export default class StoryStore {
           return acc;
         }, {} as Args),
       };
-      return passArgsFirst ? (original as ArgsStoryFn)(mapped.args, mapped) : original(mapped);
+      return passArgsFirst
+        ? (original as ArgsStoryFn)(mapped.args, mapped)
+        : (original as LegacyStoryFn)(mapped);
     };
 
     // lazily decorate the story when it's loaded