Skip to main content

zsh/extensions/
aot.rs

1//! Ahead-of-time build (`zbuild`): turn one or more shell scripts into a
2//! self-contained executable that runs with no zsh and no zshrs installed.
3//! Two modes, and they produce genuinely different artifacts:
4//!
5//! * **`zbuild`** (default, this module's trailer format below) appends the
6//!   script *source*, zstd-compressed, to a copy of the running `zshrs`
7//!   binary. At startup zshrs detects the trailer and runs every embedded
8//!   script IN INPUT ORDER as one concatenated zsh program. The launch still
9//!   parses and compiles; what you gain is a single file to ship.
10//!
11//! * **`zbuild --native`** ([`build_native`]) compiles the scripts to a fusevm
12//!   `Chunk`, lowers that to **native machine code** with Cranelift
13//!   (`fusevm::aot::compile_object` → a relocatable `.o`), and links it
14//!   against `libzsh.a` into a standalone binary. `objdump -d
15//!   fusevm_aot_entry` on the result shows the script's own control flow as
16//!   machine instructions — a `br_table` over one native block per opcode —
17//!   with no bytecode dispatch loop at run time.
18//!
19//!   Be precise about what that does and does not mean. The *shape* of the
20//!   program is native; the *work* each op does is a call into the linked zsh
21//!   runtime (`fusevm_aot_exec_op`), because a shell op is `fork`/`execve`/
22//!   glob/expansion, not arithmetic. fusevm's fully-register-native path —
23//!   where ops become real arithmetic in registers — needs
24//!   `Chunk::builtin_argc_is_arity`, which zshrs cannot set: its builtin
25//!   handlers do not treat `CallBuiltin`'s `argc` as a stack arity
26//!   (`BUILTIN_XTRACE_ARGS` is emitted with `argc = 2` and pops one value,
27//!   peeking the rest). So a shell chunk takes fusevm's threaded-native
28//!   lowering. The serialized chunk also rides along in the binary
29//!   (`fusevm_aot_chunk_blob`) — the native driver reads its ops and
30//!   constants from there.
31//!
32//! **zshrs-original infrastructure — no C source counterpart.** C zsh
33//! has the `zcompile` builtin (Src/parse.c → `bin_zcompile()`) which
34//! writes a parsed-AST `.zwc` file alongside a script for faster
35//! re-parse. That's a separate file, not a self-contained binary.
36//! AOT-ing scripts into the shell executable itself is a zshrs
37//! addition — it lets you ship a single binary that bundles its own
38//! script (think `zsh -c '...'` but compiled in).
39//!
40//! Layout (little-endian, appended to the end of a copy of the `zshrs` binary):
41//!
42//! ```text
43//!   [elf/mach-o bytes of zshrs ...]   (unchanged, still runs as `zshrs`)
44//!   [zstd-compressed payload ...]
45//!   [u64 compressed_len]
46//!   [u64 uncompressed_len]
47//!   [u32 version]
48//!   [u32 reserved (0)]
49//!   [8 bytes magic  b"ZSHRSAOT"]
50//! ```
51//!
52//! Payload v1 (single script, BACKWARD-COMPAT decoder only — new builds
53//! always emit v2 even for one input):
54//!
55//! ```text
56//!   [u32 script_name_len]
57//!   [script_name utf8]
58//!   [source bytes utf8]
59//! ```
60//!
61//! Payload v2 (ordered file list — current `zbuild` output):
62//!
63//! ```text
64//!   [u32 file_count]
65//!   for each file (file_count times, in input order):
66//!     [u32 name_len][name utf8]
67//!     [u32 source_len][source utf8]
68//! ```
69//!
70//! Files run sequentially in iteration order. zsh has no "project" concept;
71//! ordering matches `--in` argv order. Globals/functions defined by file N
72//! are visible to file N+1 (single ShellExecutor across all files).
73//!
74//! ELF (Linux) and Mach-O (macOS) loaders ignore bytes past the program-
75//! header-listed segments, so appending leaves the original `zshrs` fully
76//! runnable. macOS resulting binary is unsigned; signed-build distributors
77//! must re-codesign.
78
79use std::fs::{self, File, OpenOptions};
80use std::io::{self, Read, Seek, SeekFrom, Write};
81use std::os::unix::fs::PermissionsExt;
82use std::path::{Path, PathBuf};
83
84/// 8-byte trailer magic.
85pub const AOT_MAGIC: &[u8; 8] = b"ZSHRSAOT";
86/// Trailer format version 1: single script (legacy decode-only).
87pub const AOT_VERSION_V1: u32 = 1;
88/// Trailer format version 2: ordered file list (current build output).
89pub const AOT_VERSION_V2: u32 = 2;
90/// Fixed trailer length: `8 (cl) + 8 (ul) + 4 (ver) + 4 (rsv) + 8 (magic)`.
91pub const TRAILER_LEN: u64 = 32;
92/// `EmbeddedFile` — see fields for layout.
93#[derive(Debug, Clone)]
94pub struct EmbeddedFile {
95    /// `__FILE__` / error-reporting name (e.g. `hello.zsh`).
96    pub name: String,
97    /// UTF-8 zsh source.
98    pub source: String,
99}
100
101/// One or more embedded files, in build-order. v1 binaries decode to a
102/// 1-element vec; v2 to N elements preserving input order.
103#[derive(Debug, Clone)]
104pub struct EmbeddedFiles(pub Vec<EmbeddedFile>);
105
106fn encode_payload_v2(files: &[EmbeddedFile]) -> Vec<u8> {
107    let mut out = Vec::with_capacity(
108        64 + files
109            .iter()
110            .map(|f| f.name.len() + f.source.len() + 8)
111            .sum::<usize>(),
112    );
113    let count = u32::try_from(files.len()).expect("file count fits in u32");
114    out.extend_from_slice(&count.to_le_bytes());
115    for f in files {
116        let name_len = u32::try_from(f.name.len()).expect("name length fits in u32");
117        let src_len = u32::try_from(f.source.len()).expect("source length fits in u32");
118        out.extend_from_slice(&name_len.to_le_bytes());
119        out.extend_from_slice(f.name.as_bytes());
120        out.extend_from_slice(&src_len.to_le_bytes());
121        out.extend_from_slice(f.source.as_bytes());
122    }
123    out
124}
125
126fn decode_payload_v2(bytes: &[u8]) -> Option<EmbeddedFiles> {
127    let mut pos = 0usize;
128    if bytes.len() < 4 {
129        return None;
130    }
131    let count = u32::from_le_bytes(bytes[pos..pos + 4].try_into().ok()?) as usize;
132    pos += 4;
133    let mut out = Vec::with_capacity(count);
134    for _ in 0..count {
135        if pos + 4 > bytes.len() {
136            return None;
137        }
138        let name_len = u32::from_le_bytes(bytes[pos..pos + 4].try_into().ok()?) as usize;
139        pos += 4;
140        if pos + name_len > bytes.len() {
141            return None;
142        }
143        let name = std::str::from_utf8(&bytes[pos..pos + name_len])
144            .ok()?
145            .to_string();
146        pos += name_len;
147        if pos + 4 > bytes.len() {
148            return None;
149        }
150        let src_len = u32::from_le_bytes(bytes[pos..pos + 4].try_into().ok()?) as usize;
151        pos += 4;
152        if pos + src_len > bytes.len() {
153            return None;
154        }
155        let source = std::str::from_utf8(&bytes[pos..pos + src_len])
156            .ok()?
157            .to_string();
158        pos += src_len;
159        out.push(EmbeddedFile { name, source });
160    }
161    Some(EmbeddedFiles(out))
162}
163
164/// v1 decoder kept for backward compat: one-script payload promoted into a
165/// single-element EmbeddedFiles. Old binaries built with the previous zbuild
166/// still load.
167fn decode_payload_v1(bytes: &[u8]) -> Option<EmbeddedFiles> {
168    if bytes.len() < 4 {
169        return None;
170    }
171    let name_len = u32::from_le_bytes(bytes[0..4].try_into().ok()?) as usize;
172    if 4 + name_len > bytes.len() {
173        return None;
174    }
175    let name = std::str::from_utf8(&bytes[4..4 + name_len])
176        .ok()?
177        .to_string();
178    let source = std::str::from_utf8(&bytes[4 + name_len..])
179        .ok()?
180        .to_string();
181    Some(EmbeddedFiles(vec![EmbeddedFile { name, source }]))
182}
183
184fn build_trailer(compressed_len: u64, uncompressed_len: u64, version: u32) -> [u8; 32] {
185    let mut trailer = [0u8; 32];
186    trailer[0..8].copy_from_slice(&compressed_len.to_le_bytes());
187    trailer[8..16].copy_from_slice(&uncompressed_len.to_le_bytes());
188    trailer[16..20].copy_from_slice(&version.to_le_bytes());
189    // 20..24 reserved (zeros).
190    trailer[24..32].copy_from_slice(AOT_MAGIC);
191    trailer
192}
193
194/// Append a compressed v2 ordered-file payload to an existing file.
195/// zshrs-original — no C counterpart. C zsh's closest analog is the
196/// `zcompile` builtin in Src/parse.c which writes parsed-AST data
197/// to a separate `.zwc` file rather than appending to the binary.
198pub fn append_embedded_files(out_path: &Path, files: &[EmbeddedFile]) -> io::Result<()> {
199    let payload = encode_payload_v2(files);
200    let compressed = zstd::stream::encode_all(&payload[..], 3)?;
201    let mut f = OpenOptions::new().append(true).open(out_path)?;
202    f.write_all(&compressed)?;
203    let trailer = build_trailer(
204        compressed.len() as u64,
205        payload.len() as u64,
206        AOT_VERSION_V2,
207    );
208    f.write_all(&trailer)?;
209    f.sync_all()?;
210    Ok(())
211}
212
213/// Fast probe: read the last 32 bytes of `exe` and return embedded files
214/// in build-order if present. Decodes both v1 (legacy single-script) and
215/// v2 (current ordered list). Called at zshrs startup before arg parsing.
216/// zshrs-original — no C counterpart.
217pub fn try_load_embedded(exe: &Path) -> Option<EmbeddedFiles> {
218    let mut f = File::open(exe).ok()?;
219    let size = f.metadata().ok()?.len();
220    if size < TRAILER_LEN {
221        return None;
222    }
223    f.seek(SeekFrom::End(-(TRAILER_LEN as i64))).ok()?;
224    let mut trailer = [0u8; TRAILER_LEN as usize];
225    f.read_exact(&mut trailer).ok()?;
226    if &trailer[24..32] != AOT_MAGIC {
227        return None;
228    }
229    let compressed_len = u64::from_le_bytes(trailer[0..8].try_into().ok()?);
230    let uncompressed_len = u64::from_le_bytes(trailer[8..16].try_into().ok()?);
231    let version = u32::from_le_bytes(trailer[16..20].try_into().ok()?);
232    if compressed_len == 0 || compressed_len > size - TRAILER_LEN {
233        return None;
234    }
235    let payload_start = size - TRAILER_LEN - compressed_len;
236    f.seek(SeekFrom::Start(payload_start)).ok()?;
237    let mut compressed = vec![0u8; compressed_len as usize];
238    f.read_exact(&mut compressed).ok()?;
239    let payload = zstd::stream::decode_all(&compressed[..]).ok()?;
240    if payload.len() != uncompressed_len as usize {
241        return None;
242    }
243    match version {
244        AOT_VERSION_V1 => decode_payload_v1(&payload),
245        AOT_VERSION_V2 => decode_payload_v2(&payload),
246        _ => None,
247    }
248}
249
250#[cfg(unix)]
251fn set_executable(path: &Path) {
252    if let Ok(meta) = fs::metadata(path) {
253        let mut p = meta.permissions();
254        p.set_mode(p.mode() | 0o111);
255        let _ = fs::set_permissions(path, p);
256    }
257}
258
259#[cfg(not(unix))]
260fn set_executable(_path: &Path) {}
261
262/// Copy `src` to `dst`, skipping any existing AOT trailer on `src`. Prevents
263/// nested builds from stacking trailers: building once with trailer-A then
264/// building again with trailer-B would otherwise embed both, A then B.
265fn copy_exe_without_trailer(src: &Path, dst: &Path) -> io::Result<()> {
266    let mut sf = File::open(src)?;
267    let size = sf.metadata()?.len();
268    let keep = if size >= TRAILER_LEN {
269        sf.seek(SeekFrom::End(-(TRAILER_LEN as i64)))?;
270        let mut trailer = [0u8; TRAILER_LEN as usize];
271        if sf.read_exact(&mut trailer).is_ok() && &trailer[24..32] == AOT_MAGIC {
272            let compressed_len = u64::from_le_bytes(trailer[0..8].try_into().unwrap());
273            if compressed_len > 0 && compressed_len <= size - TRAILER_LEN {
274                size - TRAILER_LEN - compressed_len
275            } else {
276                size
277            }
278        } else {
279            size
280        }
281    } else {
282        size
283    };
284    sf.seek(SeekFrom::Start(0))?;
285    let _ = fs::remove_file(dst);
286    let mut df = File::create(dst)?;
287    let mut remaining = keep;
288    let mut buf = vec![0u8; 64 * 1024];
289    while remaining > 0 {
290        let n = std::cmp::min(remaining as usize, buf.len());
291        sf.read_exact(&mut buf[..n])?;
292        df.write_all(&buf[..n])?;
293        remaining -= n as u64;
294    }
295    df.sync_all()?;
296    Ok(())
297}
298
299/// `zbuild --in A --in B --out OUT`: bake A and B into a copy of the
300/// running zshrs binary in input order, producing a self-contained AOT
301/// executable. At runtime, all embedded files run sequentially under one
302/// ShellExecutor — globals + functions from earlier files are visible
303/// to later ones.
304/// zshrs-original — no C counterpart. C zsh's `bin_zcompile()`
305/// (Src/parse.c) writes a `.zwc` cache file but doesn't bundle into
306/// the shell binary itself.
307pub fn build(script_paths: &[PathBuf], out_path: &Path) -> Result<PathBuf, String> {
308    if script_paths.is_empty() {
309        return Err("zbuild: at least one --in PATH required".to_string());
310    }
311    let mut files: Vec<EmbeddedFile> = Vec::with_capacity(script_paths.len());
312    for p in script_paths {
313        let source = fs::read_to_string(p)
314            .map_err(|e| format!("zbuild: cannot read {}: {}", p.display(), e))?;
315        let name = p
316            .file_name()
317            .and_then(|s| s.to_str())
318            .unwrap_or("script.zsh")
319            .to_string();
320        files.push(EmbeddedFile { name, source });
321    }
322    let exe = std::env::current_exe()
323        .map_err(|e| format!("zbuild: locating current executable: {}", e))?;
324    copy_exe_without_trailer(&exe, out_path).map_err(|e| {
325        format!(
326            "zbuild: copy {} -> {}: {}",
327            exe.display(),
328            out_path.display(),
329            e
330        )
331    })?;
332    append_embedded_files(out_path, &files).map_err(|e| format!("zbuild: write trailer: {}", e))?;
333    set_executable(out_path);
334    Ok(out_path.to_path_buf())
335}
336
337// ───────────────────────── Native AOT (`zbuild --native`) ──────────────────
338//
339// Unlike the source-trailer build above (which embeds the script text and
340// re-runs it interpreted at startup), the native path compiles the script to a
341// fusevm chunk, lowers that to native machine code via `fusevm::aot`
342// (Cranelift `ObjectModule` → relocatable `.o`), and links the object against
343// the zsh runtime staticlib (`libzsh.a`) into a standalone executable. The
344// script's bytecode runs as native code with no interpreter dispatch loop;
345// command execution, expansion, and cmdsubst still go through the embedded
346// `ShellExecutor` (the shell runtime is linked in).
347
348/// Frontend runtime hook invoked by `fusevm::aot::fusevm_aot_run_embedded` at
349/// startup of a native AOT binary. Stands up a leaked-`'static` `ShellExecutor`,
350/// seeds positional parameters from `argv`, installs it as the thread's
351/// `CURRENT_EXECUTOR`, and registers the zsh builtins on the run VM — so host
352/// calls from the native chunk reach a live shell exactly as the interpreter's
353/// `run_chunk` arranges. The executor lives for the process (an AOT binary runs
354/// one program and exits), so the install is intentionally permanent.
355///
356/// # Safety
357/// `vm` is the live run VM passed by the fusevm runtime; borrowed only here.
358///
359/// Behind the default-on `aot-hook` feature because the symbol is `#[no_mangle]`:
360/// it is the ONE definition fusevm's AOT runtime resolves by name, so a
361/// dependent crate that provides its own — strykelang does, for its own
362/// `--aot` — links two and the build dies with "symbol multiply defined"
363/// (under `lto = "fat"`, as a bare `failed to load bitcode`). zshrs's own
364/// `zbuild --native` links `libzsh.a` from a default build and is unaffected;
365/// a consumer that only wants the lib takes `default-features = false`.
366#[cfg(feature = "aot-hook")]
367#[no_mangle]
368pub extern "C" fn fusevm_aot_register_builtins(vm: *mut fusevm::VM) {
369    // SAFETY: the fusevm runtime hands us the live run VM for this call.
370    let vm = unsafe { &mut *vm };
371    let exec: &'static mut crate::vm_helper::ShellExecutor =
372        Box::leak(Box::new(crate::vm_helper::ShellExecutor::new()));
373    // Positional params $1.. come from argv after the program name.
374    let args: Vec<String> = std::env::args().skip(1).collect();
375    exec.set_pparams(args);
376    let last = exec.last_status();
377    crate::fusevm_bridge::register_builtins(vm);
378    vm.last_status = last;
379    // Install permanently: forget the RAII guard so CURRENT_EXECUTOR stays set
380    // for the whole native run (the process is one-shot).
381    std::mem::forget(crate::fusevm_bridge::ExecutorContext::enter(exec));
382}
383
384/// Compile zsh `source` to a fusevm chunk via the normal parse + compile path.
385fn compile_source_to_chunk(source: &str) -> Result<fusevm::Chunk, String> {
386    let program = crate::vm_helper::parse_isolated(source);
387    let chunk = crate::compile_zsh::ZshCompiler::new().compile(&program);
388    if chunk.ops.is_empty() {
389        return Err("zbuild --native: script compiled to an empty chunk".to_string());
390    }
391    Ok(chunk)
392}
393
394/// Locate the zsh runtime staticlib to link against. `ZSHRS_AOT_RUNTIME_LIB`
395/// overrides; otherwise look for `libzsh.a` beside the running executable
396/// (the dev `target/<profile>/` layout).
397fn runtime_staticlib() -> Result<PathBuf, String> {
398    if let Ok(p) = std::env::var("ZSHRS_AOT_RUNTIME_LIB") {
399        return Ok(PathBuf::from(p));
400    }
401    let exe = std::env::current_exe().map_err(|e| e.to_string())?;
402    if let Some(dir) = exe.parent() {
403        let cand = dir.join("libzsh.a");
404        if cand.exists() {
405            return Ok(cand);
406        }
407    }
408    Err("could not locate libzsh.a (set ZSHRS_AOT_RUNTIME_LIB)".to_string())
409}
410
411/// `zbuild --native --in A.zsh [--in B.zsh] --out OUT`: AOT-compile the inputs
412/// to native machine code and link a standalone executable. Multiple inputs are
413/// concatenated in order into one program, matching the source-trailer build.
414pub fn build_native(script_paths: &[PathBuf], out_path: &Path) -> Result<PathBuf, String> {
415    if script_paths.is_empty() {
416        return Err("zbuild --native: at least one --in PATH required".to_string());
417    }
418    let mut source = String::new();
419    for p in script_paths {
420        let s = fs::read_to_string(p)
421            .map_err(|e| format!("zbuild --native: cannot read {}: {}", p.display(), e))?;
422        source.push_str(&s);
423        if !source.ends_with('\n') {
424            source.push('\n');
425        }
426    }
427    let chunk = compile_source_to_chunk(&source)?;
428
429    let runtime_lib = runtime_staticlib()?;
430    if !runtime_lib.exists() {
431        return Err(format!(
432            "zbuild --native: runtime staticlib not found at {}",
433            runtime_lib.display()
434        ));
435    }
436
437    let obj = out_path.with_extension("o");
438    fusevm::aot::compile_object(&chunk, &obj).map_err(|e| format!("zbuild --native: {}", e))?;
439
440    let stub = out_path.with_extension("aot_main.c");
441    fs::write(
442        &stub,
443        b"extern long fusevm_aot_run_embedded(void);\nint main(void){return (int)fusevm_aot_run_embedded();}\n" as &[u8],
444    )
445    .map_err(|e| format!("zbuild --native: write entry stub: {}", e))?;
446
447    let mut cmd = std::process::Command::new("cc");
448    cmd.arg(&stub).arg(&obj).arg(&runtime_lib);
449    // zsh's terminfo/termcap modules need the system terminal library.
450    cmd.arg("-lncurses");
451    if !cfg!(target_os = "macos") {
452        // glibc keeps the math functions in a separate `libm`, and rustc only
453        // passes `-lm` for the crates IT links. Linking the AOT object by hand
454        // has to say so: `zsh/mathfunc` (`callmathfunc` — `acos`, `yn`, `j0`,
455        // `lgamma`, …) and p10k's `humanize_bytes` (`log`) are undefined
456        // references without it, and the Linux link failed on a page of them.
457        // macOS carries the same symbols in libSystem, which `cc` always links.
458        cmd.arg("-lm");
459    }
460    if !cfg!(any(target_os = "linux", target_os = "android")) {
461        // `ucs4tomb` (src/ported/utils.rs) takes the iconv(3) body off glibc,
462        // and its `#[link(name = "iconv")]` does not travel inside a staticlib,
463        // so the hand-built link has to name the library itself.
464        cmd.arg("-liconv");
465    }
466    if cfg!(target_os = "macos") {
467        // Frameworks pulled by zsh's transitive deps: chrono/iana-time-zone
468        // (CoreFoundation), notify/FSEvents (CoreServices), and TLS/keychain
469        // and network-config crates (Security, SystemConfiguration).
470        for fw in [
471            "CoreFoundation",
472            "CoreServices",
473            "Security",
474            "SystemConfiguration",
475        ] {
476            cmd.arg("-framework").arg(fw);
477        }
478    }
479    cmd.arg("-o").arg(out_path);
480    // The shell's SIGCHLD reaper (`waitpid(-1)`, c:Src/signals.c:285
481    // wait_for_processes) runs on any thread and will collect `cc` before this
482    // wait gets there, leaving it with ECHILD — which reached the user as
483    // `zbuild --native: invoking cc: No child processes (os error 10)` on the
484    // macOS corpus job, after the link itself had succeeded. Wait the way
485    // every other foreground spawn does; see foreground_status.
486    let status = crate::fusevm_bridge::foreground_status(&mut cmd)
487        .map_err(|e| format!("zbuild --native: invoking cc: {}", e))?;
488    let _ = fs::remove_file(&stub);
489    let _ = fs::remove_file(&obj);
490    if !status.success() {
491        return Err(format!(
492            "zbuild --native: link failed (cc exit {:?})",
493            status.code()
494        ));
495    }
496    set_executable(out_path);
497    Ok(out_path.to_path_buf())
498}
499
500#[cfg(test)]
501mod tests {
502    use super::*;
503
504    // ── `zbuild --native` gets real machine code, not a stub ──────────────
505    //
506    // fusevm decides between a register-native lowering and a threaded one (a
507    // native block per op, each calling the runtime). Both are machine code;
508    // what is NOT acceptable is the third outcome that shipped for a while — a
509    // native plan covering a single op, with a `resume` handing the entire
510    // program back to the interpreter. That produced a working binary that
511    // printed the right answer, so only an assertion on the lowering catches it.
512    //
513    // These run the REAL zsh front end (`compile_source_to_chunk`, the same call
514    // `build_native` makes), so they also fail if a compiler change reshapes a
515    // chunk into something the analysis chokes on. No linking and no `cc` — the
516    // decision is made from the chunk alone.
517
518    use fusevm::aot::{lowering_for, Lowering};
519
520    /// Ops in the chunk zshrs compiles for `src`, and the lowering it gets.
521    fn lowering_of(src: &str) -> (usize, Lowering) {
522        let chunk = compile_source_to_chunk(src).expect("compiles");
523        (chunk.ops.len(), lowering_for(&chunk))
524    }
525
526    #[test]
527    fn native_lowering_covers_the_whole_script_not_one_op() {
528        // A spread of real shell: the simplest possible command, a loop, a
529        // function, and a script with arrays and expansion. Whatever path each
530        // takes, it must not be a native plan that gives up after a couple of
531        // ops — that is the shape that silently interpreted everything.
532        for (label, src) in [
533            ("simplest", "print hello\n"),
534            (
535                "arith loop",
536                "integer i sum\nsum=0\nfor (( i = 1; i <= 10; i++ )); do (( sum += i )); done\nprint $sum\n",
537            ),
538            (
539                "function",
540                "f() { local n=$1; print $(( n * 2 )) }\nf 21\n",
541            ),
542            (
543                "arrays",
544                "a=(one two three)\nfor x in $a; do print \"item:$x\"; done\nprint ${#a} ${a[2]}\n",
545            ),
546        ] {
547            let (ops, lowering) = lowering_of(src);
548            assert!(ops > 10, "{label}: chunk is only {ops} ops");
549            match lowering {
550                // Threaded is the expected answer for shell: every op becomes a
551                // native block. Nothing to assert beyond "it was chosen".
552                Lowering::Threaded => {}
553                // If a shell chunk ever does lower natively, it must cover the
554                // program rather than deopt out of the first handful of ops.
555                Lowering::Native { covered, .. } => assert!(
556                    covered * 2 >= ops,
557                    "{label}: native plan covers {covered} of {ops} ops — \
558                     the rest is interpreted, which a threaded lowering would not be"
559                ),
560            }
561        }
562    }
563
564    #[test]
565    fn emitted_code_grows_with_the_program() {
566        // The regression's signature: `print hello` and a million-iteration loop
567        // emitted byte-identical drivers, because neither's ops were lowered.
568        //
569        // Object size alone does not catch that — an object is the driver PLUS
570        // the serialized chunk, and the chunk grows with the script either way.
571        // Subtract the blob and what is left is code: with a stub driver that
572        // remainder is a fixed ~100 bytes whatever the script does.
573        let small = compile_source_to_chunk("print hello\n").expect("compiles");
574        let mut big_src = String::new();
575        for i in 0..60 {
576            big_src.push_str(&format!("print line{i}\n"));
577        }
578        let big = compile_source_to_chunk(&big_src).expect("compiles");
579        assert!(
580            big.ops.len() > small.ops.len() * 4,
581            "test setup: {} vs {} ops",
582            big.ops.len(),
583            small.ops.len()
584        );
585
586        let dir = tempfile::tempdir().expect("tempdir");
587        let code_len = |chunk: &fusevm::Chunk, name: &str| -> i64 {
588            let path = dir.path().join(name);
589            fusevm::aot::compile_object(chunk, &path).expect("emits an object");
590            let obj = fs::metadata(&path).expect("object exists").len() as i64;
591            let blob = bincode::serialize(chunk).expect("chunk serializes").len() as i64;
592            obj - blob
593        };
594        let small_code = code_len(&small, "small.o");
595        let big_code = code_len(&big, "big.o");
596        assert!(
597            small_code > 0 && big_code > small_code * 3,
598            "code size did not follow the program: {small_code} vs {big_code} bytes \
599             (a stub driver is constant-size whatever the script)"
600        );
601    }
602
603    #[test]
604    fn zsh_builtins_are_not_arity_honest() {
605        // Why a shell chunk takes the threaded path, stated as a test so the day
606        // it stops being true is visible. `Op::CallBuiltin(id, argc)` may only be
607        // lowered inline when every handler pops exactly `argc`; zshrs's do not
608        // (BUILTIN_XTRACE_ARGS pops one and peeks the rest), so the compiler must
609        // not set the opt-in. If an audit ever makes the table arity-honest,
610        // flip this and `compile_source_to_chunk` together — never one alone.
611        let chunk = compile_source_to_chunk("print hello\n").expect("compiles");
612        assert!(
613            !chunk.builtin_argc_is_arity,
614            "the zsh compiler declared builtin argc as a stack arity; every \
615             handler in fusevm_bridge.rs must pop exactly argc before this holds"
616        );
617    }
618
619    fn mkfile(name: &str, src: &str) -> EmbeddedFile {
620        EmbeddedFile {
621            name: name.into(),
622            source: src.into(),
623        }
624    }
625
626    #[test]
627    fn build_trailer_layout_matches_spec() {
628        let t = build_trailer(0x11_22_33_44, 0xAA_BB_CC_DD, AOT_VERSION_V2);
629        assert_eq!(t.len(), TRAILER_LEN as usize);
630        // little-endian u64 compressed_len
631        assert_eq!(&t[0..8], &0x11_22_33_44u64.to_le_bytes());
632        // little-endian u64 uncompressed_len
633        assert_eq!(&t[8..16], &0xAA_BB_CC_DDu64.to_le_bytes());
634        // little-endian u32 version
635        assert_eq!(&t[16..20], &AOT_VERSION_V2.to_le_bytes());
636        // 20..24 reserved zeros
637        assert_eq!(&t[20..24], &[0u8; 4]);
638        // magic
639        assert_eq!(&t[24..32], AOT_MAGIC);
640    }
641
642    #[test]
643    fn payload_v2_roundtrip_single_file() {
644        let files = vec![mkfile("hello.zsh", "echo hi\n")];
645        let encoded = encode_payload_v2(&files);
646        let decoded = decode_payload_v2(&encoded).expect("decode v2");
647        assert_eq!(decoded.0.len(), 1);
648        assert_eq!(decoded.0[0].name, "hello.zsh");
649        assert_eq!(decoded.0[0].source, "echo hi\n");
650    }
651
652    #[test]
653    fn payload_v2_roundtrip_multiple_files_preserves_order() {
654        let files = vec![
655            mkfile("a.zsh", "echo a\n"),
656            mkfile("b.zsh", "echo b\n"),
657            mkfile("c.zsh", "echo c\n"),
658        ];
659        let encoded = encode_payload_v2(&files);
660        let decoded = decode_payload_v2(&encoded).expect("decode v2");
661        assert_eq!(decoded.0.len(), 3);
662        assert_eq!(decoded.0[0].name, "a.zsh");
663        assert_eq!(decoded.0[1].name, "b.zsh");
664        assert_eq!(decoded.0[2].name, "c.zsh");
665        assert_eq!(decoded.0[0].source, "echo a\n");
666        assert_eq!(decoded.0[1].source, "echo b\n");
667        assert_eq!(decoded.0[2].source, "echo c\n");
668    }
669
670    #[test]
671    fn payload_v2_roundtrip_zero_files() {
672        let files: Vec<EmbeddedFile> = vec![];
673        let encoded = encode_payload_v2(&files);
674        // u32 count=0 → 4 bytes of zero.
675        assert_eq!(encoded.as_slice(), &0u32.to_le_bytes());
676        let decoded = decode_payload_v2(&encoded).expect("decode zero-file v2");
677        assert!(decoded.0.is_empty());
678    }
679
680    #[test]
681    fn payload_v2_handles_empty_source() {
682        let files = vec![mkfile("empty.zsh", "")];
683        let encoded = encode_payload_v2(&files);
684        let decoded = decode_payload_v2(&encoded).expect("decode empty source");
685        assert_eq!(decoded.0[0].source, "");
686        assert_eq!(decoded.0[0].name, "empty.zsh");
687    }
688
689    #[test]
690    fn payload_v2_preserves_utf8_in_names_and_sources() {
691        let files = vec![
692            mkfile("名前.zsh", "echo こんにちは\n"),
693            mkfile("emoji-🚀.zsh", "echo $'\\xf0\\x9f\\x9a\\x80'\n"),
694        ];
695        let encoded = encode_payload_v2(&files);
696        let decoded = decode_payload_v2(&encoded).expect("decode utf8");
697        assert_eq!(decoded.0[0].name, "名前.zsh");
698        assert_eq!(decoded.0[0].source, "echo こんにちは\n");
699        assert_eq!(decoded.0[1].name, "emoji-🚀.zsh");
700    }
701
702    #[test]
703    fn payload_v2_rejects_truncated_input() {
704        let files = vec![mkfile("x.zsh", "echo x\n")];
705        let mut encoded = encode_payload_v2(&files);
706        // Drop the last byte — source ends prematurely.
707        encoded.pop();
708        assert!(decode_payload_v2(&encoded).is_none());
709    }
710
711    #[test]
712    fn payload_v2_rejects_empty_buffer() {
713        assert!(decode_payload_v2(&[]).is_none());
714    }
715
716    #[test]
717    fn payload_v2_rejects_lying_count_header() {
718        // Header claims 5 files but no body follows.
719        let mut buf = Vec::new();
720        buf.extend_from_slice(&5u32.to_le_bytes());
721        assert!(decode_payload_v2(&buf).is_none());
722    }
723
724    #[test]
725    fn payload_v1_decodes_legacy_single_script() {
726        let name = "old.zsh";
727        let source = "echo legacy\n";
728        let mut buf = Vec::new();
729        buf.extend_from_slice(&(name.len() as u32).to_le_bytes());
730        buf.extend_from_slice(name.as_bytes());
731        buf.extend_from_slice(source.as_bytes());
732        let decoded = decode_payload_v1(&buf).expect("decode v1");
733        assert_eq!(decoded.0.len(), 1);
734        assert_eq!(decoded.0[0].name, name);
735        assert_eq!(decoded.0[0].source, source);
736    }
737
738    #[test]
739    fn payload_v1_rejects_short_buffer() {
740        assert!(decode_payload_v1(&[]).is_none());
741        assert!(decode_payload_v1(&[1, 2, 3]).is_none());
742    }
743
744    #[test]
745    fn payload_v1_rejects_name_len_larger_than_buffer() {
746        let mut buf = Vec::new();
747        buf.extend_from_slice(&999u32.to_le_bytes()); // name_len = 999
748        buf.extend_from_slice(b"abc"); // only 3 bytes follow
749        assert!(decode_payload_v1(&buf).is_none());
750    }
751
752    #[test]
753    fn payload_v1_handles_empty_source_after_name() {
754        let mut buf = Vec::new();
755        let name = "x";
756        buf.extend_from_slice(&(name.len() as u32).to_le_bytes());
757        buf.extend_from_slice(name.as_bytes());
758        // No source bytes after — should still decode with empty source.
759        let decoded = decode_payload_v1(&buf).expect("decode v1 empty source");
760        assert_eq!(decoded.0[0].source, "");
761    }
762
763    #[test]
764    fn append_and_load_roundtrip_single_file() {
765        let tmp = tempfile::NamedTempFile::new().expect("temp file");
766        // Pre-populate with non-magic data simulating an existing binary.
767        std::fs::write(tmp.path(), b"FAKE-ELF-PREFIX").expect("write prefix");
768        let files = vec![mkfile("greet.zsh", "echo hello\n")];
769        append_embedded_files(tmp.path(), &files).expect("append");
770        let loaded = try_load_embedded(tmp.path()).expect("load back");
771        assert_eq!(loaded.0.len(), 1);
772        assert_eq!(loaded.0[0].name, "greet.zsh");
773        assert_eq!(loaded.0[0].source, "echo hello\n");
774    }
775
776    #[test]
777    fn append_and_load_roundtrip_multiple_files() {
778        let tmp = tempfile::NamedTempFile::new().expect("temp file");
779        std::fs::write(tmp.path(), b"PREFIX").expect("write prefix");
780        let files = vec![
781            mkfile("first.zsh", "first()  { :; }\n"),
782            mkfile("second.zsh", "first\nsecond()  { :; }\n"),
783            mkfile("third.zsh", "echo done\n"),
784        ];
785        append_embedded_files(tmp.path(), &files).expect("append");
786        let loaded = try_load_embedded(tmp.path()).expect("load back");
787        assert_eq!(loaded.0.len(), 3);
788        assert_eq!(loaded.0[0].name, "first.zsh");
789        assert_eq!(loaded.0[1].name, "second.zsh");
790        assert_eq!(loaded.0[2].name, "third.zsh");
791    }
792
793    #[test]
794    fn try_load_embedded_returns_none_without_magic() {
795        let tmp = tempfile::NamedTempFile::new().expect("temp file");
796        // Write at least TRAILER_LEN bytes but no magic in the trailing slot.
797        std::fs::write(tmp.path(), vec![0u8; 64]).expect("write");
798        assert!(try_load_embedded(tmp.path()).is_none());
799    }
800
801    #[test]
802    fn try_load_embedded_returns_none_for_small_file() {
803        let tmp = tempfile::NamedTempFile::new().expect("temp file");
804        // Smaller than TRAILER_LEN — fast bail.
805        std::fs::write(tmp.path(), b"tiny").expect("write");
806        assert!(try_load_embedded(tmp.path()).is_none());
807    }
808
809    #[test]
810    fn try_load_embedded_returns_none_for_missing_path() {
811        let path = std::path::PathBuf::from("/this/path/does/not/exist/zshrs-aot");
812        assert!(try_load_embedded(&path).is_none());
813    }
814
815    #[test]
816    fn try_load_embedded_returns_none_for_zero_compressed_len() {
817        // Forge a trailer with valid magic but compressed_len=0.
818        let tmp = tempfile::NamedTempFile::new().expect("temp file");
819        let mut data = vec![0u8; 100];
820        let trailer = build_trailer(0, 0, AOT_VERSION_V2);
821        data.extend_from_slice(&trailer);
822        std::fs::write(tmp.path(), &data).expect("write");
823        assert!(try_load_embedded(tmp.path()).is_none());
824    }
825
826    #[test]
827    fn try_load_embedded_rejects_unknown_version() {
828        // Build a valid magic + size frame, but version=99 — unsupported.
829        let tmp = tempfile::NamedTempFile::new().expect("temp file");
830        let payload = encode_payload_v2(&[mkfile("x.zsh", "y\n")]);
831        let compressed = zstd::stream::encode_all(&payload[..], 3).expect("zstd");
832        let mut data = vec![0u8; 32]; // prefix
833        data.extend_from_slice(&compressed);
834        let trailer = build_trailer(compressed.len() as u64, payload.len() as u64, 99);
835        data.extend_from_slice(&trailer);
836        std::fs::write(tmp.path(), &data).expect("write");
837        assert!(try_load_embedded(tmp.path()).is_none());
838    }
839
840    #[test]
841    fn try_load_embedded_rejects_corrupt_uncompressed_len() {
842        // Magic + version OK, but uncompressed_len lies → decoder returns None.
843        let tmp = tempfile::NamedTempFile::new().expect("temp file");
844        let payload = encode_payload_v2(&[mkfile("x.zsh", "y\n")]);
845        let compressed = zstd::stream::encode_all(&payload[..], 3).expect("zstd");
846        let mut data = vec![0u8; 32];
847        data.extend_from_slice(&compressed);
848        // Lie: claim uncompressed is one byte larger than reality.
849        let trailer = build_trailer(
850            compressed.len() as u64,
851            payload.len() as u64 + 1,
852            AOT_VERSION_V2,
853        );
854        data.extend_from_slice(&trailer);
855        std::fs::write(tmp.path(), &data).expect("write");
856        assert!(try_load_embedded(tmp.path()).is_none());
857    }
858
859    #[test]
860    fn build_rejects_empty_input_list() {
861        let out = std::path::PathBuf::from("/tmp/zshrs-aot-empty-out");
862        let res = build(&[], &out);
863        assert!(res.is_err());
864        let err = res.unwrap_err();
865        assert!(err.contains("at least one"), "got: {err}");
866    }
867}