some changes

README.md

@@ -24,7 +24,18 @@ You can see that the async chunks are loaded directly and instantly available on
 > [!CAUTION]
 > This library is experimental and not guaranteed to work in your custom setup. 
 
-## Example showing where this library goes in your application
+
+## Install
+
+```
+$ npm install vite-preload
+```
+
+***
+
+### See [./playground](./playground/) for a basic setup with preloading
+
+## Psuedo example highlighting the important parts
 
 ### `vite.config.ts`
 
@@ -98,4 +109,16 @@ function render(req, res) {
 ```
 
 > [!NOTE]
-> In a development environment you will still have some Flash Of Unstyled Content
+> In a development environment you will still have some Flash Of Unstyled Content
+
+## Further optimizations
+
+If your app knows what pages or components that should be preloaded, like the next obvious path the user will make in your user flow, it's recommended to lazy load them with something like `lazyWithPreload`.
+
+Even if you would use a `import()` call to preload the chunk, React will still suspend 
+```tsx
+import lazyWithPreload from 'react-lazy-with-preload';
+
+const Card = lazyWithPreload(() => import('./Card'));
+Card.preload();
+```

package.json

@@ -35,6 +35,8 @@
   "license": "MIT",
   "devDependencies": {
     "@rollup/plugin-typescript": "^11.1.6",
+    "@types/babel__parser": "^7.0.0",
+    "@types/babel__traverse": "^7.20.6",
     "@types/debug": "^4.1.12",
     "@types/node": "^22.3.0",
     "@types/react": "^18.3.3",

src/context.ts

@@ -1,5 +1,6 @@
 import React from 'react';
 
 export const ModuleCollectorContext = React.createContext((id: string) => {
+    // Debug. Should do nothing on the client
     console.log('useReportModule', id);
 });

src/plugin.ts

@@ -21,6 +21,7 @@ export default function preloadPlugin({
     __internal_importHelperModuleName = 'vite-preload/__internal',
 }: PluginOptions = {}): Plugin {
     const lazyImportedModules = new Set();
+    const injectedModules = new Set();
     let count = 0;
 
     return {
@@ -31,23 +32,22 @@ export default function preloadPlugin({
             return Boolean(config.build?.ssr) || command === 'serve';
         },
 
-        transform(code, id) {
+        async transform(code, id) {
             const relative = getRelativePath(id);
 
-            if (include.test(id) && code.includes('React.lazy')) {
+            const foundLazyImports = new Set<string>();
+
+            if (include.test(id) && code.includes('lazy')) {
                 const ast = parse(code, {
                     sourceType: 'module',
                     plugins: ['jsx', 'typescript'],
                 });
 
+                // Find React.lazy imported modules
                 traverse.default(ast, {
-                    CallExpression: (path) => {
-                        if (
-                            t.isMemberExpression(path.node.callee) &&
-                            path.node.callee.object.name === 'React' &&
-                            path.node.callee.property.name === 'lazy'
-                        ) {
-                            const argument = path.node.arguments[0];
+                    CallExpression: (p) => {
+                        if (isLazyExpression(p)) {
+                            const argument = p.node.arguments[0];
                             if (t.isArrowFunctionExpression(argument)) {
                                 const body = argument.body;
                                 if (
@@ -58,34 +58,31 @@ export default function preloadPlugin({
                                         body.arguments[0] as any
                                     ).value;
 
-                                    // Shitty solution but handles index imports and imports with no extensions
-                                    lazyImportedModules.add(
-                                        modulePath + '.tsx'
-                                    );
-                                    lazyImportedModules.add(
-                                        modulePath + '.jsx'
-                                    );
-                                    lazyImportedModules.add(
-                                        modulePath + '/index.tsx'
-                                    );
-                                    lazyImportedModules.add(
-                                        modulePath + '/index.jsx'
-                                    );
-                                    lazyImportedModules.add(modulePath);
-
-                                    // Process every module that includes a React.lazy.
-                                    lazyImportedModules.add(relative);
-                                    this.info(
-                                        `${relative} imports ${modulePath}`
-                                    );
+                                    foundLazyImports.add(modulePath);
                                 }
                             }
                         }
                     },
                 });
             }
 
-            if (lazyImportedModules.has(relative)) {
+            for (const importString of foundLazyImports) {
+                const relative = path.resolve(path.dirname(id), importString);
+                const resolved = await this.resolve(importString, id);
+
+                if (!resolved) {
+                    throw new Error(`Did not find imported module ${relative}`);
+                }
+
+                this.info(
+                    `imports ${path.relative(process.cwd(), resolved.id)}`
+                );
+
+                lazyImportedModules.add(resolved.id);
+                // await this.load({ id: resolved.id });
+            }
+
+            if (lazyImportedModules.has(id)) {
                 const ast = parse(code, {
                     sourceType: 'module',
                     plugins: ['jsx', 'typescript'],
@@ -126,22 +123,25 @@ export default function preloadPlugin({
                 });
 
                 if (injected) {
-                    this.info('Injected __collectModule');
+                    this.info('Injected __collectModule in React component');
                     count++;
                     const output = generate.default(ast, {}, code);
+                    injectedModules.add(id);
                     return {
                         code: output.code,
                         map: output.map,
                     };
-                } else {
-                    this.warn('Did NOT inject __collectModule');
                 }
             }
 
             return null;
         },
 
         buildEnd() {
+            const s = lazyImportedModules.difference(injectedModules);
+            for (const z of s) {
+                this.info(`${z} was not injected`);
+            }
             this.info(`${count} hook calls injected`);
         },
     };
@@ -193,14 +193,38 @@ function injectImport(ast, importHelper) {
 }
 
 function isFunctionComponent(node) {
-    // Check if it's a function declaration or arrow function expression
     return (
         t.isFunctionDeclaration(node) ||
         t.isFunctionExpression(node) ||
         t.isArrowFunctionExpression(node)
     );
 }
 
+function isLazyExpression(p) {
+    // Ensure p.node and p.node.callee exist before proceeding
+    if (!p.node || !p.node.callee) {
+        return false;
+    }
+
+    const callee = p.node.callee;
+
+    // Check if it's a React.lazy expression
+    const isReactLazy =
+        t.isMemberExpression(callee) &&
+        callee.object.name === 'React' &&
+        callee.property.name === 'lazy';
+
+    // Check if it's a standalone lazy or lazyWithPreload function call
+    const isStandaloneLazy =
+        t.isIdentifier(callee) &&
+        (callee.name === 'lazy' || callee.name === 'lazyWithPreload');
+
+    // Log the node for debugging
+
+    // Return true if any of the conditions match
+    return isReactLazy || isStandaloneLazy;
+}
+
 function getRelativePath(filePath) {
     return path.relative(process.cwd(), filePath).replace(/\\/g, '/');
 }

src/utils.ts

@@ -10,7 +10,7 @@ export function createHtmlTag({ type, href, comment }: Module) {
 
     switch (type) {
         case 'stylesheet':
-            tag += `<link rel=stylesheet crossorigin href="/${href}" nonce="%NONCE%" />`;
+            tag += `<link rel="stylesheet" crossorigin href="/${href}" nonce="%NONCE%" />`;
             break;
         case 'modulepreload':
             tag += `<link rel="modulepreload" crossorigin href="/${href}" nonce="%NONCE%" />`;