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 {
+  STORY_PRELOAD,
   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,36 @@ 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;
+        }
+
+        // eslint-disable-next-line consistent-return
+        return lookupList[index + direction][0];
+      }
+      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, {
@@ -439,12 +446,25 @@ export const init: ModuleFn = ({
       const { sourceType } = getEventMetadata(this, fullAPI);
 
       if (sourceType === 'local') {
+        const { storyId, storiesHash } = store.getState();
         const options = fullAPI.getCurrentParameter('options');
 
         if (options) {
           checkDeprecatedOptionParameters(options);
           fullAPI.setOptions(options);
         }
+
+        // create a list of related stories to be preloaded
+        const toBePreloaded = Array.from(
+          new Set([
+            api.findSiblingStoryId(storyId, storiesHash, 1, false),
+            api.findSiblingStoryId(storyId, storiesHash, -1, false),
+            api.findSiblingStoryId(storyId, storiesHash, 1, true),
+            api.findSiblingStoryId(storyId, storiesHash, -1, true),
+          ])
+        );
+
+        fullAPI.emit(STORY_PRELOAD, toBePreloaded);
       }
     });
 

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

@@ -12,6 +12,7 @@ import { EventEmitter } from 'events';
 import global from 'global';
 import { mockChannel } from '@storybook/addons';
 
+import { stopCoverage } from 'v8';
 import { getEventMetadata } from '../lib/events';
 
 import { init as initStories } from '../modules/stories';
@@ -85,8 +86,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 +718,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
+  STORY_PRELOAD = 'storyPreload',
   // 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,
+  STORY_PRELOAD,
   STORY_RENDERED,
   STORY_MISSING,
   STORY_ERRORED,

lib/preview-web/src/PreviewWeb.tsx

@@ -6,15 +6,15 @@ import Events, { IGNORED_EXCEPTION } from '@storybook/core-events';
 import { logger } from '@storybook/client-logger';
 import { addons, Channel } from '@storybook/addons';
 import { AnyFramework, StoryId, ProjectAnnotations, Args, Globals } from '@storybook/csf';
-import {
+import type {
   ModuleImportFn,
   Selection,
   Story,
-  StoryStore,
   StorySpecifier,
   StoryIndex,
   WebProjectAnnotations,
 } from '@storybook/store';
+import { StoryStore } from '@storybook/store';
 
 import { UrlStore } from './UrlStore';
 import { WebView } from './WebView';
@@ -127,6 +127,7 @@ export class PreviewWeb<TFramework extends AnyFramework> {
     this.channel.on(Events.RESET_STORY_ARGS, this.onResetArgs.bind(this));
     this.channel.on(Events.FORCE_RE_RENDER, this.onForceReRender.bind(this));
     this.channel.on(Events.FORCE_REMOUNT, this.onForceRemount.bind(this));
+    this.channel.on(Events.STORY_PRELOAD, this.onPreloadStory.bind(this));
   }
 
   getProjectAnnotationsOrRenderError(
@@ -398,6 +399,10 @@ export class PreviewWeb<TFramework extends AnyFramework> {
     await this.onUpdateArgs({ storyId, updatedArgs });
   }
 
+  async onPreloadStory(ids: string[]) {
+    await Promise.all(ids.map((id) => this.storyStore.loadStory({ storyId: id })));
+  }
+
   // ForceReRender does not include a story id, so we simply must
   // re-render all stories in case they are relevant
   async onForceReRender() {