WIP

lib/preview-web/src/DocsRender.ts

@@ -1,35 +1,60 @@
 import global from 'global';
 import { AnyFramework, StoryId, ViewMode, StoryContextForLoaders } from '@storybook/csf';
-import { Story, StoryStore, CSFFile } from '@storybook/store';
+import { Story, StoryStore, CSFFile, ModuleExports, IndexEntry } from '@storybook/store';
 import { Channel } from '@storybook/addons';
 import { DOCS_RENDERED } from '@storybook/core-events';
 
-import { Render, StoryRender } from './StoryRender';
+import { Render, RenderType } from './StoryRender';
 import type { DocsContextProps } from './types';
 
 export class DocsRender<TFramework extends AnyFramework> implements Render<TFramework> {
+  public type: RenderType = 'docs';
+
+  public id: StoryId;
+
+  private legacy: boolean;
+
+  public story?: Story<TFramework>;
+
+  public exports?: ModuleExports;
+
+  private preparing = false;
+
   private canvasElement?: HTMLElement;
 
   private context?: DocsContextProps;
 
   public disableKeyListeners = false;
 
-  static fromStoryRender<TFramework extends AnyFramework>(storyRender: StoryRender<TFramework>) {
-    const { channel, store, id, story } = storyRender;
-    return new DocsRender<TFramework>(channel, store, id, story);
-  }
-
   // eslint-disable-next-line no-useless-constructor
   constructor(
     private channel: Channel,
     private store: StoryStore<TFramework>,
-    public id: StoryId,
-    public story: Story<TFramework>
-  ) {}
+    public entry: IndexEntry
+  ) {
+    this.id = entry.id;
+    this.legacy = entry.type === 'story' || entry.legacy;
+  }
+
+  // The two story "renders" are equal and have both loaded the same story
+  isEqual(other?: Render<TFramework>) {
+    return other && this.id === other.id && this.legacy
+      ? this.story && this.story === other.story
+      : other.type === 'docs' && this.entry === (other as DocsRender<TFramework>).entry;
+  }
+
+  async prepare() {
+    this.preparing = true;
+    if (this.legacy) {
+      this.story = await this.store.loadStory({ storyId: this.id });
+    } else {
+      this.exports = await this.store.loadDocsFileById(this.id);
+    }
+    this.preparing = false;
+  }
 
-  // DocsRender doesn't prepare, it is created *from* a prepared StoryRender
   isPreparing() {
-    return false;
+    return this.preparing;
   }
 
   async renderToElement(
@@ -38,7 +63,7 @@ export class DocsRender<TFramework extends AnyFramework> implements Render<TFram
   ) {
     this.canvasElement = canvasElement;
 
-    const { id, title, name } = this.story;
+    const { id, title, name } = this.entry;
     const csfFile: CSFFile<TFramework> = await this.store.loadCSFFileByStoryId(this.id);
 
     this.context = {
@@ -63,13 +88,20 @@ export class DocsRender<TFramework extends AnyFramework> implements Render<TFram
   }
 
   async render() {
-    if (!this.story || !this.context || !this.canvasElement)
+    if (!(this.story || this.exports) || !this.context || !this.canvasElement)
       throw new Error('DocsRender not ready to render');
 
     const renderer = await import('./renderDocs');
-    renderer.renderDocs(this.story, this.context, this.canvasElement, () =>
-      this.channel.emit(DOCS_RENDERED, this.id)
-    );
+
+    if (this.legacy) {
+      renderer.renderLegacyDocs(this.story, this.context, this.canvasElement, () =>
+        this.channel.emit(DOCS_RENDERED, this.id)
+      );
+    } else {
+      renderer.renderDocs(this.exports, this.context, this.canvasElement, () =>
+        this.channel.emit(DOCS_RENDERED, this.id)
+      );
+    }
   }
 
   async rerender() {

lib/preview-web/src/PreviewWeb.tsx

@@ -39,7 +39,7 @@ export class PreviewWeb<TFramework extends AnyFramework> extends Preview<TFramew
 
   currentSelection: Selection;
 
-  currentRender: Render<TFramework>;
+  currentRender: StoryRender<TFramework> | DocsRender<TFramework>;
 
   constructor() {
     super();
@@ -227,12 +227,17 @@ export class PreviewWeb<TFramework extends AnyFramework> extends Preview<TFramew
     }
 
     const { storyId } = selection;
+    const entry = this.storyStore.storyIndex.entries[storyId];
+
+    // Docs entries cannot be rendered in 'story' viewMode.
+    // For now story entries can be rendered in docs mode.
+    const viewMode = entry.type === 'docs' ? 'docs' : selection.viewMode;
 
     const storyIdChanged = this.currentSelection?.storyId !== storyId;
-    const viewModeChanged = this.currentSelection?.viewMode !== selection.viewMode;
+    const viewModeChanged = this.currentSelection?.viewMode !== viewMode;
 
     // Show a spinner while we load the next story
-    if (selection.viewMode === 'story') {
+    if (viewMode === 'story') {
       this.view.showPreparingStory({ immediate: viewModeChanged });
     } else {
       this.view.showPreparingDocs();
@@ -252,26 +257,32 @@ export class PreviewWeb<TFramework extends AnyFramework> extends Preview<TFramew
       lastRender = null;
     }
 
-    const storyRender = new StoryRender<TFramework>(
-      this.channel,
-      this.storyStore,
-      (...args) => {
-        // At the start of renderToDOM we make the story visible (see note in WebView)
-        this.view.showStoryDuringRender();
-        return this.renderToDOM(...args);
-      },
-      this.mainStoryCallbacks(storyId),
-      storyId,
-      'story'
-    );
+    let render;
+    if (viewMode === 'story') {
+      render = new StoryRender<TFramework>(
+        this.channel,
+        this.storyStore,
+        (...args) => {
+          // At the start of renderToDOM we make the story visible (see note in WebView)
+          this.view.showStoryDuringRender();
+          return this.renderToDOM(...args);
+        },
+        this.mainStoryCallbacks(storyId),
+        storyId,
+        'story'
+      );
+    } else {
+      render = new DocsRender<TFramework>(this.channel, this.storyStore, entry);
+    }
+    console.log(render);
+
     // We need to store this right away, so if the story changes during
     // the async `.prepare()` below, we can (potentially) cancel it
     this.currentSelection = selection;
-    // Note this may be replaced by a docsRender after preparing
-    this.currentRender = storyRender;
+    this.currentRender = render;
 
     try {
-      await storyRender.prepare();
+      await render.prepare();
     } catch (err) {
       if (err !== PREPARE_ABORTED) {
         // We are about to render an error so make sure the previous story is
@@ -281,11 +292,10 @@ export class PreviewWeb<TFramework extends AnyFramework> extends Preview<TFramew
       }
       return;
     }
-    const implementationChanged = !storyIdChanged && !storyRender.isEqual(lastRender);
+    const implementationChanged = !storyIdChanged && !render.isEqual(lastRender);
 
-    if (persistedArgs) this.storyStore.args.updateFromPersisted(storyRender.story, persistedArgs);
-
-    const { parameters, initialArgs, argTypes, args } = storyRender.context();
+    if (persistedArgs && entry.type === 'story')
+      this.storyStore.args.updateFromPersisted(render.story, persistedArgs);
 
     // Don't re-render the story if nothing has changed to justify it
     if (lastRender && !storyIdChanged && !implementationChanged && !viewModeChanged) {
@@ -304,32 +314,38 @@ export class PreviewWeb<TFramework extends AnyFramework> extends Preview<TFramew
       this.channel.emit(Events.STORY_CHANGED, storyId);
     }
 
-    if (global.FEATURES?.storyStoreV7) {
-      this.channel.emit(Events.STORY_PREPARED, {
-        id: storyId,
-        parameters,
-        initialArgs,
-        argTypes,
-        args,
-      });
-    }
+    if (render.type === 'story') {
+      const { parameters, initialArgs, argTypes, args } = (
+        render as StoryRender<TFramework>
+      ).context();
+      if (global.FEATURES?.storyStoreV7) {
+        this.channel.emit(Events.STORY_PREPARED, {
+          id: storyId,
+          parameters,
+          initialArgs,
+          argTypes,
+          args,
+        });
+      }
 
-    // For v6 mode / compatibility
-    // If the implementation changed, or args were persisted, the args may have changed,
-    // and the STORY_PREPARED event above may not be respected.
-    if (implementationChanged || persistedArgs) {
-      this.channel.emit(Events.STORY_ARGS_UPDATED, { storyId, args });
+      // For v6 mode / compatibility
+      // If the implementation changed, or args were persisted, the args may have changed,
+      // and the STORY_PREPARED event above may not be respected.
+      if (implementationChanged || persistedArgs) {
+        this.channel.emit(Events.STORY_ARGS_UPDATED, { storyId, args });
+      }
     }
 
-    if (selection.viewMode === 'docs' || parameters.docsOnly) {
-      this.currentRender = DocsRender.fromStoryRender<TFramework>(storyRender);
+    if (viewMode === 'docs') {
       this.currentRender.renderToElement(
         this.view.prepareForDocs(),
         this.renderStoryToElement.bind(this)
       );
     } else {
-      this.storyRenders.push(storyRender);
-      this.currentRender.renderToElement(this.view.prepareForStory(storyRender.story));
+      this.storyRenders.push(render as StoryRender<TFramework>);
+      (this.currentRender as StoryRender<TFramework>).renderToElement(
+        this.view.prepareForStory(render.story)
+      );
     }
   }
 

lib/preview-web/src/StoryRender.ts

@@ -40,7 +40,9 @@ export type RenderContextCallbacks<TFramework extends AnyFramework> = Pick<
 
 export const PREPARE_ABORTED = new Error('prepareAborted');
 
+export type RenderType = 'story' | 'docs';
 export interface Render<TFramework extends AnyFramework> {
+  type: RenderType;
   id: StoryId;
   story?: Story<TFramework>;
   isPreparing: () => boolean;
@@ -50,6 +52,8 @@ export interface Render<TFramework extends AnyFramework> {
 }
 
 export class StoryRender<TFramework extends AnyFramework> implements Render<TFramework> {
+  public type: RenderType = 'story';
+
   public story?: Story<TFramework>;
 
   public phase?: RenderPhase;

lib/preview-web/src/renderDocs.tsx

@@ -1,21 +1,30 @@
-import React, { ComponentType } from 'react';
+import React, { ComponentType, ReactElement } from 'react';
 import ReactDOM from 'react-dom';
 import { AnyFramework } from '@storybook/csf';
-import { Story } from '@storybook/store';
+import { ModuleExports, Story } from '@storybook/store';
 
 import { DocsContextProps } from './types';
 import { NoDocs } from './NoDocs';
 
-export function renderDocs<TFramework extends AnyFramework>(
+export function renderLegacyDocs<TFramework extends AnyFramework>(
   story: Story<TFramework>,
   docsContext: DocsContextProps<TFramework>,
   element: HTMLElement,
   callback: () => void
 ) {
-  return renderDocsAsync(story, docsContext, element).then(callback);
+  return renderLegacyDocsAsync(story, docsContext, element).then(callback);
 }
 
-async function renderDocsAsync<TFramework extends AnyFramework>(
+export function renderDocs<TFramework extends AnyFramework>(
+  exports: ModuleExports,
+  docsContext: DocsContextProps<TFramework>,
+  element: HTMLElement,
+  callback: () => void
+) {
+  return renderDocsAsync(exports, docsContext, element).then(callback);
+}
+
+async function renderLegacyDocsAsync<TFramework extends AnyFramework>(
   story: Story<TFramework>,
   docsContext: DocsContextProps<TFramework>,
   element: HTMLElement
@@ -45,6 +54,28 @@ async function renderDocsAsync<TFramework extends AnyFramework>(
   });
 }
 
+async function renderDocsAsync<TFramework extends AnyFramework>(
+  exports: ModuleExports,
+  docsContext: DocsContextProps<TFramework>,
+  element: HTMLElement
+) {
+  // FIXME -- is this at all correct?
+  const DocsContainer = ({ children }: { children: ReactElement }) => <>{children}</>;
+
+  const Page = exports.default;
+
+  // FIXME -- do we need to set a key as above?
+  const docsElement = (
+    <DocsContainer>
+      <Page />
+    </DocsContainer>
+  );
+
+  await new Promise<void>((resolve) => {
+    ReactDOM.render(docsElement, element, resolve);
+  });
+}
+
 export function unmountDocs(element: HTMLElement) {
   ReactDOM.unmountComponentAtNode(element);
 }

lib/store/src/StoryIndexStore.ts

@@ -1,8 +1,16 @@
 import dedent from 'ts-dedent';
 import { Channel } from '@storybook/addons';
 import type { StoryId } from '@storybook/csf';
+import memoize from 'memoizerific';
 
-import type { StorySpecifier, StoryIndex, IndexEntry } from './types';
+import type { StorySpecifier, StoryIndex, IndexEntry, Path } from './types';
+
+const getImportPathMap = memoize(1)((entries: StoryIndex['entries']) =>
+  Object.values(entries).reduce((acc, entry) => {
+    acc[entry.importPath] = acc[entry.importPath] || entry;
+    return acc;
+  }, {} as Record<Path, IndexEntry>)
+);
 
 export class StoryIndexStore {
   channel: Channel;
@@ -50,4 +58,8 @@ export class StoryIndexStore {
 
     return storyEntry;
   }
+
+  importPathToEntry(importPath: Path): IndexEntry {
+    return getImportPathMap(this.entries)[importPath];
+  }
 }

lib/store/src/StoryStore.ts

@@ -31,6 +31,7 @@ import type {
   StoryIndexEntry,
   V2CompatIndexEntry,
   StoryIndexV3,
+  ModuleExports,
 } from './types';
 import { HooksContext } from './hooks';
 
@@ -120,6 +121,24 @@ export class StoryStore<TFramework extends AnyFramework> {
     if (this.cachedCSFFiles) await this.cacheAllCSFFiles();
   }
 
+  // FIXME: does this need to be load
+  async loadDocsFileById(docsId: StoryId): Promise<ModuleExports> {
+    const entry = this.storyIndex.storyIdToEntry(docsId);
+    if (entry.type !== 'docs') throw new Error(`Cannot load docs file for id ${docsId}`);
+
+    const { importPath, storiesImports } = entry;
+
+    const [docsImport] = await Promise.all([
+      this.importFn(importPath),
+      ...storiesImports.map((storyImportPath) => {
+        const firstStoryEntry = this.storyIndex.importPathToEntry(storyImportPath);
+        return this.loadCSFFileByStoryId(firstStoryEntry.id);
+      }),
+    ]);
+
+    return docsImport;
+  }
+
   // To load a single CSF file to service a story we need to look up the importPath in the index
   loadCSFFileByStoryId(storyId: StoryId): PromiseLike<CSFFile<TFramework>> {
     const { importPath, title } = this.storyIndex.storyIdToEntry(storyId);