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