rllvm 0.4.2

//! Genera interfaces for the compiler wrapper

use std::{
    collections::HashSet,
    ffi::OsStr,
    path::{Path, PathBuf},
};

use crate::{
    arg_parser::{
        CompileMode, CompilerArgsInfo, universal_build_architectures, without_dependency_flags,
    },
    cache,
    compiler_wrapper::llvm::{lto_marker, marker},
    config::try_rllvm_config,
    constants::DEFAULT_LINK_OUTPUT_FILENAME,
    diagnostics::print_warning,
    error::Error,
    lto::{LtoFlavour, LtoMode, save_temps_flag, user_requested_save_temps},
    utils::{
        embed_bitcode_filepath_to_object_file, execute_command_for_status,
        extract_bitcode_filepaths_from_object_file, has_fat_lto_bitcode, is_bitcode_file,
        recorded_bitcode_filepath,
    },
};

/// Compiler type
#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
pub enum CompilerKind {
    /// Clang
    #[default]
    Clang,
    /// Clang++
    ClangXX,
}

/// What `save-temps` mode contributes to a link, and what it collects after.
#[derive(Debug, Default)]
pub struct SaveTempsPlan {
    /// Extra arguments appended to the link command.
    pub extra_args: Vec<String>,
    /// The link output whose merged module to collect, if any.
    pub collect_from: Option<PathBuf>,
    /// Whether rllvm added the flag, and so owns the cleanup.
    pub cleanup: bool,
    /// Keeps the marker object's staging directory alive until the link
    /// runs; the directory (and its contents) are removed once this drops.
    _marker_dir: Option<tempfile::TempDir>,
}

/// A general interface that wraps different compilers
pub trait CompilerWrapper {
    /// Obtain the name of the wrapper
    fn name(&self) -> &str;

    /// Obtain the path to the wrapped compiler
    fn wrapped_compiler(&self) -> &Path;

    /// Obtain the compiler kind
    fn compiler_kind(&self) -> &CompilerKind;

    /// Set the wrapper arguments parsing a command line set of arguments
    fn parse_args<S>(&mut self, args: &[S]) -> Result<&'_ mut Self, Error>
    where
        S: AsRef<str>;

    /// Obtain the argument information
    fn args(&self) -> &CompilerArgsInfo;

    /// Command to run the compiler
    fn command(&self) -> Result<Vec<String>, Error> {
        let args_info = self.args();
        let compiler_filepath = self.wrapped_compiler();
        let mut args = vec![compiler_filepath.to_string_lossy().into_owned()];

        // Append LTO LDFLAGS
        if args_info.input_files().is_empty() && !args_info.link_args().is_empty() {
            // Linking
            if args_info.is_lto() {
                // Add LTO LDFLAGS
                if let Some(lto_ldflags) = try_rllvm_config()?.lto_ldflags() {
                    args.extend(lto_ldflags.iter().cloned());
                }
            }
        }

        // Append given arguments
        args.extend(args_info.input_args().iter().cloned());

        // Remove forbidden flags
        if !args_info.forbidden_flags().is_empty() {
            let forbidden_flags_set: HashSet<String> =
                HashSet::from_iter(args_info.forbidden_flags().iter().cloned());

            // Report every dropped flag once, so the user knows the resulting
            // binary differs from the one their command asked for
            let mut removed_flags: Vec<&str> =
                forbidden_flags_set.iter().map(String::as_str).collect();
            removed_flags.sort_unstable();
            let message = format!(
                "Removed the following flag(s) from the compilation, as they are incompatible with bitcode generation: {}",
                removed_flags.join(", ")
            );

            // Deliberately not a `tracing::warn!`: the default log level is
            // ERROR, so a log record would be invisible to exactly the
            // non-interactive build-system runs that most need to know the
            // produced binary differs from the one they asked for. This goes
            // straight to stderr, where compiler diagnostics belong
            print_warning(&message);

            args.retain(|x| !forbidden_flags_set.contains(x));
        }

        Ok(args)
    }

    /// Silences the compiler wrapper output
    fn silence(&mut self, value: bool) -> &'_ mut Self;

    /// Returns `true` if `silence` was called with `true`
    fn is_silent(&self) -> bool;

    /// Decide what `save-temps` mode contributes to this invocation.
    ///
    /// Only an LTO link qualifies. ThinLTO never builds a whole-program
    /// module, so it warns and contributes nothing rather than failing a build
    /// over a mode the user set globally.
    /// The error for a build naming more than one `-arch`.
    ///
    /// Measured: clang fails the bitcode compile with "cannot use 'ir' output
    /// with multiple -arch options", and under `save-temps` the marker object
    /// comes out universal and the embedding step reports only "Unsupported
    /// file format". Neither names the cause. `lto_mode = "marker"` is not an
    /// escape -- the bitcode compile fails there too, and it fails without
    /// `-flto` at all -- so the message must not suggest one.
    fn universal_build_error(architectures: &[String]) -> Error {
        Error::UnsupportedBinaryFormat(format!(
            "rllvm cannot produce bitcode for a universal build: clang rejects \
             `-emit-llvm` with multiple `-arch` options ({}). Build one architecture \
             at a time and extract from each per-architecture binary; `rllvm-get-bc` \
             cannot read a `lipo`-combined universal binary either.",
            architectures.join(", ")
        ))
    }

    fn save_temps_plan(&self) -> Result<SaveTempsPlan, Error> {
        let args = self.args();
        if try_rllvm_config()?.lto_mode()? != LtoMode::SaveTemps
            || !matches!(args.mode(), CompileMode::LTO)
        {
            return Ok(SaveTempsPlan::default());
        }

        if args.lto_flavour() == Some(LtoFlavour::Thin) {
            print_warning(
                "ThinLTO builds no whole-program module, so lto_mode = \"save-temps\" has \
                 nothing to collect. Use lto_mode = \"marker\" for ThinLTO builds.",
            );
            return Ok(SaveTempsPlan::default());
        }

        let output_filename = match args.output_filename() {
            "" => DEFAULT_LINK_OUTPUT_FILENAME,
            name => name,
        };
        let output = PathBuf::from(output_filename);
        // Made absolute (not canonicalized: the link has not produced it yet,
        // so the path does not exist for `canonicalize` to resolve). A bare
        // relative `-o prog` otherwise breaks two ways below: `<output>.rllvm.bc`
        // is relative too, and embedding a relative bitcode path requires the
        // file to already exist; and after the link, an output with no
        // directory component makes `Path::parent` return `Some("")` rather
        // than `None`, so a naive fallback to `.` never triggers and
        // `collect_saved_module` fails to read the (nonexistent) empty path.
        let output = if output.is_absolute() {
            output
        } else {
            std::env::current_dir()?.join(output)
        };

        // A user who asked for save-temps owns the artifacts, so rllvm neither
        // adds the flag twice nor deletes what it did not create.
        let user_asked = user_requested_save_temps(args.input_args());

        let mut extra_args = vec![];
        if !user_asked {
            extra_args.push(save_temps_flag(cfg!(target_vendor = "apple")).to_string());
        }

        // Staged in its own temporary directory rather than next to the
        // output: the marker source and object are never removed by the
        // compile that produces them, and the output directory is not ours to
        // litter.
        let marker_dir = tempfile::tempdir()?;
        let bitcode = PathBuf::from(format!("{}.rllvm.bc", output.display()));
        // Built with this wrapper's own compiler and the user's compile
        // arguments, so the marker matches the link's target. A host-native
        // marker is not a link error: `-arch x86_64` on an arm64 host makes
        // ld64 warn and carry on, and the finished binary then names nothing.
        let architectures = universal_build_architectures(args.compile_args());
        if architectures.len() > 1 {
            return Err(Self::universal_build_error(&architectures));
        }

        let marker = marker::build_marker_object(
            &bitcode,
            marker_dir.path(),
            self.wrapped_compiler(),
            *self.compiler_kind(),
            args.compile_args(),
        )?;
        extra_args.push(marker.to_string_lossy().into_owned());

        Ok(SaveTempsPlan {
            extra_args,
            collect_from: Some(output),
            cleanup: !user_asked,
            _marker_dir: Some(marker_dir),
        })
    }

    /// Execute the given command with extra arguments appended.
    fn build_target_with(&self, extra_args: &[String]) -> Result<Option<i32>, Error> {
        let mut args = self.command()?;
        args.extend(extra_args.iter().cloned());
        let mode = self.args().mode();

        self.execute_command(&args, mode)
    }

    /// Run the compiler
    fn run(&mut self) -> Result<Option<i32>, Error> {
        // `save-temps` works around the link rather than after a compile: the
        // module it wants is one the linker produces, and the marker naming
        // that module has to be among the link's inputs.
        let plan = self.save_temps_plan()?;

        if let Some(code) = self.build_target_with(&plan.extra_args)?
            && code != 0
        {
            return Ok(Some(code));
        }

        if let Some(output) = plan.collect_from {
            let module = lto_marker::collect_saved_module(&output, plan.cleanup)?;
            // The module exists; nothing so far proves the binary names it.
            // A marker built for the wrong target, a dead-stripped section, or
            // a linker that dropped the input all leave a successful-looking
            // build that `rllvm-get-bc` reads nothing out of -- which is the
            // failure this mode exists to fix.
            let recorded = extract_bitcode_filepaths_from_object_file(&output)?;
            let expected = PathBuf::from(recorded_bitcode_filepath(&module)?);
            // Anything else in the section was put there by objects compiled
            // under `marker`, and `rllvm-get-bc` would merge those translation
            // units a second time -- as "symbol multiply defined", from a tool
            // and a command far away from the mode mismatch that caused it.
            let per_unit: Vec<_> = recorded.iter().filter(|path| **path != expected).collect();
            if !per_unit.is_empty() {
                print_warning(&format!(
                    "{output:?} records {} per-unit bitcode path(s) as well as the \
                     collected module {module:?}. Those objects were compiled under \
                     lto_mode = \"marker\" while this link ran under \"save-temps\", so \
                     extraction will merge those translation units twice. Build and link \
                     under one mode.",
                    per_unit.len()
                ));
            }

            if !recorded.contains(&expected) {
                return Err(Error::MissingFile(format!(
                    "The LTO link produced {module:?}, but {output:?} does not record it \
                     (recorded: {recorded:?}). The marker object naming the module never \
                     reached the linked output's bitcode-path section, so extraction would \
                     find nothing."
                )));
            }
            return Ok(Some(0));
        }

        if self.args().is_bitcode_generation_skipped()? {
            return Ok(Some(0));
        }

        self.generate_bitcode_files_and_embed_filepaths()
    }

    fn execute_command<S>(&self, args: &[S], mode: CompileMode) -> Result<Option<i32>, Error>
    where
        S: AsRef<OsStr> + std::fmt::Debug,
    {
        if !self.is_silent() {
            tracing::debug!("[{:?}] args={:?}", mode, args);
        }
        if args.is_empty() {
            return Err(Error::InvalidArguments(
                "The number of arguments cannot be 0".into(),
            ));
        }
        let status = execute_command_for_status(args[0].as_ref(), &args[1..])?;
        if !self.is_silent() {
            tracing::debug!("[{:?}] exit_status={}", mode, status);
        }

        if !status.success() {
            return Err(Error::ExecutionFailure(format!(
                "Failed to execute the command: args={:?}, exit_status={}",
                args, status
            )));
        }

        Ok(status.code())
    }

    /// Execute the given command and build the target
    fn build_target(&self) -> Result<Option<i32>, Error> {
        self.build_target_with(&[])
    }

    /// Generate bitcode files for all input files
    fn generate_bitcode_files_and_embed_filepaths(&self) -> Result<Option<i32>, Error> {
        let config = try_rllvm_config()?;
        let is_compile_only = self.args().is_compile_only();
        let artifact_filepaths = self.args().artifact_filepaths()?;

        // Determine if caching is enabled
        let caching_enabled = cache::is_cache_enabled(config.cache_enabled());
        let cache_directory = if caching_enabled {
            match cache::cache_dir(config.cache_dir().map(|p| p.as_path())) {
                Ok(dir) => Some(dir),
                Err(err) => {
                    tracing::warn!(
                        "Failed to initialize cache directory, caching disabled: {}",
                        err
                    );
                    None
                }
            }
        } else {
            None
        };

        let mut object_filepaths = vec![];
        for (src_filepath, object_filepath, bitcode_filepath) in artifact_filepaths {
            if !is_compile_only {
                // We need to explicitly build the intermediate object file
                self.build_object_file(&src_filepath, &object_filepath)?;

                // Collect all intermediate object files
                object_filepaths.push(object_filepath.clone());
            }

            let src_bitcode_filepath = if src_filepath.extension().is_some_and(|x| x == "bc") {
                // The source file is a bitcode; therefore, we do not need to
                // generate the bitcode and directly use the source file
                src_filepath
            } else if let Some(ref cache_dir) = cache_directory {
                // Caching is enabled — check for a cache hit. The manifest key
                // identifies the command; the content key adds everything the
                // last compile of that command read, which is what a header
                // edit has to invalidate.
                let manifest_key = cache::manifest_key(
                    &src_filepath,
                    self.args().compile_args(),
                    config.bitcode_generation_flags(),
                    self.wrapped_compiler(),
                );
                let depfile = cache::cached_depfile_path(cache_dir, manifest_key);
                let cached = match cache::content_key(manifest_key, &depfile) {
                    Some(key) => cache::cache_lookup(cache_dir, &src_filepath, key),
                    // No closure recorded yet: the first build of this command.
                    None => {
                        cache::record_miss(&src_filepath);
                        None
                    }
                };

                if let Some(cached_path) = cached {
                    // Cache hit — copy cached bitcode to expected output location
                    std::fs::copy(&cached_path, &bitcode_filepath).map_err(|err| {
                        tracing::error!(
                            "Failed to copy cached bitcode {:?} to {:?}: {}",
                            cached_path,
                            bitcode_filepath,
                            err
                        );
                        err
                    })?;
                    bitcode_filepath
                } else {
                    // Cache miss — generate the bitcode, recording what it
                    // read so the next build can key on that closure.
                    if let Some(code) = self.generate_bitcode_file_with_depfile(
                        &src_filepath,
                        &bitcode_filepath,
                        Some(&depfile),
                    )? && code != 0
                    {
                        return Ok(Some(code));
                    }
                    // Keyed on the closure the compile just reported, not on
                    // whatever a previous build left behind.
                    match cache::content_key(manifest_key, &depfile) {
                        Some(key) => {
                            if let Err(err) =
                                cache::cache_store(cache_dir, &src_filepath, key, &bitcode_filepath)
                            {
                                tracing::warn!("Failed to store bitcode in cache: {}", err);
                            }
                        }
                        None => tracing::warn!(
                            "No usable dependency file at {:?}; not caching {:?}",
                            depfile,
                            src_filepath
                        ),
                    }
                    bitcode_filepath
                }
            } else {
                // No caching — generate the bitcode
                if let Some(code) = self.generate_bitcode_file(&src_filepath, &bitcode_filepath)?
                    && code != 0
                {
                    return Ok(Some(code));
                }
                bitcode_filepath
            };

            // Under `-flto` the artifact is a bitcode module with no section
            // header to patch. Dispatch on content rather than on the flag:
            // `-ffat-lto-objects` produces a real object despite `-flto`, and
            // takes the ordinary path with no extra code.
            if is_bitcode_file(&object_filepath)? {
                lto_marker::inject_marker(
                    &object_filepath,
                    &src_bitcode_filepath,
                    self.args().compile_args(),
                    self.wrapped_compiler(),
                    *self.compiler_kind(),
                )?;
            } else {
                embed_bitcode_filepath_to_object_file(
                    &src_bitcode_filepath,
                    &object_filepath,
                    None,
                )?;

                // A fat LTO object carries the bitcode alongside the machine
                // code, and the section just embedded lives only in the latter.
                // GNU ld's plugin generates code from the bitcode and drops the
                // rest of the object, so the path has to be in both halves.
                //
                // The flag check comes first only to keep the content check off
                // the common path: `has_fat_lto_bitcode` reads and parses the
                // whole object, and `-ffat-lto-objects` emits nothing without
                // `-flto`, so a build that never asked for LTO cannot produce
                // one. Which half exists is still decided by content.
                if self.args().lto_flavour().is_some() && has_fat_lto_bitcode(&object_filepath)? {
                    lto_marker::inject_marker_into_fat_object(
                        &object_filepath,
                        &src_bitcode_filepath,
                        self.args().compile_args(),
                        self.wrapped_compiler(),
                        *self.compiler_kind(),
                    )?;
                }
            }
        }

        // Log cache statistics if caching was used
        if cache_directory.is_some() {
            cache::log_cache_stats();
        }

        // In compile-only mode the wrapped compiler already produced the final
        // object file and there is nothing left to link. The same holds when no
        // intermediate objects were built, in which case a link step would
        // invoke the compiler with no inputs at all.
        if is_compile_only || object_filepaths.is_empty() {
            return Ok(Some(0));
        }

        // Without an explicit `-o` the compiler wrote its default output, and that
        // is the file we must relink over. `output_filename` is only populated
        // when `-o` is parsed, so it is empty here -- and `PathBuf::from("")`
        // canonicalises to ENOENT. That surfaced as autoconf's "C compiler cannot
        // create executables" on its very first probe, which looks nothing like a
        // wrapper bug. CMake always passes `-o`, so this hid behind CMake builds.
        let output_filename = match self.args().output_filename() {
            "" => DEFAULT_LINK_OUTPUT_FILENAME,
            name => name,
        };
        let output_filepath = PathBuf::from(output_filename).canonicalize()?;
        self.link_object_files(&object_filepaths, output_filepath)
    }

    /// Generate bitcode file for one input file
    fn generate_bitcode_file<P>(
        &self,
        src_filepath: P,
        bitcode_filepath: P,
    ) -> Result<Option<i32>, Error>
    where
        P: AsRef<Path>,
    {
        self.generate_bitcode_file_with_depfile(src_filepath, bitcode_filepath, None)
    }

    /// Generate bitcode for one input file, optionally recording what it read.
    ///
    /// `depfile` is the cache's own dependency file, inside the cache
    /// directory. `-MD` costs almost nothing during a compile that is
    /// happening anyway, and it is what lets the next build know the include
    /// closure without running the preprocessor itself.
    fn generate_bitcode_file_with_depfile<P>(
        &self,
        src_filepath: P,
        bitcode_filepath: P,
        depfile: Option<&Path>,
    ) -> Result<Option<i32>, Error>
    where
        P: AsRef<Path>,
    {
        let src_filepath = src_filepath.as_ref();
        let bitcode_filepath = bitcode_filepath.as_ref();
        let compiler_filepath = self.wrapped_compiler();

        // Checked here rather than once up front so the user's own build still
        // runs first: rllvm reports what it could not do, without deciding for
        // the build system whether the object it asked for gets produced.
        let architectures = universal_build_architectures(self.args().compile_args());
        if architectures.len() > 1 {
            return Err(Self::universal_build_error(&architectures));
        }

        let mut args = vec![compiler_filepath.to_string_lossy().into_owned()];
        // Not `compile_args()` verbatim: the user's command already wrote the
        // dependency file, and this compile would overwrite it with one that
        // names the `.bc`.
        args.extend(without_dependency_flags(self.args().compile_args()));
        // Add bitcode generation flags
        if let Some(bitcode_generation_flags) = try_rllvm_config()?.bitcode_generation_flags() {
            args.extend(bitcode_generation_flags.iter().cloned());
        }
        // Ours, not the user's: their `-M*` flags were stripped above, and
        // this file is written inside the cache directory.
        if let Some(depfile) = depfile {
            args.extend_from_slice(&[
                "-MD".to_string(),
                "-MF".to_string(),
                depfile.to_string_lossy().into_owned(),
            ]);
        }
        args.extend_from_slice(&[
            "-emit-llvm".to_string(),
            "-c".to_string(),
            "-o".to_string(),
            bitcode_filepath.to_string_lossy().into_owned(),
            src_filepath.to_string_lossy().into_owned(),
        ]);

        let mode = CompileMode::BitcodeGeneration;

        self.execute_command(&args, mode)
    }

    /// Execute the command and build the object file
    fn build_object_file<P>(
        &self,
        src_filepath: P,
        object_filepath: P,
    ) -> Result<Option<i32>, Error>
    where
        P: AsRef<Path>,
    {
        let src_filepath = src_filepath.as_ref();
        let object_filepath = object_filepath.as_ref();
        let wrapped_compiler = self.wrapped_compiler();

        let mut args = vec![wrapped_compiler.to_string_lossy().into_owned()];
        // Same reason as in `generate_bitcode_file`: this intermediate object
        // must not become the last writer of the user's dependency file.
        args.extend(without_dependency_flags(self.args().compile_args()));
        args.extend_from_slice(&[
            "-c".to_string(),
            "-o".to_string(),
            object_filepath.to_string_lossy().into_owned(),
            src_filepath.to_string_lossy().into_owned(),
        ]);

        let mode = CompileMode::Compiling;

        self.execute_command(&args, mode)
    }

    fn link_object_files<P>(
        &self,
        object_filepaths: &[P],
        output_filepath: P,
    ) -> Result<Option<i32>, Error>
    where
        P: AsRef<Path>,
    {
        let output_filepath = output_filepath.as_ref();
        let wrapped_compiler = self.wrapped_compiler();

        let mut args = vec![wrapped_compiler.to_string_lossy().into_owned()];
        if self.args().is_lto() {
            // Add LTO LDFLAGS
            if let Some(lto_ldflags) = try_rllvm_config()?.lto_ldflags() {
                args.extend(lto_ldflags.iter().cloned());
            }
        }
        // Link arguments
        args.extend(self.args().link_args().iter().cloned());
        // Output
        args.extend_from_slice(&[
            "-o".to_string(),
            output_filepath.to_string_lossy().into_owned(),
        ]);
        // Input object files
        args.extend(
            object_filepaths
                .iter()
                .map(|x| x.as_ref().to_string_lossy().into_owned()),
        );

        // Mode
        let mode = CompileMode::Linking;

        self.execute_command(&args, mode)
    }
}

/// A general interface for the compiler wrapper builder
pub trait CompilerWrapperBuilder {
    type OutputType;

    /// Build the compiler wrapper
    fn build(&self) -> Result<Self::OutputType, Error>;

    /// Set the compiler name
    #[must_use]
    fn name(self, name: &str) -> Self;

    /// Set the compiler kind
    #[must_use]
    fn compiler_kind(self, compiler_kind: CompilerKind) -> Self;

    /// Set the wrapped compiler path
    fn wrapped_compiler<P>(self, wrapped_compiler: P) -> Self
    where
        P: AsRef<Path>;

    /// Set the silence flag
    fn silence(self, value: bool) -> Self;
}