Skip to main content

numcodecs_wasm_builder/
main.rs

1#![expect(missing_docs)] // FIXME
2#![allow(clippy::multiple_crate_versions)] // windows_i686_gnullvm
3
4use std::{
5    collections::HashMap,
6    env, fs, io,
7    path::{Path, PathBuf},
8    process::Command,
9    str::FromStr,
10};
11
12use clap::Parser;
13use semver::Version;
14
15#[derive(Parser, Debug)]
16#[command()]
17struct Args {
18    /// Name of the numcodecs codec crate to compile
19    #[arg(name = "crate", long)]
20    crate_: String,
21
22    /// Version of the numcodecs codec crate to compile
23    #[arg(long)]
24    version: Version,
25
26    /// Path to the codec type to export, without the leading crate name
27    #[arg(long)]
28    codec: String,
29
30    /// Path to which the wasm file is output
31    #[arg(long, short)]
32    output: PathBuf,
33
34    /// Compile the local crate instead of the published one
35    #[arg(long)]
36    local: bool,
37
38    /// Compile the crate with debug information enabled
39    #[arg(long)]
40    debug: bool,
41
42    /// Enable verbose logging while compiling the crate
43    #[arg(long)]
44    verbose: bool,
45}
46
47fn main() -> io::Result<()> {
48    let args = Args::parse();
49
50    let scratch_dir = scratch::path(concat!(
51        env!("CARGO_PKG_NAME"),
52        "-",
53        env!("CARGO_PKG_VERSION"),
54    ));
55    eprintln!("scratch_dir={}", scratch_dir.display());
56
57    let target_dir = scratch_dir.join("target");
58    eprintln!("target_dir={}", target_dir.display());
59    eprintln!("creating {}", target_dir.display());
60    fs::create_dir_all(&target_dir)?;
61
62    let crate_dir = create_codec_wasm_component_crate(
63        &scratch_dir,
64        &args.crate_,
65        &args.version,
66        &args.codec,
67        args.local,
68    )?;
69    copy_buildenv_to_crate(&crate_dir)?;
70
71    let nix_env = NixEnv::new(&crate_dir)?;
72
73    let wasm = build_wasm_codec(
74        &nix_env,
75        &target_dir,
76        &crate_dir,
77        &format!("{}-wasm", args.crate_),
78        args.debug,
79        args.verbose,
80    )?;
81    let wasm = optimize_wasm_codec(&wasm, &nix_env, args.debug)?;
82    let wasm = adapt_wasi_snapshot_to_preview2(&wasm)?;
83
84    fs::copy(wasm, args.output)?;
85
86    Ok(())
87}
88
89fn create_codec_wasm_component_crate(
90    scratch_dir: &Path,
91    crate_: &str,
92    version: &Version,
93    codec: &str,
94    local: bool,
95) -> io::Result<PathBuf> {
96    let crate_dir = scratch_dir.join(format!("{crate_}-wasm-{version}"));
97    eprintln!("crate_dir={}", crate_dir.display());
98    eprintln!("creating {}", crate_dir.display());
99    if crate_dir.exists() {
100        fs::remove_dir_all(&crate_dir)?;
101    }
102    fs::create_dir_all(&crate_dir)?;
103
104    let (numcodecs_wasm_logging_path, numcodecs_wasm_guest_path, numcodecs_my_codec_path) = if local
105    {
106        let numcodecs = Path::new(env!("CARGO_MANIFEST_DIR")).join("..").join("..");
107        eprintln!("looking for local workspace in {}", numcodecs.display());
108        let numcodecs = numcodecs.canonicalize()?;
109        let numcodecs_wasm_logging = numcodecs.join("crates").join("numcodecs-wasm-logging");
110        let numcodecs_wasm_guest = numcodecs.join("crates").join("numcodecs-wasm-guest");
111        let numcodecs_my_codec = numcodecs
112            .join("codecs")
113            .join(crate_.strip_prefix("numcodecs-").unwrap_or(crate_));
114        (
115            format!(r#" path = "{}","#, numcodecs_wasm_logging.display()),
116            format!(r#" path = "{}","#, numcodecs_wasm_guest.display()),
117            format!(r#" path = "{}","#, numcodecs_my_codec.display()),
118        )
119    } else {
120        (String::new(), String::new(), String::new())
121    };
122
123    fs::write(
124        crate_dir.join("Cargo.toml"),
125        format!(
126            r#"
127[workspace]
128
129[package]
130name = "{crate_}-wasm"
131version = "{version}"
132edition = "2024"
133
134# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
135
136[dependencies]
137numcodecs-wasm-logging = {{ version = "0.2",{numcodecs_wasm_logging_path} default-features = false }}
138numcodecs-wasm-guest = {{ version = "0.3",{numcodecs_wasm_guest_path} default-features = false }}
139numcodecs-my-codec = {{ package = "{crate_}", version = "{version}",{numcodecs_my_codec_path} default-features = false }}
140    "#
141        ),
142    )?;
143
144    fs::create_dir_all(crate_dir.join("src"))?;
145
146    fs::write(
147        crate_dir.join("src").join("lib.rs"),
148        format!(
149            "
150#![cfg_attr(not(test), no_main)]
151
152numcodecs_wasm_guest::export_codec!(
153    numcodecs_wasm_logging::LoggingCodec<numcodecs_my_codec::{codec}>
154);
155    "
156        ),
157    )?;
158
159    Ok(crate_dir)
160}
161
162fn copy_buildenv_to_crate(crate_dir: &Path) -> io::Result<()> {
163    fs::write(
164        crate_dir.join("flake.nix"),
165        include_str!("../buildenv/flake.nix"),
166    )?;
167    fs::write(
168        crate_dir.join("flake.lock"),
169        include_str!("../buildenv/flake.lock"),
170    )?;
171
172    fs::write(
173        crate_dir.join("include.h"),
174        include_str!("../buildenv/include.h"),
175    )?;
176    fs::write(
177        crate_dir.join("include.hpp"),
178        include_str!("../buildenv/include.hpp"),
179    )?;
180
181    fs::write(
182        crate_dir.join("rust-toolchain"),
183        include_str!("../buildenv/rust-toolchain"),
184    )?;
185
186    Ok(())
187}
188
189struct NixEnv {
190    llvm_version: String,
191    ar: PathBuf,
192    clang: PathBuf,
193    libclang: PathBuf,
194    lld: PathBuf,
195    nm: PathBuf,
196    ranlib: PathBuf,
197    strip: PathBuf,
198    objdump: PathBuf,
199    dlltool: PathBuf,
200    wasi_sysroot: PathBuf,
201    libclang_rt: PathBuf,
202    wasm_opt: PathBuf,
203    pkg_config: PathBuf,
204    #[expect(dead_code)]
205    python3: PathBuf,
206}
207
208impl NixEnv {
209    pub fn new(flake_parent_dir: &Path) -> io::Result<Self> {
210        fn try_read_env<T: FromStr<Err: std::error::Error>>(
211            env: &HashMap<&str, &str>,
212            key: &str,
213        ) -> Result<T, io::Error> {
214            let Some(var) = env.get(key).copied() else {
215                return Err(io::Error::new(
216                    io::ErrorKind::InvalidData,
217                    format!("missing flake env key: {key}"),
218                ));
219            };
220
221            T::from_str(var).map_err(|err| {
222                io::Error::new(
223                    io::ErrorKind::InvalidData,
224                    format!("invalid flake env variable {key}={var}: {err}"),
225                )
226            })
227        }
228
229        let mut env = Command::new("nix");
230        env.current_dir(flake_parent_dir);
231        env.arg("develop");
232        // env.arg("--store");
233        // env.arg(nix_store_path);
234        env.arg("path:.");
235        env.arg("--no-update-lock-file");
236        env.arg("--ignore-environment");
237        env.arg("--command");
238        env.arg("env");
239        eprintln!("executing {env:?}");
240        let env = env.output()?;
241        eprintln!(
242            "{}\n{}",
243            String::from_utf8_lossy(&env.stdout),
244            String::from_utf8_lossy(&env.stderr)
245        );
246        let env = std::str::from_utf8(&env.stdout).map_err(|err| {
247            io::Error::new(
248                io::ErrorKind::InvalidData,
249                format!("invalid flake env output: {err}"),
250            )
251        })?;
252        let env = env
253            .lines()
254            .filter_map(|line| line.split_once('='))
255            .collect::<HashMap<_, _>>();
256
257        Ok(Self {
258            llvm_version: try_read_env(&env, "MY_LLVM_VERSION")?,
259            ar: try_read_env(&env, "MY_AR")?,
260            clang: try_read_env(&env, "MY_CLANG")?,
261            libclang: try_read_env(&env, "MY_LIBCLANG")?,
262            lld: try_read_env(&env, "MY_LLD")?,
263            nm: try_read_env(&env, "MY_NM")?,
264            ranlib: try_read_env(&env, "MY_RANLIB")?,
265            strip: try_read_env(&env, "MY_STRIP")?,
266            objdump: try_read_env(&env, "MY_OBJDUMP")?,
267            dlltool: try_read_env(&env, "MY_DLLTOOL")?,
268            wasi_sysroot: try_read_env(&env, "MY_WASI_SYSROOT")?,
269            libclang_rt: try_read_env(&env, "MY_LIBCLANG_RT")?,
270            wasm_opt: try_read_env(&env, "MY_WASM_OPT")?,
271            pkg_config: try_read_env(&env, "MY_PKG_CONFIG")?,
272            python3: try_read_env(&env, "MY_PYTHON3")?,
273        })
274    }
275}
276
277#[expect(clippy::too_many_lines)]
278fn configure_cargo_cmd(
279    nix_env: &NixEnv,
280    target_dir: &Path,
281    crate_dir: &Path,
282    debug: bool,
283) -> Command {
284    let NixEnv {
285        llvm_version,
286        ar,
287        clang,
288        libclang,
289        lld,
290        nm,
291        ranlib,
292        strip,
293        objdump,
294        dlltool,
295        wasi_sysroot,
296        libclang_rt,
297        pkg_config,
298        ..
299    } = nix_env;
300
301    let mut cmd = Command::new("nix");
302    cmd.current_dir(crate_dir);
303    cmd.arg("develop");
304    // cmd.arg("--store");
305    // cmd.arg(nix_store_path);
306    cmd.arg("--no-update-lock-file");
307    cmd.arg("--ignore-environment");
308    cmd.arg("path:.");
309    cmd.arg("--command");
310    cmd.arg("env");
311    cmd.arg("GMP_MPFR_SYS_CACHE=");
312    cmd.arg(format!("CC={clang}", clang = clang.join("clang").display()));
313    cmd.arg(format!(
314        "CXX={clang}",
315        clang = clang.join("clang++").display()
316    ));
317    cmd.arg(format!("LD={lld}", lld = lld.join("lld").display()));
318    cmd.arg(format!("LLD={lld}", lld = lld.join("lld").display()));
319    cmd.arg(format!("AR={ar}", ar = ar.display()));
320    cmd.arg(format!("NM={nm}", nm = nm.display()));
321    cmd.arg(format!("RANLIB={ranlib}", ranlib = ranlib.display()));
322    cmd.arg(format!("STRIP={strip}", strip = strip.display()));
323    cmd.arg(format!("OBJDUMP={objdump}", objdump = objdump.display()));
324    cmd.arg(format!("DLLTOOL={dlltool}", dlltool = dlltool.display()));
325    cmd.arg(format!(
326        "PKG_CONFIG={pkg_config}",
327        pkg_config = pkg_config.display()
328    ));
329    cmd.arg(format!(
330        "LIBCLANG_PATH={libclang}",
331        libclang = libclang.display()
332    ));
333    cmd.arg(format!(
334        "CFLAGS=--target=wasm32-wasip1 -nodefaultlibs -resource-dir {resource_dir} \
335         --sysroot={wasi_sysroot} -isystem {clang_include} -isystem {wasi32_wasi_include} \
336         -isystem {include} -B {lld} -D_WASI_EMULATED_PROCESS_CLOCKS -D_WASI_EMULATED_SIGNAL \
337         -include {c_include_path} -O3 {debug} \
338         -DHAVE_STRNLEN=1 -DHAVE_MEMSET=1 -DHAVE_RAISE=1",
339        resource_dir = libclang.join("clang").join(llvm_version).display(),
340        wasi_sysroot = wasi_sysroot.display(),
341        clang_include = libclang
342            .join("clang")
343            .join(llvm_version)
344            .join("include")
345            .display(),
346        wasi32_wasi_include = wasi_sysroot.join("include").join("wasm32-wasip1").display(),
347        include = wasi_sysroot.join("include").display(),
348        lld = lld.display(),
349        c_include_path = crate_dir.join("include.h").display(),
350        debug = if debug { "-g" } else { "" },
351    ));
352    cmd.arg(format!(
353        "CXXFLAGS=--target=wasm32-wasip1 -nodefaultlibs -resource-dir {resource_dir} \
354         --sysroot={wasi_sysroot} -isystem {wasm32_wasi_cxx_include} -isystem {cxx_include} \
355         -isystem {clang_include} -isystem {wasi32_wasi_include} -isystem {include} -B {lld} \
356         -D_WASI_EMULATED_PROCESS_CLOCKS -include {cpp_include_path} -O3 {debug} \
357         -fwasm-exceptions -mllvm -wasm-use-legacy-eh=false",
358        resource_dir = libclang.join("clang").join(llvm_version).display(),
359        wasi_sysroot = wasi_sysroot.display(),
360        wasm32_wasi_cxx_include = wasi_sysroot
361            .join("include")
362            .join("wasm32-wasip1")
363            .join("c++")
364            .join("v1")
365            .display(),
366        cxx_include = wasi_sysroot
367            .join("include")
368            .join("c++")
369            .join("v1")
370            .display(),
371        clang_include = libclang
372            .join("clang")
373            .join(llvm_version)
374            .join("include")
375            .display(),
376        wasi32_wasi_include = wasi_sysroot.join("include").join("wasm32-wasip1").display(),
377        include = wasi_sysroot.join("include").display(),
378        lld = lld.display(),
379        cpp_include_path = crate_dir.join("include.hpp").display(),
380        debug = if debug { "-g" } else { "" },
381    ));
382    cmd.arg(format!(
383        "BINDGEN_EXTRA_CLANG_ARGS=--target=wasm32-wasip1 -nodefaultlibs -resource-dir \
384         {resource_dir} --sysroot={wasi_sysroot} -isystem {wasm32_wasi_cxx_include} -isystem \
385         {cxx_include} -isystem {clang_include} -isystem {wasi32_wasi_include} -isystem {include} \
386         -B {lld} -D_WASI_EMULATED_PROCESS_CLOCKS -fvisibility=default",
387        resource_dir = libclang.join("clang").join(llvm_version).display(),
388        wasi_sysroot = wasi_sysroot.display(),
389        wasm32_wasi_cxx_include = wasi_sysroot
390            .join("include")
391            .join("wasm32-wasip1")
392            .join("c++")
393            .join("v1")
394            .display(),
395        cxx_include = wasi_sysroot
396            .join("include")
397            .join("c++")
398            .join("v1")
399            .display(),
400        clang_include = libclang
401            .join("clang")
402            .join(llvm_version)
403            .join("include")
404            .display(),
405        wasi32_wasi_include = wasi_sysroot.join("include").join("wasm32-wasip1").display(),
406        include = wasi_sysroot.join("include").display(),
407        lld = lld.display(),
408    ));
409    cmd.arg("CXXSTDLIB=c++");
410    // disable default flags from cc
411    cmd.arg("CRATE_CC_NO_DEFAULTS=1");
412    cmd.arg(format!(
413        "LDFLAGS=-lc -lwasi-emulated-process-clocks -lwasi-emulated-signal \
414        -L{libclang_rt} -lclang_rt.builtins -lunwind -lc++ -lc++abi",
415        libclang_rt = libclang_rt.join("wasm32-unknown-wasip1").display(),
416    ));
417    cmd.arg("PKG_CONFIG_PATH=\"\"");
418    cmd.arg(format!(
419        "PKG_CONFIG_LIBDIR={pkg_config_lib}:{pkg_config_share}",
420        pkg_config_lib = wasi_sysroot.join("lib").join("pkgconfig").display(),
421        pkg_config_share = wasi_sysroot.join("share").join("pkgconfig").display()
422    ));
423    cmd.arg(format!(
424        "PKG_CONFIG_SYSROOT_DIR={wasi_sysroot}",
425        wasi_sysroot = wasi_sysroot.display()
426    ));
427    cmd.arg(format!(
428        "RUSTFLAGS=-C panic=abort {debug} \
429        -C link-arg=-L{wasm32_wasi_lib} \
430        -C link-arg=-L{libclang_rt} -C link-arg=-lclang_rt.builtins \
431        -C link-arg=-lunwind -C link-arg=-lc++ -C link-arg=-lc++abi \
432        -C llvm-args=-wasm-use-legacy-eh=false",
433        debug = if debug { "-g" } else { "-C strip=symbols" },
434        wasm32_wasi_lib = wasi_sysroot.join("lib").join("wasm32-wasip1").display(),
435        libclang_rt = libclang_rt.join("wasm32-unknown-wasip1").display(),
436    ));
437    cmd.arg(format!(
438        "CARGO_TARGET_DIR={target_dir}",
439        target_dir = target_dir.display()
440    ));
441
442    // we don't need nightly Rust features but need to compile std with immediate
443    // panic abort instead of compiling with nightly, we fake it and forbid the
444    // unstable_features lint
445    cmd.arg("RUSTC_BOOTSTRAP=1");
446
447    cmd.arg("cargo");
448
449    cmd
450}
451
452fn build_wasm_codec(
453    nix_env: &NixEnv,
454    target_dir: &Path,
455    crate_dir: &Path,
456    crate_name: &str,
457    debug: bool,
458    verbose: bool,
459) -> io::Result<PathBuf> {
460    let mut cmd = configure_cargo_cmd(nix_env, target_dir, crate_dir, debug);
461    cmd.arg("rustc")
462        .arg("--crate-type=cdylib")
463        .arg("-Z")
464        .arg("build-std=std,panic_abort")
465        .arg("-Z")
466        .arg("build-std-features=panic_immediate_abort")
467        .arg("--release")
468        .arg("--target=wasm32-wasip1");
469
470    if verbose {
471        cmd.arg("-vv");
472    }
473
474    eprintln!("executing {cmd:?}");
475
476    let status = cmd.status()?;
477    if !status.success() {
478        return Err(io::Error::other(format!("cargo exited with code {status}")));
479    }
480
481    Ok(target_dir
482        .join("wasm32-wasip1")
483        .join("release")
484        .join(crate_name.replace('-', "_"))
485        .with_extension("wasm"))
486}
487
488fn optimize_wasm_codec(wasm: &Path, nix_env: &NixEnv, debug: bool) -> io::Result<PathBuf> {
489    let NixEnv { wasm_opt, .. } = nix_env;
490
491    let opt_out = wasm.with_extension("opt.wasm");
492
493    let mut cmd = Command::new(wasm_opt);
494
495    cmd.arg("--enable-sign-ext")
496        .arg("--disable-threads")
497        .arg("--enable-mutable-globals")
498        .arg("--enable-nontrapping-float-to-int")
499        .arg("--enable-simd")
500        .arg("--enable-bulk-memory")
501        .arg("--enable-exception-handling")
502        .arg("--disable-tail-call")
503        .arg("--enable-reference-types")
504        .arg("--enable-multivalue")
505        .arg("--disable-gc")
506        .arg("--disable-memory64")
507        .arg("--disable-relaxed-simd")
508        .arg("--disable-extended-const")
509        .arg("--disable-strings")
510        .arg("--disable-multimemory")
511        .arg(if debug { "-g" } else { "--strip-debug" });
512
513    // FIXME: https://github.com/WebAssembly/binaryen/issues/8325
514    cmd.arg("-O3").arg("-o").arg(&opt_out).arg(wasm);
515
516    eprintln!("executing {cmd:?}");
517
518    let status = cmd.status()?;
519    if !status.success() {
520        return Err(io::Error::other(format!(
521            "wasm-opt exited with code {status}"
522        )));
523    }
524
525    Ok(opt_out)
526}
527
528fn adapt_wasi_snapshot_to_preview2(wasm: &Path) -> io::Result<PathBuf> {
529    let wasm_preview2 = wasm.with_extension("preview2.wasm");
530
531    eprintln!("reading from {}", wasm.display());
532    let wasm = fs::read(wasm)?;
533
534    let mut encoder = wit_component::ComponentEncoder::default()
535        .module(&wasm)
536        .map_err(|err| {
537            io::Error::other(
538                // FIXME: better error reporting in the build script
539                format!("wit_component::ComponentEncoder::module failed: {err:#}"),
540            )
541        })?
542        .adapter(
543            wasi_preview1_component_adapter_provider::WASI_SNAPSHOT_PREVIEW1_ADAPTER_NAME,
544            wasi_preview1_component_adapter_provider::WASI_SNAPSHOT_PREVIEW1_REACTOR_ADAPTER,
545        )
546        .map_err(|err| {
547            io::Error::other(
548                // FIXME: better error reporting in the build script
549                format!("wit_component::ComponentEncoder::adapter failed: {err:#}"),
550            )
551        })?;
552
553    let wasm = encoder.encode().map_err(|err| {
554        io::Error::other(
555            // FIXME: better error reporting in the build script
556            format!("wit_component::ComponentEncoder::encode failed: {err:#}"),
557        )
558    })?;
559
560    eprintln!("writing to {}", wasm_preview2.display());
561    fs::write(&wasm_preview2, wasm)?;
562
563    Ok(wasm_preview2)
564}