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