From 6019388c5d0dfda693dea4c2b47ac0af7d566c64 Mon Sep 17 00:00:00 2001
From: Mateusz Kadlubowski <xeho91@pm.me>
Date: Mon, 22 Jul 2024 16:34:24 +0800
Subject: [PATCH] fix empty instance tag not being removed

---
 src/compiler/pre-transform/index.test.ts |  4 --
 src/compiler/pre-transform/index.ts      | 60 ++++++++++++++----------
 2 files changed, 35 insertions(+), 29 deletions(-)

diff --git a/src/compiler/pre-transform/index.test.ts b/src/compiler/pre-transform/index.test.ts
index 43c67ed..c5a9f3b 100644
--- a/src/compiler/pre-transform/index.test.ts
+++ b/src/compiler/pre-transform/index.test.ts
@@ -153,10 +153,6 @@ describe(codemodLegacyNodes.name, () => {
     expect(print(transformed)).toMatchInlineSnapshot(`
       "<script context="module">
       	import { defineMeta } from "@storybook/addon-svelte-csf";
-      </script>
-
-      <script>
-      	import { defineMeta } from "@storybook/addon-svelte-csf";
       </script>"
     `);
   });
diff --git a/src/compiler/pre-transform/index.ts b/src/compiler/pre-transform/index.ts
index 50425f0..1e9009a 100644
--- a/src/compiler/pre-transform/index.ts
+++ b/src/compiler/pre-transform/index.ts
@@ -1,5 +1,5 @@
 import pkg from '@storybook/addon-svelte-csf/package.json' with { type: 'json' };
-import type { ImportDeclaration, VariableDeclaration } from 'estree';
+import type { ImportDeclaration, Program, VariableDeclaration } from 'estree';
 import type { Comment, Fragment, Root, Script, SvelteNode } from 'svelte/compiler';
 
 import { transformComponentMetaToDefineMeta } from '#compiler/pre-transform/codemods/component-meta-to-define-meta';
@@ -161,54 +161,64 @@ export async function codemodLegacyNodes(params: Params): Promise<Root> {
     },
   });
 
-  const { componentIdentifierName, defineMetaFromMeta } = state;
-
   // Clean-up
   // TODO: To optimize it (stop walking on AST again)...
   // it might be possible to use `visit(path(-1))` for some cases in the previous AST walk,
   // but I haven't managed to get it to work properly
   transformedAst = walk(transformedAst, state, {
-    Root(_node, context) {
-      const { next, state } = context;
+    Root(node, context) {
+      let { fragment, instance, module, ...rest } = node;
+      const { state, visit } = context;
 
-      next(state);
+      // NOTE: At this point, we decide for walker where it should walk first instead of using `next(state)`
+      module = visit(module as Script, state) as Script;
+      instance = instance ? (visit(instance, state) as Script) : null;
+      fragment = visit(fragment, state) as Fragment;
+
+      // NOTE: Remove if body is empty
+      if (instance && instance.content.body.length === 0) {
+        instance = null;
+      }
+
+      return {
+        ...rest,
+        module,
+        instance,
+        fragment,
+      };
     },
 
     Script(node, context) {
-      const { context: scriptContext } = node;
+      let { content, context: scriptContext, ...rest } = node;
+      const { state, visit } = context;
 
-      const { next, state } = context;
+      state.currentScript = scriptContext === 'module' ? 'module' : 'instance';
+
+      content = visit(content, state) as Program;
 
-      next({ ...state, currentScript: scriptContext === 'module' ? 'module' : 'instance' });
+      return { ...rest, content, context: scriptContext };
     },
 
     Program(node, _context) {
       let { body, ...rest } = node;
       const { currentScript, pkgImportDeclaration, defineMetaFromMeta } = state;
 
-      if (pkgImportDeclaration) {
-        if (currentScript === 'module') {
-          body.unshift(pkgImportDeclaration);
-        }
-
-        if (currentScript === 'instance') {
-          body.filter((declaration) => {
-            return (
-              declaration.type === 'ImportDeclaration' && declaration.source.value === pkg.name
-            );
-          });
-        }
+      if (pkgImportDeclaration && currentScript === 'instance') {
+        body = body.filter((declaration) => {
+          return declaration.type === 'ImportDeclaration' && declaration.source.value !== pkg.name;
+        });
       }
 
-      if (defineMetaFromMeta) {
+      if (currentScript === 'module' && defineMetaFromMeta) {
         body.push(defineMetaFromMeta);
-
-        return { ...rest, body };
       }
+
+      return { ...rest, body };
     },
 
     Fragment(node, context) {
-      const { stop } = context;
+      const { state, stop } = context;
+      const { componentIdentifierName, defineMetaFromMeta } = state;
 
       if (defineMetaFromMeta) {
         let { nodes, ...rest } = node;