library/stdarch/crates/intrinsic-test/src/common/gen_rust.rs

use itertools::Itertools;
use std::process::Command;

use crate::common::argument::ArgumentList;
use crate::common::intrinsic::Intrinsic;

use super::indentation::Indentation;
use super::intrinsic_helpers::IntrinsicTypeDefinition;

// The number of times each intrinsic will be called.
pub(crate) const PASSES: u32 = 20;

macro_rules! concatln {
    ($($lines:expr),* $(,)?) => {
        concat!($( $lines, "\n" ),*)
    };
}

fn write_cargo_toml_header(w: &mut impl std::io::Write, name: &str) -> std::io::Result<()> {
    writeln!(
        w,
        concatln!(
            "[package]",
            "name = \"{name}\"",
            "version = \"{version}\"",
            "authors = [{authors}]",
            "license = \"{license}\"",
            "edition = \"2018\"",
        ),
        name = name,
        version = env!("CARGO_PKG_VERSION"),
        authors = env!("CARGO_PKG_AUTHORS")
            .split(":")
            .format_with(", ", |author, fmt| fmt(&format_args!("\"{author}\""))),
        license = env!("CARGO_PKG_LICENSE"),
    )
}

pub fn write_bin_cargo_toml(
    w: &mut impl std::io::Write,
    module_count: usize,
) -> std::io::Result<()> {
    write_cargo_toml_header(w, "intrinsic-test-programs")?;

    writeln!(w, "[dependencies]")?;
    writeln!(w, "core_arch = {{ path = \"../crates/core_arch\" }}")?;

    for i in 0..module_count {
        writeln!(w, "mod_{i} = {{ path = \"mod_{i}/\" }}")?;
    }

    Ok(())
}

pub fn write_lib_cargo_toml(w: &mut impl std::io::Write, name: &str) -> std::io::Result<()> {
    write_cargo_toml_header(w, name)?;

    writeln!(w, "[dependencies]")?;
    writeln!(w, "core_arch = {{ path = \"../../crates/core_arch\" }}")?;

    Ok(())
}

pub fn write_main_rs<'a>(
    w: &mut impl std::io::Write,
    chunk_count: usize,
    cfg: &str,
    definitions: &str,
    intrinsics: impl Iterator<Item = &'a str> + Clone,
) -> std::io::Result<()> {
    writeln!(w, "#![feature(simd_ffi)]")?;
    writeln!(w, "#![feature(f16)]")?;
    writeln!(w, "#![allow(unused)]")?;

    // Cargo will spam the logs if these warnings are not silenced.
    writeln!(w, "#![allow(non_upper_case_globals)]")?;
    writeln!(w, "#![allow(non_camel_case_types)]")?;
    writeln!(w, "#![allow(non_snake_case)]")?;

    writeln!(w, "{cfg}")?;
    writeln!(w, "{definitions}")?;

    for module in 0..chunk_count {
        writeln!(w, "use mod_{module}::*;")?;
    }

    writeln!(w, "fn main() {{")?;

    writeln!(w, "    match std::env::args().nth(1).unwrap().as_str() {{")?;

    for binary in intrinsics {
        writeln!(w, "        \"{binary}\" => run_{binary}(),")?;
    }

    writeln!(
        w,
        "        other => panic!(\"unknown intrinsic `{{}}`\", other),"
    )?;

    writeln!(w, "    }}")?;
    writeln!(w, "}}")?;

    Ok(())
}

pub fn write_lib_rs<T: IntrinsicTypeDefinition>(
    w: &mut impl std::io::Write,
    notice: &str,
    cfg: &str,
    definitions: &str,
    intrinsics: &[Intrinsic<T>],
) -> std::io::Result<()> {
    write!(w, "{notice}")?;

    writeln!(w, "#![feature(simd_ffi)]")?;
    writeln!(w, "#![feature(f16)]")?;
    writeln!(w, "#![allow(unused)]")?;

    // Cargo will spam the logs if these warnings are not silenced.
    writeln!(w, "#![allow(non_upper_case_globals)]")?;
    writeln!(w, "#![allow(non_camel_case_types)]")?;
    writeln!(w, "#![allow(non_snake_case)]")?;

    writeln!(w, "{cfg}")?;

    writeln!(w, "{definitions}")?;

    let mut seen = std::collections::HashSet::new();

    for intrinsic in intrinsics {
        for arg in &intrinsic.arguments.args {
            if !arg.has_constraint() && arg.ty.is_rust_vals_array_const() {
                let name = arg.rust_vals_array_name().to_string();

                if seen.insert(name) {
                    ArgumentList::gen_arg_rust(arg, w, Indentation::default(), PASSES)?;
                }
            }
        }
    }

    for intrinsic in intrinsics {
        crate::common::gen_rust::create_rust_test_module(w, intrinsic)?;
    }

    Ok(())
}

pub fn compile_rust_programs(toolchain: Option<&str>, target: &str, linker: Option<&str>) -> bool {
    /* If there has been a linker explicitly set from the command line then
     * we want to set it via setting it in the RUSTFLAGS*/

    // This is done because `toolchain` is None when
    // the --generate-only flag is passed
    if toolchain.is_none() {
        return true;
    }

    trace!("Building cargo command");

    let mut cargo_command = Command::new("cargo");
    cargo_command.current_dir("rust_programs");

    // Do not use the target directory of the workspace please.
    cargo_command.env("CARGO_TARGET_DIR", "target");

    if toolchain.is_some_and(|val| !val.is_empty()) {
        cargo_command.arg(toolchain.unwrap());
    }
    cargo_command.args(["build", "--target", target, "--release"]);

    let mut rust_flags = "-Cdebuginfo=0".to_string();
    if let Some(linker) = linker {
        rust_flags.push_str(" -C linker=");
        rust_flags.push_str(linker);
        rust_flags.push_str(" -C link-args=-static");

        cargo_command.env("CPPFLAGS", "-fuse-ld=lld");
    }

    cargo_command.env("RUSTFLAGS", rust_flags);

    trace!("running cargo");

    if log::log_enabled!(log::Level::Trace) {
        cargo_command.stdout(std::process::Stdio::inherit());
        cargo_command.stderr(std::process::Stdio::inherit());
    }

    let output = cargo_command.output();
    trace!("cargo is done");

    if let Ok(output) = output {
        if output.status.success() {
            true
        } else {
            error!(
                "Failed to compile code for rust intrinsics\n\nstdout:\n{}\n\nstderr:\n{}",
                std::str::from_utf8(&output.stdout).unwrap_or(""),
                std::str::from_utf8(&output.stderr).unwrap_or("")
            );
            false
        }
    } else {
        error!("Command failed: {output:#?}");
        false
    }
}

pub fn generate_rust_test_loop<T: IntrinsicTypeDefinition>(
    w: &mut impl std::io::Write,
    intrinsic: &Intrinsic<T>,
    indentation: Indentation,
    specializations: &[Vec<i32>],
    passes: u32,
) -> std::io::Result<()> {
    let intrinsic_name = &intrinsic.name;

    // Each function (and each specialization) has its own type. Erase that type with a cast.
    let mut coerce = String::from("unsafe fn(");
    for _ in intrinsic.arguments.iter().filter(|a| !a.has_constraint()) {
        coerce += "_, ";
    }
    coerce += ") -> _";

    match specializations {
        [] => {
            writeln!(w, "    let specializations = [(\"\", {intrinsic_name})];")?;
        }
        [const_args] if const_args.is_empty() => {
            writeln!(w, "    let specializations = [(\"\", {intrinsic_name})];")?;
        }
        _ => {
            writeln!(w, "    let specializations = [")?;

            for specialization in specializations {
                let mut specialization: Vec<_> =
                    specialization.iter().map(|d| d.to_string()).collect();

                let const_args = specialization.join(",");

                // The identifier is reversed.
                specialization.reverse();
                let id = specialization.join("-");

                writeln!(
                    w,
                    "        (\"-{id}\", {intrinsic_name}::<{const_args}> as {coerce}),"
                )?;
            }

            writeln!(w, "    ];")?;
        }
    }

    write!(
        w,
        concatln!(
            "    for (id, f) in specializations {{",
            "        for i in 0..{passes} {{",
            "            unsafe {{",
            "{loaded_args}",
            "                let __return_value = f({args});",
            "                println!(\"Result {{id}}-{{}}: {{:?}}\", i + 1, {return_value});",
            "            }}",
            "        }}",
            "    }}",
        ),
        loaded_args = intrinsic.arguments.load_values_rust(indentation.nest_by(4)),
        args = intrinsic.arguments.as_call_param_rust(),
        return_value = intrinsic.results.print_result_rust(),
        passes = passes,
    )
}

/// Generate the specializations (unique sequences of const-generic arguments) for this intrinsic.
fn generate_rust_specializations(
    constraints: &mut impl Iterator<Item = impl Iterator<Item = i64>>,
) -> Vec<Vec<i32>> {
    let mut specializations = vec![vec![]];

    for constraint in constraints {
        specializations = constraint
            .flat_map(|right| {
                specializations.iter().map(move |left| {
                    let mut left = left.clone();
                    left.push(i32::try_from(right).unwrap());
                    left
                })
            })
            .collect();
    }

    specializations
}

// Top-level function to create complete test program
pub fn create_rust_test_module<T: IntrinsicTypeDefinition>(
    w: &mut impl std::io::Write,
    intrinsic: &Intrinsic<T>,
) -> std::io::Result<()> {
    trace!("generating `{}`", intrinsic.name);
    let indentation = Indentation::default();

    writeln!(w, "pub fn run_{}() {{", intrinsic.name)?;

    // Define the arrays of arguments.
    let arguments = &intrinsic.arguments;
    arguments.gen_arglists_rust(w, indentation.nested(), PASSES)?;

    // Define any const generics as `const` items, then generate the actual test loop.
    let specializations = generate_rust_specializations(
        &mut arguments
            .iter()
            .filter_map(|i| i.constraint.as_ref().map(|v| v.iter())),
    );

    generate_rust_test_loop(w, intrinsic, indentation, &specializations, PASSES)?;

    writeln!(w, "}}")?;

    Ok(())
}