diffr web: pack grammar tables in the browser engine (127 MB → 25 MB)

diffr/crates/diffr-grammars/Cargo.toml

@@ -22,18 +22,22 @@ include = [
 [dependencies]
 tree-sitter-language = "0.1.7"
 
-# The grammars' C code calls ZSTD_decompress, except on wasm32, where the
-# tables are not packed.
+# The grammars' C code calls ZSTD_decompress: libzstd's, or on wasm32, where it
+# does not build, src/lib.rs's over ruzstd.
 [target.'cfg(not(target_arch = "wasm32"))'.dependencies]
 zstd-sys = { version = "2.1.0", default-features = false }
 
+[target.'cfg(target_arch = "wasm32")'.dependencies]
+ruzstd = { version = "0.7.3", default-features = false, features = ["std"] }
+
 [build-dependencies]
 cc = "1.2.48"
 object = { version = "0.40", default-features = false, features = ["read_core", "coff", "elf", "macho", "std"] }
 rayon = "1.7.0"
 serde_json = "1.0"
 tree-sitter = "0.26.10"
 tree-sitter-c = "0.24"
+wasmparser = { version = "0.259.0", default-features = false }
 zstd = { version = "0.13.3", default-features = false }
 
 # The build script reads the parser sources from these packages.

diffr/crates/diffr-grammars/build.rs

@@ -7,10 +7,10 @@
 //! function, `diffr_unpack_<symbol>`, fills the arrays, and `src/lib.rs` runs it
 //! once before the grammar's first use.
 //!
-//! A wasm32 build compiles each grammar as generated, with an empty unpack
-//! function: a page is served compressed, so packing would only cost the
-//! decompressor, and wasm object files hold their tables in a form this
-//! script does not read.
+//! A wasm32 build packs them too. Wasm copies a module's data into every
+//! instance's memory, so unpacked tables would cost each worker all ~110 MB of
+//! them; packed, a worker holds only the grammars it has parsed with. There the
+//! decompressor is `src/lib.rs`'s pure-Rust `ZSTD_decompress`.
 use object::{Object, ObjectSection, ObjectSymbol};
 use rayon::prelude::*;
 use std::collections::BTreeMap;
@@ -190,21 +190,6 @@ fn package_roots() -> BTreeMap<String, PathBuf> {
 /// fills them, and the grammar's constructor renamed to `diffr_<symbol>`, then compile it.
 fn pack(symbol: &str, source: &Path, out: &Path) {
     std::fs::create_dir_all(out).unwrap();
-    if std::env::var("CARGO_CFG_TARGET_ARCH").unwrap() == "wasm32" {
-        let grammar = format!(
-            "#define {symbol} diffr_{symbol}\n\
-             #include \"parser.c\"\n\
-             void diffr_unpack_{symbol}(void) {{}}\n"
-        );
-        std::fs::write(out.join("grammar.c"), grammar).unwrap();
-        let mut build = compiler(source, out);
-        build.file(out.join("grammar.c"));
-        if source.join("scanner.c").is_file() {
-            build.file(source.join("scanner.c"));
-        }
-        build.compile(symbol);
-        return;
-    }
     let parser = std::fs::read_to_string(source.join("parser.c")).unwrap();
     let header = std::fs::read_to_string(source.join("tree_sitter/parser.h")).unwrap();
     let tables = tables(&parser, &header);
@@ -289,6 +274,9 @@ fn table_bytes(
     .unwrap();
 
     let data = std::fs::read(object_path).unwrap();
+    if data.starts_with(b"\0asm") {
+        return wasm_table_bytes(symbol, &data, tables);
+    }
     let file = object::File::parse(&*data).unwrap();
     let prefix = if file.format() == object::BinaryFormat::MachO {
         "_"
@@ -322,6 +310,56 @@ fn table_bytes(
         .collect()
 }
 
+/// The same tables read from a wasm object file. Its `linking` section places each defined
+/// data symbol in a data segment, at an offset and with a size, so no size globals are needed.
+fn wasm_table_bytes(symbol: &str, data: &[u8], tables: &[Table]) -> Vec<Vec<u8>> {
+    let mut segments = Vec::new();
+    let mut symbols = BTreeMap::new();
+    for payload in wasmparser::Parser::new(0).parse_all(data) {
+        match payload.unwrap() {
+            wasmparser::Payload::DataSection(reader) => {
+                for segment in reader {
+                    segments.push(segment.unwrap().data);
+                }
+            }
+            wasmparser::Payload::CustomSection(section) if section.name() == "linking" => {
+                let linking = wasmparser::LinkingSectionReader::new(wasmparser::BinaryReader::new(
+                    section.data(),
+                    section.data_offset(),
+                ))
+                .unwrap();
+                for subsection in linking {
+                    let wasmparser::Linking::SymbolTable(table) = subsection.unwrap() else {
+                        continue;
+                    };
+                    for info in table {
+                        if let wasmparser::SymbolInfo::Data {
+                            name,
+                            symbol: Some(defined),
+                            ..
+                        } = info.unwrap()
+                        {
+                            symbols.insert(name, defined);
+                        }
+                    }
+                }
+            }
+            _ => {}
+        }
+    }
+    tables
+        .iter()
+        .map(|Table { name, .. }| {
+            let defined = symbols
+                .get(name)
+                .unwrap_or_else(|| panic!("{name} is missing from the {symbol} tables object"));
+            let start = usize::try_from(defined.offset).unwrap();
+            let end = start + usize::try_from(defined.size).unwrap();
+            segments[usize::try_from(defined.index).unwrap()][start..end].to_vec()
+        })
+        .collect()
+}
+
 fn compiler(source: &Path, out: &Path) -> cc::Build {
     let mut build = cc::Build::new();
     build.include(source).warnings(false).out_dir(out);

diffr/crates/diffr-grammars/src/lib.rs

@@ -4,10 +4,36 @@
 //! The first call unpacks that grammar's tables.
 use std::sync::Once;
 use tree_sitter_language::LanguageFn;
-// Links libzstd for the grammars' C code; wasm32 grammars are not packed.
+// Links libzstd for the grammars' C code; wasm32 has the pure-Rust decoder below.
 #[cfg(not(target_arch = "wasm32"))]
 use zstd_sys as _;
 
+/// libzstd's one-shot decompression, which the grammars' unpack functions call, for wasm32,
+/// where libzstd does not build. Returns the bytes written, or `usize::MAX` on a bad frame.
+///
+/// # Safety
+///
+/// `dst` must be writable for `dst_capacity` bytes and `src` readable for `src_size`.
+#[cfg(target_arch = "wasm32")]
+#[no_mangle]
+pub unsafe extern "C" fn ZSTD_decompress(
+    dst: *mut u8,
+    dst_capacity: usize,
+    src: *const u8,
+    src_size: usize,
+) -> usize {
+    // Safety: the caller passes a table's own storage and its packed bytes, with their sizes.
+    let (output, input) = unsafe {
+        (
+            std::slice::from_raw_parts_mut(dst, dst_capacity),
+            std::slice::from_raw_parts(src, src_size),
+        )
+    };
+    ruzstd::FrameDecoder::new()
+        .decode_all(input, output)
+        .unwrap_or(usize::MAX)
+}
+
 macro_rules! grammar {
     ($name:ident, $unpack:ident, $constructor:ident) => {
         pub const $name: LanguageFn = {