> ## Documentation Index
> Fetch the complete documentation index at: https://mintlify.com/mercuryworkshop/scramjet/llms.txt
> Use this file to discover all available pages before exploring further.

# WebAssembly integration

> Understanding Scramjet's Rust-powered WASM rewriter for high-performance JavaScript transformation

Scramjet's JavaScript rewriter is implemented in Rust and compiled to WebAssembly for maximum performance. This architecture enables AST-level transformations with near-native speed.

## Architecture overview

The WASM rewriter is a Cargo workspace with multiple components:

```
rewriter/
├── Cargo.toml          # Workspace configuration
├── js/                 # JavaScript rewriter logic
├── wasm/               # WASM bindings
├── transform/          # AST transformations
└── native/             # Native binary for testing
```

### Key dependencies

```toml theme={null}
[workspace.dependencies]
oxc = { version = "0.77.2", features = ["ast_visit"] }

[profile.release]
opt-level = 3
lto = true              # Link-time optimization
codegen-units = 1       # Single codegen unit for better optimization
panic = "abort"         # Smaller binary size
```

<Info>
  Scramjet uses [OXC](https://oxc-project.github.io/), a blazingly fast JavaScript parser written in Rust, achieving performance comparable to native compilers.
</Info>

## WASM module structure

### Initialization

The WASM module is loaded and initialized in `src/shared/rewriters/wasm.ts`:

```typescript theme={null}
import { initSync, Rewriter } from "../../../rewriter/wasm/out/wasm.js";

let wasm_u8: Uint8Array;

// Load WASM binary (embedded or fetched)
if (REWRITERWASM) {
  // Embedded at build time
  wasm_u8 = Uint8Array.from(atob(REWRITERWASM), c => c.charCodeAt(0));
} else if (self.WASM) {
  // Injected by service worker
  wasm_u8 = Uint8Array.from(atob(self.WASM), c => c.charCodeAt(0));
}

function initWasm() {
  // Validate WASM magic bytes
  const MAGIC = "\0asm".split("").map(x => x.charCodeAt(0));
  if (![...wasm_u8.slice(0, 4)].every((x, i) => x === MAGIC[i])) {
    throw new Error("Invalid WASM binary");
  }
  
  initSync({
    module: new WebAssembly.Module(wasm_u8),
  });
}
```

<Warning>
  The WASM module must be loaded before any JavaScript rewriting occurs. Scramjet automatically handles this in service workers via `asyncSetWasm()`.
</Warning>

### Service worker WASM loading

In service workers, WASM is fetched dynamically:

```typescript theme={null}
export async function asyncSetWasm() {
  const buf = await fetch(config.files.wasm).then(r => r.arrayBuffer());
  wasm_u8 = new Uint8Array(buf);
}

// Called during service worker initialization
await scramjet.loadConfig();
// -> setConfig(config);
// -> await asyncSetWasm();
```

## Rewriter class

The WASM `Rewriter` class exposes a JavaScript-friendly API:

```rust theme={null}
// Simplified from rewriter/wasm/src/lib.rs
#[wasm_bindgen]
pub struct Rewriter {
    alloc: Allocator,
    scramjet: Object,
    js: JsRewriter,
}

#[wasm_bindgen]
impl Rewriter {
    #[wasm_bindgen(constructor)]
    pub fn new(scramjet: Object) -> Result<Self> {
        Ok(Self {
            alloc: Allocator::default(),
            js: create_js(&scramjet)?,
            scramjet,
        })
    }
    
    #[wasm_bindgen]
    pub fn rewrite_js(
        &mut self,
        js: String,
        base: String,
        url: String,
        module: bool,
    ) -> Result<JsRewriterOutput> {
        let flags = get_js_flags(&self.scramjet, base, module)?;
        let out = self.js.rewrite(&self.alloc, &js, flags)?;
        
        let ret = create_js_output(out, url, js);
        self.alloc.reset(); // Reuse allocator
        ret
    }
    
    #[wasm_bindgen]
    pub fn rewrite_js_bytes(
        &mut self,
        js: Vec<u8>,
        base: String,
        url: String,
        module: bool,
    ) -> Result<JsRewriterOutput> {
        // SAFETY: Assumes valid UTF-8
        let js = unsafe { String::from_utf8_unchecked(js) };
        self.rewrite_js(js, base, url, module)
    }
}
```

### TypeScript bindings

The Rust code generates TypeScript definitions:

```typescript theme={null}
export class Rewriter {
  constructor(scramjet: object);
  
  rewrite_js(
    js: string,
    base: string,
    url: string,
    module: boolean
  ): JsRewriterOutput;
  
  rewrite_js_bytes(
    js: Uint8Array,
    base: string,
    url: string,
    module: boolean
  ): JsRewriterOutput;
}

export interface JsRewriterOutput {
  js: Uint8Array;
  map: Uint8Array;
  scramtag: string;
  errors: string[];
}
```

## Rewriter configuration

The `Rewriter` constructor accepts a configuration object:

```typescript theme={null}
const rewriter = new Rewriter({
  config: scramjetConfig,
  shared: {
    rewrite: {
      htmlRules,
      rewriteUrl,
      rewriteCss,
      rewriteJs,
      getHtmlInjectCode(cookieStore, foundHead) {
        let inject = getInjectScripts(
          cookieStore,
          src => `<script src="${src}"></script>`
        ).join("");
        return foundHead ? `<head>${inject}</head>` : inject;
      },
    },
  },
  flagEnabled,
  codec: {
    encode: codecEncode,
    decode: codecDecode,
  },
});
```

### Shared functions

The rewriter calls back into JavaScript for certain operations:

* **URL rewriting**: `rewriteUrl(url, meta)`
* **CSS rewriting**: `rewriteCss(css, meta)` (for `<style>` tags)
* **JS rewriting**: `rewriteJs(js, url, meta)` (for nested scripts)
* **HTML injection**: `getHtmlInjectCode(cookieStore, foundHead)`

<Accordion title="Why callbacks instead of pure Rust?">
  While implementing everything in Rust would be faster, JavaScript callbacks provide:

  1. **Flexibility**: Configuration (codec, flags) can be changed without recompiling WASM
  2. **Code sharing**: URL/CSS rewriting logic is shared between client and worker contexts
  3. **Smaller binary**: Avoiding duplicate implementations reduces WASM size
  4. **Maintainability**: Complex logic (HTML parsing) stays in TypeScript where it's easier to debug
</Accordion>

## AST transformation

The Rust rewriter performs AST-level transformations using OXC:

```rust theme={null}
use oxc::{
    allocator::Allocator,
    ast_visit::Visit,
    parser::{Parser, ParseOptions},
    span::SourceType,
};

pub fn rewrite<'alloc>(
    &self,
    alloc: &'alloc Allocator,
    js: &str,
    flags: Flags,
) -> Result<RewriteResult<'alloc>> {
    let source_type = SourceType::unambiguous()
        .with_javascript(true)
        .with_module(flags.is_module)
        .with_standard(true);
    
    let parsed = Parser::new(alloc, js, source_type)
        .with_options(ParseOptions {
            allow_v8_intrinsics: true,
            allow_return_outside_function: true,
            ..Default::default()
        })
        .parse();
    
    if parsed.panicked {
        return Err(RewriterError::OxcPanicked(format_errors(&parsed.errors)));
    }
    
    let mut visitor = Visitor {
        alloc,
        config: &self.cfg,
        rewriter: &self.url,
        flags,
        jschanges,
        error: None,
    };
    
    visitor.visit_program(&parsed.program);
    
    let changed = visitor.jschanges.perform(js, &self.cfg, &visitor.flags)?;
    
    Ok(RewriteResult {
        js: changed.source,
        sourcemap: changed.map,
        errors: parsed.errors,
        flags: visitor.flags,
    })
}
```

### Visitor pattern

The visitor walks the AST and collects transformations:

```rust theme={null}
impl<'a> Visit<'a> for Visitor<'a> {
    fn visit_call_expression(&mut self, expr: &CallExpression<'a>) {
        // Intercept: fetch(url)
        if is_fetch_call(expr) {
            let url_arg = &expr.arguments[0];
            self.rewrite_url_argument(url_arg);
        }
    }
    
    fn visit_new_expression(&mut self, expr: &NewExpression<'a>) {
        // Intercept: new WebSocket(url)
        if is_websocket_constructor(expr) {
            let url_arg = &expr.arguments[0];
            self.rewrite_url_argument(url_arg);
        }
    }
    
    fn visit_member_expression(&mut self, expr: &MemberExpression<'a>) {
        // Intercept: location.href
        if is_location_href(expr) {
            self.wrap_location_access(expr);
        }
    }
}
```

## Building the WASM module

Scramjet includes a build script for the WASM module:

```bash theme={null}
# rewriter/wasm/build.sh
cd rewriter/wasm/
wasm-pack build --target web --out-dir out
cd ../..
```

### Build configuration

For production builds, Scramjet uses aggressive optimizations:

```toml theme={null}
[profile.release]
opt-level = 3        # Maximum optimization
lto = true           # Link-time optimization
codegen-units = 1    # Single codegen unit
panic = "abort"      # Smaller binary (no unwinding)
```

<Info>
  The production WASM binary is approximately 500KB gzipped, which is loaded once and cached by the browser.
</Info>

## Sourcemap generation

The rewriter generates sourcemaps for debugging:

```typescript theme={null}
if (flagEnabled("sourcemaps", meta.base)) {
  const pushmap = globalThis[config.globals.pushsourcemapfn];
  if (pushmap) {
    // Direct function call (faster)
    pushmap(Array.from(res.map), res.tag);
  } else {
    // Inline sourcemap as code
    const sourcemapfn = `${config.globals.pushsourcemapfn}([${res.map.join(",")}], "${res.tag}");`;
    
    // Preserve "use strict" directive
    const strictMode = /^\s*(['"])use strict\1;?/;
    if (strictMode.test(newjs)) {
      newjs = newjs.replace(strictMode, `$&\n${sourcemapfn}`);
    } else {
      newjs = `${sourcemapfn}\n${newjs}`;
    }
  }
}
```

Sourcemaps enable:

* Accurate stack traces in DevTools
* Breakpoint debugging in original code
* Proper error line numbers

## Error handling

### Parse errors

The rewriter collects parse errors without failing:

```typescript theme={null}
const result = rewriteJs(code, url, meta, isModule);

if (flagEnabled("rewriterLogs", meta.base)) {
  for (const error of result.errors) {
    console.error("oxc parse error", error);
  }
}
```

### Panic recovery

If OXC panics, the rewriter catches it:

```rust theme={null}
if parsed.panicked {
    use std::fmt::Write;
    
    let mut errors = String::new();
    for error in parsed.errors {
        writeln!(errors, "{error}")?;
    }
    return Err(RewriterError::OxcPanicked(errors));
}
```

In TypeScript:

```typescript theme={null}
try {
  return rewriteJs(code, url, meta, isModule);
} catch (err) {
  console.warn("Failed rewriting JS:", err.message);
  
  if (flagEnabled("allowInvalidJs", meta.base)) {
    return code; // Return original
  }
  throw err;
}
```

## Performance optimizations

### Rewriter pooling

Scramjet maintains a pool of rewriter instances:

```typescript theme={null}
let rewriters: Array<{ rewriter: Rewriter; inUse: boolean }> = [];

export function getRewriter(meta: URLMeta): [Rewriter, () => void] {
  initWasm();
  
  let obj = rewriters.find(x => !x.inUse);
  
  if (!obj) {
    if (flagEnabled("rewriterLogs", meta.base)) {
      console.log(`Creating new rewriter, ${rewriters.length} already exist`);
    }
    
    const rewriter = new Rewriter({ config, shared, flagEnabled, codec });
    obj = { rewriter, inUse: false };
    rewriters.push(obj);
  }
  
  obj.inUse = true;
  return [obj.rewriter, () => (obj.inUse = false)];
}
```

### Usage

```typescript theme={null}
function rewriteJsWasm(
  input: string | Uint8Array,
  source: string | null,
  meta: URLMeta,
  module: boolean
): RewriterResult {
  let [rewriter, release] = getRewriter(meta);
  
  try {
    const out = typeof input === "string"
      ? rewriter.rewrite_js(input, meta.base.href, source, module)
      : rewriter.rewrite_js_bytes(input, meta.base.href, source, module);
    
    return {
      js: typeof input === "string" ? textDecoder.decode(out.js) : out.js,
      map: out.map,
      tag: out.scramtag,
      errors: out.errors,
    };
  } finally {
    release(); // Return to pool
  }
}
```

<Note>
  Always release the rewriter back to the pool using the `release()` function to avoid memory leaks and rewriter exhaustion.
</Note>

## Testing the rewriter

Scramjet includes a native test runner:

```bash theme={null}
# Build and run native binary
cd rewriter/native
cargo run --release
```

This allows testing the rewriter without compiling to WASM, which is useful for:

* Debugging Rust code
* Running benchmarks
* Integration tests

## Debugging tips

### Enable rewriter logs

```typescript theme={null}
import { flagEnabled } from "@/shared";

if (flagEnabled("rewriterLogs", meta.base)) {
  console.log("Rewriter created");
  console.log("Parse errors:", result.errors);
}
```

### Inspect WASM binary

```typescript theme={null}
function initWasm() {
  const MAGIC = [0x00, 0x61, 0x73, 0x6D]; // "\0asm"
  
  console.log("WASM size:", wasm_u8.length, "bytes");
  console.log("Magic bytes:", [...wasm_u8.slice(0, 4)]);
  
  if (![...wasm_u8.slice(0, 4)].every((x, i) => x === MAGIC[i])) {
    console.error("Invalid WASM:", textDecoder.decode(wasm_u8.slice(0, 100)));
    throw new Error("Invalid WASM binary");
  }
}
```

### Compare input/output

```typescript theme={null}
const original = code;
const rewritten = rewriteJs(code, url, meta, false);

console.log("Original:", original);
console.log("Rewritten:", rewritten);
console.log("Size change:", rewritten.length - original.length, "bytes");
```