Core: Add story preloading to optimize lazy compilation

lib/api/src/lib/stories.ts

@@ -1,3 +1,4 @@
+import memoize from 'memoizerific';
 import React from 'react';
 import deprecate from 'util-deprecate';
 import dedent from 'ts-dedent';
@@ -327,3 +328,17 @@ export function isStory(item: Item): item is Story {
   }
   return false;
 }
+
+export const getComponentLookupList = memoize(1)((hash: StoriesHash) => {
+  return Object.entries(hash).reduce((acc, i) => {
+    const value = i[1];
+    if (value.isComponent) {
+      acc.push([...i[1].children]);
+    }
+    return acc;
+  }, [] as StoryId[][]);
+});
+
+export const getStoriesLookupList = memoize(1)((hash: StoriesHash) => {
+  return Object.keys(hash).filter((k) => !(hash[k].children || Array.isArray(hash[k])));
+});

lib/api/src/modules/stories.ts

@@ -1,6 +1,7 @@
 import global from 'global';
 import { toId, sanitize } from '@storybook/csf';
 import {
+  PRELOAD_STORIES,
   STORY_PREPARED,
   UPDATE_STORY_ARGS,
   RESET_STORY_ARGS,
@@ -22,7 +23,10 @@ import {
   isStory,
   isRoot,
   transformStoryIndexToStoriesHash,
+  getComponentLookupList,
+  getStoriesLookupList,
 } from '../lib/stories';
+
 import type {
   StoriesHash,
   Story,
@@ -77,6 +81,12 @@ export interface SubAPI {
   updateStoryArgs(story: Story, newArgs: Args): void;
   resetStoryArgs: (story: Story, argNames?: string[]) => void;
   findLeafStoryId(StoriesHash: StoriesHash, storyId: StoryId): StoryId;
+  findSiblingStoryId(
+    storyId: StoryId,
+    hash: StoriesHash,
+    direction: Direction,
+    toSiblingGroup: boolean // when true, skip over leafs within the same group
+  ): StoryId;
   fetchStoryList: () => Promise<void>;
   setStoryList: (storyList: StoryIndex) => Promise<void>;
   updateStory: (storyId: StoryId, update: StoryUpdate, ref?: ComposedRef) => Promise<void>;
@@ -187,26 +197,7 @@ export const init: ModuleFn = ({
       }
 
       const hash = refId ? refs[refId].stories || {} : storiesHash;
-
-      const lookupList = Object.entries(hash).reduce((acc, i) => {
-        const value = i[1];
-        if (value.isComponent) {
-          acc.push([...i[1].children]);
-        }
-        return acc;
-      }, []);
-
-      const index = lookupList.findIndex((i) => i.includes(storyId));
-
-      // cannot navigate beyond fist or last
-      if (index === lookupList.length - 1 && direction > 0) {
-        return;
-      }
-      if (index === 0 && direction < 0) {
-        return;
-      }
-
-      const result = lookupList[index + direction][0];
+      const result = api.findSiblingStoryId(storyId, hash, direction, true);
 
       if (result) {
         api.selectStory(result, undefined, { ref: refId });
@@ -227,21 +218,7 @@ export const init: ModuleFn = ({
       }
 
       const hash = story.refId ? refs[story.refId].stories : storiesHash;
-
-      const lookupList = Object.keys(hash).filter(
-        (k) => !(hash[k].children || Array.isArray(hash[k]))
-      );
-      const index = lookupList.indexOf(storyId);
-
-      // cannot navigate beyond fist or last
-      if (index === lookupList.length - 1 && direction > 0) {
-        return;
-      }
-      if (index === 0 && direction < 0) {
-        return;
-      }
-
-      const result = lookupList[index + direction];
+      const result = api.findSiblingStoryId(storyId, hash, direction, false);
 
       if (result) {
         api.selectStory(result, undefined, { ref: refId });
@@ -332,6 +309,39 @@ export const init: ModuleFn = ({
       const childStoryId = storiesHash[storyId].children[0];
       return api.findLeafStoryId(storiesHash, childStoryId);
     },
+    findSiblingStoryId(storyId, hash, direction, toSiblingGroup) {
+      if (toSiblingGroup) {
+        const lookupList = getComponentLookupList(hash);
+        const index = lookupList.findIndex((i) => i.includes(storyId));
+
+        // cannot navigate beyond fist or last
+        if (index === lookupList.length - 1 && direction > 0) {
+          return;
+        }
+        if (index === 0 && direction < 0) {
+          return;
+        }
+
+        if (lookupList[index + direction]) {
+          // eslint-disable-next-line consistent-return
+          return lookupList[index + direction][0];
+        }
+        return;
+      }
+      const lookupList = getStoriesLookupList(hash);
+      const index = lookupList.indexOf(storyId);
+
+      // cannot navigate beyond fist or last
+      if (index === lookupList.length - 1 && direction > 0) {
+        return;
+      }
+      if (index === 0 && direction < 0) {
+        return;
+      }
+
+      // eslint-disable-next-line consistent-return
+      return lookupList[index + direction];
+    },
     updateStoryArgs: (story, updatedArgs) => {
       const { id: storyId, refId } = story;
       fullAPI.emit(UPDATE_STORY_ARGS, {
@@ -448,6 +458,36 @@ export const init: ModuleFn = ({
       }
     });
 
+    fullAPI.on(STORY_PREPARED, function handler({ id, ...update }) {
+      const { ref, sourceType } = getEventMetadata(this, fullAPI);
+      fullAPI.updateStory(id, { ...update, prepared: true }, ref);
+
+      if (!ref) {
+        if (!store.getState().hasCalledSetOptions) {
+          const { options } = update.parameters;
+          checkDeprecatedOptionParameters(options);
+          fullAPI.setOptions(options);
+          store.setState({ hasCalledSetOptions: true });
+        }
+      } else {
+        fullAPI.updateRef(ref.id, { ready: true });
+      }
+
+      if (sourceType === 'local') {
+        const { storyId, storiesHash } = store.getState();
+
+        // create a list of related stories to be preloaded
+        const toBePreloaded = Array.from(
+          new Set([
+            api.findSiblingStoryId(storyId, storiesHash, 1, true),
+            api.findSiblingStoryId(storyId, storiesHash, -1, true),
+          ])
+        ).filter(Boolean);
+
+        fullAPI.emit(PRELOAD_STORIES, toBePreloaded);
+      }
+    });
+
     fullAPI.on(SET_STORIES, function handler(data: SetStoriesPayload) {
       const { ref } = getEventMetadata(this, fullAPI);
       const stories = data.v ? denormalizeStoryParameters(data) : data.stories;
@@ -489,22 +529,6 @@ export const init: ModuleFn = ({
       }
     );
 
-    fullAPI.on(STORY_PREPARED, function handler({ id, ...update }) {
-      const { ref } = getEventMetadata(this, fullAPI);
-      fullAPI.updateStory(id, { ...update, prepared: true }, ref);
-
-      if (!ref) {
-        if (!store.getState().hasCalledSetOptions) {
-          const { options } = update.parameters;
-          checkDeprecatedOptionParameters(options);
-          fullAPI.setOptions(options);
-          store.setState({ hasCalledSetOptions: true });
-        }
-      } else {
-        fullAPI.updateRef(ref.id, { ready: true });
-      }
-    });
-
     fullAPI.on(
       STORY_ARGS_UPDATED,
       function handleStoryArgsUpdated({ storyId, args }: { storyId: StoryId; args: Args }) {

lib/api/src/tests/stories.test.js

@@ -85,8 +85,22 @@ describe('stories API', () => {
   });
   const parameters = {};
   const storiesHash = {
-    'a--1': { kind: 'a', name: '1', parameters, path: 'a--1', id: 'a--1', args: {} },
-    'a--2': { kind: 'a', name: '2', parameters, path: 'a--2', id: 'a--2', args: {} },
+    'a--1': {
+      kind: 'a',
+      name: '1',
+      parameters,
+      path: 'a--1',
+      id: 'a--1',
+      args: {},
+    },
+    'a--2': {
+      kind: 'a',
+      name: '2',
+      parameters,
+      path: 'a--2',
+      id: 'a--2',
+      args: {},
+    },
     'b-c--1': {
       kind: 'b/c',
       name: '1',
@@ -703,6 +717,38 @@ describe('stories API', () => {
     });
   });
 
+  describe('findSiblingStoryId', () => {
+    it('works forward', () => {
+      const navigate = jest.fn();
+      const store = createMockStore();
+
+      const storyId = 'a--1';
+      const {
+        api: { setStories, findSiblingStoryId },
+        state,
+      } = initStories({ store, navigate, storyId, viewMode: 'story', provider });
+      store.setState(state);
+      setStories(storiesHash);
+
+      const result = findSiblingStoryId(storyId, storiesHash, 1, false);
+      expect(result).toBe('a--2');
+    });
+    it('works forward toSiblingGroup', () => {
+      const navigate = jest.fn();
+      const store = createMockStore();
+
+      const storyId = 'a--1';
+      const {
+        api: { setStories, findSiblingStoryId },
+        state,
+      } = initStories({ store, navigate, storyId, viewMode: 'story', provider });
+      store.setState(state);
+      setStories(storiesHash);
+
+      const result = findSiblingStoryId(storyId, store.getState().storiesHash, 1, true);
+      expect(result).toBe('b-c--1');
+    });
+  });
   describe('jumpToComponent', () => {
     it('works forward', () => {
       const navigate = jest.fn();

lib/core-events/src/index.ts

@@ -16,6 +16,8 @@ enum events {
   FORCE_RE_RENDER = 'forceReRender',
   // Force the current story to re-render from scratch, with its initial args
   FORCE_REMOUNT = 'forceRemount',
+  // Request the story has been loaded into the store, ahead of time, before it's actually
+  PRELOAD_STORIES = 'preloadStories',
   // The story has been loaded into the store, we have parameters/args/etc
   STORY_PREPARED = 'storyPrepared',
   // The next 6 events are emitted by the StoryRenderer when rendering the current story
@@ -71,6 +73,7 @@ export const {
   STORY_PREPARED,
   STORY_CHANGED,
   STORY_UNCHANGED,
+  PRELOAD_STORIES,
   STORY_RENDERED,
   STORY_MISSING,
   STORY_ERRORED,

lib/preview-web/src/PreviewWeb.test.ts

@@ -1022,6 +1022,18 @@ describe('PreviewWeb', () => {
     });
   });
 
+  describe('onPreloadStories', () => {
+    it('loads stories', async () => {
+      document.location.search = '?id=component-one--a&viewMode=docs';
+      const preview = await createAndRenderPreview();
+      await waitForRender();
+
+      importFn.mockClear();
+      await preview.onPreloadStories(['component-two--c']);
+      expect(importFn).toHaveBeenCalledWith('./src/ComponentTwo.stories.js');
+    });
+  });
+
   describe('onResetArgs', () => {
     it('emits STORY_ARGS_UPDATED', async () => {
       document.location.search = '?id=component-one--a';

lib/preview-web/src/PreviewWeb.tsx

@@ -4,7 +4,7 @@ import global from 'global';
 import Events, { IGNORED_EXCEPTION } from '@storybook/core-events';
 import { logger } from '@storybook/client-logger';
 import { AnyFramework, StoryId, ProjectAnnotations, Args, Globals } from '@storybook/csf';
-import {
+import type {
   ModuleImportFn,
   Selection,
   Story,
@@ -66,6 +66,7 @@ export class PreviewWeb<TFramework extends AnyFramework> extends Preview<TFramew
 
     this.channel.on(Events.SET_CURRENT_STORY, this.onSetCurrentStory.bind(this));
     this.channel.on(Events.UPDATE_QUERY_PARAMS, this.onUpdateQueryParams.bind(this));
+    this.channel.on(Events.PRELOAD_STORIES, this.onPreloadStories.bind(this));
   }
 
   initializeWithProjectAnnotations(projectAnnotations: WebProjectAnnotations<TFramework>) {
@@ -208,6 +209,10 @@ export class PreviewWeb<TFramework extends AnyFramework> extends Preview<TFramew
     if (this.currentRender instanceof DocsRender) await this.currentRender.rerender();
   }
 
+  async onPreloadStories(ids: string[]) {
+    await Promise.all(ids.map((id) => this.storyStore.loadStory({ storyId: id })));
+  }
+
   // RENDERING
 
   // We can either have: