From 1550e08fbfa6b731dc84b85bba8e0285eaa2401f Mon Sep 17 00:00:00 2001
From: Jeppe Reinhold <jeppe@reinhold.is>
Date: Mon, 3 Jun 2024 22:56:06 +0200
Subject: [PATCH] use function declarations to get storiesFunctionDeclarations
 instead of componentName

---
 src/compiler/plugin.ts                        | 11 -----
 ...eate-variable-from-runtime-stories-call.ts |  8 ++--
 src/compiler/transform/create-appendix.ts     |  7 ++-
 src/parser/extract/compiled/nodes.ts          | 48 +++++++++----------
 4 files changed, 29 insertions(+), 45 deletions(-)

diff --git a/src/compiler/plugin.ts b/src/compiler/plugin.ts
index 01c7a6b6..286e3203 100644
--- a/src/compiler/plugin.ts
+++ b/src/compiler/plugin.ts
@@ -19,12 +19,10 @@ import { removeExportDefault } from './transform/remove-export-default.js';
 import { insertDefineMetaJSDocCommentAsDescription } from './transform/define-meta/description.js';
 import { destructureMetaFromDefineMeta } from './transform/define-meta/destructure-meta.js';
 import { insertStoryHTMLCommentAsDescription } from './transform/Story/description.js';
-import { moveSourceAttributeToParameters } from './transform/Story/source.js';
 import { getSvelteAST } from '../parser/ast.js';
 import { extractStoriesNodesFromExportDefaultFn } from '../parser/extract/compiled/stories.js';
 import { extractCompiledASTNodes } from '../parser/extract/compiled/nodes.js';
 import { extractSvelteASTNodes } from '../parser/extract/svelte/nodes.js';
-import { getComponentName } from '../utils/get-component-name.js';
 
 export async function plugin(): Promise<Plugin> {
   const [{ createFilter }, { loadSvelteConfig }] = await Promise.all([
@@ -45,14 +43,6 @@ export async function plugin(): Promise<Plugin> {
       const compiledAST = this.parse(compiledCode);
       let magicCompiledCode = new MagicString(compiledCode);
 
-      const componentName = getComponentName(id);
-
-      if (!componentName) {
-        // TODO: make error message more user friendly
-        // what happened, how to fix
-        throw new Error(`Failed to extract component name from filename: ${id}`);
-      }
-
       let rawCode = fs.readFileSync(id).toString();
 
       if (svelteConfig?.preprocess) {
@@ -116,7 +106,6 @@ export async function plugin(): Promise<Plugin> {
         filename: id,
       });
       await createAppendix({
-        componentName,
         code: magicCompiledCode,
         nodes: {
           compiled: compiledNodes,
diff --git a/src/compiler/transform/appendix/create-variable-from-runtime-stories-call.ts b/src/compiler/transform/appendix/create-variable-from-runtime-stories-call.ts
index a6efa044..4541ac1d 100644
--- a/src/compiler/transform/appendix/create-variable-from-runtime-stories-call.ts
+++ b/src/compiler/transform/appendix/create-variable-from-runtime-stories-call.ts
@@ -1,15 +1,15 @@
-import type { VariableDeclaration } from 'estree';
+import type { FunctionDeclaration, VariableDeclaration } from 'estree';
 
 import type { getMetaIdentifier } from '../../../parser/analyse/meta/identifier.js';
 
 interface Params {
-  componentName: string;
+  storiesFunctionDeclaration: FunctionDeclaration;
   metaIdentifier: ReturnType<typeof getMetaIdentifier>;
   filename: string;
 }
 
 export function createVariableFromRuntimeStoriesCall(params: Params): VariableDeclaration {
-  const { componentName, metaIdentifier } = params;
+  const { storiesFunctionDeclaration, metaIdentifier } = params;
 
   return {
     type: 'VariableDeclaration',
@@ -32,7 +32,7 @@ export function createVariableFromRuntimeStoriesCall(params: Params): VariableDe
           arguments: [
             {
               type: 'Identifier',
-              name: componentName,
+              name: storiesFunctionDeclaration.id.name,
             },
             metaIdentifier,
           ],
diff --git a/src/compiler/transform/create-appendix.ts b/src/compiler/transform/create-appendix.ts
index d0da81e9..3260e78c 100644
--- a/src/compiler/transform/create-appendix.ts
+++ b/src/compiler/transform/create-appendix.ts
@@ -13,7 +13,6 @@ import type { SvelteASTNodes } from '../../parser/extract/svelte/nodes.js';
 import { getStoriesIdentifiers } from '../../parser/analyse/Story/attributes/identifiers.js';
 
 interface Params {
-  componentName: string;
   code: MagicString;
   nodes: {
     compiled: CompiledASTNodes;
@@ -23,9 +22,9 @@ interface Params {
 }
 
 export async function createAppendix(params: Params) {
-  const { componentName, code, nodes, filename } = params;
+  const { code, nodes, filename } = params;
   const { compiled, svelte } = nodes;
-  const { defineMetaVariableDeclaration } = compiled;
+  const { defineMetaVariableDeclaration, storiesFunctionDeclaration } = compiled;
 
   const storyIdentifiers = await getStoriesIdentifiers({ nodes: svelte, filename });
   const metaIdentifier = getMetaIdentifier({
@@ -33,7 +32,7 @@ export async function createAppendix(params: Params) {
     filename,
   });
   const variableFromRuntimeStoriesCall = createVariableFromRuntimeStoriesCall({
-    componentName,
+    storiesFunctionDeclaration,
     metaIdentifier,
     filename,
   });
diff --git a/src/parser/extract/compiled/nodes.ts b/src/parser/extract/compiled/nodes.ts
index dfc917a4..7b104eaf 100644
--- a/src/parser/extract/compiled/nodes.ts
+++ b/src/parser/extract/compiled/nodes.ts
@@ -56,7 +56,9 @@ export async function extractCompiledASTNodes(params: Params): Promise<CompiledA
   const { walk } = await import('zimmerframe');
 
   const { ast, filename } = params;
-  const state: Partial<CompiledASTNodes> = {};
+  const state: Partial<CompiledASTNodes> & {
+    potentialStoriesFunctionDeclaration: FunctionDeclaration[];
+  } = { potentialStoriesFunctionDeclaration: [] };
   const visitors: Visitors<Node | Comment, typeof state> = {
     ImportDeclaration(node, { state, visit }) {
       const { source, specifiers } = node;
@@ -106,27 +108,28 @@ export async function extractCompiledASTNodes(params: Params): Promise<CompiledA
       }
     },
 
+    FunctionDeclaration(node, { state }) {
+      state.potentialStoriesFunctionDeclaration.push(node);
+    },
+    
     ExportDefaultDeclaration(node, { state }) {
       state.exportDefault = node;
-
-      // WARN: This may be confusing.
-      // In the `NODE_ENV="production"` the export default is different.
-      // Identifier to a FunctionDeclaration.
-      if (
-        process.env.NODE_ENV === 'production' &&
-        node.declaration.type === 'FunctionDeclaration' &&
-        isStoriesComponentFn(node.declaration as FunctionDeclaration)
-      ) {
+      if (node.declaration.type === 'FunctionDeclaration') {
+        /*
+        In production, Svelte will compile the component to:
+        export default COMPONENT_NAME () {...}
+        */
         state.storiesFunctionDeclaration = node.declaration as FunctionDeclaration;
-      }
-    },
-
-    FunctionDeclaration(node, { state }) {
-      // WARN: This may be confusing.
-      // In the `NODE_ENV="development"` the export default is different.
-      // A `FunctionDeclaration`
-      if (isStoriesComponentFn(node)) {
-        state.storiesFunctionDeclaration = node;
+      } else if (node.declaration.type === 'Identifier') {
+        /*
+        In development, Svelte will compile the component to:
+        function COMPONENT_NAME () {...}
+        export default COMPONENT_NAME;
+        */
+        const { name } = node.declaration as Identifier;
+        state.storiesFunctionDeclaration = state.potentialStoriesFunctionDeclaration?.find(
+          (potential) => potential.id.name === name
+        );
       }
     },
   };
@@ -179,10 +182,3 @@ export async function extractCompiledASTNodes(params: Params): Promise<CompiledA
     storiesFunctionDeclaration,
   };
 }
-
-/**
- *:The main component function of those stories file _(`*.stories.svelte`)_ will always end up with `.stories`.
- * @see {@link "file://./../../../utils/get-component-name.ts"}
- */
-const isStoriesComponentFn = (fnDeclaration: FunctionDeclaration) =>
-  fnDeclaration.id?.name.endsWith('.stories');