arceos-rust 0.1.2

// Package name of the `lib` directory
const LIB_NAME: &str = "arceos_rust_interface";

use std::{
    collections::BTreeSet,
    env, fs,
    path::{Path, PathBuf},
    process::Command,
    time::UNIX_EPOCH,
};

fn main() {
    println!(
        "cargo:warning=Running build script for ArceOS rust library. Time: {:?}",
        std::time::SystemTime::now()
    );
    let out_dir = PathBuf::from(env::var("OUT_DIR").unwrap());

    let manifest_dir = PathBuf::from(env::var("CARGO_MANIFEST_DIR").unwrap());
    let lib_dir = manifest_dir.join("lib");
    let features = env::var("CARGO_CFG_FEATURE").unwrap_or_default();
    let passthrough_features = env::var("ARCEOS_RUST_FEATURES").unwrap_or_default();
    let feature_list = resolve_nested_features(&features, &passthrough_features);

    // generate configuration file
    let config_path = get_config_path(&manifest_dir, &out_dir, has_dynamic_platform(&feature_list));
    if let Some(config_path) = &config_path {
        println!("cargo:warning=config path: {}", config_path.display());
    }
    // build the ArceOS library
    let artifact_path = compile_project(
        &lib_dir,
        &out_dir,
        config_path.as_deref(),
        &feature_list,
        get_log_level(&features),
    );
    // rename symbols to avoid conflicts
    let lib_file = artifact_path.join(format!("lib{}.a", LIB_NAME));

    // Auto-generate rename list by extracting all ___rustc symbols
    let rename_list = generate_rename_list(&lib_file);
    rename_symbols(&lib_file, &rename_list);
    emit_linker_script_search_paths(&artifact_path, has_dynamic_platform(&feature_list));
    println!("cargo:rustc-link-lib=static={}", LIB_NAME);

    // Trick: specify a non-existent path to always trigger a rebuild
    // See https://doc.rust-lang.org/cargo/faq.html#why-is-cargo-rebuilding-my-code
    println!("cargo:rerun-if-changed=always");
}

fn emit_linker_script_search_paths(artifact_path: &Path, plat_dyn: bool) {
    println!("cargo:rustc-link-search=native={}", artifact_path.display());

    if !plat_dyn {
        // Use the final linker script generated by ax-runtime.
        let linker_script_path = find_runtime_linker_script(artifact_path);
        let linker_search_dirs = find_linker_search_dirs(artifact_path);

        println!(
            "cargo:warning=Linker script path: {}",
            linker_script_path.display()
        );

        for dir in linker_search_dirs {
            println!("cargo:warning=Linker script search dir: {}", dir.display());
            println!("cargo:rustc-link-search=native={}", dir.display());
        }
        return;
    }

    let scripts = dynamic_linker_scripts(artifact_path);
    for script in &scripts {
        if !script.is_file() {
            panic!(
                "Linker script not found at {}. Expected from inner ArceOS build.",
                script.display()
            );
        }
        println!("cargo:warning=Linker script path: {}", script.display());
        println!(
            "cargo:rustc-link-search=native={}",
            script.parent().unwrap().display()
        );
    }
}

fn dynamic_linker_scripts(artifact_path: &Path) -> Vec<PathBuf> {
    vec![
        find_runtime_linker_script(artifact_path),
        required_build_output_file(artifact_path, "axplat.x"),
        required_build_output_file(artifact_path, "link.x"),
        required_build_output_file(artifact_path, "someboot.x"),
    ]
}

fn required_build_output_file(artifact_path: &Path, file_name: &str) -> PathBuf {
    find_build_output_file(artifact_path, file_name).unwrap_or_else(|| {
        panic!(
            "Linker script {file_name} was not found under {}",
            artifact_path.display()
        )
    })
}

fn find_build_output_file(artifact_path: &Path, file_name: &str) -> Option<PathBuf> {
    let build_dir = artifact_path.join("build");
    walkdir::WalkDir::new(build_dir)
        .max_depth(3)
        .into_iter()
        .filter_map(Result::ok)
        .map(|entry| entry.into_path())
        .find(|path| {
            path.is_file() && path.file_name().and_then(|name| name.to_str()) == Some(file_name)
        })
}

fn get_config_path(manifest_dir: &Path, out_dir: &Path, plat_dyn: bool) -> Option<PathBuf> {
    if let Ok(path) = env::var("ARCEOS_RUST_CONFIG")
        && !path.trim().is_empty()
    {
        return Some(PathBuf::from(path));
    }

    if plat_dyn {
        return None;
    }

    Some(generate_config(manifest_dir, out_dir))
}

fn generate_config(manifest_dir: &Path, out_dir: &Path) -> PathBuf {
    let template = manifest_dir.join("defconfig.toml");
    let arch = get_arch();
    let platform = default_static_platform().unwrap_or_else(|| {
        panic!("ARCEOS_RUST_PLATFORM_CONFIG is required for static {arch} std builds")
    });
    let platform_config_path = get_platform_config_path(platform);
    let out_config_path = out_dir.join("axconfig.toml");

    let command = Command::new("axconfig-gen")
        .arg(&template)
        .arg(platform_config_path)
        .arg("-w")
        .arg(format!(r#"arch="{}""#, arch))
        .arg("-w")
        .arg(format!(r#"platform="{}""#, platform))
        .arg("-o")
        .arg(&out_config_path)
        .status()
        .expect("Failed to generate configuration file.");

    if !command.success() {
        panic!("Failed to generate configuration file.");
    }

    out_config_path
}

fn get_platform_config_path(platform: &str) -> PathBuf {
    if let Ok(path) = env::var("ARCEOS_RUST_PLATFORM_CONFIG") {
        return PathBuf::from(path);
    }

    // Fallback for direct upstream usage without axbuild.
    let output = Command::new(cargo())
        .arg("axplat")
        .arg("info")
        .arg(format!("ax-plat-{}", platform))
        .arg("-c")
        .output()
        .expect("Failed to get platform config path.");

    if !output.status.success() {
        println!(
            "cargo:warning=axplat output: {} {}",
            String::from_utf8_lossy(&output.stdout),
            String::from_utf8_lossy(&output.stderr)
        );
        panic!("Failed to get platform config path.");
    }

    PathBuf::from(String::from_utf8_lossy(&output.stdout).trim())
}

fn compile_project(
    lib_dir: &PathBuf,
    out_dir: &PathBuf,
    config_path: Option<&Path>,
    feature_list: &[String],
    log_level: &str,
) -> PathBuf {
    let profile = env::var("PROFILE").unwrap();
    let is_debug = profile == "debug";
    let arch = get_arch();
    let target = get_target(&arch);
    let plat_dyn = has_dynamic_platform(feature_list);
    let feature_list = feature_list.join(" ");

    let mut command = Command::new(cargo());
    command.env("AX_TARGET", target);
    command.env("AX_MODE", profile);
    if let Some(config_path) = config_path {
        command.env("AX_CONFIG_PATH", config_path);
    }
    if plat_dyn {
        append_inner_rustflag(&mut command, "-Crelocation-model=pic");
    }
    command.env("AX_LOG", log_level);
    if env::var("AX_IP").is_err() {
        command.env("AX_IP", "10.0.2.15");
    }
    if env::var("AX_GW").is_err() {
        command.env("AX_GW", "10.0.2.2");
    }
    command
        .current_dir(lib_dir)
        .arg("build")
        .arg("--target-dir")
        .arg(out_dir)
        .arg("--target")
        .arg(target)
        .arg("--no-default-features");
    if !feature_list.is_empty() {
        command.arg("--features").arg(feature_list);
    }
    if !is_debug {
        command.arg("--release");
    }
    println!(
        "cargo:warning=FATURES are {}",
        env::var("CARGO_CFG_FEATURE").unwrap_or("none".to_string())
    );
    println!("cargo:warning=command: {:?}", command);

    let status = command.status().expect("Failed to build ArceOS library.");

    if !status.success() {
        panic!("Failed to build ArceOS library.");
    }

    let build_type = if is_debug { "debug" } else { "release" };
    out_dir.join(target).join(build_type)
}

fn find_runtime_linker_script(artifact_path: &Path) -> PathBuf {
    latest_out_dir_with_script(&artifact_path.join("build"), "ax-runtime-", "linker.x")
        .join("linker.x")
}

fn find_linker_search_dirs(artifact_path: &Path) -> BTreeSet<PathBuf> {
    let build_dir = artifact_path.join("build");
    find_linker_search_dirs_for_owner(&build_dir, platform_linker_owner_prefix())
}

fn find_linker_search_dirs_for_owner(
    build_dir: &Path,
    platform_owner_prefix: &str,
) -> BTreeSet<PathBuf> {
    let mut dirs = BTreeSet::new();
    dirs.insert(latest_out_dir_with_script(
        build_dir,
        "ax-runtime-",
        "linker.x",
    ));
    dirs.insert(latest_out_dir_with_script(
        build_dir,
        platform_owner_prefix,
        "axplat.x",
    ));
    dirs
}

fn latest_out_dir_with_script(
    build_dir: &Path,
    package_prefix: &str,
    script_name: &str,
) -> PathBuf {
    let mut candidates = Vec::new();
    let entries = fs::read_dir(build_dir)
        .unwrap_or_else(|err| panic!("failed to read {}: {err}", build_dir.display()));
    for entry in entries {
        let path = entry
            .unwrap_or_else(|err| {
                panic!("failed to read entry under {}: {err}", build_dir.display())
            })
            .path();
        if path
            .file_name()
            .and_then(|name| name.to_str())
            .is_some_and(|name| name.starts_with(package_prefix))
        {
            let out_dir = path.join("out");
            let script_path = out_dir.join(script_name);
            if script_path.is_file() {
                let modified = script_path
                    .metadata()
                    .and_then(|metadata| metadata.modified())
                    .unwrap_or(UNIX_EPOCH);
                candidates.push((modified, out_dir));
            }
        }
    }

    candidates.sort_by(|(left_time, left_path), (right_time, right_path)| {
        right_time
            .cmp(left_time)
            .then_with(|| left_path.cmp(right_path))
    });
    candidates
        .into_iter()
        .map(|(_, path)| path)
        .next()
        .unwrap_or_else(|| {
            panic!(
                "Linker script {script_name} not found under {}. Expected from {package_prefix}* \
                 build.",
                build_dir
                    .join(format!("{package_prefix}*/out/{script_name}"))
                    .display()
            )
        })
}

fn platform_linker_owner_prefix() -> &'static str {
    platform_linker_owner_prefix_for(&selected_platform(), has_passthrough_feature("plat-dyn"))
}

fn platform_linker_owner_prefix_for(platform: &str, is_dynamic_platform: bool) -> &'static str {
    if is_dynamic_platform {
        return "axplat-dyn-";
    }

    match platform {
        "loongarch64-qemu-virt" => "ax-plat-loongarch64-qemu-virt-",
        "x86-qemu-q35" => "ax-plat-x86-qemu-q35-",
        _ => "ax-hal-",
    }
}

fn selected_platform() -> String {
    passthrough_features()
        .into_iter()
        .find_map(|feature| {
            feature
                .strip_prefix("ax-hal/")
                .filter(|platform| !matches!(*platform, "defplat" | "myplat" | "plat-dyn"))
                .map(ToOwned::to_owned)
        })
        .or_else(|| default_static_platform().map(ToOwned::to_owned))
        .unwrap_or_else(|| {
            panic!(
                "ARCEOS_RUST_FEATURES must include an explicit static platform feature for {}",
                get_arch()
            )
        })
}

fn has_passthrough_feature(feature_name: &str) -> bool {
    passthrough_features()
        .iter()
        .any(|feature| feature == feature_name || feature.ends_with(&format!("/{feature_name}")))
}

fn passthrough_features() -> Vec<String> {
    env::var("ARCEOS_RUST_FEATURES")
        .unwrap_or_default()
        .split(|ch: char| ch == ',' || ch.is_whitespace())
        .map(str::trim)
        .filter(|feature| !feature.is_empty())
        .map(map_nested_feature)
        .collect()
}

fn cargo() -> String {
    env::var("CARGO").unwrap()
}

fn append_inner_rustflag(command: &mut Command, flag: &str) {
    if let Ok(existing) = env::var("CARGO_ENCODED_RUSTFLAGS")
        && !existing.is_empty()
    {
        command.env("CARGO_ENCODED_RUSTFLAGS", format!("{existing}\x1f{flag}"));
        return;
    }

    if let Ok(existing) = env::var("RUSTFLAGS")
        && !existing.is_empty()
    {
        command.env("RUSTFLAGS", format!("{existing} {flag}"));
        return;
    }

    command.env("CARGO_ENCODED_RUSTFLAGS", flag);
}

fn generate_rename_list(lib_path: &Path) -> PathBuf {
    // Use rust-nm to extract all symbols from the library
    let nm_output = Command::new("rust-nm")
        .arg(lib_path)
        .output()
        .expect("Failed to run rust-nm. Please ensure llvm-tools-preview is installed.");

    if !nm_output.status.success() {
        panic!(
            "rust-nm failed:\n{}",
            String::from_utf8_lossy(&nm_output.stderr)
        );
    }

    let symbols_output = String::from_utf8_lossy(&nm_output.stdout);

    // Extract all unique symbols containing "___rustc" and generate rename pairs
    let mut rename_pairs = std::collections::HashSet::new();
    for line in symbols_output.lines() {
        // nm output format: <address> <type> <symbol_name>
        // Extract the last field (symbol name)
        if let Some(symbol) = line.split_whitespace().last()
            && symbol.contains("___rustc")
        {
            // Rename by appending "_1" suffix
            rename_pairs.insert((symbol.to_string(), format!("{}_1", symbol)));
        }
    }

    // Write rename list to a temporary file
    let rename_list_path = lib_path.parent().unwrap().join("symbol_rename_auto.txt");
    let mut file = std::fs::File::create(&rename_list_path)
        .expect("Failed to create auto-generated rename list file");

    use std::io::Write;
    for (old_symbol, new_symbol) in &rename_pairs {
        writeln!(file, "{} {}", old_symbol, new_symbol)
            .expect("Failed to write to rename list file");
    }

    println!(
        "cargo:warning=Auto-generated {} symbol rename rules",
        rename_pairs.len()
    );

    rename_list_path
}

fn rename_symbols(lib_path: &Path, rename_list: &Path) {
    let output = Command::new("rust-objcopy")
        .arg("--redefine-syms")
        .arg(rename_list)
        .arg(lib_path)
        .output();

    match output {
        Ok(output) if output.status.success() => {}
        Ok(output) => panic!(
            "Failed to rename symbols with rust-objcopy (exit: {}).\nstdout:\n{}\nstderr:\n{}",
            output
                .status
                .code()
                .map_or_else(|| "signal".to_string(), |c| c.to_string()),
            String::from_utf8_lossy(&output.stdout),
            String::from_utf8_lossy(&output.stderr),
        ),
        Err(_) => panic!(
            "Failed to run rust-objcopy. Please install required tools with:\n  rustup component \
             add llvm-tools-preview\n  cargo install cargo-binutils"
        ),
    }
}

fn get_arch() -> String {
    env::var("CARGO_CFG_TARGET_ARCH").unwrap()
}

fn default_static_platform_for_arch(arch: &str) -> Option<&'static str> {
    match arch {
        "x86_64" => Some("x86-pc"),
        "aarch64" => None,
        "riscv64" => None,
        "loongarch64" => Some("loongarch64-qemu-virt"),
        _ => panic!("Unsupported architecture: {}", arch),
    }
}

fn get_target(arch: &str) -> &'static str {
    match arch {
        "x86_64" => "x86_64-unknown-none",
        "aarch64" => "aarch64-unknown-none-softfloat",
        "riscv64" => "riscv64gc-unknown-none-elf",
        "loongarch64" => "loongarch64-unknown-none-softfloat",
        _ => panic!("Unsupported architecture: {}", arch),
    }
}

fn default_static_platform() -> Option<&'static str> {
    let arch = get_arch();
    default_static_platform_for_arch(&arch)
}

fn get_log_level(feature_list: &str) -> &str {
    let mut level = "off";
    for feature in feature_list.split(',') {
        if let Some(stripped) = feature.strip_prefix("log-level-") {
            level = stripped;
        }
    }
    level
}

fn map_nested_feature(feature: &str) -> String {
    if let Some(rest) = feature.strip_prefix("ax_hal/") {
        format!("ax-hal/{rest}")
    } else if let Some(rest) = feature.strip_prefix("ax_driver/") {
        format!("ax-driver/{rest}")
    } else {
        feature.to_string()
    }
}

fn resolve_nested_features(features: &str, passthrough_features: &str) -> Vec<String> {
    let mut feature_list = features
        .split(',')
        .chain(passthrough_features.split(','))
        .map(str::trim)
        .filter(|feature| !feature.is_empty())
        .map(map_nested_feature)
        .collect::<Vec<_>>();

    if !has_platform_feature(&feature_list) {
        feature_list.push("defplat".to_string());
    }
    feature_list.sort();
    feature_list.dedup();
    feature_list
}

fn has_platform_feature(features: &[String]) -> bool {
    features.iter().any(|feature| {
        matches!(
            feature.as_str(),
            "default" | "defplat" | "myplat" | "plat-dyn"
        ) || feature.starts_with("ax-hal/")
    })
}

fn has_dynamic_platform(features: &[String]) -> bool {
    features.iter().any(|feature| feature == "plat-dyn")
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn single_feature_builds_with_default_platform() {
        assert_eq!(
            resolve_nested_features("alloc", ""),
            vec!["alloc".to_string(), "defplat".to_string()]
        );
    }

    #[test]
    fn default_feature_keeps_inner_defaults() {
        assert_eq!(resolve_nested_features("default", ""), vec!["default"]);
    }

    #[test]
    fn explicit_platform_is_preserved_without_defplat() {
        assert_eq!(
            resolve_nested_features("alloc", "ax_hal/x86-pc"),
            vec!["alloc".to_string(), "ax-hal/x86-pc".to_string()]
        );
    }

    #[test]
    fn explicit_myplat_is_preserved_without_defplat() {
        assert_eq!(
            resolve_nested_features("alloc,myplat", ""),
            vec!["alloc".to_string(), "myplat".to_string()]
        );
    }

    #[test]
    fn explicit_dynamic_platform_is_preserved_without_defplat() {
        assert_eq!(
            resolve_nested_features("alloc,plat-dyn", ""),
            vec!["alloc".to_string(), "plat-dyn".to_string()]
        );
    }

    #[test]
    fn loongarch_platform_linker_is_platform_owned() {
        assert_eq!(
            platform_linker_owner_prefix_for("loongarch64-qemu-virt", false),
            "ax-plat-loongarch64-qemu-virt-"
        );
    }

    #[test]
    fn loongarch_default_platform_linker_is_platform_owned() {
        let platform = default_static_platform_for_arch("loongarch64").unwrap();

        assert_eq!(
            platform_linker_owner_prefix_for(platform, false),
            "ax-plat-loongarch64-qemu-virt-"
        );
    }

    #[test]
    fn generic_static_platform_linker_is_axhal_owned() {
        assert_eq!(
            platform_linker_owner_prefix_for("aarch64-qemu-virt", false),
            "ax-hal-"
        );
    }

    #[test]
    fn dynamic_platform_linker_is_axplat_dyn_owned() {
        assert_eq!(
            platform_linker_owner_prefix_for("loongarch64-qemu-virt", true),
            "axplat-dyn-"
        );
    }

    #[test]
    fn static_aarch64_has_no_deleted_default_platform() {
        assert_eq!(default_static_platform_for_arch("aarch64"), None);
    }

    #[test]
    fn inner_dynamic_build_uses_pic_codegen() {
        let mut command = Command::new("cargo");

        append_inner_rustflag(&mut command, "-Crelocation-model=pic");

        assert_eq!(
            command
                .get_envs()
                .find(|(key, _)| key == "CARGO_ENCODED_RUSTFLAGS"),
            Some((
                std::ffi::OsStr::new("CARGO_ENCODED_RUSTFLAGS"),
                Some(std::ffi::OsStr::new("-Crelocation-model=pic"))
            ))
        );
    }

    #[test]
    fn linker_search_dirs_include_runtime_and_platform_out_dirs() {
        let root = unique_temp_dir("arceos-rust-linker-search");
        let build_dir = root.join("build");
        let runtime_out = build_dir.join("ax-runtime-123/out");
        let platform_out = build_dir.join("ax-plat-loongarch64-qemu-virt-456/out");
        std::fs::create_dir_all(&runtime_out).unwrap();
        std::fs::create_dir_all(&platform_out).unwrap();
        std::fs::write(runtime_out.join("linker.x"), "").unwrap();
        std::fs::write(platform_out.join("axplat.x"), "").unwrap();

        let dirs = find_linker_search_dirs_for_owner(&build_dir, "ax-plat-loongarch64-qemu-virt-");

        assert!(dirs.contains(&runtime_out));
        assert!(dirs.contains(&platform_out));
        std::fs::remove_dir_all(root).unwrap();
    }

    #[test]
    fn dynamic_linker_scripts_include_runtime_entrypoint_and_includes() {
        let root = unique_temp_dir("arceos-rust-dyn-linker-search");
        let build_dir = root.join("build");
        let runtime_out = build_dir.join("ax-runtime-123/out");
        let axplat_out = build_dir.join("axplat-dyn-456/out");
        let somehal_out = build_dir.join("somehal-789/out");
        let someboot_out = build_dir.join("someboot-abc/out");
        std::fs::create_dir_all(&runtime_out).unwrap();
        std::fs::create_dir_all(&axplat_out).unwrap();
        std::fs::create_dir_all(&somehal_out).unwrap();
        std::fs::create_dir_all(&someboot_out).unwrap();
        std::fs::write(runtime_out.join("linker.x"), "").unwrap();
        std::fs::write(axplat_out.join("axplat.x"), "").unwrap();
        std::fs::write(somehal_out.join("link.x"), "").unwrap();
        std::fs::write(someboot_out.join("someboot.x"), "").unwrap();

        let scripts = dynamic_linker_scripts(&root);

        assert_eq!(scripts[0], runtime_out.join("linker.x"));
        assert!(scripts.contains(&axplat_out.join("axplat.x")));
        assert!(scripts.contains(&somehal_out.join("link.x")));
        assert!(scripts.contains(&someboot_out.join("someboot.x")));
        std::fs::remove_dir_all(root).unwrap();
    }

    fn unique_temp_dir(prefix: &str) -> PathBuf {
        let nanos = std::time::SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_nanos();
        env::temp_dir().join(format!("{prefix}-{}-{nanos}", std::process::id()))
    }
}