cargo-compiler-interrupts 4.0.1

//! Implementation of `cargo-build-ci`.

use std::path::{Path, PathBuf};
use std::process::Output;
use std::str::FromStr;
use std::sync::mpsc::{Receiver, Sender};
use std::sync::{mpsc, Arc, Mutex};
use std::vec::IntoIter;

use anyhow::{bail, Context};
use cargo_util::{paths, ProcessBuilder, ProcessError};
use clap::Parser;
use colored::Colorize;
use crossbeam_utils::thread;
use indicatif::{ProgressBar, ProgressStyle};
use tracing::{debug, info, warn, Level};

use crate::args::BuildArgs;
use crate::cargo::{Cargo, Linker};
use crate::config::Config;
use crate::error::Error;
use crate::llvm::{LlvmToolchain, LlvmUtility};
use crate::paths::PathExt;
use crate::{llvm, util, CIResult, BUILD_CI_BIN_NAME};

/// Default pre-optimization passes for Compiler Interrupts.
const DEFAULT_OPT_PASSES: [&str; 6] = [
    "--postdomtree",
    "--mem2reg",
    "--indvars",
    "--loop-simplify",
    "--branch-prob",
    "--scalar-evolution",
];

/// State of the stage.
#[derive(Debug)]
enum State {
    /// Started.
    Started,
    /// Finished.
    Finished,
}

/// Stage of the integration.
#[derive(Debug)]
enum Stage {
    /// Integrating.
    Integrating(State),
    /// Static compiling.
    StaticCompiling(State),
    /// Linking
    Linking(State),
    /// Crate is skipped.
    Skipped,
    /// Error occurred.
    Error(String),
}

/// Shared context of the integration.
#[derive(Debug)]
struct IntegrationContext {
    /// Name of the crate.
    crate_name: Arc<String>,
    /// Current stage.
    stage: Stage,
}

/// Main routine for `cargo-build-ci`.
pub fn exec() -> CIResult<()> {
    let args = if std::env::args().next().unwrap_or_default() == BUILD_CI_BIN_NAME {
        BuildArgs::parse()
    } else {
        BuildArgs::parse_from(std::env::args().skip(1))
    };

    util::init_logger(&args.log_level)?;
    util::set_current_workspace_root_dir()?;

    let config = Config::load()?;
    let toolchain = llvm::toolchain()?;

    _exec(&config, &args, &toolchain)
}

/// Core routine for `cargo-build-ci`.
fn _exec(config: &Config, args: &BuildArgs, toolchain: &LlvmToolchain) -> CIResult<()> {
    if !config.library_path.is_file() {
        bail!(Error::LibraryNotInstalled);
    }

    if args.debug {
        warn!("Debugging mode is enabled");
    }

    let mut cargo = Cargo::with_args(args.cargo_args.clone());
    cargo.build()?;

    let time = std::time::Instant::now();

    let target_dir = cargo.target_dir;

    let llvm_predicate = |path: &PathBuf| -> bool {
        let file_stem = path.file_stem().unwrap_or_default();
        let extension = path.extension().unwrap_or_default();
        file_stem.contains("rcgu") && !file_stem.contains("-ci") && extension == "ll"
    };

    // *.rcgu.ll are intermediate files generated by `rustc -C save-temps`
    let mut llvm_ir_files = target_dir.join("deps").read_dir(llvm_predicate)?;
    llvm_ir_files.append(&mut target_dir.join("examples").read_dir(llvm_predicate)?);

    // parse cargo build output to get the linker invocation
    let linkers = cargo.linkers;

    // total length of the process bar
    let length = llvm_ir_files.len() * 2 + linkers.len() + 1;

    let llvm_ir_iter = Arc::new(Mutex::new(llvm_ir_files.into_iter()));
    let linker_iter = Arc::new(Mutex::new(linkers.into_iter()));

    thread::scope(move |s| -> CIResult<()> {
        let timestamp = chrono::Local::now().format("%y%m%dT%H%M%S").to_string();
        let mut path = Config::dir()?;
        path.push(format!("CI-{}.log", timestamp));

        let verify = |results: Vec<CIResult<()>>| -> CIResult<()> {
            let mut ok = true;
            for result in results {
                if let Err(error) = result {
                    paths::append(&path, error.to_string().as_ref())?;
                    ok = false;
                }
            }

            if !ok {
                bail!(format!(
                    "Consider filing an issue report on \
                        https://github.com/bitslab/CompilerInterrupts \
                        with the LLVM IR file and log attached.\n\
                        Path to the log: {}",
                    path.display(),
                ));
            }

            Ok(())
        };

        // communication between the progress bar thread and other threads
        let (tx, rx) = mpsc::channel::<IntegrationContext>();

        // number of threads based on number of logical cores in CPU
        let num_cpus = num_cpus::get();

        // progress bar rendering
        let pb_thread =
            s.spawn(move |_| -> CIResult<()> { progress_bar(rx, length as u64, &args.log_level) });

        // integration
        let mut threads = Vec::new();
        for _ in 0..num_cpus {
            let tx = tx.clone();
            let files = Arc::clone(&llvm_ir_iter);
            let thread =
                s.spawn(move |_| -> CIResult<()> { integrate(config, args, toolchain, tx, files) });
            threads.push(thread);
        }

        let mut results = Vec::new();
        for thread in threads {
            let result = thread.join().expect("integration thread panicked");
            results.push(result);
        }
        verify(results)?;

        // linking
        let mut threads = Vec::new();
        for _ in 0..num_cpus {
            let tx = tx.clone();
            let linkers = Arc::clone(&linker_iter);
            let thread = s.spawn(move |_| -> CIResult<()> { link(toolchain, tx, linkers) });
            threads.push(thread);
        }

        let mut results = Vec::new();
        for thread in threads {
            let result = thread.join().expect("linking thread panicked");
            results.push(result);
        }
        verify(results)?;

        drop(tx);

        pb_thread
            .join()
            .expect("progress bar thread panicked")
            .context("progress bar failed")?;

        Ok(())
    })
    .expect("main scoped thread panicked")?;

    println!(
        "{:>12} integrated {} target(s) in {}",
        "Finished".green().bold(),
        length,
        util::human_duration(time.elapsed())
    );

    Ok(())
}

/// Handle the progress bar rendering.
fn progress_bar(rx: Receiver<IntegrationContext>, len: u64, log_level: &String) -> CIResult<()> {
    let log_level = Level::from_str(&log_level)?;
    // progress bar
    let pb = if log_level <= Level::WARN {
        ProgressBar::new(len)
    } else {
        ProgressBar::hidden()
    };
    pb.set_prefix("Building");

    let mut names: Vec<String> = Vec::new();
    let mut error = false;

    while let Ok(integration) = rx.recv() {
        if error {
            // halt updating status until rx closed
            continue;
        }

        let name = integration.crate_name;
        let status_line =
            |status: &str| -> String { format!("{:>12} {}", status.green().bold(), name) };
        let mut remove = |name: &String| -> CIResult<()> {
            let idx = names
                .iter()
                .position(|e| e == name)
                .context("failed to find crate name to be removed")?;
            names.remove(idx);
            Ok(())
        };
        let llc_name = format!("{}(llc)", name);
        let ld_name = format!("{}(bin)", name);

        use Stage::*;
        use State::*;
        match integration.stage {
            Integrating(state) => match state {
                Started => {
                    pb.println(status_line("Integrating"));
                    pb.inc(1);
                    names.insert(0, name.to_string());
                }
                Finished => remove(&name)?,
            },
            StaticCompiling(state) => match state {
                Started => {
                    pb.inc(1);
                    names.insert(0, llc_name);
                }
                Finished => remove(&llc_name)?,
            },
            Linking(state) => match state {
                Started => {
                    pb.println(status_line("Linking"));
                    pb.inc(1);
                    names.insert(0, ld_name);
                }
                Finished => remove(&ld_name)?,
            },
            Skipped => {
                // redundant to print `compiler_interrupts` status as it is always skipped
                if *name != "compiler_interrupts" {
                    pb.println(status_line("Skipped"));
                }
                pb.inc(1);
            }
            Error(_) => {
                pb.finish_and_clear();
                println!(
                    "{:>12} Compiler Interrupts integration has unexpectedly failed",
                    "Error".red().bold(),
                );
                println!(
                    "{:>12} Waiting for other jobs to finish",
                    "Warning".yellow().bold()
                );

                // we must not prematurely close the channel
                // channel must live until all threads are done sending signals
                error = true;
                continue;
            }
        }

        // progress bar message
        let term_size = terminal_size::terminal_size()
            .map(|(w, h)| (w.0.into(), h.0.into()))
            .unwrap_or((80, 24));
        let prefix_size = if term_size.0 > 80 { 50 } else { 20 };
        let template = if term_size.0 > 80 {
            "{prefix:>12.cyan.bold} [{bar:27}] {pos}/{len}: {wide_msg}"
        } else {
            "{prefix:>12.cyan.bold} {pos}/{len}: {wide_msg}"
        };
        let mut message = String::new();
        let mut iter = names.iter();
        let first = match iter.next() {
            Some(first) => first,
            None => "",
        };
        message.push_str(first);
        for name in iter {
            message.push_str(", ");
            // truncate the message if it is too lengthy
            if message.len() + name.len() < term_size.0 - prefix_size - /* padding */ 15 {
                message.push_str(name);
            } else {
                message.push_str("...");
                break;
            }
        }
        pb.set_style(ProgressStyle::with_template(template)?.progress_chars("=> "));
        pb.set_message(message);
    }

    if !error {
        pb.inc(1);
        pb.finish_and_clear();
    }

    Ok(())
}

/// Handle the integration process.
fn integrate(
    config: &Config,
    args: &BuildArgs,
    toolchain: &LlvmToolchain,
    tx: Sender<IntegrationContext>,
    files: Arc<Mutex<IntoIter<PathBuf>>>,
) -> CIResult<()> {
    loop {
        let file = files.lock().expect("failed to acquire lock").next();
        if let Some(file) = file {
            let mut integrate = true;
            let crate_name = Arc::new(crate_name(&file)?);
            let ci_file = file.append_suffix("ci")?;

            // `nm -jU` displays defined symbol names
            let output = LlvmUtility::NameMangling
                .process_builder(toolchain)
                .arg("-jU")
                .arg(file.with_extension("o"))
                .exec_with_output()?;
            let stdout = String::from_utf8(output.stdout)?;
            if stdout.contains("intvActionHook") {
                // skip the crate that has CI symbols defined
                integrate = false;
            }

            if let Some(skip_crates) = &args.skip_crates {
                for skip_crate in skip_crates {
                    if skip_crate.replace('-', "_").contains(&*crate_name) {
                        // skip the given crates
                        integrate = false;
                        break;
                    }
                }
            }

            if integrate {
                info!("integrating: {}", file.display());
                tx.send(IntegrationContext {
                    crate_name: Arc::clone(&crate_name),
                    stage: Stage::Integrating(State::Started),
                })?;

                // `opt` runs the integration
                let mut opt = LlvmUtility::Optimizer.process_builder(toolchain);
                opt.args(&[
                    "-S",
                    "--enable-new-pm=0",
                    "--load",
                    &config.library_path.to_string()?,
                    "--logicalclock",
                ]);
                opt.args(&DEFAULT_OPT_PASSES);
                opt.args(&config.library_args);
                opt.arg(&file);
                opt.arg("-o");
                opt.arg(&ci_file);
                // debug!("opt: opt {:#?}", opt.get_args());
                let output = opt.exec_with_output();
                handle_output(&tx, output, &ci_file)?;

                tx.send(IntegrationContext {
                    crate_name: Arc::clone(&crate_name),
                    stage: Stage::Integrating(State::Finished),
                })?;
            } else {
                info!("integration skipped: {}", file.display());
                tx.send(IntegrationContext {
                    crate_name: Arc::clone(&crate_name),
                    stage: Stage::Skipped,
                })?;
                paths::copy(&file, &ci_file)?;
            }

            // `llc` transforms integrated IR bitcode to object file
            debug!("run llc on: {}", ci_file.display());
            tx.send(IntegrationContext {
                crate_name: Arc::clone(&crate_name),
                stage: Stage::StaticCompiling(State::Started),
            })?;

            let mut llc = LlvmUtility::StaticCompiler.process_builder(toolchain);
            llc.arg("-filetype=obj");
            llc.arg(&ci_file);

            // fixes mismatch relocation symbols on linux
            if cfg!(target_os = "linux") {
                llc.arg("-code-model=large");
            }

            let output = llc.exec_with_output();
            handle_output(&tx, output, &ci_file)?;

            tx.send(IntegrationContext {
                crate_name: Arc::clone(&crate_name),
                stage: Stage::StaticCompiling(State::Finished),
            })?;
        } else {
            break;
        }
    }

    Ok(())
}

/// Handle the linking process.
fn link(
    toolchain: &LlvmToolchain,
    tx: Sender<IntegrationContext>,
    linkers: Arc<Mutex<IntoIter<Linker>>>,
) -> CIResult<()> {
    loop {
        let linker = linkers.lock().expect("failed to acquire lock").next();
        if let Some(mut linker) = linker {
            if linker
                .args
                .input_files
                .iter()
                .any(|e| e.contains("build_script_build"))
            {
                debug!("linking skipped: {}", linker.args.output_file);
                continue;
            }

            let output_file = linker.args.output_file.clone();
            let _crate_name = crate_name(&output_file)?;
            let crate_name = Arc::new(_crate_name.clone());
            info!("linking: {}", crate_name);

            tx.send(IntegrationContext {
                crate_name: Arc::clone(&crate_name),
                stage: Stage::Linking(State::Started),
            })?;

            for file in &mut linker.args.input_files {
                if !file.contains("deps") {
                    continue;
                }

                let output = LlvmUtility::NameMangling
                    .process_builder(toolchain)
                    .arg("-jU")
                    .arg(&file)
                    .exec_with_output()?;
                let stdout = String::from_utf8(output.stdout)?;
                if stdout.contains("__rust_alloc") {
                    // skip the object file contains the symbol for memory allocator
                    debug!("found allocator shim: {}", file);
                } else {
                    *file = file.append_suffix("ci")?.to_string()?;
                }
            }

            // make a copy and replace *.o with *-ci.o in the rlib files
            for file in &mut linker.args.rlib_files {
                if !file.contains("deps") {
                    continue;
                }

                debug!("original rlib: {}", file);
                let ci_file = file.append_suffix("ci")?;
                paths::copy(&file, &ci_file)?;

                debug!("replacing object file for rlib: {}", ci_file.display());
                // list all object files inside rlib
                let output = LlvmUtility::Archiver
                    .process_builder(toolchain)
                    .arg("-t")
                    .arg(&ci_file)
                    .exec_with_output()?;
                let stdout = String::from_utf8(output.stdout)?;
                if let Some(rcgu_obj_file_name) = stdout
                    .lines()
                    .find(|e| e.contains("rcgu") && !e.contains("-ci"))
                {
                    let rcgu_obj_file = ci_file.parent()?.join(rcgu_obj_file_name);
                    let rcgu_obj_ci_file = rcgu_obj_file.append_suffix("ci")?;

                    // replace *.o with *-ci.o
                    LlvmUtility::Archiver
                        .process_builder(toolchain)
                        .arg("-rb")
                        .arg(&rcgu_obj_file)
                        .arg(&ci_file)
                        .arg(&rcgu_obj_ci_file)
                        .exec_with_output()?;

                    // delete old *.o
                    LlvmUtility::Archiver
                        .process_builder(toolchain)
                        .arg("-d")
                        .arg(&ci_file)
                        .arg(&rcgu_obj_file)
                        .exec_with_output()?;
                }

                *file = ci_file.to_string()?;
            }

            let output_ci_file = output_file.append_suffix("ci")?.to_string()?;
            linker.args.output_file = output_ci_file.clone();

            // execute the linker
            debug!("linker: {:#?}", linker);
            let mut builder = ProcessBuilder::new(&linker.program);
            builder.args(&linker.args.build());
            let output = builder.exec_with_output();
            handle_output(&tx, output, &output_ci_file)?;

            // hard link the CI-integrated binary file to the parent directory
            let link_file = output_file
                .parent()?
                .parent()?
                .join(_crate_name.append_suffix("ci")?);
            debug!(?output_file);
            debug!(?link_file);
            paths::link_or_copy(&output_file, &link_file)?;

            tx.send(IntegrationContext {
                crate_name: Arc::clone(&crate_name),
                stage: Stage::Linking(State::Finished),
            })?;
        } else {
            break;
        }
    }

    Ok(())
}

/// Handle output from the process and validate output file.
fn handle_output<P: AsRef<Path>>(
    tx: &Sender<IntegrationContext>,
    output: anyhow::Result<Output>,
    output_file: P,
) -> CIResult<()> {
    let output_file = output_file.as_ref();
    let crate_name = Arc::new(crate_name(output_file)?);
    match output {
        Ok(output) => {
            if !output_file.is_file() {
                // output file does not exist
                let stderr = String::from_utf8(output.stderr.clone())?;

                tx.send(IntegrationContext {
                    crate_name: Arc::clone(&crate_name),
                    stage: Stage::Error(String::new()),
                })?;

                bail!(
                    "process returned success but output file does not exist\n\
                    process: {:#?}\n\
                    expected file: {}\n\
                    --- stderr\n{}",
                    output,
                    output_file.display(),
                    stderr
                );
            }

            Ok(())
        }
        Err(err) => {
            let proc_err = err
                .downcast_ref::<ProcessError>()
                .context("failed to downcast to ProcessError")?;

            tx.send(IntegrationContext {
                crate_name: Arc::clone(&crate_name),
                stage: Stage::Error(String::new()),
            })?;

            bail!(ToString::to_string(&proc_err.desc));
        }
    }
}

/// Get the binary name from path.
fn crate_name<P: AsRef<Path>>(path: P) -> CIResult<String> {
    Ok(path
        .file_stem()?
        .split('.')
        .next()
        .context("invalid crate name, expected '.'")?
        .split('-')
        .next()
        .context("invalid crate name, expected '-'")?
        .to_string())
}