Improve bundling of web-viewer package

Cargo.toml

@@ -240,6 +240,8 @@ url = "2.3"
 uuid = "1.1"
 vec1 = "1.8"
 walkdir = "2.0"
+# NOTE: `rerun_js/web-viewer/build-wasm.mjs` is HIGHLY sensitive to changes in `wasm-bindgen`.
+#       Whenever updating `wasm-bindgen`, make sure that the build script still works.
 wasm-bindgen = "0.2.89"
 wasm-bindgen-cli-support = "0.2.89"
 wasm-bindgen-futures = "0.4.33"

crates/re_dev_tools/src/build_web_viewer/lib.rs

@@ -49,6 +49,20 @@ impl Profile {
 pub enum Target {
     Browser,
     Module,
+
+    /// Custom target meant for post-processing inside `rerun_js`.
+    NoModulesBase,
+}
+
+impl argh::FromArgValue for Target {
+    fn from_arg_value(value: &str) -> Result<Self, String> {
+        match value {
+            "browser" => Ok(Self::Browser),
+            "module" => Ok(Self::Module),
+            "no-modules-base" => Ok(Self::NoModulesBase),
+            _ => Err(format!("Unknown target: {value}")),
+        }
+    }
 }
 
 /// Build `re_viewer` as Wasm, generate .js bindings for it, and place it all into the `build_dir` folder.
@@ -158,6 +172,10 @@ pub fn build(
         match target {
             Target::Browser => bindgen_cmd.no_modules(true)?.typescript(false),
             Target::Module => bindgen_cmd.no_modules(false)?.typescript(true),
+            Target::NoModulesBase => bindgen_cmd
+                .no_modules(true)?
+                .reference_types(true)
+                .typescript(true),
         };
         if let Err(err) = bindgen_cmd.generate(build_dir.as_str()) {
             if err
@@ -191,7 +209,13 @@ pub fn build(
         // to get wasm-opt:  apt/brew/dnf install binaryen
         let mut cmd = std::process::Command::new("wasm-opt");
 
-        let mut args = vec![wasm_path.as_str(), "-O2", "--output", wasm_path.as_str()];
+        let mut args = vec![
+            wasm_path.as_str(),
+            "-O2",
+            "--output",
+            wasm_path.as_str(),
+            "--enable-reference-types",
+        ];
         if debug_symbols {
             args.push("-g");
         }

crates/re_dev_tools/src/build_web_viewer/mod.rs

@@ -27,9 +27,9 @@ pub struct Args {
     #[argh(switch, short = 'g')]
     debug_symbols: bool,
 
-    /// if set, will build the module target instead of the browser target.
-    #[argh(switch)]
-    module: bool,
+    /// target to build for.
+    #[argh(option, short = 't', long = "target", default = "Target::Browser")]
+    target: Target,
 
     /// set the output directory. This is a path relative to the cargo workspace root.
     #[argh(option, short = 'o', long = "out")]
@@ -47,12 +47,7 @@ pub fn main(args: Args) -> anyhow::Result<()> {
         ));
     };
 
-    let target = if args.module {
-        Target::Module
-    } else {
-        Target::Browser
-    };
     let build_dir = args.build_dir.unwrap_or_else(default_build_dir);
 
-    build(profile, args.debug_symbols, target, &build_dir)
+    build(profile, args.debug_symbols, args.target, &build_dir)
 }

rerun_js/web-viewer/.gitignore

@@ -7,3 +7,5 @@ index.d.ts
 index.d.ts.map
 index.js
 index.js.map
+inlined.js
+inlined.d.ts

rerun_js/web-viewer/build-wasm.mjs

@@ -0,0 +1,97 @@
+// Script responsible for building the wasm and transforming the JS bindings for the web viewer.
+
+import * as child_process from "node:child_process";
+import { fileURLToPath } from "node:url";
+import * as path from "node:path";
+import * as fs from "node:fs";
+import * as util from "node:util";
+
+const __filename = path.resolve(fileURLToPath(import.meta.url));
+const __dirname = path.dirname(__filename);
+
+const exec = (cmd) => {
+  console.log(cmd);
+  child_process.execSync(cmd, { cwd: __dirname, stdio: "inherit" });
+};
+
+function wasm(mode) {
+  switch (mode) {
+    case "debug": {
+      return exec(
+        "cargo run -p re_dev_tools -- build-web-viewer --debug --target no-modules-base -o rerun_js/web-viewer",
+      );
+    }
+    case "release": {
+      return exec(
+        "cargo run -p re_dev_tools -- build-web-viewer --release -g --target no-modules-base -o rerun_js/web-viewer",
+      );
+    }
+    default:
+      throw new Error(`Unknown mode: ${mode}`);
+  }
+}
+
+child_process.execSync(
+  "cargo run -p re_dev_tools -- build-web-viewer --debug --target no-modules-base -o rerun_js/web-viewer",
+  { cwd: __dirname, stdio: "inherit" },
+);
+
+function script() {
+  let code = fs.readFileSync(path.join(__dirname, "re_viewer.js"), "utf-8");
+
+  // this transforms the module, wrapping it in a default-exported function.
+  // calling the function produces a new "instance" of the module, because
+  // all of the globals are scoped to the function, and become closure state
+  // for any functions that reference them within the module.
+  //
+  // we do this so that we don't leak globals across web viewer instantiations.
+  //
+  // this is HIGHLY sensitive to the exact output of `wasm-bindgen`, so if
+  // the output changes, this will need to be updated.
+
+  const start = `let wasm_bindgen;
+(function() {`;
+  const end = `wasm_bindgen = Object.assign(__wbg_init, { initSync }, __exports);
+
+})();`;
+  code = code.replace(start, "").replace(end, "");
+
+  code = `
+export default function() {
+${code}
+return Object.assign(__wbg_init, { initSync }, __exports);
+}
+`;
+
+  fs.writeFileSync(path.join(__dirname, "re_viewer.js"), code);
+}
+
+function types() {
+  let code = fs.readFileSync(path.join(__dirname, "re_viewer.d.ts"), "utf-8");
+
+  // this transformation just re-exports WebHandle and adds a default export inside the `.d.ts` file
+
+  code = `
+${code}
+export type WebHandle = wasm_bindgen.WebHandle;
+export default function(): wasm_bindgen;
+`;
+
+  fs.writeFileSync(path.join(__dirname, "re_viewer.d.ts"), code);
+}
+
+const args = util.parseArgs({
+  options: {
+    mode: {
+      type: "string",
+    },
+  },
+});
+
+if (!args.values.mode) {
+  throw new Error("Missing required argument: mode");
+}
+
+wasm(args.values.mode);
+script();
+types();

rerun_js/web-viewer/bundle.mjs

@@ -0,0 +1,55 @@
+// Script responsible for taking the generated Wasm/JS, and transpiled TS
+// and producing a single file with everything inlined.
+
+import { fileURLToPath } from "node:url";
+import * as path from "node:path";
+import * as fs from "node:fs";
+import * as util from "node:util";
+
+const __filename = path.resolve(fileURLToPath(import.meta.url));
+const __dirname = path.dirname(__filename);
+
+const wasm = fs.readFileSync(path.join(__dirname, "re_viewer_bg.wasm"));
+const js = fs.readFileSync(path.join(__dirname, "re_viewer.js"), "utf-8");
+const index = fs.readFileSync(path.join(__dirname, "index.js"), "utf-8");
+
+const INLINE_MARKER = "/*<INLINE-MARKER>*/";
+
+/** @param {Buffer} buffer */
+function buffer_to_data_url(buffer) {
+  return `data:application/wasm;base64,${buffer.toString("base64")}`;
+}
+
+async function data_url_to_buffer(dataUrl) {
+  const response = await fetch(dataUrl);
+  return response.arrayBuffer();
+}
+
+const inlined_js = js.replace("export default function", "return function");
+
+const inlined_code = `
+async function fetch_viewer_js() {
+  ${inlined_js}
+}
+
+async function fetch_viewer_wasm() {
+  ${data_url_to_buffer.toString()}
+  const dataUrl = ${JSON.stringify(buffer_to_data_url(wasm))};
+  const buffer = await data_url_to_buffer(dataUrl);
+  return new Response(buffer, { "headers": { "Content-Type": "application/wasm" } });
+}
+`;
+
+// replace INLINE_MARKER, inclusive
+const inline_start = index.indexOf(INLINE_MARKER);
+const inline_end =
+  index.indexOf(INLINE_MARKER, inline_start + 1) + INLINE_MARKER.length;
+
+const bundle =
+  index.substring(0, inline_start) + inlined_code + index.substring(inline_end);
+
+fs.writeFileSync(path.join(__dirname, "inlined.js"), bundle);
+fs.copyFileSync(
+  path.join(__dirname, "index.d.ts"),
+  path.join(__dirname, "inlined.d.ts"),
+);