ReScriptify Server and implement asset loading

.gitignore

@@ -10,4 +10,5 @@ src/__generated__/**
 src/**/*.mjs
 router/**/*.mjs
 cli/**/*.mjs
-_release
+Server.mjs
+_release

RescriptRelayVitePlugin.mjs

@@ -1,8 +1,15 @@
 import fs from "fs";
+import fsPromised from "fs/promises";
 import path from "path";
 import readline from "readline";
+import MagicString from "magic-string";
+import { normalizePath } from 'vite'
 import { runCli } from "./cli/RescriptRelayRouterCli__Commands.mjs";
 
+/**
+ * @typedef {import("vite").ResolvedConfig} ResolvedConfig
+ */
+
 // Expected to run in vite.config.js folder, right next to bsconfig.
 let cwd = process.cwd();
 
@@ -68,15 +75,24 @@ export let rescriptRelayVitePlugin = ({
   autoScaffoldRenderers = true,
   deleteRemoved = true,
 } = {}) => {
+  // The watcher for the ReScript Relay Router CLI.
   let watcher;
-  let outputCount;
-  let manifest = {};
+  // An in-memory copy of the ssr-manifest.json for bundle manipulation.
+  let ssrManifest = {};
+  // The resolved Vite config to ensure we do what the rest of Vite does.
+  /** @type ResolvedConfig */
   let config;
 
   return {
     name: "rescript-relay",
+    /**
+     * Workaround until we can upgrade to Vite 3.
+     *
+     * Remove manualChunks if this is SSR, since it doesn't work in SSR mode.
+     * See https://github.com/vitejs/vite/issues/8836
+     */
     config(userConfig) {
-      // Remove manualChunks if this is SSR, since it doesn't work in SSR mode.
+      //
       if (
         Boolean(userConfig.build.ssr) &&
         userConfig.build?.rollupOptions?.output?.manualChunks != null
@@ -86,9 +102,17 @@ export let rescriptRelayVitePlugin = ({
 
       return userConfig;
     },
+    /**
+     * @param {ResolvedConfig} resolvedConfig
+     */
     configResolved(resolvedConfig) {
-      config = resolvedConfig;
-      outputCount = 0;
+      config = resolvedConfig
+      // For the server build in SSR we read the client manifest from disk.
+      if (config.build.ssr) {
+        // TODO: This relies on the client and server paths being next to eachother. Perhaps add config?
+        // TODO: SSR Manifest name is configurable in Vite and may be different.
+        ssrManifest = JSON.parse(fs.readFileSync(path.resolve(config.build.outDir, "../client/ssr-manifest.json"), 'utf-8'));
+      }
     },
     buildStart() {
       // Run single generate in prod
@@ -145,148 +169,148 @@ export let rescriptRelayVitePlugin = ({
         }
       }
     },
-
-    // This below takes care of inlining chunk names into the source JS, so we
-    // can use those chunk names to preload code via script tags on the
-    // client.
-
-    // This first pass is only relevant in dev. It will replace all
-    // `__$rescriptChunkName__: "ModuleName" entries with the full file
-    // location of "ModuleName", so that we can preload the full module via
-    // the file name. This is handled differently in production, where we need
-    // to do a full lookup of what bundle the target file is in, which we
-    // can't do until all bundles have been rendered.
-    async transform(code) {
-      if (config.mode === "production") {
-        return null;
-      }
-
+    // Transforms the magic object property's value `__$rescriptChunkName__` from `ModuleName` (without extension)
+    // into the actual path for the compiled asset.
+    async transform(code, id) {
+      // The __$rescriptChunkName__ is a non-public identifier used to bridge the gap
+      // between the ReScript and JavaScript world. It's public API is `chunk` within
+      // ReScript and it's not intended to be used from a non-ReScript codebase.
       if (!code.startsWith("// Generated by ReScript")) {
         return null;
       }
 
-      let didReplace = false;
-
-      let newCode = await replaceAsync(
+      const transformedCode = await replaceAsyncWithMagicString(
         code,
-        /__\$rescriptChunkName__: "([A-Za-z0-9_]+)"/gm,
-        async (_match, moduleName) => {
-          if (moduleName != null) {
-            let loc = await findGeneratedModule(moduleName);
-
-            if (loc != null) {
-              // This transforms the found loc to a URL relative to the project
-              // root. That's also what vite uses internally as URL for src
-              // assets.
-              let locRelativeToProjectRoot = loc.replace(config.root, "");
-              didReplace = true;
-              // TODO: Source maps
-              return `__$rescriptChunkName__: "${locRelativeToProjectRoot}"`;
+        /__\$rescriptChunkName__:\s*"([A-Za-z0-9_]+)"/gm,
+        async (fullMatch, moduleId) => {
+          if (moduleId != null && moduleId !== "") {
+            let resolved = await findGeneratedModule(moduleId);
+            if (resolved != null) {
+              // The location of findGeneratedModule is an absolute URL but we
+              // want the URL relative to the project root. That's also what
+              // vite uses internally as URL for src assets.
+              resolved = resolved.replace(config.root, "");
+              return `__$rescriptChunkName__: "${resolved}"`;
             }
+            console.warn(`Could not resolve Rescript Module '${moduleId}' for match '${fullMatch}'.`);
           }
+          else {
+            console.warn(`Tried to resolve ReScript module to path but match '${fullMatch}' didn't contain a moduleId.`);
+          }
+
+          return fullMatch;
         }
       );
 
-      if (didReplace) {
-        return newCode;
+      if (!transformedCode.hasChanged()) {
+        return null;
       }
 
-      return null;
-    },
+      const sourceMap = transformedCode.generateMap({
+        source: id,
+        file: `${id}.map`,
+      });
 
-    // This second step runs only in prod, and will track all bundle chunks,
-    // and map them to whatever ReScript modules they contain. This helps us
-    // in two ways:
-    //
-    // 1) It allows us to emit a specialized manifest we can use during SSR to
-    //    push preloading of code assets we'll know we'll need ASAP
-    //
-    // 2) It lets us use the same information to inline bundle chunk location
-    //    strings of ReScript modules in the prod client build, which we then
-    //    use to preload code via script tags on the client.
-    generateBundle(_, bundle) {
-      for (let file in bundle) {
-        let chunk = bundle[file];
-        if (chunk.type === "chunk" && chunk.isDynamicEntry) {
-          manifest[chunk.name] = chunk.fileName;
-        }
+      return {
+        code: transformedCode.toString(),
+        map: sourceMap.toString(),
+      };
+    },
+    // In addition to the transform from ReScript module name to JS file.
+    // In production we want to change the JS file name to the corresponding chunk that contains the compiled JS.
+    // This is similar to what Rollup does for us for `import` statements.
+    // We start out by creating a lookup table of JS files to output assets.
+    // This is copied from vite/packages/vite/src/node/ssr/ssrManifestPlugin.ts but does not track CSS files.
+    generateBundle(_options, bundle) {
+      // We only have to collect the ssr-manifest during client bundling.
+      // For SSR it's just read from disk.
+      if (config.build.ssr) {
+        return;
       }
-
-      // This below is mostly copied from the internal Vite manifest plugin.
-      outputCount++;
-      let output = config.build.rollupOptions?.output;
-      let outputLength = Array.isArray(output) ? output.length : 1;
-      if (outputCount >= outputLength) {
-        if (config.mode === "production" && !Boolean(config.build.ssr)) {
-          this.emitFile({
-            fileName: "res-manifest.json",
-            type: "asset",
-            source: JSON.stringify(manifest, null, 2),
-          });
+      for (const file in bundle) {
+        const chunk = bundle[file]
+        if (chunk.type === 'chunk') {
+          for (const id in chunk.modules) {
+            const normalizedId = normalizePath(path.relative(config.root, id))
+            const mappedChunks =
+              ssrManifest[normalizedId] ?? (ssrManifest[normalizedId] = [])
+            if (!chunk.isEntry) {
+              mappedChunks.push(config.base + chunk.fileName)
+            }
+            chunk.viteMetadata.importedAssets.forEach((file) => {
+              mappedChunks.push(config.base + file)
+            })
+          }
         }
       }
     },
-    // This final step runs when all bundles have been written. It does the
-    // final swap of ReScript module names to a loadable src string pointing
-    // to the bundle where that ReScript module is included. This runs as a
-    // post-build step, and needs to run as that because we need the full
-    // build to finish before we have enough information to do the replace.
+    // We can't do the gathering of chunk names at the same time but must complete all of that
+    // before we can do the replacement so we know we replace all. Therefore we do this in
+    // writeBundle which also only runs in production like generateBundle.
     writeBundle(outConfig, bundle) {
-      // TODO: We might be able to optimize this (skip reading through all of
-      // the generated code) by tracking what ReScript module chunk strings
-      // are included in what bundles, and only replacing in those bundles,
-      // skipping the rest.
-      Object.entries(bundle).forEach(([_bundleName, bundleContents]) => {
-        let code = bundleContents.code;
-
-        if (code != null) {
-          let didReplace = false;
-          // TODO: Source maps
-          let newCode = code.replace(
-            /__\$rescriptChunkName__:"([A-Za-z0-9_]+)"/gm,
-            (match, moduleName) => {
-              didReplace = true;
-
-              return `__$rescriptChunkName__:"${
-                manifest[moduleName] ?? match
-              }"`;
+      Object.entries(bundle).forEach(async ([_bundleName, bundleContents]) => {
+        const code = bundleContents.code;
+        if (typeof code === "undefined") {
+          return;
+        }
+        const transformedCode = await replaceAsyncWithMagicString(
+          code,
+          /__\$rescriptChunkName__:\s*"\/([A-Za-z0-9_\/\.]+)"/gm,
+          (fullMatch, jsUrl) => {
+            if (jsUrl != null && jsUrl !== "") {
+              let chunk = (ssrManifest[jsUrl] ?? [])[0] ?? null;
+              if (chunk !== null) {
+                return `__$rescriptChunkName__:"${chunk}"`;
+              }
+              console.warn(`Could not find chunk path for '${jsUrl}' for match '${fullMatch}'.`);
+            }
+            else {
+              console.warn(`Tried to rewrite compiled path to chunk but match '${fullMatch}' didn't contain a compiled path.`);
             }
-          );
 
-          if (didReplace) {
-            fs.writeFileSync(
-              path.resolve(outConfig.dir, bundleContents.fileName),
-              newCode
-            );
+            return fullMatch;
           }
+        );
+
+        if (transformedCode.hasChanged()) {
+          await fsPromised.writeFile(
+            path.resolve(outConfig.dir, bundleContents.fileName),
+            transformedCode.toString()
+          );
         }
       });
-    },
+    }
   };
 };
 
-// Copied from some async replacer lib.
-function replaceAsync(string, searchValue, replacer) {
+/**
+ * Performs a string replace with an async replacer function returning a source map.
+ *
+ * Takes the following steps:
+ * 1. Run fake pass of `replace`, collect values from `replacer` calls
+ * 2. Resolve them with `Promise.all`
+ * 3. Create a 'MagicString' (using the magic-string package).
+ * 4. Run `replace` with resolved values
+ */
+function replaceAsyncWithMagicString(string, searchValue, replacer) {
+  if (typeof replacer !== "function") {
+    throw new Error("Must provide a replacer function, otherwise just call replace directly.");
+  }
   try {
-    if (typeof replacer === "function") {
-      // 1. Run fake pass of `replace`, collect values from `replacer` calls
-      // 2. Resolve them with `Promise.all`
-      // 3. Run `replace` with resolved values
-      var values = [];
-      String.prototype.replace.call(string, searchValue, function () {
-        values.push(replacer.apply(undefined, arguments));
-        return "";
-      });
-      return Promise.all(values).then(function (resolvedValues) {
-        return String.prototype.replace.call(string, searchValue, function () {
-          return resolvedValues.shift();
-        });
-      });
-    } else {
-      return Promise.resolve(
-        String.prototype.replace.call(string, searchValue, replacer)
+    var values = [];
+    String.prototype.replace.call(string, searchValue, function () {
+      values.push(replacer.apply(undefined, arguments));
+      return "";
+    });
+    let mapTrackingString = new MagicString(string)
+    return Promise.all(values).then(function (resolvedValues) {
+      // Call replace again, this time on the string that tracks a sourcemap.
+      // We use the replacerFunction so each occurrence can be replaced by the
+      // previously resolved value for that index.
+      return mapTrackingString.replace(searchValue,
+        () => resolvedValues.shift()
       );
-    }
+    });
   } catch (error) {
     return Promise.reject(error);
   }