Fix treeshaking

Author: lxsmnsyc

packages/start/src/directives/bubble-function-declaration.ts

@@ -18,7 +18,7 @@ export function bubbleFunctionDeclaration(path: babel.NodePath<t.FunctionDeclara
       ]),
     );
     path.scope.registerDeclaration(tmp);
-    tmp.skip();
+    // tmp.skip();
     if (path.parentPath.isExportNamedDeclaration()) {
       path.parentPath.replaceWith(
         t.exportNamedDeclaration(undefined, [t.exportSpecifier(decl.id, decl.id)]),

packages/start/src/directives/plugin.ts

@@ -8,6 +8,7 @@ import { getImportIdentifier } from "./get-import-identifier.ts";
 import { getRootStatementPath } from "./get-root-statement-path.ts";
 import { isStatementTopLevel } from "./is-statement-top-level.ts";
 import { isPathValid, unwrapPath } from "./paths.ts";
+import { removeUnusedVariables } from "./remove-unused-variables.ts";
 import type { ImportDefinition } from "./types.ts";
 
 export interface StateContext {
@@ -124,6 +125,8 @@ function transformFunction(
       ]),
     );
   }
+
+  path.scope.crawl();
 }
 
 function traceBinding(path: babel.NodePath, name: string): Binding | undefined {
@@ -166,23 +169,6 @@ function transformBindingForServer(ctx: StateContext, binding: Binding) {
   }
 }
 
-function treeshake(path: babel.NodePath, name: string): void {
-  const binding = path.scope.getBinding(name);
-
-  if (!(binding && binding.references + binding.constantViolations.length > 0)) {
-    if (isPathValid(path.parentPath, t.isImportDeclaration)) {
-      const parent = path.parentPath;
-      if (parent.node.specifiers.length === 1) {
-        parent.remove();
-      } else {
-        path.remove();
-      }
-    } else {
-      path.remove();
-    }
-  }
-}
-
 interface State extends babel.PluginPass {
   opts: StateContext;
 }
@@ -377,20 +363,11 @@ export function directivesPlugin(): babel.PluginObj<State> {
               transformFunction(ctx.opts, path, false);
             },
           });
-
           program.scope.crawl();
-          // Tree-shaking
-          program.traverse({
-            ImportDefaultSpecifier(path) {
-              treeshake(path, path.node.local.name);
-            },
-            ImportNamespaceSpecifier(path) {
-              treeshake(path, path.node.local.name);
-            },
-            ImportSpecifier(path) {
-              treeshake(path, path.node.local.name);
-            },
-          });
+
+          if (ctx.opts.count > 0) {
+            removeUnusedVariables(program);
+          }
         }
       },
     },

packages/start/src/directives/remove-unused-variables.ts

@@ -0,0 +1,40 @@
+import type * as babel from "@babel/core";
+import type * as t from "@babel/types";
+
+export function removeUnusedVariables(program: babel.NodePath<t.Program>) {
+  // TODO(Alexis):
+  // This implementation is simple but slow
+  // We repeat removing unused variables from each pass
+  // until no potential unused variables are left.
+  // There might be a simpler implementation.
+  let dirty = true;
+
+  while (dirty) {
+    dirty = false;
+    program.traverse({
+      BindingIdentifier(path) {
+        const binding = path.scope.getBinding(path.node.name);
+
+        if (binding) {
+          switch (binding.kind) {
+            case "const":
+            case "let":
+            case "var":
+            case "hoisted":
+            case "module":
+              if (binding.references === 0 && !binding.path.removed) {
+                binding.path.remove();
+                dirty = true;
+              }
+              break;
+            case "local":
+            case "param":
+            case "unknown":
+              break;
+          }
+        }
+      },
+    });
+    program.scope.crawl();
+  }
+}