zisk-build 1.1.0-alpha

use crate::{
    command::create_command, utils::cargo_rerun_if_changed, BuildArgs, HELPER_TARGET_SUBDIR,
    ZISK_TARGET,
};
use cargo_metadata::camino::Utf8PathBuf;
use std::{
    io::{BufRead, BufReader},
    path::PathBuf,
    process::{exit, Command, Stdio},
    thread,
};
use zisk_rom_setup::{
    assembly_files_exist, gen_assembly, get_assembly_file_paths, get_output_path,
};

use anyhow::{Context, Result};

fn should_skip_guest_build() -> bool {
    if std::env::var("SKIP_GUEST_BUILD").is_ok() {
        return true;
    }
    // cargo clippy sets RUSTC_WORKSPACE_WRAPPER to the clippy-driver binary.
    // Cross-compiling the guest during a clippy run is pointless and would fail
    // without the ZisK toolchain installed, so skip it.
    if std::env::var("RUSTC_WORKSPACE_WRAPPER").map(|v| v.contains("clippy")).unwrap_or(false) {
        return true;
    }
    false
}

// Helper for building a ZisK program.
pub(crate) fn build_program_internal(path: &str, args: Option<BuildArgs>) {
    // Always declare the cfg so rustc doesn't warn about it being unexpected in the host crate.
    println!("cargo:rustc-check-cfg=cfg(zisk_skip_guest_build)");

    if should_skip_guest_build() {
        println!("cargo:rustc-cfg=zisk_skip_guest_build");
        return;
    }

    // Get the root package name and metadata.
    let program_dir = std::path::Path::new(path);
    let metadata_file = program_dir.join("Cargo.toml");
    let mut metadata_cmd = cargo_metadata::MetadataCommand::new();
    let metadata = metadata_cmd.manifest_path(metadata_file).exec().unwrap();

    // Activate the build command if the dependencies change.
    cargo_rerun_if_changed(&metadata, program_dir);

    // Build the program with the given arguments.
    let path_output = if let Some(args) = &args {
        execute_build_program(args, Some(program_dir.to_path_buf()))
    } else {
        // Detect the host's build profile and use it for the guest program
        let profile = std::env::var("PROFILE").unwrap_or_else(|_| "debug".to_string());
        let default_args = BuildArgs { release: profile == "release", ..Default::default() };

        execute_build_program(&default_args, Some(program_dir.to_path_buf()))
    };
    if let Err(err) = path_output {
        panic!("Failed to build ZisK program: {err:#}.");
    }

    if let Err(err) = crate::aggregation::process_aggregations(program_dir) {
        panic!("Failed to process ZisK aggregation definitions: {err:#}.");
    }
}

pub fn execute_build_program(
    args: &BuildArgs,
    program_dir: Option<PathBuf>,
) -> Result<Vec<(String, Utf8PathBuf)>> {
    println!("cargo:rustc-check-cfg=cfg(zisk_skip_guest_build)");

    if should_skip_guest_build() {
        println!("cargo:rustc-cfg=zisk_skip_guest_build");
        return Ok(vec![]);
    }

    // If the program directory is not specified, use the current directory.
    let program_dir = program_dir
        .unwrap_or_else(|| std::env::current_dir().expect("Failed to get current directory."));
    let program_dir: Utf8PathBuf =
        program_dir.try_into().expect("Failed to convert PathBuf to Utf8PathBuf");

    // Get the program metadata.
    let program_metadata_file = program_dir.join("Cargo.toml");
    let mut program_metadata_cmd = cargo_metadata::MetadataCommand::new();
    let program_metadata = program_metadata_cmd.manifest_path(program_metadata_file).exec()?;

    // Get the command corresponding to Docker or local build.
    let mut cmd = create_command(args, &program_dir, &program_metadata)?;

    // Guest rustflags + linker script; keep the temp file alive until cargo
    // finishes. Env rustflags are NOT inherited: in this build-script context
    // they are the host build's flags, injected by the outer cargo.
    let target_features =
        crate::target_features_from_features(args.features.as_deref(), args.all_features);
    let _linker_script = crate::apply_guest_rustflags(
        &mut cmd,
        Some(program_dir.as_std_path()),
        false,
        &target_features,
    )?;

    let target_elf_paths = generate_elf_paths(&program_metadata, Some(args))?;

    if target_elf_paths.len() > 1 && args.elf_name.is_some() {
        anyhow::bail!("--elf-name is not supported when --output-directory is used and multiple ELFs are built.");
    }

    execute_command(cmd)?;

    // Generate assembly for all ELF files (only if not already generated)
    let asm = args.asm.unwrap_or(false);
    let hints = args.hints.unwrap_or(false);

    let output_path = get_output_path(&None)?;
    for (_, elf_path) in target_elf_paths.iter() {
        let elf_path_std = elf_path.as_std_path();

        let assembly_exists = assembly_files_exist(elf_path_std, &output_path, hints)?;
        let hints_marker = output_path.join(format!(
            "{}.assembly_hints",
            elf_path_std.file_name().unwrap().to_string_lossy()
        ));
        let new_value = if hints { "on" } else { "off" };

        let hints_changed = match std::fs::read_to_string(&hints_marker) {
            Ok(prev) => prev != new_value,
            Err(_) => true,
        };

        if asm && (!assembly_exists || hints_changed) {
            gen_assembly(elf_path_std, &None, hints, true)?;
            std::fs::write(&hints_marker, new_value)?;
        }

        // Tell cargo to rerun if any assembly file is deleted
        let assembly_files = get_assembly_file_paths(elf_path_std, &output_path, hints)?;
        for asm_file in assembly_files {
            println!("cargo:rerun-if-changed={}", asm_file.display());
        }
    }

    if let Some(output_directory) = &args.output_directory {
        // The path to the output directory, maybe relative or absolute.
        let output_directory = PathBuf::from(output_directory);

        // Ensure the output directory is a directory. If it doesnt exist, this is false.
        if output_directory.is_file() {
            anyhow::bail!("--output-directory is a file.");
        }

        // Ensure the output directory exists.
        std::fs::create_dir_all(&output_directory)?;

        // Copy the ELF file to the output directory.
        for (_, elf_path) in target_elf_paths.iter() {
            let elf_path = elf_path.to_path_buf();
            let elf_name = elf_path.file_name().expect("ELF path has a file name");
            let output_path = output_directory.join(args.elf_name.as_deref().unwrap_or(elf_name));

            std::fs::copy(&elf_path, &output_path)?;
        }
    }

    print_elf_paths_cargo_directives(&target_elf_paths, hints);

    Ok(target_elf_paths)
}

/// Execute the command and handle the output depending on the context.
pub(crate) fn execute_command(mut command: Command) -> Result<()> {
    // Add necessary tags for stdout and stderr from the command.
    let mut child = command
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .spawn()
        .context("failed to spawn command")?;
    let stdout = BufReader::new(child.stdout.take().unwrap());
    let stderr = BufReader::new(child.stderr.take().unwrap());

    // Add prefix to the output of the process depending on the context.
    let msg = "[ZisK] ";

    // Pipe stdout and stderr to the parent process with [docker] prefix
    let stdout_handle = thread::spawn(move || {
        stdout.lines().for_each(|line| {
            println!("{} {}", msg, line.unwrap());
        });
    });
    stderr.lines().for_each(|line| {
        eprintln!("{} {}", msg, line.unwrap());
    });
    stdout_handle.join().unwrap();

    // Wait for the child process to finish and check the result.
    let result = child.wait()?;
    if !result.success() {
        // Error message is already printed by cargo.
        exit(result.code().unwrap_or(1))
    }
    Ok(())
}

/// Collects the list of targets that would be built and their output ELF file paths.
pub fn generate_elf_paths(
    metadata: &cargo_metadata::Metadata,
    args: Option<&BuildArgs>,
) -> Result<Vec<(String, Utf8PathBuf)>> {
    let profile = args.map(|v| if v.release { "release" } else { "debug" }).unwrap_or("debug");
    let mut target_elf_paths = vec![];

    let packages_to_iterate = match args {
        Some(args) if !args.packages.is_empty() => args
            .packages
            .iter()
            .map(|wanted_package| {
                metadata
                    .packages
                    .iter()
                    .find(|p| p.name == *wanted_package)
                    .ok_or_else(|| anyhow::anyhow!("cannot find package named {wanted_package}"))
                    .map(|p| p.id.clone())
            })
            .collect::<Result<Vec<_>>>()?,
        _ => {
            if let Some(root_package) = metadata.root_package() {
                vec![root_package.id.clone()]
            } else {
                metadata.workspace_default_members.to_vec()
            }
        }
    };

    for program_crate in packages_to_iterate {
        let program = metadata
            .packages
            .iter()
            .find(|p| p.id == program_crate)
            .ok_or_else(|| anyhow::anyhow!("cannot find package for {program_crate}"))?;

        for bin_target in program.targets.iter().filter(|t| {
            t.kind.contains(&cargo_metadata::TargetKind::Bin)
                && t.crate_types.contains(&cargo_metadata::CrateType::Bin)
        }) {
            if let Some(args) = args {
                if !args.binaries.is_empty() && !args.binaries.contains(&bin_target.name) {
                    continue;
                }
            }

            let elf_path = metadata
                .target_directory
                .join(HELPER_TARGET_SUBDIR)
                .join(ZISK_TARGET)
                .join(profile)
                .join(&bin_target.name);

            target_elf_paths.push((bin_target.name.to_owned(), elf_path));
        }
    }

    Ok(target_elf_paths)
}

fn print_elf_paths_cargo_directives(target_elf_paths: &[(String, Utf8PathBuf)], hints: bool) {
    for (target_name, elf_path) in target_elf_paths.iter() {
        // Only set env var if the ELF file actually exists
        if elf_path.exists() {
            println!("cargo:rustc-env=ZISK_ELF_{target_name}={elf_path}");
            if hints {
                println!("cargo:rustc-env=ZISK_ELF_{target_name}_WITH_HINTS=1");
            }

            // Compute and emit blake3 hash of the ELF file
            if let Ok(elf_bytes) = std::fs::read(elf_path) {
                let hash = blake3::hash(&elf_bytes).to_hex();
                println!("cargo:rustc-env=ZISK_ELF_HASH_{target_name}={hash}");
            }
        }
    }
}