rllvm 0.2.0

//! Genera interfaces for the compiler wrapper

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

use crate::{
    arg_parser::{CompileMode, CompilerArgsInfo},
    cache,
    config::try_rllvm_config,
    constants::DEFAULT_LINK_OUTPUT_FILENAME,
    diagnostics::print_warning,
    error::Error,
    utils::{embed_bitcode_filepath_to_object_file, execute_command_for_status},
};

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

/// 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;

    /// Run the compiler
    fn run(&mut self) -> Result<Option<i32>, Error> {
        if let Some(code) = self.build_target()?
            && code != 0
        {
            return Ok(Some(code));
        }
        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> {
        let args = self.command()?;
        let mode = self.args().mode();

        self.execute_command(&args, mode)
    }

    /// 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
            };

            // Embed the path of the bitcode to the corresponding object file
            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;
}