fix: enhance Vite plugin to track type usage and refresh project on hot updates

Author: HichemTab-tech

src/ts-transformer.ts

@@ -1,11 +1,23 @@
 import { Project, SyntaxKind } from "ts-morph";
 
-const project = new Project({
-    tsConfigFilePath: "tsconfig.json",
-});
+export const fileToTypes = new Map<string, Set<string>>();
+
+export const typeToFile = new Map<string, string>();
+
+let project = null as unknown as Project;
+
+export const refreshProject = () => {
+    project = new Project({
+        tsConfigFilePath: "tsconfig.json",
+    });
+}
+
+refreshProject();
 
 export function transform(code: string, filePath: string): string {
 
+    const usedTypes = new Set<string>();
+
     let sourceFile = project.getSourceFile(filePath);
 
     if (!sourceFile) {
@@ -16,6 +28,20 @@ export function transform(code: string, filePath: string): string {
         sourceFile.replaceWithText(code);
     }
 
+    sourceFile.getDescendantsOfKind(SyntaxKind.TypeReference).forEach(ref => {
+        const typeName = ref.getTypeName().getText();
+        usedTypes.add(typeName);
+
+        // BONUS: resolve where this type is defined
+        const decl = ref.getType().getSymbol()?.getDeclarations()?.[0];
+        const source = decl?.getSourceFile()?.getFilePath();
+        if (source) {
+            typeToFile.set(typeName, source);
+        }
+    });
+
+    fileToTypes.set(filePath, usedTypes);
+
     // Figure out what the “createPicker” import is called (alias or direct)
     let createPickerAlias: string | null = null;
     for (const importDecl of sourceFile.getImportDeclarations()) {

src/vite-plugin.ts

@@ -1,5 +1,5 @@
 import { PluginOption } from "vite";
-import { transform } from "./ts-transformer.js";
+import {fileToTypes, refreshProject, transform, typeToFile} from "./ts-transformer.js";
 
 // noinspection JSUnusedGlobalSymbols
 /**
@@ -27,6 +27,42 @@ export default function TsRuntimePickerVitePlugin(): PluginOption {
 
             return null;
         },
+
+        handleHotUpdate: ({server, file}) => {
+            const affectedFiles = [];
+
+            // Loop over all files that used types
+            for (const [userFile, usedTypes] of fileToTypes.entries()) {
+                for (const typeName of usedTypes) {
+                    if (typeToFile.get(typeName) === file) {
+                        affectedFiles.push(userFile);
+                        break;
+                    }
+                }
+            }
+
+            if (affectedFiles.length) {
+                refreshProject();
+
+                for (const file of affectedFiles) {
+                    const mod = server.moduleGraph.getModuleById(file);
+                    if (mod) {
+                        server.moduleGraph.invalidateModule(mod);
+                        server.ws.send({
+                            type: 'update',
+                            updates: [
+                                {
+                                    type: 'js-update',
+                                    path: mod.url, // relative to root
+                                    acceptedPath: mod.url,
+                                    timestamp: Date.now()
+                                }
+                            ]
+                        });
+                    }
+                }
+            }
+        }
     } as PluginOption;
 }