use page_size=1 for instrumented components

Makefile

@@ -4,8 +4,8 @@ all: build-components build-cli
 build-cli:
 	cargo build --all-features --release
 build-components:
-	cargo build -p debug --target wasm32-wasip2 --release
-	cargo build -p recorder --target wasm32-wasip2 --release
+	RUSTFLAGS="-C link-arg=--page-size=1 -C link-arg=--adapt=assets/wasi_snapshot_preview1.reactor.one_byte_page.wasm" cargo build -p debug --target wasm32-wasip2 --release
+	RUSTFLAGS="-C link-arg=--page-size=1 -C link-arg=--adapt=assets/wasi_snapshot_preview1.reactor.one_byte_page.wasm" cargo build -p recorder --target wasm32-wasip2 --release
 	cp target/wasm32-wasip2/release/debug.wasm assets/debug.wasm
 	cp target/wasm32-wasip2/release/recorder.wasm assets/recorder.wasm
 
@@ -25,7 +25,7 @@ test-record:
 	$(MAKE) run-record WASM=tests/rust.wasm
 	# test the same trace with a different wasm replay
 	target/release/proxy-component instrument -m replay tests/rust.debug.wasm
-	wasmtime --invoke 'start()' composed.wasm < trace.out
+	wasmtime -W custom-page-sizes=y --invoke 'start()' composed.wasm < trace.out
 	# build-only test
 	target/release/proxy-component instrument -m record tests/calculator.wasm
 	target/release/proxy-component instrument -m replay tests/calculator.wasm
@@ -40,13 +40,13 @@ test-dialog:
 
 run-fuzz:
 	target/release/proxy-component instrument -m fuzz $(WASM)
-	wasmtime --invoke 'start()' composed.wasm
+	wasmtime -W custom-page-sizes=y --invoke 'start()' composed.wasm
 
 run-record:
 	target/release/proxy-component instrument -m record $(WASM)
 	$(MAKE) run-viceroy URL=localhost:7676
 	target/release/proxy-component instrument -m replay $(WASM)
-	wasmtime --invoke 'start()' composed.wasm < trace.out
+	wasmtime -W custom-page-sizes=y --invoke 'start()' composed.wasm < trace.out
 	# test host replay
 	target/release/proxy-component instrument -m replay --use-host-recorder $(WASM)
 	target/release/proxy-component run composed.wasm --invoke 'start()' --trace trace.out

assets/proj_cargo.toml

@@ -11,3 +11,4 @@ path = "lib.rs"
 wit-bindgen.workspace = true
 wasm-wave.workspace = true
 arbitrary.workspace = true
+talc.workspace = true

assets/workspace_cargo.toml

@@ -7,6 +7,7 @@ wit-bindgen = { version = "0.62.0", default-features = false, features = ["bitfl
 wasm-wave = { git = "https://github.com/chenyan2002/wasm-tools.git", branch = "extend-wave", version = "0.239.0", default-features = false }
 #wasm-wave = { path = "/Users/chenyan/src/bytecodealliance/wasm-tools/crates/wasm-wave", default-features = false }
 arbitrary = "1.4.2"
+talc = "5.0.3"
 
 [profile.release]
 codegen-units = 1

components/debug/Cargo.toml

@@ -8,4 +8,5 @@ crate-type = ["cdylib"]
 
 [dependencies]
 wit-bindgen.workspace = true
-getrandom = "0.3"
+getrandom = "0.3"
+talc = "5.0.3"

components/debug/src/lib.rs

@@ -1,3 +1,12 @@
+#[cfg(all(not(target_feature = "atomics"), target_family = "wasm"))]
+#[global_allocator]
+static TALC: talc::wasm::WasmArenaTalc = {
+    use core::mem::MaybeUninit;
+    static mut MEMORY: [MaybeUninit<u8>; 0x80000] = [MaybeUninit::uninit(); 0x80000];
+    // SAFETY: the memory for MEMORY is never modified externally. It's the allocator's.
+    unsafe { talc::wasm::new_wasm_arena_allocator(&raw mut MEMORY) }
+};
+
 mod bindings {
     wit_bindgen::generate!({
         path: "../../assets/util.wit",

components/recorder/Cargo.toml

@@ -10,3 +10,4 @@ crate-type = ["cdylib"]
 trace = { path = "../../crates/trace" }
 serde_json.workspace = true
 wit-bindgen.workspace = true
+talc = "5.0.3"

components/recorder/src/lib.rs

@@ -1,3 +1,12 @@
+#[cfg(all(not(target_feature = "atomics"), target_family = "wasm"))]
+#[global_allocator]
+static TALC: talc::wasm::WasmArenaTalc = {
+    use core::mem::MaybeUninit;
+    static mut MEMORY: [MaybeUninit<u8>; 0x80000] = [MaybeUninit::uninit(); 0x80000];
+    // SAFETY: the memory for MEMORY is never modified externally. It's the allocator's.
+    unsafe { talc::wasm::new_wasm_arena_allocator(&raw mut MEMORY) }
+};
+
 mod bindings {
     wit_bindgen::generate!({
         path: "../../assets/recorder.wit",

src/codegen/mod.rs

@@ -22,6 +22,8 @@ pub struct GenerateArgs {
     /// The path to the output file.
     #[arg(short, long, default_value("lib.rs"))]
     pub output_file: PathBuf,
+    #[arg(long)]
+    pub use_custom_allocator: bool,
 }
 #[derive(clap::ValueEnum, clap::Parser, Clone)]
 pub enum GenerateMode {
@@ -84,7 +86,19 @@ impl GenerateArgs {
         let traits = trait_generator.generate();
         drop(trait_generator);
         state.output.extend(traits);
-        let file = state.into_output_file();
+        let mut file = state.into_output_file();
+        if self.use_custom_allocator {
+            file.items.push(parse_quote! {
+                #[cfg(all(not(target_feature = "atomics"), target_family = "wasm"))]
+                #[global_allocator]
+                static TALC: talc::wasm::WasmArenaTalc = {
+                    use core::mem::MaybeUninit;
+                    static mut MEMORY: [MaybeUninit<u8>; 0x80000] = [MaybeUninit::uninit(); 0x80000];
+                    // SAFETY: the memory for MEMORY is never modified externally. It's the allocator's.
+                    unsafe { talc::wasm::new_wasm_arena_allocator(&raw mut MEMORY) }
+                };
+            });
+        }
         let output = prettyplease::unparse(&file);
         std::fs::write(&self.output_file, output)?;
         Ok(())

src/instrument.rs

@@ -18,10 +18,16 @@ pub struct InstrumentArgs {
     /// Whether to use the host recorder implementation or link the recorder component
     #[arg(long)]
     pub use_host_recorder: bool,
+    /// Whether to use a one-byte page size for the wasm component
+    #[arg(long, default_value_t = true)]
+    pub one_byte_page_size: bool,
 }
 
+// TODO: Add one-byte and regular page size wasm binaries
 const DEBUG_WASM: &[u8] = include_bytes!("../assets/debug.wasm");
 const RECORDER_WASM: &[u8] = include_bytes!("../assets/recorder.wasm");
+const WASI_ADAPTER_ONE_BYTE_PAGE: &[u8] =
+    include_bytes!("../assets/wasi_snapshot_preview1.reactor.one_byte_page.wasm");
 
 pub fn run(args: InstrumentArgs) -> Result<()> {
     if args.use_host_recorder && !matches!(args.mode, Mode::Record | Mode::Replay) {
@@ -57,20 +63,47 @@ pub fn run(args: InstrumentArgs) -> Result<()> {
     }
 
     // 5. Generate Rust binding for both import and export interface
-    bindgen(&tmp_dir, &wit_dir, &args.mode, "imports", "record_imports")?;
-    bindgen(&tmp_dir, &wit_dir, &args.mode, "exports", "record_exports")?;
+    bindgen(
+        &tmp_dir,
+        &wit_dir,
+        &args.mode,
+        "imports",
+        "record_imports",
+        args.one_byte_page_size,
+    )?;
+    bindgen(
+        &tmp_dir,
+        &wit_dir,
+        &args.mode,
+        "exports",
+        "record_exports",
+        args.one_byte_page_size,
+    )?;
     // 6. cargo build
     let mut cmd = Command::new("cargo");
+    if args.one_byte_page_size {
+        cmd.env("RUSTFLAGS", "-C link-arg=--page-size=1");
+    }
     cmd.arg("build")
         .arg("--target=wasm32-unknown-unknown")
         .current_dir(tmp_dir.as_path());
     let status = cmd.status()?;
     assert!(status.success());
 
-    let exports_wasm_path =
-        component_new(&tmp_dir, &wit_dir, "exports", "debug/record_exports.wasm")?;
-    let imports_wasm_path =
-        component_new(&tmp_dir, &wit_dir, "imports", "debug/record_imports.wasm")?;
+    let exports_wasm_path = component_new(
+        &tmp_dir,
+        &wit_dir,
+        "exports",
+        "debug/record_exports.wasm",
+        args.one_byte_page_size,
+    )?;
+    let imports_wasm_path = component_new(
+        &tmp_dir,
+        &wit_dir,
+        "imports",
+        "debug/record_imports.wasm",
+        args.one_byte_page_size,
+    )?;
     // 7. run wac
     opts.generate_wac(&imports_wasm_path, &exports_wasm_path, &wit_dir)?;
     let output_file = "composed.wasm";
@@ -122,6 +155,7 @@ fn bindgen(
     mode: &Mode,
     world_name: &str,
     dest_name: &str,
+    use_custom_allocator: bool,
 ) -> Result<()> {
     let out_dir = tmp_dir.join(dest_name);
     crate::util::generate_bindings(wit_dir, world_name, &out_dir)?;
@@ -136,6 +170,7 @@ fn bindgen(
         bindings: binding_file.clone(),
         output_file: out_dir.join("lib.rs"),
         mode: codegen_mode,
+        use_custom_allocator,
     };
     codegen_opt.generate()?;
     fs::rename(&binding_file, out_dir.join("bindings.rs"))?;
@@ -146,6 +181,7 @@ fn component_new(
     wit_dir: &Path,
     world_name: &str,
     wasm_file: &str,
+    one_byte_page_size: bool,
 ) -> Result<PathBuf> {
     let wasm_path = tmp_dir
         .join("target/wasm32-unknown-unknown/")
@@ -161,8 +197,12 @@ fn component_new(
         wit_component::StringEncoding::UTF8,
     )?;
     // create component from the embedded module
-    let component = ComponentEncoder::default()
-        .module(&wasm)?
+    let mut encoder = ComponentEncoder::default();
+    encoder.module(&wasm)?;
+    if one_byte_page_size {
+        encoder.adapter("wasi_snapshot_preview1", WASI_ADAPTER_ONE_BYTE_PAGE)?;
+    }
+    let component = encoder
         .encode()
         .context("failed to encode a component from module")?;
     fs::write(&wasm_path, component)?;

src/run.rs

@@ -74,6 +74,7 @@ pub fn run(args: RunArgs) -> anyhow::Result<()> {
     let mut config = Config::new();
     config
         .consume_fuel(true)
+        .wasm_custom_page_sizes(true)
         //.debug_info(true)
         .wasm_backtrace_details(WasmBacktraceDetails::Enable);
     let engine = Engine::new(&config)?;