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

packages/fs-router-vite-plugin/tests/generator.spec.ts

@@ -5,7 +5,7 @@ import { describe, it, expect } from 'vitest'
 import { routeGenerator } from '../src/generator'
 
 describe('generator works', async () => {
-  const folderNames = await fs.readdir(process.cwd() + '/tests/generator')
+  const folderNames = await fs.readdir(`${process.cwd()}/tests/generator`)
 
   it.each(folderNames)(
     'should wire-up the routes for a "%s" tree',

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
     },
   }

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

@@ -1,22 +1,25 @@
-import type {
-  Identifier,
-  ImportDeclaration,
-  ImportSpecifier,
+import {
+  isIdentifier,
+  isImportDeclaration,
+  isImportSpecifier,
 } from '@babel/types'
 
 import { TUONO_MAIN_PACKAGE, TUONO_DYNAMIC_FN_ID } from './constants'
 
-export const isTuonoDynamicFnImported = (
-  path: babel.NodePath<ImportSpecifier>,
-): boolean => {
-  if ((path.node.imported as Identifier).name !== TUONO_DYNAMIC_FN_ID) {
-    return false
-  }
-  if (
-    (path.parentPath.node as ImportDeclaration).source.value !==
-    TUONO_MAIN_PACKAGE
-  ) {
-    return false
-  }
-  return true
+/**
+ * By a given AST Node path returns true if the path involves an import specifier
+ * importing {@link TUONO_DYNAMIC_FN_ID}
+ */
+export const isTuonoDynamicFnImported = (path: babel.NodePath): boolean => {
+  // If the node isn't an import declaration there is no need to process it
+  if (!isImportDeclaration(path.parentPath?.node)) return false
+
+  // if the import doesn't import from 'tuono' we don't need to process it
+  if (path.parentPath.node.source.value !== TUONO_MAIN_PACKAGE) return false
+
+  // ensure that we are processing an import specifier
+  if (!isImportSpecifier(path.node)) return false
+
+  // finally check if the imported item is `TUONO_DYNAMIC_FN_ID`
+  return isIdentifier(path.node.imported, { name: TUONO_DYNAMIC_FN_ID })
 }

packages/lazy-fn-vite-plugin/tests/sources/dynamic-only/client.expected.tsx

@@ -0,0 +1,9 @@
+import React, { Suspense, type JSX } from "react";
+import { __tuono__internal__lazyLoadComponent as dynamic } from "tuono";
+const DynamicComponent = dynamic(() => import("../components/DynamicComponent"));
+const Loading = (): JSX.Element => <>Loading</>;
+export default function IndexPage(): JSX.Element {
+  return <Suspense fallback={<Loading />}>
+			<DynamicComponent />
+		</Suspense>;
+}

packages/lazy-fn-vite-plugin/tests/sources/dynamic-only/server.expected.tsx

@@ -0,0 +1,8 @@
+import React, { Suspense, type JSX } from "react";
+import DynamicComponent from "../components/DynamicComponent";
+const Loading = (): JSX.Element => <>Loading</>;
+export default function IndexPage(): JSX.Element {
+  return <Suspense fallback={<Loading />}>
+			<DynamicComponent />
+		</Suspense>;
+}

packages/lazy-fn-vite-plugin/tests/sources/dynamic-only/source.tsx

@@ -0,0 +1,14 @@
+import React, { Suspense, type JSX } from "react";
+import { dynamic } from "tuono";
+
+const DynamicComponent = dynamic(() => import("../components/DynamicComponent"))
+
+const Loading = (): JSX.Element => <>Loading</>
+
+export default function IndexPage(): JSX.Element {
+	return (
+		<Suspense fallback={<Loading />}>
+			<DynamicComponent />
+		</Suspense>
+	);
+}

packages/lazy-fn-vite-plugin/tests/transpileSource.test.ts

@@ -1,5 +1,5 @@
 import fs from 'node:fs/promises'
-import path from 'node:path'
+import os from 'node:os'
 
 import { it, expect, describe } from 'vitest'
 import type { Plugin } from 'vite'
@@ -17,35 +17,48 @@ function getTransform(): (...args: Parameters<ViteTransformHandler>) => string {
   return LazyLoadingPlugin().transform as never
 }
 
-describe('"dynamic" fn', async () => {
+describe('"dynamic" sources', async () => {
   const folderNames = await fs.readdir(`${process.cwd()}/tests/sources`)
 
-  it.each(folderNames)(
-    'should correctly build the "%s" dynamic fn',
-    async (folderName) => {
-      const testDirPath = `${process.cwd()}/tests/sources/${folderName}`
+  describe.each(folderNames)('%s', async (folderName) => {
+    const testDirPath = `${process.cwd()}/tests/sources/${folderName}`
 
-      const source = await fs.readFile(
-        path.join(testDirPath, 'source.tsx'),
-        'utf-8',
-      )
+    const sourceRaw = await fs.readFile(`${testDirPath}/source.tsx`, 'utf-8')
+    /**
+     * When adding `packages/lazy-fn-vite-plugin/tests/sources/dynamic-only` only
+     * the test involving that fixture were broken on Windows... but not the one in the other fixtures:
+     * - packages/lazy-fn-vite-plugin/tests/sources/vanilla
+     * - packages/lazy-fn-vite-plugin/tests/sources/external-dynamic
+     *
+     * Awkwardly this doesn't happen on  `packages/fs-router-vite-plugin/tests/generator.spec.ts`
+     *
+     * Too much pain and sadness to investigate this right now.
+     * Might worth creating an utility function in the future if this happens again
+     */
+    const source = sourceRaw.replace(new RegExp(os.EOL, 'g'), '\n')
 
+    it('should generate file for client', async () => {
       const pluginTransform = getTransform()
       const clientBundle = pluginTransform(source, 'id')
-      const serverBundle = pluginTransform(source, 'id', { ssr: true })
 
       const expectedClientSrc = `${testDirPath}/client.expected.tsx`
-      const expectedServerSrc = `${testDirPath}/server.expected.tsx`
 
       await expect(clientBundle).toMatchFileSnapshot(
         expectedClientSrc,
         `${testDirPath} client build should be equal to ${expectedClientSrc}`,
       )
+    })
+
+    it('should generate file for server', async () => {
+      const pluginTransform = getTransform()
+      const serverBundle = pluginTransform(source, 'id', { ssr: true })
+
+      const expectedServerSrc = `${testDirPath}/server.expected.tsx`
 
       await expect(serverBundle).toMatchFileSnapshot(
         expectedServerSrc,
         `${testDirPath} server build should be equal to ${expectedServerSrc}`,
       )
-    },
-  )
+    })
+  })
 })