rllvm 0.4.1

//! Link-time optimization support.
//!
//! With `-flto` the compiler writes a bitcode module where an object file
//! belongs, so there is no section header to record a bitcode path in. Two
//! mechanisms answer that, and [`LtoMode`] chooses between them.

use serde::{Deserialize, Serialize};

use crate::constants::{DARWIN_SECTION_NAME, DARWIN_SEGMENT_NAME, ELF_SECTION_NAME};
use crate::error::Error;

/// How the wrappers handle a build that enables link-time optimization.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
#[serde(rename_all = "kebab-case")]
pub enum LtoMode {
    /// Compile a marker module naming each translation unit's bitcode and
    /// merge it into the LTO object. Works on every linker and both LTO
    /// flavours, and costs one extra compile per translation unit.
    #[default]
    Marker,
    /// Ask the linker to keep the module its own LTO pipeline merged. Full
    /// LTO only, and the link has to go through the wrapper.
    SaveTemps,
    /// Generate nothing, and say so. The behaviour before #96.
    Skip,
}

impl std::fmt::Display for LtoMode {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            LtoMode::Marker => write!(f, "marker"),
            LtoMode::SaveTemps => write!(f, "save-temps"),
            LtoMode::Skip => write!(f, "skip"),
        }
    }
}

impl std::str::FromStr for LtoMode {
    type Err = Error;

    fn from_str(value: &str) -> Result<Self, Self::Err> {
        match value {
            "marker" => Ok(Self::Marker),
            "save-temps" => Ok(Self::SaveTemps),
            "skip" => Ok(Self::Skip),
            other => Err(Error::ConfigError(format!(
                "Unknown lto_mode {other:?}; expected one of: marker, save-temps, skip"
            ))),
        }
    }
}

/// Which flavour of link-time optimization the command line asked for.
///
/// The distinction is not cosmetic: full LTO merges every module into one,
/// which the linker can be asked to save, and ThinLTO deliberately never
/// builds such a module.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum LtoFlavour {
    /// `-flto`, `-flto=full`.
    Full,
    /// `-flto=thin`.
    Thin,
}

/// Escape a path for an assembler string literal that is itself written as a
/// C string literal.
///
/// The result is interpolated into `__asm__("....ascii \"<path>\"...")`, so it
/// passes two decoders: the C compiler reads the literal and hands the
/// assembler what is left. Escaping for one layer only puts a bare `\` or `"`
/// in front of the assembler -- "invalid escape sequence" for the first, an
/// operand that ends early for the second -- and a path containing `\n` would
/// reach the assembler as a real newline, splitting one entry into two garbage
/// paths.
///
/// So each layer gets its own escape: the assembler needs `\\` to emit one
/// backslash, which the C source spells `\\\\`, and `\"` to emit one quote,
/// which the C source spells `\\\"`.
fn escape_for_assembler(path: &str) -> String {
    // Backslashes first: escaping quotes first would then re-escape the
    // backslashes this step introduces.
    path.replace('\\', r"\\\\").replace('"', r#"\\\""#)
}

/// Build the marker translation unit that records `recorded_path`.
///
/// The unit is module-level assembly rather than a `used` global on purpose. A
/// global's string literal carries a NUL terminator that corrupts the
/// newline-separated list, and a `used` const global lands in an *alloc*
/// section while rllvm's `llvm-objcopy` section is not, which makes `ld.bfd`
/// emit two sections with the same name. The directive writes exact bytes with
/// exact flags, so both linkers produce one section.
///
/// The target picks the directive: a bitcode object has no binary format to
/// read, and inspecting the host would be wrong under cross-compilation.
/// Exactly one of `__MACH__` and `__ELF__` is defined for the targets rllvm
/// supports, and both follow `--target=`.
///
/// `recorded_path` is the entry as it must appear in the section, already
/// resolved against `bitcode_root` -- the marker records exactly what every
/// other writer records, so one binary never mixes absolute and relative
/// entries.
pub fn marker_source(recorded_path: &str) -> String {
    let path = escape_for_assembler(recorded_path);
    format!(
        r#"/* Generated by rllvm. Records the bitcode path for an LTO object. */
#if defined(__MACH__)
__asm__(".section {DARWIN_SEGMENT_NAME},{DARWIN_SECTION_NAME},regular,no_dead_strip\n.ascii \"{path}\\n\"\n.previous");
#elif defined(__ELF__)
__asm__(".section {ELF_SECTION_NAME},\"\",@progbits\n.ascii \"{path}\\n\"\n.previous");
#else
#error "rllvm: -flto bitcode extraction supports ELF and Mach-O only; set lto_mode = \"skip\""
#endif
"#
    )
}

/// The linker flag that makes an LTO link keep its intermediate modules.
///
/// ld64 has its own spelling. Every ELF linker accepts gold's, including lld:
/// `-Wl,-plugin-opt=save-temps` was measured producing byte-identical file
/// names under ld.lld and ld.bfd.
pub fn save_temps_flag(target_is_darwin: bool) -> &'static str {
    if target_is_darwin {
        "-Wl,-save-temps"
    } else {
        "-Wl,-plugin-opt=save-temps"
    }
}

/// Whether `filename` is the merged module an LTO link saved for `output_name`.
///
/// The last stage before codegen, on both linker families: it is the module
/// the linker actually generated code from, and it is the only stage that
/// exists everywhere. ld64's earliest module is already internalized, while
/// lld and bfd expose a pre-internalize `preopt`, so no earlier stage is
/// comparable across linkers.
pub fn is_saved_module(output_name: &str, filename: &str, target_is_darwin: bool) -> bool {
    if target_is_darwin {
        filename == format!("{output_name}.lto.opt.bc")
    } else {
        // `<output>.<partition>.5.precodegen.bc`, one partition by default and
        // more under `--lto-partitions`. A plain prefix-and-suffix match would
        // also accept `prog.debug.0.5.precodegen.bc` for output `prog`, which
        // collides with a sibling output that happens to share a prefix, so
        // the partition segment is required to be digits and nothing else.
        let Some(rest) = filename.strip_prefix(&format!("{output_name}.")) else {
            return false;
        };
        let Some(partition) = rest.strip_suffix(".5.precodegen.bc") else {
            return false;
        };
        !partition.is_empty() && partition.chars().all(|c| c.is_ascii_digit())
    }
}

/// Whether `filename` is a save-temps artifact this link produced.
///
/// Used to clean up after a link rllvm added the flag to. The patterns are the
/// measured ones and deliberately nothing wider: this deletes files, and a
/// user file that merely starts with the output's name must survive.
pub fn is_save_temps_artifact(output_name: &str, filename: &str) -> bool {
    let Some(rest) = filename.strip_prefix(&format!("{output_name}.")) else {
        return false;
    };
    // The module rllvm keeps, under the name rllvm chose.
    if rest == "rllvm.bc" {
        return false;
    }
    if matches!(
        rest,
        "lto.bc" | "lto.opt.bc" | "lto.o" | "index.bc" | "index.dot" | "resolution.txt"
    ) {
        return true;
    }
    if rest
        .strip_prefix("lto.o")
        .is_some_and(|n| !n.is_empty() && n.chars().all(|c| c.is_ascii_digit()))
    {
        return true;
    }

    // `<output>.<task>.<stage>.<name>.bc` and `<output>.<N>.thinlto.o`: a bare
    // suffix match would also accept `debug.0.0.preopt.bc` for output `prog`,
    // which collides with a sibling output that happens to share a prefix, so
    // every numeric segment is required to be digits and nothing else.
    let segments: Vec<&str> = rest.split('.').collect();
    match segments.as_slice() {
        [task, stage, name, "bc"] => {
            is_digits(task)
                && is_digits(stage)
                && matches!(
                    *name,
                    "preopt" | "promote" | "internalize" | "import" | "opt" | "precodegen"
                )
        }
        [n, "thinlto", "o"] => is_digits(n),
        _ => false,
    }
}

/// Whether `s` is non-empty and every character is an ASCII digit.
fn is_digits(s: &str) -> bool {
    !s.is_empty() && s.chars().all(|c| c.is_ascii_digit())
}

/// Whether the user's own link arguments already ask the linker to keep its
/// save-temps modules.
///
/// A user who did this owns the artifacts: rllvm must not add the flag a
/// second time, and must not delete files it did not create. Missing a
/// spelling deletes linker intermediates the user explicitly asked to keep.
///
/// Every route to the linker is flattened into one token list first. `-Wl,`
/// groups several linker options behind one comma-separated argument --
/// `-Wl,-O2,-save-temps` is a common `LDFLAGS` shape that an exact-argument
/// match misses entirely -- and the driver also forwards single options
/// through `-Xlinker <option>` and `-Xlinker=<option>`. Within the list,
/// gold/lld accept `-plugin-opt` and its value as two separate tokens as well
/// as joined with `=`, and both one and two leading dashes throughout.
pub fn user_requested_save_temps(args: &[String]) -> bool {
    let mut tokens: Vec<&str> = vec![];
    let mut args = args.iter();
    while let Some(arg) = args.next() {
        if let Some(rest) = arg.strip_prefix("-Wl,") {
            tokens.extend(rest.split(','));
        } else if let Some(option) = arg.strip_prefix("-Xlinker=") {
            tokens.push(option);
        } else if arg == "-Xlinker"
            && let Some(option) = args.next()
        {
            tokens.push(option);
        }
    }

    tokens.iter().enumerate().any(|(i, token)| {
        matches!(
            *token,
            "-save-temps" | "--save-temps" | "-plugin-opt=save-temps" | "--plugin-opt=save-temps"
        ) || (matches!(*token, "-plugin-opt" | "--plugin-opt")
            && tokens.get(i + 1) == Some(&"save-temps"))
    })
}

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

    #[test]
    fn lto_mode_defaults_to_marker() {
        assert_eq!(LtoMode::default(), LtoMode::Marker);
    }

    #[test]
    fn lto_mode_round_trips_through_its_documented_spelling() {
        for mode in [LtoMode::Marker, LtoMode::SaveTemps, LtoMode::Skip] {
            let spelled = mode.to_string();
            assert_eq!(spelled.parse::<LtoMode>().unwrap(), mode, "{spelled}");
        }
    }

    #[test]
    fn unknown_lto_mode_is_an_error_not_a_fallback() {
        // A typo must stop the build, not silently disable extraction.
        let err = "marker ".parse::<LtoMode>().unwrap_err();
        assert!(matches!(err, Error::ConfigError(_)));
        assert!(err.to_string().contains("save-temps"), "{err}");
    }

    /// Decode one layer of string-literal escaping, the way the C compiler
    /// reads the `__asm__` operand and the assembler then reads the `.ascii`
    /// operand. Both accept the same escapes for the characters at issue.
    fn decode_one_layer(text: &str) -> String {
        let mut out = String::new();
        let mut chars = text.chars();
        while let Some(c) = chars.next() {
            if c != '\\' {
                out.push(c);
                continue;
            }
            match chars.next() {
                Some('n') => out.push('\n'),
                Some(escaped) => out.push(escaped),
                None => out.push('\\'),
            }
        }
        out
    }

    #[test]
    fn assembler_escaping_survives_both_decoders() {
        // The escaped text is written into a C string literal inside
        // `__asm__(...)`, so the C compiler decodes it once and the assembler
        // decodes what is left a second time. Asserting on the C source text
        // instead would pin whichever form the code happens to produce; what
        // has to hold is that the bytes reaching the section are the path.
        for path in [
            r#"/tmp/we"ird\path.bc"#,
            r"/tmp/od\d/a.bc",
            // A literal backslash-n: decoded one layer short it becomes a real
            // newline, which is the section's record separator.
            r"/tmp/new\nline.bc",
            "/tmp/plain.bc",
        ] {
            let escaped = escape_for_assembler(path);
            assert_eq!(
                decode_one_layer(&decode_one_layer(&escaped)),
                path,
                "escaped as {escaped:?}"
            );
        }
    }

    #[test]
    fn marker_source_terminates_the_entry_without_a_nul() {
        // The linker concatenates these sections and the entries are
        // newline-separated. `.asciz` would append a NUL between entries and
        // corrupt the list -- measured, not theorised.
        let source = marker_source("/tmp/a.bc");
        assert!(source.contains(r"\\n"), "no newline terminator: {source}");
        assert!(
            !source.contains(".asciz"),
            "must not NUL-terminate: {source}"
        );
    }

    #[test]
    fn marker_source_covers_both_supported_formats() {
        let source = marker_source("/tmp/a.bc");
        assert!(
            source.contains("__RLLVM,__rllvm_bc,regular,no_dead_strip"),
            "{source}"
        );
        assert!(source.contains(".rllvm_bc,\\\"\\\",@progbits"), "{source}");
        assert!(
            source.contains("#error"),
            "unsupported formats must not compile"
        );
    }

    /// A vector shaped like what `cc`/cargo actually pass: real compile
    /// flags interleaved with the dependency-generation flags that must not
    /// reach the marker compile.

    #[test]
    fn save_temps_flag_matches_the_linker_family() {
        // ld64 has its own spelling; every ELF linker accepts gold's, and
        // ld.lld and ld.bfd were measured producing identical file names.
        assert_eq!(save_temps_flag(true), "-Wl,-save-temps");
        assert_eq!(save_temps_flag(false), "-Wl,-plugin-opt=save-temps");
    }

    #[test]
    fn saved_module_is_recognised_per_linker() {
        assert!(is_saved_module("prog", "prog.lto.opt.bc", true));
        assert!(!is_saved_module("prog", "prog.lto.bc", true));

        assert!(is_saved_module("prog", "prog.0.5.precodegen.bc", false));
        assert!(!is_saved_module("prog", "prog.0.4.opt.bc", false));
        assert!(!is_saved_module("other", "prog.0.5.precodegen.bc", false));
    }

    #[test]
    fn saved_module_does_not_collide_with_a_sibling_output_sharing_a_prefix() {
        // `prog` and `prog.debug` in the same directory: a plain prefix-and-
        // suffix match would let `prog`'s link claim `prog.debug`'s module.
        assert!(!is_saved_module(
            "prog",
            "prog.debug.0.5.precodegen.bc",
            false
        ));
        assert!(is_saved_module(
            "prog.debug",
            "prog.debug.0.5.precodegen.bc",
            false
        ));
    }

    #[test]
    fn cleanup_spares_the_module_rllvm_keeps_and_the_user_s_files() {
        // Deleting is destructive, so the patterns are the measured ones and
        // nothing wider.
        for litter in [
            "prog.lto.bc",
            "prog.lto.o",
            "prog.lto.o1",
            "prog.0.0.preopt.bc",
            "prog.0.2.internalize.bc",
            "prog.index.bc",
            "prog.index.dot",
            "prog.resolution.txt",
            "prog.0.thinlto.o",
        ] {
            assert!(is_save_temps_artifact("prog", litter), "{litter}");
        }
        for keep in ["prog.rllvm.bc", "prog.bc", "prog", "prog.c", "progress.bc"] {
            assert!(!is_save_temps_artifact("prog", keep), "{keep}");
        }
    }

    #[test]
    fn cleanup_does_not_collide_with_a_sibling_output_sharing_a_prefix() {
        // `prog` and `prog.debug` in the same directory: a bare suffix match
        // would let `prog`'s cleanup delete `prog.debug`'s own intermediates.
        for artifact in ["prog.debug.0.0.preopt.bc", "prog.debug.0.thinlto.o"] {
            assert!(!is_save_temps_artifact("prog", artifact), "{artifact}");
        }
        // The same files are still recognised as litter for their own output.
        for artifact in ["prog.debug.0.0.preopt.bc", "prog.debug.0.thinlto.o"] {
            assert!(is_save_temps_artifact("prog.debug", artifact), "{artifact}");
        }
    }

    #[test]
    fn user_requested_save_temps_recognises_every_measured_spelling() {
        for args in [
            vec!["-Wl,-save-temps".to_string()],
            // A common `LDFLAGS` shape: several linker options behind one
            // comma-separated `-Wl,` argument.
            vec!["-Wl,-O2,-save-temps".to_string()],
            vec!["-Wl,--save-temps".to_string()],
            vec!["-Wl,-plugin-opt=save-temps".to_string()],
            // LLD documents the two-dash spelling of the same option.
            vec!["-Wl,--plugin-opt=save-temps".to_string()],
            // gold/lld also accept the option and its value as two separate
            // comma-joined tokens rather than joined with `=`.
            vec!["-Wl,-plugin-opt,save-temps".to_string()],
            vec!["-Wl,--plugin-opt,save-temps".to_string()],
            // The driver forwards a single option to the linker either way.
            vec!["-Xlinker".to_string(), "-save-temps".to_string()],
            vec!["-Xlinker=-save-temps".to_string()],
            vec![
                "-Xlinker".to_string(),
                "--plugin-opt=save-temps".to_string(),
            ],
            vec![
                "-Xlinker".to_string(),
                "-plugin-opt".to_string(),
                "-Xlinker".to_string(),
                "save-temps".to_string(),
            ],
        ] {
            assert!(user_requested_save_temps(&args), "{args:?}");
        }
    }

    #[test]
    fn user_requested_save_temps_is_false_without_the_flag() {
        let args = ["-flto", "-O2", "-Wl,-dead_strip"]
            .into_iter()
            .map(String::from)
            .collect::<Vec<String>>();
        assert!(!user_requested_save_temps(&args));

        // `-Xlinker` consumes its own value, and a neighbouring option is not
        // a request: matching one would make rllvm skip both its flag and its
        // cleanup for a link that keeps nothing.
        for args in [
            vec!["-Xlinker".to_string(), "-O2".to_string()],
            vec!["-Wl,-plugin-opt=thinlto".to_string()],
        ] {
            assert!(!user_requested_save_temps(&args), "{args:?}");
        }
    }
}