From 460fc9f061f0957923b6df789644de712cbade23 Mon Sep 17 00:00:00 2001
From: xeho91 <xeho91@pm.me>
Date: Sat, 6 Jul 2024 13:34:16 +0800
Subject: [PATCH] Provide usage test & fix runtime issues

---
 src/compiler/plugins.ts                       |  4 +-
 .../codemods/template-to-snippet.ts           | 45 +++++++++++++------
 src/compiler/pre-transform/extractor.ts       |  5 ++-
 src/compiler/pre-transform/index.ts           |  7 ++-
 src/preset.ts                                 |  3 +-
 src/runtime/Template.svelte                   | 34 +++++---------
 src/runtime/contexts/renderer.svelte.ts       |  6 +--
 tests/stories/LegacyTemplate.stories.svelte   | 44 ++++++++++++++++++
 tests/stories/LegacyTemplate.svelte           | 27 +++++++++++
 9 files changed, 126 insertions(+), 49 deletions(-)
 create mode 100644 tests/stories/LegacyTemplate.stories.svelte
 create mode 100644 tests/stories/LegacyTemplate.svelte

diff --git a/src/compiler/plugins.ts b/src/compiler/plugins.ts
index 0377334a..02da8a73 100644
--- a/src/compiler/plugins.ts
+++ b/src/compiler/plugins.ts
@@ -53,9 +53,9 @@ export async function preTransformPlugin(): Promise<Plugin> {
       let magicLegacyCode = new MagicString(legacyCode);
 
       const svelteAST = getSvelteAST({ code: legacyCode, filename: id });
-      const legacyNodes = extractLegacyNodes(svelteAST);
+      const legacyNodes = await extractLegacyNodes(svelteAST);
 
-      codemodLegacyNodes({
+      await codemodLegacyNodes({
         code: magicLegacyCode,
         legacyNodes,
       });
diff --git a/src/compiler/pre-transform/codemods/template-to-snippet.ts b/src/compiler/pre-transform/codemods/template-to-snippet.ts
index e13c0f98..1a3675a3 100644
--- a/src/compiler/pre-transform/codemods/template-to-snippet.ts
+++ b/src/compiler/pre-transform/codemods/template-to-snippet.ts
@@ -1,5 +1,5 @@
 import { getStringValueFromAttribute } from '#parser/analyse/story/attributes';
-import type { Attribute, Component, SnippetBlock } from 'svelte/compiler';
+import type { Attribute, Component, LetDirective, SnippetBlock } from 'svelte/compiler';
 
 /**
  *
@@ -40,24 +40,43 @@ export function transformTemplateToSnippet(component: Component): SnippetBlock {
     component,
   });
 
+  const letDirectiveArgs = attributes.find((attr) => {
+    if (attr.type === 'LetDirective') {
+      return attr.name === 'args';
+    }
+    // Will TypeScript 5.5 handle type inference for this one better? 🤔
+  }) as LetDirective | undefined;
+
+  const letDirectiveContext = attributes.find((attr) => {
+    if (attr.type === 'LetDirective') {
+      return attr.name === 'context';
+    }
+    // Will TypeScript 5.5 handle type inference for this one better? 🤔
+  }) as LetDirective | undefined;
+
+  let parameters: SnippetBlock['parameters'] = [];
+
+  if (letDirectiveArgs) {
+    parameters.push({
+      type: 'Identifier',
+      name: 'args',
+    });
+  }
+
+  if (letDirectiveContext) {
+    parameters.push({
+      type: 'Identifier',
+      name: 'context',
+    });
+  }
+
   return {
     type: 'SnippetBlock',
     expression: {
       type: 'Identifier',
       name: id ?? 'children',
     },
-    // WARN: I suspect at this point, it doesn't matter if user used one of directives `let:args` and `let:context`.
-    // W provide both parameters, just in case.
-    parameters: [
-      {
-        type: 'Identifier',
-        name: 'args',
-      },
-      {
-        type: 'Identifier',
-        name: 'storyContext',
-      },
-    ],
+    parameters,
     body: fragment,
     // NOTE: Those are useless, but I want TypeScript to 🤫
     start,
diff --git a/src/compiler/pre-transform/extractor.ts b/src/compiler/pre-transform/extractor.ts
index 1978ead1..5386d361 100644
--- a/src/compiler/pre-transform/extractor.ts
+++ b/src/compiler/pre-transform/extractor.ts
@@ -1,5 +1,4 @@
 import type { Component, SvelteNode } from 'svelte/compiler';
-import { walk } from 'zimmerframe';
 
 interface Results {
   componentsTemplate: Component[];
@@ -8,7 +7,9 @@ interface Results {
 /**
  * Extract legacy AST nodes, for usage with codemods into modern syntax.
  */
-export function extractLegacyNodes(parsed: SvelteNode): Results {
+export async function extractLegacyNodes(parsed: SvelteNode): Promise<Results> {
+  const { walk } = await import('zimmerframe');
+
   const state: {
     componentsTemplate: Component[];
   } = {
diff --git a/src/compiler/pre-transform/index.ts b/src/compiler/pre-transform/index.ts
index f1926df9..b254eee9 100644
--- a/src/compiler/pre-transform/index.ts
+++ b/src/compiler/pre-transform/index.ts
@@ -1,15 +1,14 @@
-import { print } from 'svelte-ast-print';
-
 import type { extractLegacyNodes } from '#compiler/pre-transform/extractor';
 import { transformTemplateToSnippet } from './codemods/template-to-snippet';
 import type MagicString from 'magic-string';
 
 interface Params {
   code: MagicString;
-  legacyNodes: ReturnType<typeof extractLegacyNodes>;
+  legacyNodes: Awaited<ReturnType<typeof extractLegacyNodes>>;
 }
 
-export function codemodLegacyNodes(params: Params): void {
+export async function codemodLegacyNodes(params: Params): Promise<void> {
+  const { print } = await import('svelte-ast-print');
   const { legacyNodes, code } = params;
   const { componentsTemplate } = legacyNodes;
 
diff --git a/src/preset.ts b/src/preset.ts
index 1d03a58d..a156cbc0 100644
--- a/src/preset.ts
+++ b/src/preset.ts
@@ -7,8 +7,9 @@ export const viteFinal: StorybookConfig['viteFinal'] = async (config, options) =
   return {
     ...config,
     plugins: [
+      /** TODO: Is this the place for `options.supportLegacy`? */
+      preTransformPlugin(),
       ...(config.plugins ?? []),
-      /** TODO: Is this the place for `options.supportLegacy`? */ preTransformPlugin(),
       postTransformPlugin(),
     ],
   };
diff --git a/src/runtime/Template.svelte b/src/runtime/Template.svelte
index 4e48249f..a8b96459 100644
--- a/src/runtime/Template.svelte
+++ b/src/runtime/Template.svelte
@@ -1,29 +1,15 @@
 <script lang="ts">
-  import type { Snippet } from 'svelte';
+  /**
+   * This is a legacy template, this functionality is just a "mock".
+   * To allow user still have typing experience
+   * Vite pre-transform hook does codemod where this component gets transformed into Svelte v5 SnippetBlock.
+   */
 
-  import { useStoryRenderer, type StoryRendererContext } from './contexts/renderer.svelte';
+  import { type StoryRendererContext } from './contexts/renderer.svelte';
 
-  interface Props {
-    /**
-     * The content of template to be rendered inside the `<Story />`
-     */
-    children: Snippet<
-      [
-        //
-        StoryRendererContext['args'],
-        StoryRendererContext['storyContext'],
-      ]
-    >;
-    /**
-     * Unique id of template - to be used in `<Story templateId={id} />`
-     * @default id
-     */
-    id?: string;
-  }
-
-  const renderer = useStoryRenderer();
-
-  let { children, id = 'default' }: Props = $props();
+  let id: string = 'default';
+  let args: StoryRendererContext['storyContext'];
+  let context: StoryRendererContext['args'];
 </script>
 
-{@render children(renderer.args, renderer.storyContext)}
+<slot {context} {args} />
diff --git a/src/runtime/contexts/renderer.svelte.ts b/src/runtime/contexts/renderer.svelte.ts
index 6648ae97..067dae07 100644
--- a/src/runtime/contexts/renderer.svelte.ts
+++ b/src/runtime/contexts/renderer.svelte.ts
@@ -39,9 +39,9 @@ function buildContext<TOverrideArgs extends Args, TCmp extends Cmp, TMeta extend
 }
 
 export type StoryRendererContext<
-  TOverrideArgs extends Args,
-  TCmp extends Cmp,
-  TMeta extends Meta<TCmp>,
+  TOverrideArgs extends Args = Args,
+  TCmp extends Cmp = Cmp,
+  TMeta extends Meta<TCmp> = Meta<TCmp>,
 > = ReturnType<typeof buildContext<TOverrideArgs, TCmp, TMeta>>;
 
 function createStoryRendererContext<
diff --git a/tests/stories/LegacyTemplate.stories.svelte b/tests/stories/LegacyTemplate.stories.svelte
new file mode 100644
index 00000000..a5615ccf
--- /dev/null
+++ b/tests/stories/LegacyTemplate.stories.svelte
@@ -0,0 +1,44 @@
+<script context="module" lang="ts">
+  import { Template, defineMeta } from '@storybook/addon-svelte-csf';
+
+  import LegacyTemplate from './LegacyTemplate.svelte';
+
+  /**
+   * Description set explicitly in the comment above `defineMeta`.
+   *
+   * Multiline supported. And also Markdown syntax:
+   *
+   * * **Bold**,
+   * * _Italic_,
+   * * `Code`.
+   */
+  const { Story } = defineMeta({
+    title: 'LegacyTemplate',
+    component: LegacyTemplate,
+    tags: ['autodocs'],
+  });
+</script>
+
+<script lang="ts">
+  let count = $state(0);
+
+  function handleClick() {
+    count += 1;
+  }
+</script>
+
+<Template let:args let:context>
+  <p>Using default template</p>
+  <LegacyTemplate {...args} onclick={handleClick} />
+</Template>
+
+<Template id="rounded" let:args>
+  <p>Using rounded template</p>
+  <LegacyTemplate {...args} onclick={handleClick} rounded />
+</Template>
+
+<!-- TODO: Need to restore legacy Story first -->
+<Story name="Default" {children} />
+
+<!-- TODO: Need to restore legacy Story first -->
+<Story name="Rounded" children={rounded} />
diff --git a/tests/stories/LegacyTemplate.svelte b/tests/stories/LegacyTemplate.svelte
new file mode 100644
index 00000000..20bd6185
--- /dev/null
+++ b/tests/stories/LegacyTemplate.svelte
@@ -0,0 +1,27 @@
+<script lang="ts">
+  import type { HTMLAttributes } from 'svelte/elements';
+
+  interface Props extends HTMLAttributes<HTMLButtonElement> {
+    rounded?: boolean;
+  }
+
+  let { rounded = false, ...restProps }: Props = $props();
+</script>
+
+<button class="button" class:rounded {...restProps}>
+  <strong>{rounded ? 'Round' : 'Square'} corners</strong>
+  <hr />
+</button>
+
+<style>
+  .rounded {
+    border-radius: 35px;
+  }
+
+  .button {
+    border: 3px solid;
+    padding: 10px 20px;
+    background-color: white;
+    outline: none;
+  }
+</style>