//! This crate provides code coverage support for no_std and embedded programs.
//!
//! # Usage
//!
//! Note: This crate requires a recent nightly compiler (2020-04-20 or later).
//!
//! 1. Install the `cargo-minicov` tool:
//!
//! ```sh
//! cargo install cargo-minicov
//! ```
//!
//! 2. Ensure that the following environment variables are set up:
//!
//! ```sh
//! export CARGO_INCREMENTAL=0
//! export RUSTFLAGS="-Zprofile -Zno-profiler-runtime -Copt-level=0 -Clink-dead-code -Coverflow-checks=off"
//! ```
//!
//! Note that the use of these flags may cause build-dependencies and proc
//! macros to fail to compile. This can be worked around by explicitly
//! specifying a target when invoking cargo:
//!
//! ```sh
//! # Fails to compile
//! cargo build
//!
//! # Works
//! cargo build --target x86_64-unknown-linux-gnu
//! ```
//!
//!
//! 3. Add the `minicov` crate as a dependency to your program:
//!
//! ```toml
//! [dependencies]
//! minicov = "0.1"
//! ```
//!
//! 4. (optional) The profiling instrumentation generated by LLVM relies on
//!    global constructors to work. This will generally work out of the box on
//!    most systems. However if your program runs on bare metal then you may
//!    need to do this yourself. minicov provides a helper function for this:
//!
//! ```ignore
//! unsafe {
//!     minicov::run_static_constructors();
//! }
//! ```
//!
//! WARNING: Make sure you don't call static constructors again if your runtime
//! has already done this for you. Doing so is undefined behavior and may result
//! in crashes and/or data corruption.
//!
//! 5. Before your program exits, call `minicov::capture_coverage` which returns
//!    a `Vec<u8>` and dump its contents to a file:
//!
//! ```ignore
//! fn main() {
//!     // ...
//!
//!     let coverage = minicov::capture_coverage().unwrap();
//!     std::fs::write("output.minicov", coverage).unwrap();
//! }
//! ```
//!
//! If your program is running on a different system than your build system then
//! you will need to transfer this file back to your build system.
//!
//! 6. After your program finishes running, use the `cargo minicov` command to
//!    generate GCOV .gcda files from the captured coverage:
//!
//! ```sh
//! cargo minicov output.minicov
//! ```
//!
//! You can pass multiple input files to `cargo minicov`, or even concatenate
//! multiple input files into a single file:
//!
//! ```sh
//! cargo minicov a.minicov b.minicov
//!
//! # Or
//!
//! cat a.minicov b.minicov > combined.minicov
//! cargo minicov combined.minicov
//! ```
//!
//! 7. Use your favorite GCOV-compatible coverage tool (e.g. [grcov]) to
//!    process the .gcda files.
//!
//! ```sh
//! grcov ./target/x86_64-unknown-linux-gnu/debug/ -s . -t html --llvm --branch --ignore-not-existing -o ./target/debug/coverage/
//! ```
//!
//! [grcov]: https://github.com/mozilla/grcov

#![no_std]
#![warn(missing_docs)]
#![warn(rust_2018_idioms)]

extern crate alloc;

use alloc::vec::Vec;
use core::fmt;
use cstr_core::{c_char, CStr};
use postcard::flavors::AllocVec;
use serde::ser::{SerializeSeq, Serializer};
use serde::Serialize;
use spinning_top::Spinlock;

/// Unfortunately LLVM's GCOV implementation relies on global state, so we have
/// to use this ugly hack.
///
/// This is only accessed while the REGISTERED_FNS lock is held.
static mut CALLBACK: Option<*mut dyn FnMut(CovEvent)> = None;

/// These point to compiler-generated functions which will invoke the various
/// `llvm_gcda_*` functions below to actually emit the coverage data.
static REGISTERED_FNS: Spinlock<Vec<extern "C" fn()>> = Spinlock::new(Vec::new());

/// Magic bytes at the start of the data to avoid deserializing garbage.
const MAGIC: [u8; 8] = *b"minicov\0";

/// File format version.
const VERSION: [u8; 4] = *b"v01\0";

/// Recording of a call to a `llvm_gcda_*` function.
///
/// We basically just capture the parameters to the functions and serialize them.
#[derive(Serialize)]
enum CovEvent {
    StartFile {
        orig_filename: &'static [u8],
        version: [c_char; 4],
        checksum: u32,
    },
    EmitFunction {
        ident: u32,
        function_name: &'static [u8],
        func_checksum: u32,
        use_extra_checksum: u8,
        cfg_checksum: u32,
    },
    EmitArcs {
        counters: Counters,
    },
    SummaryInfo,
    EndFile,
    End,
}

/// Custom implementation of `Serialize` for the counters. It can be
/// deserialized as a `Vec<u64>` later.
struct Counters {
    num_counters: usize,
    counters: *mut u64,
}

impl Serialize for Counters {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        let mut serializer = serializer.serialize_seq(Some(self.num_counters))?;
        for i in 0..self.num_counters {
            // Use a volatile read since other threads may be concurrently
            // modifying the pointers.
            let val = unsafe { self.counters.add(i).read_volatile() };
            serializer.serialize_element(&val)?;
        }
        serializer.end()
    }
}

/// LLVM generates global constructors that call this function.
#[no_mangle]
unsafe extern "C" fn llvm_gcov_init(_writeout: extern "C" fn(), flush: extern "C" fn()) {
    REGISTERED_FNS.lock().push(flush);
}

#[no_mangle]
unsafe extern "C" fn llvm_gcda_start_file(
    orig_filename: *const c_char,
    version: *const [c_char; 4],
    checksum: u32,
) {
    if let Some(callback) = CALLBACK {
        // Exclude the null byte so that deserialization matches that of CString.
        let orig_filename = CStr::from_ptr(orig_filename).to_bytes();

        (*callback)(CovEvent::StartFile {
            orig_filename,
            version: *version,
            checksum,
        });
    }
}

#[no_mangle]
unsafe extern "C" fn llvm_gcda_emit_function(
    ident: u32,
    function_name: *const c_char,
    func_checksum: u32,
    use_extra_checksum: u8,
    cfg_checksum: u32,
) {
    if let Some(callback) = CALLBACK {
        // Exclude the null byte so that deserialization matches that of CString.
        let function_name = CStr::from_ptr(function_name).to_bytes();

        (*callback)(CovEvent::EmitFunction {
            ident,
            function_name,
            func_checksum,
            use_extra_checksum,
            cfg_checksum,
        });
    }
}

#[no_mangle]
unsafe extern "C" fn llvm_gcda_emit_arcs(num_counters: u32, counters: *mut u64) {
    if let Some(callback) = CALLBACK {
        let counters = Counters {
            num_counters: num_counters as usize,
            counters,
        };
        (*callback)(CovEvent::EmitArcs { counters });
    }
}

#[no_mangle]
unsafe extern "C" fn llvm_gcda_summary_info() {
    if let Some(callback) = CALLBACK {
        (*callback)(CovEvent::SummaryInfo);
    }
}

#[no_mangle]
unsafe extern "C" fn llvm_gcda_end_file() {
    if let Some(callback) = CALLBACK {
        (*callback)(CovEvent::EndFile);
    }
}

/// Error type returned if no coverage data is found.
///
/// This typically happens if you didn't compile your code with `-Zprofile` or
/// if you didn't run the static constructors to register the coverage handlers.
#[derive(Copy, Clone, Debug)]
pub struct NoCoverageData;
impl fmt::Display for NoCoverageData {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str("no coverage data found")
    }
}

/// Captures the coverage data for the current program and returns it as a
/// binary blob.
///
/// The resulting blobs should be passed to the `cargo minicov` tool on the
/// system used to build your program. This will generate .gcda files in your
/// build directory which can be parsed by any GCOV-compatible tool such as
/// `lcov` or `grcov`.
///
/// This function can be called multiple times. The coverage counters are reset
/// after each call, so you will need to merge the results of each call with the
/// `cargo minicov` tool.
pub fn capture_coverage() -> Result<Vec<u8>, NoCoverageData> {
    // Prepend some identifying information about the byte stream.
    let mut out = Vec::from(MAGIC);
    out.extend_from_slice(&VERSION);

    // We need to cheat with the lifetime here, but it's fine since this is the
    // only place that (indirectly) invokes the llvm_gcda_* functions.
    let vec_ptr = &mut out as *mut Vec<u8>;
    let mut callback = move |event: CovEvent| unsafe {
        let out = vec_ptr.replace(Vec::new());
        let out = postcard::serialize_with_flavor(&event, AllocVec(out)).unwrap();
        *vec_ptr = out;
    };

    // Invoke all the registered functions
    unsafe {
        let registered_fns = REGISTERED_FNS.lock();
        CALLBACK = Some(&mut callback as *mut dyn FnMut(CovEvent));
        for f in registered_fns.iter() {
            f();
        }
        CALLBACK = None;
    }

    if out.len() == MAGIC.len() + VERSION.len() {
        Err(NoCoverageData)
    } else {
        // Append a marker to indicate the end of the data stream
        let out = postcard::serialize_with_flavor(&CovEvent::End, AllocVec(out)).unwrap();
        Ok(out)
    }
}

/// Resets all coverage counters in the program to zero.
///
/// This function should be called after a process forks to avoid recording
/// coverage data for the parent process twice.
///
/// Note that calls to [`capture_coverage`] will implicitly reset the counters.
pub fn reset_coverage_counters() {
    // LLVM's GCOV interface doesn't support __gcov_reset, but we can emulate it
    // by doing the equivalent of __gcov_flush and discarding the results.
    unsafe {
        CALLBACK = None;

        let registered_fns = REGISTERED_FNS.lock();
        for f in registered_fns.iter() {
            f();
        }
    }
}

/// Helper function to run the static constructors if your runtime doesn't do it
/// for you.
///
/// This is necessary to register the LLVM-generated coverage callbacks with
/// minicov.
///
/// # Safety
///
/// This must only be called **once** at program startup. Do not call this
/// function if your runtime environment (e.g. libc) does this for you.
///
/// The global allocator must be available to use when this function is called.
pub unsafe fn run_static_constructors() {
    extern "C" {
        static mut __init_array_start: [extern "C" fn(); 0];
        static mut __init_array_end: [extern "C" fn(); 0];
    }

    let mut ptr = __init_array_start.as_ptr();
    let end = __init_array_end.as_ptr();
    while ptr != end {
        (*ptr)();
        ptr = ptr.add(1);
    }
}