Using mmap to Read Files and Write to Files Directly (#4)

Author: Spxg

TODO.md

@@ -9,6 +9,7 @@
 
 # High Priority
 
+- E2E testing for the linker. Should also ensure deterministic output.
 - Add a `disable_hygiene` attribute to `#[js_sys]` to reduce the compile-time of the output to an
   absolute minimum. This can avoid all `interpolate`s.
 - Escape namespaces and function names if they are not valid JS identifiers.
@@ -33,6 +34,7 @@
   - Store them next to the output file.
   - Pass an environment variable from a `build.rs` pointing to the target folder and go from there.
     This seems to have failed. No build script instruction can reach the linker on Wasm.
+- Memory-mapped file reading should lock files to make it safe.
 
 # Medium Priority
 
@@ -41,8 +43,6 @@
 - Provide an absolutely minimal allocator.
 - The `js_sys` proc-macro should remove the `extern "C" { ... }` part of the input on error to avoid
   triggering the `unsafe` requirement downstream.
-- Optimize linker file interactions by using memory mapped files instead of reading and writing
-  everything into memory.
 
 [Emscripten's]:
 	https://github.com/emscripten-core/emscripten/blob/28bcb86466a273859b8adb43cb167b97e05e145d/src/lib/libstrings.js
@@ -85,6 +85,7 @@ This is a list of upstream issues that could make our lives significantly easier
 - `TextDe/Encoder` could support `SharedArrayBuffer`s:
   - [Chrome Bug](https://issues.chromium.org/issues/40102463)
   - [Firefox Bug](https://bugzilla.mozilla.org/show_bug.cgi?id=1561594)
+- `wasm-encoder` `io::Write` support: [bytecodealliance/wasm-tools#778]
 
 [llvm/llvm-project#136594]: https://github.com/llvm/llvm-project/issues/136594
 [rust-lang/rust#136382]: https://github.com/rust-lang/rust/issues/136382
@@ -100,3 +101,4 @@ This is a list of upstream issues that could make our lives significantly easier
 [rust-lang/rfcs#3834]: https://github.com/rust-lang/rfcs/pull/3834
 [rust-lang/rust#136096]: https://github.com/rust-lang/rust/issues/136096
 [rust-lang/rust#133508]: https://github.com/rust-lang/rust/issues/133508
+[bytecodealliance/wasm-tools#778]: https://github.com/bytecodealliance/wasm-tools/issues/778

host/Cargo.toml

@@ -14,6 +14,7 @@ indoc = "2"
 itertools = { version = "0.14", default-features = false }
 js-bindgen-ld-shared = { path = "ld-shared" }
 js-bindgen-macro-shared = { path = "macro-shared" }
+memmap2 = "0.9"
 object = { version = "0.38", default-features = false, features = ["archive", "read_core", "std"] }
 prettyplease = "0.2"
 proc-macro2 = { version = "1", default-features = false }
@@ -27,6 +28,8 @@ wasm-encoder = { version = "0.244", default-features = false, features = ["wasmp
 wasmparser = { version = "0.244", default-features = false }
 
 [patch.crates-io]
+# Detect unsupported platforms at run-time: https://github.com/RazrFalcon/memmap2-rs/issues/160.
+memmap2 = { git = "https://github.com/daxpedda/memmap2-rs", branch = "stable-unsupported" }
 # Proc-macro bug: https://github.com/oli-obk/ui_test/issues/361.
 ui_test = { git = "https://github.com/daxpedda/ui_test", branch = "js-bindgen" }
 

host/js-sys-macro/src/tests/mod.rs

@@ -1,7 +1,7 @@
 mod function;
 
 use std::ffi::OsStr;
-use std::io::Cursor;
+use std::io::{self, Cursor};
 use std::path::Path;
 use std::process::Command;
 use std::{env, fs};
@@ -182,6 +182,7 @@ fn inner(tmp: &Path, source: &str) -> Result<(String, String)> {
 								js_bindgen_ld_shared::assembly_to_object(
 									OsStr::new("wasm32"),
 									assembly,
+									&mut io::sink(),
 								)?;
 							}
 							Payload::CustomSection(c)

host/ld-shared/Cargo.toml

@@ -4,6 +4,7 @@ name = "js-bindgen-ld-shared"
 rust-version = "1.91"
 
 [dependencies]
+memmap2 = { workspace = true }
 object = { workspace = true }
 wasmparser = { workspace = true }
 

host/ld-shared/src/lib.rs

@@ -1,16 +1,22 @@
 use std::ffi::OsStr;
 use std::fmt::{self, Debug, Formatter};
-use std::fs;
-use std::io::{Error, Write};
+use std::fs::File;
+use std::io::{self, Error, ErrorKind, Read, Write};
+use std::ops::Deref;
 use std::path::Path;
 use std::process::{Command, Stdio};
 
+use memmap2::Mmap;
 use object::read::archive::ArchiveFile;
 use wasmparser::CustomSectionReader;
 
 /// Currently this simply passes the LLVM s-format assembly to `llvm-mc` to
 /// convert to an object file the linker can consume.
-pub fn assembly_to_object(arch_str: &OsStr, assembly: &[u8]) -> Result<Vec<u8>, Error> {
+pub fn assembly_to_object(
+	arch_str: &OsStr,
+	assembly: &[u8],
+	output: &mut dyn Write,
+) -> Result<(), Error> {
 	let mut child = Command::new("llvm-mc")
 		.arg(format!("-arch={}", arch_str.display()))
 		// In the future we will switch to something supporting auto-detection.
@@ -27,41 +33,54 @@ pub fn assembly_to_object(arch_str: &OsStr, assembly: &[u8]) -> Result<Vec<u8>,
 		.ok_or_else(|| Error::other("`llvm-mc` process should have `stdin`"))?;
 	stdin.write_all(assembly)?;
 
-	let output = child.wait_with_output()?;
+	let status = child.wait()?;
 
-	if output.status.success() {
-		Ok(output.stdout)
+	let mut child_stdout = child
+		.stdout
+		.ok_or_else(|| Error::other("`llvm-mc` process should have `stdout`"))?;
+
+	if status.success() {
+		io::copy(&mut child_stdout, output)?;
+		Ok(())
 	} else {
 		eprintln!(
 			"------ llvm-mc input -------\n{}",
 			String::from_utf8_lossy(assembly)
 		);
 
-		if !output.stdout.is_empty() {
+		let mut stdout = Vec::new();
+		child_stdout.read_to_end(&mut stdout)?;
+
+		if !stdout.is_empty() {
 			eprintln!(
 				"------ llvm-mc stdout ------\n{}",
-				String::from_utf8_lossy(&output.stdout)
+				String::from_utf8_lossy(&stdout)
 			);
 
-			if !output.stdout.ends_with(b"\n") {
+			if !stdout.ends_with(b"\n") {
 				eprintln!();
 			}
 		}
 
-		if !output.stderr.is_empty() {
+		let mut stderr = Vec::new();
+		child
+			.stderr
+			.ok_or_else(|| Error::other("`llvm-mc` process should have `stderr`"))?
+			.read_to_end(&mut stderr)?;
+
+		if !stderr.is_empty() {
 			eprintln!(
 				"------ llvm-mc stderr ------\n{}",
-				String::from_utf8_lossy(&output.stderr)
+				String::from_utf8_lossy(&stderr)
 			);
 
-			if !output.stderr.ends_with(b"\n") {
+			if !stderr.ends_with(b"\n") {
 				eprintln!();
 			}
 		}
 
 		Err(Error::other(format!(
-			"`llvm-mc` process failed with status: {}",
-			output.status
+			"`llvm-mc` process failed with status: {status}"
 		)))
 	}
 }
@@ -73,7 +92,7 @@ pub fn ld_input_parser<E>(
 	// We found a UNIX archive.
 	if input.as_encoded_bytes().ends_with(b".rlib") {
 		let archive_path = Path::new(&input);
-		let archive_data = match fs::read(archive_path) {
+		let archive_data = match ReadFile::new(archive_path) {
 			Ok(archive_data) => archive_data,
 			Err(error) => {
 				eprintln!(
@@ -130,7 +149,7 @@ pub fn ld_input_parser<E>(
 		}
 	} else if input.as_encoded_bytes().ends_with(b".o") {
 		let object_path = Path::new(&input);
-		let object = match fs::read(object_path) {
+		let object = match ReadFile::new(object_path) {
 			Ok(object) => object,
 			Err(error) => {
 				eprintln!(
@@ -147,6 +166,42 @@ pub fn ld_input_parser<E>(
 	Ok(())
 }
 
+pub struct ReadFile(ReadInner);
+
+enum ReadInner {
+	Mmap(Mmap),
+	File(Vec<u8>),
+}
+
+impl ReadFile {
+	pub fn new(path: &Path) -> Result<Self, Error> {
+		let mut file = File::open(path)?;
+		// SAFETY: the file is not mutated while the mapping is in use.
+		let result = unsafe { Mmap::map(&file) };
+
+		match result {
+			Ok(mmap) => Ok(Self(ReadInner::Mmap(mmap))),
+			Err(error) if matches!(error.kind(), ErrorKind::Unsupported) => {
+				let mut output = Vec::new();
+				file.read_to_end(&mut output)?;
+				Ok(Self(ReadInner::File(output)))
+			}
+			Err(error) => Err(error),
+		}
+	}
+}
+
+impl Deref for ReadFile {
+	type Target = [u8];
+
+	fn deref(&self) -> &Self::Target {
+		match &self.0 {
+			ReadInner::Mmap(mmap) => mmap.deref(),
+			ReadInner::File(data) => data.as_slice(),
+		}
+	}
+}
+
 #[derive(Clone)]
 pub struct CustomSectionParser<'cs> {
 	name: &'cs str,

host/ld/src/main.rs

@@ -3,14 +3,15 @@ mod wasm_ld;
 use std::borrow::Cow;
 use std::convert::Infallible;
 use std::ffi::{OsStr, OsString};
-use std::fmt::Write;
+use std::fs::File;
+use std::io::{BufWriter, Write};
 use std::path::Path;
 use std::process::{self, Command};
 use std::{env, fs};
 
 use hashbrown::{HashMap, HashSet};
 use itertools::{Itertools, Position};
-use js_bindgen_ld_shared::CustomSectionParser;
+use js_bindgen_ld_shared::{CustomSectionParser, ReadFile};
 use wasm_encoder::{EntityType, ImportSection, Module, RawSection, Section};
 use wasmparser::{Encoding, Parser, Payload, TypeRef};
 
@@ -148,11 +149,14 @@ fn process_object(
 				// ensures freshness:
 				// https://doc.rust-lang.org/1.92.0/nightly-rustc/cargo/core/compiler/fingerprint/index.html#fingerprints-and-unithashs
 				if !asm_path.exists() {
-					let asm_object = js_bindgen_ld_shared::assembly_to_object(arch_str, assembly)
+					let mut asm_file = BufWriter::new(
+						File::create(&asm_path).expect("output assembly object should be writable"),
+					);
+
+					js_bindgen_ld_shared::assembly_to_object(arch_str, assembly, &mut asm_file)
 						.expect("compiling assembly should be valid");
 
-					fs::write(&asm_path, asm_object)
-						.expect("writing assembly object file should succeed");
+					asm_file.into_inner().unwrap().sync_all().unwrap();
 				}
 
 				add_args.push(asm_path.into());
@@ -166,7 +170,7 @@ fn post_processing(output_path: &Path, main_memory: MainMemory<'_>) {
 	// Unfortunately we don't receive the final output path adjustments Cargo makes.
 	// So for the JS file we just figure it out ourselves.
 	let package = env::var_os("CARGO_CRATE_NAME").expect("`CARGO_CRATE_NAME` should be present");
-	let wasm_input = fs::read(output_path).expect("output file should be readable");
+	let wasm_input = ReadFile::new(output_path).expect("output file should be readable");
 	let mut wasm_output = Vec::new();
 
 	let mut found_import: HashMap<&str, HashMap<&str, &str>> = HashMap::new();
@@ -396,12 +400,15 @@ fn post_processing(output_path: &Path, main_memory: MainMemory<'_>) {
 		"missing JS embed: {expected_embed:?}"
 	);
 
-	fs::write(output_path, wasm_output).expect("output Wasm file should be writable");
-
-	let mut js_output = String::new();
+	let mut js_output = BufWriter::new(
+		File::create(output_path.with_file_name(package).with_extension("js"))
+			.expect("output JS file should be writable"),
+	);
 
 	// Create our `WebAssembly.Memory`.
-	js_output.push_str("const memory = new WebAssembly.Memory({ ");
+	js_output
+		.write_all(b"const memory = new WebAssembly.Memory({ ")
+		.unwrap();
 
 	if memory.memory64 {
 		write!(js_output, "initial: {}n", memory.initial).unwrap();
@@ -418,18 +425,18 @@ fn post_processing(output_path: &Path, main_memory: MainMemory<'_>) {
 	}
 
 	if memory.memory64 {
-		js_output.push_str(", address: 'i64'");
+		js_output.write_all(b", address: 'i64'").unwrap();
 	}
 
 	if memory.shared {
-		js_output.push_str(", shared: true");
+		js_output.write_all(b", shared: true").unwrap();
 	}
 
-	js_output.push_str(" })\n\n");
+	js_output.write_all(b" })\n\n").unwrap();
 
 	// Output requested embedded JS.
 	if !found_embed.is_empty() {
-		js_output.push_str("const jsEmbed = {\n");
+		js_output.write_all(b"const jsEmbed = {\n").unwrap();
 
 		for (package, embeds) in found_embed {
 			writeln!(js_output, "\t{package}: {{").unwrap();
@@ -438,25 +445,27 @@ fn post_processing(output_path: &Path, main_memory: MainMemory<'_>) {
 				write!(js_output, "\t\t\"{name}\": ").unwrap();
 
 				for (position, line) in js.lines().with_position() {
-					js_output.push_str(line);
+					js_output.write_all(line.as_bytes()).unwrap();
 
 					if let Position::First | Position::Middle = position {
-						js_output.push_str("\n\t\t");
+						js_output.write_all(b"\n\t\t").unwrap();
 					}
 				}
 
-				js_output.push_str(",\n");
+				js_output.write_all(b",\n").unwrap();
 			}
 
-			js_output.push_str("\t},\n");
+			js_output.write_all(b"\t},\n").unwrap();
 		}
 
-		js_output.push_str("}\n\n");
+		js_output.write_all(b"}\n\n").unwrap();
 	}
 
 	// Create our `importObject`.
-	js_output.push_str("export const importObject = {\n");
-	js_output.push_str("\tjs_bindgen: { memory },\n");
+	js_output
+		.write_all(b"export const importObject = {\n")
+		.unwrap();
+	js_output.write_all(b"\tjs_bindgen: { memory },\n").unwrap();
 
 	for (module, names) in found_import {
 		writeln!(js_output, "\t{module}: {{").unwrap();
@@ -465,24 +474,27 @@ fn post_processing(output_path: &Path, main_memory: MainMemory<'_>) {
 			write!(js_output, "\t\t\"{name}\": ").unwrap();
 
 			for (position, line) in js.lines().with_position() {
-				js_output.push_str(line);
+				js_output.write_all(line.as_bytes()).unwrap();
 
 				if let Position::First | Position::Middle = position {
-					js_output.push_str("\n\t\t");
+					js_output.write_all(b"\n\t\t").unwrap();
 				}
 			}
 
-			js_output.push_str(",\n");
+			js_output.write_all(b",\n").unwrap();
 		}
 
-		js_output.push_str("\t},\n");
+		js_output.write_all(b"\t},\n").unwrap();
 	}
 
-	js_output.push_str("}\n");
+	js_output.write_all(b"}\n").unwrap();
 
-	fs::write(
-		output_path.with_file_name(package).with_extension("js"),
-		js_output,
-	)
-	.expect("output JS file should be writable");
+	js_output.into_inner().unwrap().sync_all().unwrap();
+
+	// We could write into the file directly, but `wasm-encoder` doesn't support
+	// `io::Write`: https://github.com/bytecodealliance/wasm-tools/issues/778.
+	//
+	// When it does, we should rename the old file and write to a new file. This way
+	// we can keep parsing and writing at the same time without allocating memory.
+	fs::write(output_path, wasm_output).expect("output Wasm file should be writable");
 }