rllvm 0.3.0

//! Genera interfaces for the compiler wrapper

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

use crate::{
    arg_parser::{CompileMode, CompilerArgsInfo},
    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, 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.
    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 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)?);
            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
                let cache_key = cache::compute_cache_key(
                    &src_filepath,
                    self.args().compile_args(),
                    config.bitcode_generation_flags(),
                )?;

                if let Some(cached_path) = cache::cache_lookup(cache_dir, &src_filepath, cache_key)
                {
                    // 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 bitcode and store in cache
                    if let Some(code) =
                        self.generate_bitcode_file(&src_filepath, &bitcode_filepath)?
                        && code != 0
                    {
                        return Ok(Some(code));
                    }
                    if let Err(err) =
                        cache::cache_store(cache_dir, &src_filepath, cache_key, &bitcode_filepath)
                    {
                        tracing::warn!("Failed to store bitcode in cache: {}", err);
                    }
                    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,
                )?;
            }
        }

        // 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>,
    {
        let src_filepath = src_filepath.as_ref();
        let bitcode_filepath = bitcode_filepath.as_ref();
        let compiler_filepath = self.wrapped_compiler();

        let mut args = vec![compiler_filepath.to_string_lossy().into_owned()];
        args.extend(self.args().compile_args().iter().cloned());
        // Add bitcode generation flags
        if let Some(bitcode_generation_flags) = try_rllvm_config()?.bitcode_generation_flags() {
            args.extend(bitcode_generation_flags.iter().cloned());
        }
        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()];
        args.extend(self.args().compile_args().iter().cloned());
        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;
}