fix(packages/lazy-fn-vite-plugin): make dynamic function detection stricter (#223)

packages/lazy-fn-vite-plugin/src/index.ts

@@ -1,13 +1,7 @@
-import * as babel from '@babel/core'
-import * as t from '@babel/types'
-import type { Plugin } from 'vite'
-import type { PluginItem } from '@babel/core'
-import type {
-  Identifier,
-  CallExpression,
-  ArrowFunctionExpression,
-  StringLiteral,
-} from '@babel/types'
+import { transformSync } from '@babel/core'
+import type { PluginItem as BabelPluginItem } from '@babel/core'
+import * as BabelTypes from '@babel/types'
+import type { Plugin as VitePlugin, Rollup } from 'vite'
 
 import {
   TUONO_MAIN_PACKAGE,
@@ -20,11 +14,24 @@ import { isTuonoDynamicFnImported } from './utils'
  * [SERVER build]
  * This plugin just removes the `dynamic` imported function from any tuono import
  */
-const RemoveTuonoLazyImport: PluginItem = {
+const RemoveTuonoLazyImport: BabelPluginItem = {
   name: 'remove-tuono-lazy-import-plugin',
   visitor: {
-    ImportSpecifier: (path) => {
-      if (isTuonoDynamicFnImported(path)) {
+    ImportDeclaration: (path) => {
+      const importNode = path.node
+      if (importNode.source.value !== TUONO_MAIN_PACKAGE) return
+
+      path.traverse({
+        ImportSpecifier: (importSpecifierPath) => {
+          if (isTuonoDynamicFnImported(importSpecifierPath)) {
+            importSpecifierPath.remove()
+          }
+        },
+      })
+
+      // If there are no specifiers left after traverse
+      // remove the import to avoid unwanted side effects
+      if (importNode.specifiers.length === 0) {
         path.remove()
       }
     },
@@ -35,34 +42,54 @@ const RemoveTuonoLazyImport: PluginItem = {
  * [CLIENT build]
  * This plugin replace the `dynamic` function with the `__tuono__internal__lazyLoadComponent` one
  */
-const ReplaceTuonoLazyImport: PluginItem = {
-  name: 'remove-tuono-lazy-import-plugin',
+const ReplaceTuonoLazyImport: BabelPluginItem = {
+  name: 'replace-tuono-lazy-import-plugin',
   visitor: {
     ImportSpecifier: (path) => {
-      if (isTuonoDynamicFnImported(path)) {
-        ;(path.node.imported as Identifier).name = TUONO_LAZY_FN_ID
+      if (
+        BabelTypes.isIdentifier(path.node.imported) &&
+        isTuonoDynamicFnImported(path)
+      ) {
+        path.node.imported.name = TUONO_LAZY_FN_ID
       }
     },
   },
 }
 
 const turnLazyIntoStatic = {
-  VariableDeclaration: (path: babel.NodePath<t.VariableDeclaration>): void => {
-    path.node.declarations.forEach((el) => {
-      const init = el.init as CallExpression
-      if ((init.callee as Identifier).name === TUONO_DYNAMIC_FN_ID) {
-        const importName = (el.id as Identifier).name
-        const importPath = (
-          (
-            (init.arguments[0] as ArrowFunctionExpression)
-              .body as CallExpression
-          ).arguments[0] as StringLiteral
-        ).value
+  VariableDeclaration: (
+    path: babel.NodePath<BabelTypes.VariableDeclaration>,
+  ): void => {
+    path.node.declarations.forEach((variableDeclarationNode) => {
+      const init = variableDeclarationNode.init
+
+      if (
+        BabelTypes.isCallExpression(init) &&
+        // ensures that the method call is `TUONO_DYNAMIC_FN_ID`
+        BabelTypes.isIdentifier(init.callee, { name: TUONO_DYNAMIC_FN_ID }) &&
+        // import name must be an identifier
+        BabelTypes.isIdentifier(variableDeclarationNode.id) &&
+        // check that the first function parameter is an arrow function
+        BabelTypes.isArrowFunctionExpression(init.arguments[0])
+      ) {
+        const cmpImportFn = init.arguments[0]
+
+        // ensures that the first parameter is a call expression (may be a block statement)
+        if (!BabelTypes.isCallExpression(cmpImportFn.body)) return
+        // ensures that the first parameter is a string literal (the import path)
+        if (!BabelTypes.isStringLiteral(cmpImportFn.body.arguments[0])) return
+
+        const importName = variableDeclarationNode.id.name
+        const importPath = cmpImportFn.body.arguments[0].value
 
         if (importName && importPath) {
-          const importDeclaration = t.importDeclaration(
-            [t.importDefaultSpecifier(t.identifier(importName))],
-            t.stringLiteral(importPath),
+          const importDeclaration = BabelTypes.importDeclaration(
+            [
+              BabelTypes.importDefaultSpecifier(
+                BabelTypes.identifier(importName),
+              ),
+            ],
+            BabelTypes.stringLiteral(importPath),
           )
 
           path.replaceWith(importDeclaration)
@@ -76,7 +103,7 @@ const turnLazyIntoStatic = {
  * [SERVER build]
  * This plugin statically imports the lazy loaded components
  */
-const TurnLazyIntoStaticImport: PluginItem = {
+const TurnLazyIntoStaticImport: BabelPluginItem = {
   name: 'turn-lazy-into-static-import-plugin',
   visitor: {
     Program: (path) => {
@@ -91,28 +118,38 @@ const TurnLazyIntoStaticImport: PluginItem = {
   },
 }
 
-export function LazyLoadingPlugin(): Plugin {
+export function LazyLoadingPlugin(): VitePlugin {
   return {
     name: 'vite-plugin-tuono-lazy-loading',
     enforce: 'pre',
-    transform(code, _id, opts): string | undefined | null {
+    transform(code, _id, opts): Rollup.TransformResult {
+      /**
+       * @todo we should exclude non tsx files from this transformation
+       *       this might benefit build time avoiding running `includes` on non-tsx files.
+       *       This can be executed using `_id` parameter
+       *       which is the filepath that is being processed
+       */
+
       if (
         code.includes(TUONO_DYNAMIC_FN_ID) &&
         code.includes(TUONO_MAIN_PACKAGE)
       ) {
-        const res = babel.transformSync(code, {
-          plugins: [
-            ['@babel/plugin-syntax-jsx', {}],
-            ['@babel/plugin-syntax-typescript', { isTSX: true }],
-            [!opts?.ssr ? ReplaceTuonoLazyImport : []],
-            [opts?.ssr ? RemoveTuonoLazyImport : []],
-            [opts?.ssr ? TurnLazyIntoStaticImport : []],
-          ],
-          sourceMaps: true,
-        })
+        const plugins: Array<BabelPluginItem> = [
+          ['@babel/plugin-syntax-jsx', {}],
+          ['@babel/plugin-syntax-typescript', { isTSX: true }],
+        ]
+
+        if (opts?.ssr) {
+          plugins.push(RemoveTuonoLazyImport, TurnLazyIntoStaticImport)
+        } else {
+          plugins.push(ReplaceTuonoLazyImport)
+        }
+
+        const res = transformSync(code, { plugins })
 
         return res?.code
       }
+
       return code
     },
   }