feat: enhance CI workflow to include Rust setup and tests, update rewriter logic for improved handling of export declarations

Author: lemon-mint

.github/workflows/ci.yml

@@ -15,7 +15,7 @@ concurrency:
 
 jobs:
   test:
-    name: Go, JavaScript, Puppeteer, and build checks
+    name: Go, Rust, JavaScript, Puppeteer, and build checks
     runs-on: ubuntu-24.04
     timeout-minutes: 30
 
@@ -28,22 +28,34 @@ jobs:
         with:
           go-version-file: go.mod
           cache: true
+      - name: Set up Rust
+        uses: dtolnay/rust-toolchain@stable
+        with:
+          targets: wasm32-unknown-unknown
+
 
       - name: Set up Node.js
         uses: actions/setup-node@v4
         with:
           node-version: lts/*
           cache: npm
-
       - name: Print toolchain versions
         run: |
           go version
+          rustc --version
+          cargo --version
           node --version
           npm --version
 
+      - name: Install wasm-bindgen CLI
+        run: cargo install wasm-bindgen-cli --version 0.2.122 --locked
+
       - name: Install Node dependencies
         run: npm ci
 
+      - name: Run Rust rewriter tests
+        run: cargo test --manifest-path rewriter-rs/Cargo.toml
+
       - name: Run Go tests
         run: go test ./...
 

rewriter-rs/Cargo.toml

@@ -5,7 +5,7 @@ edition = "2021"
 publish = false
 
 [lib]
-crate-type = ["cdylib"]
+crate-type = ["cdylib", "rlib"]
 
 [dependencies]
 oxc_allocator = "0.60"

rewriter-rs/src/lib.rs

@@ -59,7 +59,6 @@ struct Replacement {
 #[derive(Clone, Copy, PartialEq, Eq)]
 enum ScopeMode {
     FunctionRoot,
-    Function,
     Block,
 }
 
@@ -216,12 +215,41 @@ impl<'a> Rewriter<'a> {
             Statement::ImportDeclaration(decl) => self.add_replacement(decl.source.span, format!("{:?}", self.module_specifier(decl.source.value.as_str())), 95),
             Statement::ExportNamedDeclaration(decl) => {
                 if let Some(source) = &decl.source { self.add_replacement(source.span, format!("{:?}", self.module_specifier(source.value.as_str())), 95); }
+                if let Some(inner) = &decl.declaration { self.walk_declaration(inner); }
             }
             Statement::ExportAllDeclaration(decl) => self.add_replacement(decl.source.span, format!("{:?}", self.module_specifier(decl.source.value.as_str())), 95),
+            Statement::ExportDefaultDeclaration(decl) => self.walk_export_default(decl),
             _ => {}
         }
     }
-
+    fn walk_declaration(&mut self, decl: &Declaration<'a>) {
+        match decl {
+            Declaration::VariableDeclaration(decl) => self.walk_variable_declaration(decl),
+            Declaration::FunctionDeclaration(func) => self.walk_function(func),
+            Declaration::ClassDeclaration(class) => {
+                if let Some(id) = &class.id { self.declare(id.name.as_str()); }
+                for elem in &class.body.body {
+                    match elem {
+                        ClassElement::PropertyDefinition(prop) => if let Some(value) = &prop.value { self.walk_expression(value); },
+                        ClassElement::AccessorProperty(prop) => if let Some(value) = &prop.value { self.walk_expression(value); },
+                        _ => {}
+                    }
+                }
+            }
+            _ => {}
+        }
+    }
+    fn walk_export_default(&mut self, decl: &ExportDefaultDeclaration<'a>) {
+        match &decl.declaration {
+            ExportDefaultDeclarationKind::FunctionDeclaration(func) => self.walk_function(func),
+            ExportDefaultDeclarationKind::ClassDeclaration(class) => {
+                if let Some(id) = &class.id { self.declare(id.name.as_str()); }
+            }
+            other => {
+                if let Some(expr) = other.as_expression() { self.walk_expression(expr); }
+            }
+        }
+    }
     fn walk_function(&mut self, func: &Function<'a>) {
         if let Some(id) = &func.id { self.declare(id.name.as_str()); }
         let mut scope = HashSet::new();
@@ -614,34 +642,34 @@ impl<'a> Rewriter<'a> {
             Statement::FunctionDeclaration(func) => { if let Some(id) = &func.id { names.insert(id.name.to_string()); } }
             Statement::ClassDeclaration(class) => { if let Some(id) = &class.id { names.insert(id.name.to_string()); } }
             Statement::VariableDeclaration(decl) => {
-                if mode == ScopeMode::FunctionRoot {
-                    if decl.kind == VariableDeclarationKind::Var { for d in &decl.declarations { self.collect_binding_pattern(&d.id, names); } }
-                } else if decl.kind != VariableDeclarationKind::Var {
+                if mode == ScopeMode::Block {
+                    if decl.kind != VariableDeclarationKind::Var { for d in &decl.declarations { self.collect_binding_pattern(&d.id, names); } }
+                } else if decl.kind == VariableDeclarationKind::Var {
                     for d in &decl.declarations { self.collect_binding_pattern(&d.id, names); }
                 }
             }
-            Statement::BlockStatement(block) if mode == ScopeMode::FunctionRoot => for stmt in &block.body { self.collect_statement_bindings(stmt, ScopeMode::Function, names); },
-            Statement::IfStatement(stmt) if mode == ScopeMode::FunctionRoot => { self.collect_statement_bindings(&stmt.consequent, ScopeMode::Function, names); if let Some(alt) = &stmt.alternate { self.collect_statement_bindings(alt, ScopeMode::Function, names); } }
-            Statement::ForStatement(stmt) if mode == ScopeMode::FunctionRoot => {
+            Statement::BlockStatement(block) if mode != ScopeMode::Block => for stmt in &block.body { self.collect_statement_bindings(stmt, mode, names); },
+            Statement::IfStatement(stmt) if mode != ScopeMode::Block => { self.collect_statement_bindings(&stmt.consequent, mode, names); if let Some(alt) = &stmt.alternate { self.collect_statement_bindings(alt, mode, names); } }
+            Statement::ForStatement(stmt) if mode != ScopeMode::Block => {
                 if let Some(ForStatementInit::VariableDeclaration(decl)) = &stmt.init { if decl.kind == VariableDeclarationKind::Var { for d in &decl.declarations { self.collect_binding_pattern(&d.id, names); } } }
-                self.collect_statement_bindings(&stmt.body, ScopeMode::Function, names);
-            }
-            Statement::ForInStatement(stmt) if mode == ScopeMode::FunctionRoot => { if let ForStatementLeft::VariableDeclaration(decl) = &stmt.left { if decl.kind == VariableDeclarationKind::Var { for d in &decl.declarations { self.collect_binding_pattern(&d.id, names); } } } self.collect_statement_bindings(&stmt.body, ScopeMode::Function, names); }
-            Statement::ForOfStatement(stmt) if mode == ScopeMode::FunctionRoot => { if let ForStatementLeft::VariableDeclaration(decl) = &stmt.left { if decl.kind == VariableDeclarationKind::Var { for d in &decl.declarations { self.collect_binding_pattern(&d.id, names); } } } self.collect_statement_bindings(&stmt.body, ScopeMode::Function, names); }
-            Statement::WhileStatement(stmt) if mode == ScopeMode::FunctionRoot => self.collect_statement_bindings(&stmt.body, ScopeMode::Function, names),
-            Statement::DoWhileStatement(stmt) if mode == ScopeMode::FunctionRoot => self.collect_statement_bindings(&stmt.body, ScopeMode::Function, names),
-            Statement::LabeledStatement(stmt) if mode == ScopeMode::FunctionRoot => self.collect_statement_bindings(&stmt.body, ScopeMode::Function, names),
-            Statement::SwitchStatement(stmt) if mode == ScopeMode::FunctionRoot => for case in &stmt.cases { for child in &case.consequent { self.collect_statement_bindings(child, ScopeMode::Function, names); } },
-            Statement::TryStatement(stmt) if mode == ScopeMode::FunctionRoot => {
-                self.collect_block_bindings(&stmt.block, names);
-                if let Some(handler) = &stmt.handler { self.collect_block_bindings(&handler.body, names); }
-                if let Some(finalizer) = &stmt.finalizer { self.collect_block_bindings(finalizer, names); }
+                self.collect_statement_bindings(&stmt.body, mode, names);
+            }
+            Statement::ForInStatement(stmt) if mode != ScopeMode::Block => { if let ForStatementLeft::VariableDeclaration(decl) = &stmt.left { if decl.kind == VariableDeclarationKind::Var { for d in &decl.declarations { self.collect_binding_pattern(&d.id, names); } } } self.collect_statement_bindings(&stmt.body, mode, names); }
+            Statement::ForOfStatement(stmt) if mode != ScopeMode::Block => { if let ForStatementLeft::VariableDeclaration(decl) = &stmt.left { if decl.kind == VariableDeclarationKind::Var { for d in &decl.declarations { self.collect_binding_pattern(&d.id, names); } } } self.collect_statement_bindings(&stmt.body, mode, names); }
+            Statement::WhileStatement(stmt) if mode != ScopeMode::Block => self.collect_statement_bindings(&stmt.body, mode, names),
+            Statement::DoWhileStatement(stmt) if mode != ScopeMode::Block => self.collect_statement_bindings(&stmt.body, mode, names),
+            Statement::LabeledStatement(stmt) if mode != ScopeMode::Block => self.collect_statement_bindings(&stmt.body, mode, names),
+            Statement::SwitchStatement(stmt) if mode != ScopeMode::Block => for case in &stmt.cases { for child in &case.consequent { self.collect_statement_bindings(child, mode, names); } },
+            Statement::TryStatement(stmt) if mode != ScopeMode::Block => {
+                self.collect_block_bindings(&stmt.block, names, mode);
+                if let Some(handler) = &stmt.handler { self.collect_block_bindings(&handler.body, names, mode); }
+                if let Some(finalizer) = &stmt.finalizer { self.collect_block_bindings(finalizer, names, mode); }
             }
             _ => {}
         }
     }
-    fn collect_block_bindings(&self, block: &BlockStatement<'a>, names: &mut HashSet<String>) {
-        for stmt in &block.body { self.collect_statement_bindings(stmt, ScopeMode::Function, names); }
+    fn collect_block_bindings(&self, block: &BlockStatement<'a>, names: &mut HashSet<String>, mode: ScopeMode) {
+        for stmt in &block.body { self.collect_statement_bindings(stmt, mode, names); }
     }
 
     fn collect_binding_pattern(&self, pattern: &BindingPattern<'a>, names: &mut HashSet<String>) {
@@ -801,6 +829,9 @@ fn join_url(base: &str, raw: &str) -> String {
     }
     let prefix = match base.rfind('/') { Some(i) => &base[..=i], None => base };
     let mut parts: Vec<&str> = prefix.split('/').collect();
+    if parts.last() == Some(&"") {
+        parts.pop();
+    }
     for part in raw.split('/') {
         match part {
             "." => {}
@@ -832,3 +863,58 @@ fn hex(v: u8) -> char {
         _ => (b'A' + (v - 10)) as char,
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn rewrite_ok(source: &str, kind: &str, target_url: &str) -> String {
+        let out = rewrite_script(source, kind, target_url, "/zp/");
+        assert!(out.ok, "rewrite failed: {}", out.error);
+        out.code
+    }
+
+    #[test]
+    fn rewrites_virtualized_globals_and_shadowing() {
+        let code = rewrite_ok(
+            "function f(x) { if (x) { var location = { href: 'local' }; } return location.href; }\nwindow.location.hash += '-tail';\ndocument.defaultView.location.href;",
+            "classic",
+            "https://example.com/app.js",
+        );
+        assert!(code.contains("return location.href;"));
+        assert!(code.contains("__zp_assign(__zp_get(__zp_get(globalThis,\"window\"),\"location\"),\"hash\""));
+        assert!(code.contains("__zp_get(__zp_get(globalThis,\"document\"),\"defaultView\")"));
+        assert!(!code.contains("return __zp_get(globalThis,\"location\")"));
+    }
+
+    #[test]
+    fn rewrites_module_urls_and_dynamic_imports() {
+        let code = rewrite_ok(
+            "import './dep.js'; export async function load(name) { await import('./chunks/' + name + '.js'); return new URL('/worker-fixture.js', import.meta.url).href; }",
+            "module",
+            "https://example.com/assets/main.js",
+        );
+        assert!(code.contains("import \"/zp/api/script?kind=module&u=https%3A%2F%2Fexample.com%2Fassets%2Fdep.js\";"));
+        assert!(code.contains("__zp_module_url('./chunks/' + name + '.js',\"https://example.com/assets/main.js\")"));
+        assert!(code.contains("\"https://example.com/assets/main.js\""));
+    }
+
+    #[test]
+    fn rewrites_calls_and_constructors() {
+        let code = rewrite_ok(
+            "window.location = '/next'; const ws = new WebSocket('/ws', ['chat']); Object.getOwnPropertyDescriptor(window, 'location');",
+            "classic",
+            "https://example.com/app.js",
+        );
+        assert!(code.contains("__zp_set(__zp_get(globalThis,\"window\"),\"location\",'/next')"));
+        assert!(code.contains("__zp_construct(__zp_get(globalThis,\"WebSocket\"),['/ws',['chat']])"));
+        assert!(code.contains("__zp_call(Object,\"getOwnPropertyDescriptor\",[__zp_get(globalThis,\"window\"),'location'])"));
+    }
+
+    #[test]
+    fn parse_failures_return_error() {
+        let out = rewrite_script("if (", "classic", "https://example.com/app.js", "/zp/");
+        assert!(!out.ok);
+        assert_eq!(out.error, "PARSE_FAILED");
+    }
+}

test/js/rewriter.test.js

@@ -7,11 +7,34 @@ const path = require('node:path');
 const vm = require('node:vm');
 const oxc = require('@oxc-parser/wasm');
 
+let builtRustAssetPromise = null;
+
 async function loadRewriter() {
   if (!globalThis.ZPRewriter) vm.runInThisContext(fs.readFileSync('web/js-rewriter.js', 'utf8'), { filename: 'web/js-rewriter.js' });
   await globalThis.ZPRewriter.init({ parser: oxc });
   return globalThis.ZPRewriter;
 }
+
+function loadBuiltRustContext() {
+  if (!builtRustAssetPromise) {
+    builtRustAssetPromise = (async () => {
+      const outDir = fs.mkdtempSync(path.join(os.tmpdir(), 'zp-rust-unit-'));
+      const result = childProcess.spawnSync('node', ['scripts/build.mjs', '--web-only', '--out', outDir], {
+        cwd: path.resolve(__dirname, '../..'),
+        encoding: 'utf8',
+      });
+      if (result.status !== 0) {
+        throw new Error(`node scripts/build.mjs --web-only --out ${outDir} failed\nstdout:\n${result.stdout}\nstderr:\n${result.stderr}`);
+      }
+      const ctx = { console, atob, btoa, TextEncoder, TextDecoder, Uint8Array, WebAssembly, FinalizationRegistry, URL, globalThis: null };
+      ctx.globalThis = ctx;
+      vm.createContext(ctx);
+      vm.runInContext(fs.readFileSync(path.join(outDir, 'web', 'rust-rewriter.js'), 'utf8'), ctx, { filename: 'rust-rewriter.js' });
+      return ctx;
+    })();
+  }
+  return builtRustAssetPromise;
+}
 test('rewriter prefers Rust engine when available', async () => {
   vm.runInThisContext(fs.readFileSync('web/js-rewriter.js', 'utf8'), { filename: 'web/js-rewriter.js' });
   globalThis.ZPRustRewriter = { rewriteScript(source, kind, targetUrl, controlPrefix) { return { ok: true, code: `/*rust:${kind}:${targetUrl}:${controlPrefix}*/`, error: '' }; } };
@@ -20,19 +43,46 @@ test('rewriter prefers Rust engine when available', async () => {
   assert.equal(out.ok, true);
   assert.equal(out.code, '/*rust:classic:https://example.com/app.js:/zp/*/');
 });
-test('built Rust rewriter asset rewrites live code paths', () => {
-  const outDir = fs.mkdtempSync(path.join(os.tmpdir(), 'zp-rust-unit-'));
-  childProcess.execFileSync('node', ['scripts/build.mjs', '--web-only', '--out', outDir], { cwd: path.resolve(__dirname, '../..'), stdio: 'ignore' });
-  const ctx = { console, atob, btoa, TextEncoder, TextDecoder, Uint8Array, WebAssembly, FinalizationRegistry, URL, globalThis: null };
-  ctx.globalThis = ctx;
-  vm.createContext(ctx);
-  vm.runInContext(fs.readFileSync(path.join(outDir, 'web', 'rust-rewriter.js'), 'utf8'), ctx, { filename: 'rust-rewriter.js' });
-  const assign = ctx.ZPRustRewriter.rewriteScript('window.location.hash += \"-tail\";', 'classic', 'https://example.com/app.js', '/zp/');
-  const meta = ctx.ZPRustRewriter.rewriteScript('new URL(\"/worker-fixture.js\", import.meta.url).href;', 'module', 'https://example.com/module-worker.js', '/zp/');
+test('built Rust rewriter asset rewrites live code paths', async () => {
+  const ctx = await loadBuiltRustContext();
+  const assign = ctx.ZPRustRewriter.rewriteScript('window.location.hash += "-tail"; document.defaultView.location.href;', 'classic', 'https://example.com/app.js', '/zp/');
+  const meta = ctx.ZPRustRewriter.rewriteScript('new URL("/worker-fixture.js", import.meta.url).href;', 'module', 'https://example.com/module-worker.js', '/zp/');
+  const dynamic = ctx.ZPRustRewriter.rewriteScript('export async function load(name) { return import("./chunks/" + name + ".js"); }', 'module', 'https://example.com/assets/main.js', '/zp/');
   assert.equal(assign.ok, true);
-  assert.ok(assign.code.includes('__zp_assign(__zp_get(__zp_get(globalThis,\"window\"),\"location\"),\"hash\"'));
+  assert.ok(assign.code.includes('__zp_assign(__zp_get(__zp_get(globalThis,"window"),"location"),"hash"'));
+  assert.ok(assign.code.includes('__zp_get(__zp_get(globalThis,"document"),"defaultView")'));
   assert.equal(meta.ok, true);
-  assert.ok(meta.code.includes('\"https://example.com/module-worker.js\"'));
+  assert.ok(meta.code.includes('https://example.com/module-worker.js'));
+  assert.equal(dynamic.ok, true);
+  assert.ok(dynamic.code.includes('__zp_module_url('));
+  assert.ok(dynamic.code.includes('https://example.com/assets/main.js'));
+});
+
+test('built Rust rewriter asset reports parse failures', async () => {
+  const ctx = await loadBuiltRustContext();
+  const out = ctx.ZPRustRewriter.rewriteScript('if (', 'classic', 'https://example.com/app.js', '/zp/');
+  assert.equal(out.ok, false);
+  assert.equal(out.error, 'PARSE_FAILED');
+});
+
+test('JS wrapper falls back to JS engine for import-map and event-handler paths', async () => {
+  const ctx = await loadBuiltRustContext();
+  globalThis.ZPRustRewriter = ctx.ZPRustRewriter;
+  const rewriter = await loadRewriter();
+  const handler = rewriter.rewriteScript('return location.href', { kind: 'event-handler', targetUrl: 'https://example.com/' });
+  assert.equal(handler.ok, true);
+  assert.match(handler.code, /__zp_runEvent/);
+  const fnBody = rewriter.rewriteScript('return location.href;', { kind: 'function', targetUrl: 'https://example.com/' });
+  assert.equal(fnBody.ok, true);
+  assert.ok(fnBody.code.includes('__zp_get(globalThis,"location")'));
+  const importMap = rewriter.rewriteScript('import React from \"react\";', {
+    kind: 'module',
+    targetUrl: 'https://example.com/app.js',
+    importMap: { imports: { react: 'https://cdn.example/react.js' } },
+  });
+  delete globalThis.ZPRustRewriter;
+  assert.equal(importMap.ok, true);
+  assert.ok(importMap.code.includes('https%3A%2F%2Fcdn.example%2Freact.js'));
 });
 
 test('OXC rewriter virtualizes dangerous globals without rewriting local bindings', async () => {

web/js-rewriter.js

@@ -57,7 +57,8 @@
     source = String(source || '');
     const kind = normalizeKind(options.scriptKind || options.kind);
     const rust = globalThis.ZPRustRewriter;
-    if (rust && typeof rust.rewriteScript === 'function') {
+    const rustEligible = kind !== 'event-handler' && kind !== 'function' && !options.importMap;
+    if (rustEligible && rust && typeof rust.rewriteScript === 'function') {
       try {
         const out = rust.rewriteScript(source, kind, options.url || options.targetUrl || '', options.controlPrefix || globalThis.ZP && globalThis.ZP.CONTROL_PREFIX || '/zp/');
         if (out && out.ok && typeof out.code === 'string') return ok(out.code, []);
@@ -69,7 +70,7 @@
     if (kind === 'function') return rewriteFunctionBody(source, options);
     const parsed = parse(source, kind === 'module' ? 'module' : 'script', options.url || options.targetUrl || 'target.js');
     if (!parsed.ok) return parsed;
-    const rewritten = rewriteProgram(source, parsed.program, { module: kind === 'module', targetUrl: options.url || options.targetUrl || '' });
+    const rewritten = rewriteProgram(source, parsed.program, { module: kind === 'module', targetUrl: options.url || options.targetUrl || '', importMap: options.importMap, controlPrefix: options.controlPrefix });
     return ok(rewritten.code, parsed.diagnostics.concat(rewritten.diagnostics));
   }