cargo_difftests/

difftest.rs

/*
 *        Copyright (c) 2023-2024 Dinu Blanovschi
 *
 *    Licensed under the Apache License, Version 2.0 (the "License");
 *    you may not use this file except in compliance with the License.
 *    You may obtain a copy of the License at
 *
 *        https://www.apache.org/licenses/LICENSE-2.0
 *
 *    Unless required by applicable law or agreed to in writing, software
 *    distributed under the License is distributed on an "AS IS" BASIS,
 *    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *    See the License for the specific language governing permissions and
 *    limitations under the License.
 */

//! Holds the [`Difftest`] struct and related functions.

use std::ffi::OsStr;
use std::fs;
use std::path::{Path, PathBuf};
use std::process::Command;

use cargo_difftests_core::CoreTestDesc;
use log::{debug, info, warn};

use crate::analysis::AnalysisContext;
use crate::index_data::{IndexDataCompilerConfig, TestIndex};
use crate::{analysis_data, DifftestsError, DifftestsResult};

/// A single difftest.
///
/// Points to the files that were generated by the call to
/// `cargo_difftests_testclient::init`.
#[derive(serde::Serialize, serde::Deserialize, Debug)]
pub struct Difftest {
    pub(crate) dir: PathBuf,
    pub(crate) test_binary_path: PathBuf,
    pub(crate) test_name_path: PathBuf,
    pub(crate) profraws: Vec<PathBuf>,
    pub(crate) self_json: Option<PathBuf>,
    pub(crate) test_run_time: std::time::SystemTime,
    pub(crate) profdata_file: Option<PathBuf>,
    pub(crate) index_data: Option<PathBuf>,

    pub(crate) cleaned: bool,
}

impl Difftest {
    const CLEANED_FILE_NAME: &'static str = "cargo_difftests_cleaned";

    /// Get the directory to the difftest.
    ///
    /// Note that you should not modify the contents
    /// of this directory in any way.
    pub fn dir(&self) -> &Path {
        &self.dir
    }

    /// Tests whether the difftest has a test index
    /// file attached.
    pub fn has_index(&self) -> bool {
        self.index_data.is_some()
    }

    fn read_test_binary_path(&self) -> DifftestsResult<PathBuf> {
        let s = fs::read_to_string(&self.test_binary_path)?;
        Ok(PathBuf::from(s))
    }

    pub fn test_info(&self) -> DifftestsResult<TestInfo> {
        let test_name = std::fs::read_to_string(&self.test_name_path)?;
        let test_binary = self.read_test_binary_path()?;

        let extra_desc = self.load_test_desc()?;

        Ok(TestInfo {
            test_name,
            test_binary,
            extra_desc,
        })
    }

    /// Reads the [`TestIndex`] data associated with the [`Difftest`].
    ///
    /// Returns `Ok(None)` if the [`Difftest`] does not have an associated index file.
    pub fn read_index_data(&self) -> DifftestsResult<Option<TestIndex>> {
        let Some(index_data) = &self.index_data else {
            return Ok(None);
        };

        TestIndex::read_from_file(index_data).map(Some)
    }

    pub(crate) fn test_run_time(&self) -> std::time::SystemTime {
        self.test_run_time
    }

    /// Loads the `self.json` file from the test directory, and parses it into
    /// a [`CoreTestDesc`] value.
    pub fn load_test_desc(&self) -> DifftestsResult<Option<CoreTestDesc>> {
        if let Some(self_json) = self.self_json.as_ref() {
            let s = fs::read_to_string(self_json)?;
            let desc = serde_json::from_str(&s)
                .map_err(|e| DifftestsError::Json(e, Some(self_json.clone())))?;
            Ok(desc)
        } else {
            Ok(None)
        }
    }

    /// Merges the `.profraw` files into a `.profdata` file, via `llvm-profdata merge`.
    pub fn merge_profraw_files_into_profdata(&mut self, force: bool) -> DifftestsResult<()> {
        if self.cleaned {
            return Err(DifftestsError::DifftestCleaned);
        }

        if self.profdata_file.is_some() && !force {
            return Ok(());
        }

        merge_profraws(self)?;

        self.profdata_file = Some(self.out_profdata_path());

        Ok(())
    }

    /// Cleans up the profiling data from the [`Difftest`], reducing
    /// the size on-disk substantially, but all the profiling data
    /// gets lost in the process.
    ///
    /// This is used by the `--index-strategy=always-and-clean` flag, which keeps
    /// the test index data around, even if the profiling data is deleted.
    pub fn clean(&mut self) -> DifftestsResult<()> {
        fn clean_file(f: &mut Option<PathBuf>) -> DifftestsResult<()> {
            if let Some(f) = f {
                fs::remove_file(f)?;
            }

            *f = None;
            Ok(())
        }

        clean_file(&mut self.profdata_file)?;

        for profraw in self.profraws.drain(..) {
            fs::remove_file(profraw)?;
        }

        fs::write(self.dir.join(Self::CLEANED_FILE_NAME), b"")?;

        self.cleaned = true;

        Ok(())
    }

    /// Checks whether the [`Difftest`] has been cleaned.
    pub fn was_cleaned(&self) -> bool {
        self.cleaned
    }

    /// Checks whether the [`Difftest`] has the `.profdata` file.
    pub fn has_profdata(&self) -> bool {
        if self.cleaned {
            return false;
        }

        self.profdata_file.is_some()
    }

    /// Tries to discover a [`Difftest`] from the given directory, with the
    /// given [`DiscoverIndexPathResolver`] to resolve the index path.
    pub fn discover_from(
        dir: PathBuf,
        index_resolver: Option<&DiscoverIndexPathResolver>,
    ) -> DifftestsResult<Self> {
        let test_name_path = dir.join(cargo_difftests_core::CARGO_DIFFTESTS_TEST_NAME_FILENAME);

        if !test_name_path.exists() {
            return Err(DifftestsError::IO(std::io::Error::new(
                std::io::ErrorKind::NotFound,
                format!(
                    "test name file does not exist: {}",
                    test_name_path.display()
                ),
            )));
        }

        discover_difftest_from_tempdir(dir, test_name_path, index_resolver)
    }

    /// Reads the `.profdata` profiling data file, to be able to use
    /// it for analysis.
    ///
    /// This function should be ran after
    /// [`Difftest::merge_profraw_files_into_profdata`].
    pub fn export_profdata(
        &self,
        config: ExportProfdataConfig,
    ) -> DifftestsResult<analysis_data::CoverageData> {
        assert!(self.has_profdata());

        let ExportProfdataConfig {
            ignore_registry_files,
            mut other_binaries,
        } = config;

        for other_binary in &mut other_binaries {
            if !other_binary.is_absolute() {
                use path_absolutize::Absolutize;
                *other_binary = other_binary.absolutize()?.into_owned();
            }
        }

        let r = export_profdata_file(
            &ProfdataExportableWrapper {
                difftest: self,
                other_bins: other_binaries,
            },
            ignore_registry_files,
            ExportProfdataAction::Read,
        )?;
        let r = r.coverage_data();
        Ok(r)
    }

    /// Starts the analysis of the exported `.json` profiling data file.
    ///
    /// See the [`AnalysisContext`] type and the [`analysis`](crate::analysis)
    /// module for how to perform the analysis.
    ///
    /// This function should be ran after
    /// [`Difftest::merge_profraw_files_into_profdata`].
    pub fn start_analysis(
        &mut self,
        config: ExportProfdataConfig,
    ) -> DifftestsResult<AnalysisContext<'_>> {
        info!("Starting analysis...");
        let profdata = self.export_profdata(config)?;

        Ok(AnalysisContext::new(self, profdata))
    }

    /// Compiles the exported `.json` profiling data file into a
    /// [`TestIndex`] object, but does not save it on-disk.
    ///
    /// To save it on-disk, use [`TestIndex::write_to_file`].
    pub fn compile_test_index_data(
        &self,
        config: ExportProfdataConfig,
        index_data_compiler_config: IndexDataCompilerConfig,
    ) -> DifftestsResult<TestIndex> {
        info!("Compiling test index data...");

        let profdata = self.export_profdata(config)?;
        let test_index_data = TestIndex::index(self, profdata, index_data_compiler_config)?;

        info!("Done compiling test index data.");
        Ok(test_index_data)
    }
}

/// The configuration to use when exporting a `.profdata` file
/// into a `.json` file.
#[derive(Clone)]
pub struct ExportProfdataConfig {
    /// Whether to ignore files from the cargo registry.
    pub ignore_registry_files: bool,
    /// Other binaries to include in the export.
    ///
    /// By default, only the test binary is included (via the [`CoreTestDesc`]'s
    /// `bin_path` field), but if the test has spawned some other child process
    /// (as often happens when testing binaries), and they
    /// were profiled, the paths to those binaries should
    /// be passed here.
    pub other_binaries: Vec<PathBuf>,
}

/// A resolver for test index data file paths.
///
/// # Examples
///
/// ```
/// # use std::path::{Path, PathBuf};
/// # use cargo_difftests::difftest::DiscoverIndexPathResolver;
///
/// let resolver = DiscoverIndexPathResolver::Remap {from: "foo".into(), to: "bar".into()};
///
/// assert_eq!(resolver.resolve(Path::new("foo/bar/baz")), Some(PathBuf::from("bar/bar/baz")));
/// assert_eq!(resolver.resolve(Path::new("bar/baz")), None);
/// ```
pub enum DiscoverIndexPathResolver {
    /// Remaps the index path from the given `from` path to the given `to` path.
    Remap {
        /// The path to strip from the index path.
        from: PathBuf,
        /// The path to append to the stripped index path.
        to: PathBuf,
    },
    /// A custom remapping function.
    Custom {
        /// The remapping function.
        f: Box<dyn Fn(&Path) -> Option<PathBuf>>,
    },
}

impl DiscoverIndexPathResolver {
    /// Resolves the index path from the given [`Difftest`] directory.
    pub fn resolve(&self, p: &Path) -> Option<PathBuf> {
        match self {
            DiscoverIndexPathResolver::Remap { from, to } => {
                let p = p.strip_prefix(from).ok()?;
                Some(to.join(p))
            }
            DiscoverIndexPathResolver::Custom { f } => f(p),
        }
    }
}

fn discover_difftest_from_tempdir(
    dir: PathBuf,
    test_name_path: PathBuf,
    index_resolver: Option<&DiscoverIndexPathResolver>,
) -> DifftestsResult<Difftest> {
    let self_json = dir.join(cargo_difftests_core::CARGO_DIFFTESTS_SELF_JSON_FILENAME);

    let self_json = self_json.exists().then_some(self_json);

    if !test_name_path.exists() {
        return Err(DifftestsError::IO(std::io::Error::new(
            std::io::ErrorKind::NotFound,
            format!(
                "test name file does not exist: {}",
                test_name_path.display()
            ),
        )));
    }

    let cargo_difftests_version = dir.join(cargo_difftests_core::CARGO_DIFFTESTS_VERSION_FILENAME);

    if !cargo_difftests_version.exists() {
        return Err(DifftestsError::CargoDifftestsVersionDoesNotExist(
            cargo_difftests_version,
        ));
    }

    let version = fs::read_to_string(&cargo_difftests_version)?;

    if version != env!("CARGO_PKG_VERSION") {
        return Err(DifftestsError::CargoDifftestsVersionMismatch(
            version,
            env!("CARGO_PKG_VERSION").to_owned(),
        ));
    }

    let test_binary_path = dir.join(cargo_difftests_core::CARGO_DIFFTESTS_TEST_BINARY_FILENAME);

    if !test_binary_path.exists() {
        return Err(DifftestsError::IO(std::io::Error::new(
            std::io::ErrorKind::NotFound,
            format!(
                "test binary file does not exist: {}",
                test_binary_path.display()
            ),
        )));
    }

    let test_run = test_binary_path.metadata()?.modified()?;

    let mut profraws = Vec::new();

    let mut profdata_file = None;

    let mut cleaned = false;

    for e in dir.read_dir()? {
        let e = e?;
        let p = e.path();

        if !p.is_file() {
            continue;
        }

        let file_name = p.file_name();
        let ext = p.extension();

        if ext == Some(OsStr::new("profraw")) {
            profraws.push(p);
            continue;
        }

        if ext == Some(OsStr::new("profdata")) {
            if profdata_file.is_none() {
                profdata_file = Some(p);
            } else {
                warn!(
                    "multiple profdata files found in difftest directory: {}",
                    dir.display()
                );
                warn!("ignoring: {}", p.display());
            }
            continue;
        }

        if file_name == Some(OsStr::new(Difftest::CLEANED_FILE_NAME)) {
            cleaned = true;
        }
    }

    let index_data = 'index_data: {
        let index_data = index_resolver.and_then(|resolver| resolver.resolve(&dir));

        if let Some(ind) = &index_data {
            if !ind.exists() {
                debug!("index data file does not exist: {}", ind.display());
                break 'index_data None;
            } else if !ind.is_file() {
                debug!("index data file is not a file: {}", ind.display());
                break 'index_data None;
            }

            if ind.metadata()?.modified()? < test_run {
                warn!("index data file is older than test run");
                break 'index_data None;
            }
        }

        index_data
    };

    Ok(Difftest {
        dir,
        test_binary_path,
        test_name_path,
        profraws,
        self_json,
        test_run_time: test_run,
        profdata_file,
        index_data,
        cleaned,
    })
}

fn discover_difftests_to_vec(
    dir: &Path,
    discovered: &mut Vec<Difftest>,
    ignore_incompatible: bool,
    index_resolver: Option<&DiscoverIndexPathResolver>,
) -> DifftestsResult {
    let test_name_path = dir.join(cargo_difftests_core::CARGO_DIFFTESTS_TEST_NAME_FILENAME);
    if test_name_path.exists() && test_name_path.is_file() {
        let r = discover_difftest_from_tempdir(dir.to_path_buf(), test_name_path, index_resolver);

        if let Err(DifftestsError::CargoDifftestsVersionMismatch(_, _)) = r {
            if ignore_incompatible {
                return Ok(());
            }
        }

        discovered.push(r?);
        return Ok(());
    }

    for entry in fs::read_dir(dir)? {
        let entry = entry?;
        let path = entry.path();
        if path.is_dir() {
            discover_difftests_to_vec(&path, discovered, ignore_incompatible, index_resolver)?;
        }
    }

    Ok(())
}

/// Discovers all the [`Difftest`]s in the given directory,
/// optionally using the `index_resolver` to resolve the index paths,
/// and ignoring incompatible [`Difftest`] directories if `ignore_incompatible`
/// is true.
pub fn discover_difftests(
    dir: &Path,
    ignore_incompatible: bool,
    index_resolver: Option<&DiscoverIndexPathResolver>,
) -> DifftestsResult<Vec<Difftest>> {
    let mut discovered = Vec::new();

    discover_difftests_to_vec(dir, &mut discovered, ignore_incompatible, index_resolver)?;

    Ok(discovered)
}

pub trait ProfrawsMergeable {
    fn list_profraws(&self) -> impl Iterator<Item = &Path>;
    fn out_profdata_path(&self) -> PathBuf;
}

impl ProfrawsMergeable for Difftest {
    fn list_profraws(&self) -> impl Iterator<Item = &Path> {
        self.profraws.iter().map(PathBuf::as_path)
    }

    fn out_profdata_path(&self) -> PathBuf {
        const OUT_FILE_NAME: &str = "merged.profdata";

        self.dir.join(OUT_FILE_NAME)
    }
}

pub fn merge_profraws<T: ProfrawsMergeable>(d: &T) -> DifftestsResult<()> {
    debug!(
        "Merging profraw files into {}...",
        d.out_profdata_path().display()
    );

    let p = d.out_profdata_path();

    let mut cmd = Command::new("rust-profdata");

    cmd.arg("merge")
        .arg("-sparse")
        .args(d.list_profraws())
        .arg("-o")
        .arg(&p);

    let status = cmd.status()?;

    if !status.success() {
        return Err(DifftestsError::ProcessFailed {
            name: "rust-profdata",
        });
    }

    Ok(())
}

pub trait ProfDataExportable {
    fn profdata_path(&self) -> &Path;
    fn main_bin_path(&self) -> DifftestsResult<PathBuf>;
    fn other_bins(&self) -> impl Iterator<Item = &Path>;
}

struct ProfdataExportableWrapper<'a> {
    difftest: &'a Difftest,
    other_bins: Vec<PathBuf>,
}

impl<'a> ProfDataExportable for ProfdataExportableWrapper<'a> {
    fn profdata_path(&self) -> &Path {
        self.difftest.profdata_file.as_ref().unwrap()
    }

    fn main_bin_path(&self) -> DifftestsResult<PathBuf> {
        self.difftest.read_test_binary_path()
    }

    fn other_bins(&self) -> impl Iterator<Item = &Path> {
        self.other_bins.iter().map(PathBuf::as_path)
    }
}

#[derive(Debug)]
pub enum ExportProfdataAction {
    Store(PathBuf),
    Read,
}

#[derive(Debug)]
pub enum ExportProfdataActionResult {
    Store,
    Read(analysis_data::CoverageData),
}

impl ExportProfdataActionResult {
    pub fn coverage_data(self) -> analysis_data::CoverageData {
        match self {
            ExportProfdataActionResult::Store => unreachable!(),
            ExportProfdataActionResult::Read(r) => r,
        }
    }
}

pub fn export_profdata_file(
    d: &impl ProfDataExportable,
    ignore_registry_files: bool,
    action: ExportProfdataAction,
) -> DifftestsResult<ExportProfdataActionResult> {
    debug!(
        "Exporting profdata file from {}...",
        d.profdata_path().display()
    );

    let mut cmd = Command::new("rust-cov");

    cmd.arg("export")
        .arg("-instr-profile")
        .arg(d.profdata_path())
        .arg(d.main_bin_path()?)
        .args(
            d.other_bins()
                .flat_map(|it| [OsStr::new("--object"), it.as_os_str()]),
        );

    #[cfg(not(windows))]
    const REGISTRY_FILES_REGEX: &str = r#"/\.cargo/registry"#;

    #[cfg(windows)]
    const REGISTRY_FILES_REGEX: &str = r#"\\\.cargo\\registry"#;

    if ignore_registry_files {
        cmd.arg("-ignore-filename-regex").arg(REGISTRY_FILES_REGEX);
    }

    match &action {
        ExportProfdataAction::Store(p) => {
            cmd.stdout(fs::File::create(p)?);
        }
        ExportProfdataAction::Read => {
            cmd.stdout(std::process::Stdio::piped());
        }
    }

    debug!("Running: {:?}", cmd);

    let mut process = cmd.spawn()?;

    let r = match &action {
        ExportProfdataAction::Store(_p) => {
            let status = process.wait()?;

            if !status.success() {
                return Err(DifftestsError::ProcessFailed { name: "rust-cov" });
            }

            ExportProfdataActionResult::Store
        }
        ExportProfdataAction::Read => {
            let r = {
                let stdout = process.stdout.as_mut().unwrap();
                let stdout = std::io::BufReader::new(stdout);

                serde_json::from_reader(stdout)
                    .map_err(|e| DifftestsError::Json(e, Some(d.profdata_path().to_path_buf())))?
            };

            let status = process.wait()?;

            if !status.success() {
                return Err(DifftestsError::ProcessFailed { name: "rust-cov" });
            }

            ExportProfdataActionResult::Read(r)
        }
    };

    Ok(r)
}

#[derive(serde::Serialize, serde::Deserialize, Debug, Clone)]
pub struct TestInfo {
    pub test_name: String,
    pub test_binary: PathBuf,

    pub extra_desc: Option<CoreTestDesc>,
}