criterion-polyglot 0.1.0

use std::{borrow::Cow, env, fs, io, path::PathBuf, process::Command, ffi::OsString, ops::Index};

use lazy_static::lazy_static;
use tempfile::{TempDir, TempPath, NamedTempFile};

use crate::BenchSpec;

const BEGIN_TMPL_VAR: &str = "@\"";
const BEGIN_TMPL_VAR_LEN: usize = BEGIN_TMPL_VAR.len();
const END_TMPL_VAR: &str = "\"@";
const END_TMPL_VAR_LEN: usize = END_TMPL_VAR.len();

const TEMP_PREFIX: &str = concat!(env!("CARGO_PKG_NAME"), "__");
// Use a Cow<str> to make using the `+` operator possible
const ENVVAR_PREFIX: Cow<str> = Cow::Borrowed("CRITERION_");

// Choose the most "specific" temporary directory, based
// on what is available in the build environment. Order
// of preference:
//
// 1. `CARGO_TARGET_TMPDIR` - supplied by Cargo when building benchmark binaries, but not
//    available to this crate at build time; synthesized from std::env::current_exe()
// 2. System default temporary directory; may be affected by `noexec` restrictions,
//    causing benchmarks to fail
lazy_static! {
    static ref CARGO_TARGET_TMPDIR: PathBuf = {
        let exe_path = env::current_exe().ok();
        if let Some(mut target_tmp_path) = exe_path {
            // Remove binary filename
            target_tmp_path.pop();
            // Check if we're inside target/{target}/deps
            if target_tmp_path.ends_with("deps") {
                target_tmp_path.pop();
            }
            // Check if we're inside target/{target} and go up
            // a level to match CARGO_TARGET_TMPDIR
            if target_tmp_path.parent().map_or(false, |p| p.ends_with("target")) {
                target_tmp_path.pop();
            }
            // Make a tmpdir inside the target dir if it doesn't exist
            target_tmp_path.push("tmp");
            if !target_tmp_path.try_exists().expect("Can't check existence of target tmpdir") {
                fs::create_dir(&target_tmp_path).expect("Can't create target tmpdir");
            }
            target_tmp_path
        } else {
            env::temp_dir()
        }
    };
}

pub fn make_tempfile<'a, I>(extension: &str, content: I) -> TempPath
where
    I: IntoIterator<Item = &'a str>,
{
    use io::Write;

    let mut file = make_tempfile_writer(extension);
    for segment in content {
        file.write_all(segment.as_bytes()).unwrap();
    }

    // into_temp_path() closes and flushes the file
    file.into_temp_path()
}

pub fn make_spec_tempfile(extension: &str, harness: &str, spec: &BenchSpec) -> TempPath {
    let mut file = make_tempfile_writer(extension);
    process_harness(&mut file, harness, spec)
        .expect("failed writing harness tempfile");
    file.into_temp_path()
}

pub fn make_empty_tempfile(extension: &str) -> TempPath {
    make_tempfile(extension, [])
}

pub fn make_tempfile_writer(extension: &str) -> NamedTempFile {
    tempfile::Builder::new()
        .prefix(TEMP_PREFIX)
        .suffix(extension)
        .tempfile_in(&*CARGO_TARGET_TMPDIR)
        .unwrap()
}

pub fn make_tempdir() -> TempDir {
    tempfile::Builder::new()
        .prefix(TEMP_PREFIX)
        .tempdir_in(&*CARGO_TARGET_TMPDIR)
        .unwrap()
}

pub fn format_command(command: &Command) -> String {
    use std::fmt::Write;

    let mut command_line = command.get_program().to_string_lossy().into_owned();

    let args = command.get_args().map(|a| a.to_string_lossy());
    for arg in args {
        write!(command_line, " {}", shell_quote(&arg)).unwrap();
    }

    command_line
}

pub fn reindent(s: &str, target_indent: usize) -> Cow<str> {
    use std::fmt::Write;

    let non_blank_lines = s.lines().filter(|l| !l.trim_start_matches(|c| char::is_ascii_whitespace(&c)).is_empty());
    let mut remove_indent = non_blank_lines.map(|l| l.len() - l.trim_start_matches(' ').len()).min().unwrap_or_default();
    let add_indent = if remove_indent >= target_indent {
        remove_indent -= target_indent;
        0
    } else {
        target_indent
    };
    if remove_indent > 0 || add_indent > 0 {
        let mut outdented_s = " ".repeat(add_indent);
        for line in s.lines() {
            let outdented_line = if line.len() > remove_indent {
                &line[remove_indent..]
            } else { line };
            writeln!(outdented_s, "{}", outdented_line).unwrap();
        }
        outdented_s.into()
    } else {
        s.into()
    }
}

pub fn shell_quote(s: &str) -> Cow<str> {
    // Adapted from python's shlex.quote()
    fn is_shell_unsafe(s: &str) -> bool {
        const SAFE_SHELL_CHARS: &str = "@%+=:,./-_";
        fn is_char_unsafe(c: char) -> bool {
            let is_safe = SAFE_SHELL_CHARS.contains(c) || char::is_ascii_alphanumeric(&c);
            !is_safe
        }
        s.contains(is_char_unsafe)
    }

    let mut s = Cow::from(s);
    if is_shell_unsafe(&s) {
        let s_mut = s.to_mut();
        if s_mut.contains('\'') {
            *s_mut = s_mut.replace('\'', "'\"'\"'");
        }
        s_mut.insert(0, '\'');
        s_mut.push('\'');
    }

    s
}

pub fn get_criterion_env_var(name: &str, default: &str) -> OsString {
    // Internal: no lowercase characters in environment variable name
    debug_assert!(name.chars().all(|c| !c.is_ascii_lowercase()));
    let var_name = ENVVAR_PREFIX + name;
    env::var_os(&*var_name)
        .unwrap_or_else(|| OsString::from(default))
}

/// PANICS: if `vars` argument doesn't contain a variable used in `template` (but template
/// contents are entirely under this crate's control, so this is actually desirable because it will
/// cause tests to fail if templates contain invalid variable names)
pub fn process_harness<'a, W, V>(out: &mut W, template: &'a str, vars: &V) -> io::Result<()>
where
    W: io::Write,
    V: Index<&'a str, Output = str>
{
    let mut remainder = template;
    while let Some(var_start) = remainder.find(BEGIN_TMPL_VAR) {
        // Extract the variable name from between the delimiters, and return an error if no ending
        // delimiter is found
        let var_name_start = var_start + BEGIN_TMPL_VAR_LEN;
        let var_name_len = remainder[var_name_start..].find(END_TMPL_VAR)
            .ok_or_else(|| io::Error::from(io::ErrorKind::InvalidInput))?;
        let var_name_end = var_name_start + var_name_len;
        let var_name = &remainder[var_name_start..var_name_end];
        // Look up the variable value (may panic)
        let var_value = &vars[var_name];

        // Count leading spaces on the line containing the variable
        let line_until_var = remainder[..var_start].rsplit('\n').next().unwrap();
        let var_indent = line_until_var.len() - line_until_var.trim_start_matches(' ').len();
        // All the text before the variable, excluding the variable-containing line's
        // leading whitespace, because it will be added back by `reindent()`.
        let preceding = &remainder[..(var_start - var_indent)];
        let value_indented = reindent(var_value, var_indent);
        write!(out, "{preceding}{value_indented}")?;

        let var_end = var_name_end + END_TMPL_VAR_LEN;
        remainder = &remainder[var_end..];
    }
    write!(out, "{remainder}")
}