lazaret/crates/lazaret-ignored-workspace/src/bin/lazaret-workspace-benchmark.rs

use std::env;
use std::ffi::OsString;
use std::fs;
use std::io::{self, Read};
use std::os::unix::ffi::{OsStrExt, OsStringExt};
use std::os::unix::fs::{FileTypeExt, MetadataExt};
use std::path::{Path, PathBuf};

use rustix::fs::{AtFlags, StatxFlags, statx};
use sha2::{Digest, Sha256};

const FIXTURE_CONTEXT: &[u8] = b"lazaret-ignored-workspace-fixture-v1\0";

fn main() {
    if let Err(error) = run() {
        eprintln!("lazaret workspace benchmark refused: {error}");
        std::process::exit(1);
    }
}

fn run() -> io::Result<()> {
    let arguments = env::args_os().skip(1).collect::<Vec<_>>();
    let root = parse_metadata_arguments(&arguments)?;
    let result = measure_metadata(&root)?;
    println!(
        "{}",
        serde_json::json!({
            "content_bytes": result.content_bytes,
            "fixture_identity": hex(&result.fixture_identity),
            "regular_files": result.regular_files,
        })
    );
    Ok(())
}

fn parse_metadata_arguments(arguments: &[OsString]) -> io::Result<PathBuf> {
    if arguments.len() != 6
        || arguments[0] != "metadata"
        || arguments[1] != "--root"
        || arguments[3] != "--bytewise-order"
        || arguments[4] != "--statx-mask"
        || arguments[5] != "basic-stats"
    {
        return Err(io::Error::other(
            "usage: lazaret-workspace-benchmark metadata --root ROOT --bytewise-order --statx-mask basic-stats",
        ));
    }
    Ok(PathBuf::from(&arguments[2]))
}

fn measure_metadata(root: &Path) -> io::Result<MetadataResult> {
    let mut entries = Vec::new();
    collect_entries(root, Path::new(""), &mut entries)?;
    entries.sort_by(|left, right| {
        left.relative
            .as_os_str()
            .as_bytes()
            .cmp(right.relative.as_os_str().as_bytes())
    });

    let mut stream = Vec::from(FIXTURE_CONTEXT);
    stream.extend_from_slice(
        &u32::try_from(entries.len())
            .map_err(|_| io::Error::other("fixture contains too many entries"))?
            .to_be_bytes(),
    );
    let mut regular_files = 0_u64;
    let mut content_bytes = 0_u64;
    for entry in entries {
        let path_bytes = entry.relative.as_os_str().as_bytes();
        stream.extend_from_slice(
            &u32::try_from(path_bytes.len())
                .map_err(|_| io::Error::other("fixture path is too long"))?
                .to_be_bytes(),
        );
        stream.extend_from_slice(path_bytes);
        stream.push(entry.kind);
        stream.extend_from_slice(&(entry.metadata.mode() & 0o7777).to_be_bytes());
        let payload_length = u64::try_from(entry.payload.len())
            .map_err(|_| io::Error::other("fixture payload is too large"))?;
        stream.extend_from_slice(&payload_length.to_be_bytes());
        stream.extend_from_slice(&Sha256::digest(&entry.payload));
        if entry.kind == 2 {
            regular_files = regular_files.saturating_add(1);
            content_bytes = content_bytes.saturating_add(payload_length);
        }
    }
    Ok(MetadataResult {
        regular_files,
        content_bytes,
        fixture_identity: Sha256::digest(stream).into(),
    })
}

fn collect_entries(
    root: &Path,
    relative: &Path,
    entries: &mut Vec<FixtureEntry>,
) -> io::Result<()> {
    let mut children = fs::read_dir(root.join(relative))?
        .map(|entry| entry.map(|entry| entry.file_name()))
        .collect::<io::Result<Vec<_>>>()?;
    children.sort_by(|left, right| left.as_bytes().cmp(right.as_bytes()));
    for name in children {
        let child_relative = relative.join(name);
        let path = root.join(&child_relative);
        let metadata = fs::symlink_metadata(&path)?;
        statx(
            rustix::fs::CWD,
            &path,
            AtFlags::SYMLINK_NOFOLLOW,
            StatxFlags::BASIC_STATS,
        )?;
        let file_type = metadata.file_type();
        let (kind, payload) = if file_type.is_dir() {
            (1, Vec::new())
        } else if file_type.is_file() {
            let mut payload = Vec::new();
            fs::File::open(&path)?.read_to_end(&mut payload)?;
            (2, payload)
        } else if file_type.is_symlink() {
            (3, fs::read_link(&path)?.into_os_string().into_vec())
        } else if file_type.is_socket()
            || file_type.is_fifo()
            || file_type.is_block_device()
            || file_type.is_char_device()
        {
            return Err(io::Error::other(
                "fixture contains an unsupported entry kind",
            ));
        } else {
            return Err(io::Error::other("fixture entry kind is unknown"));
        };
        entries.push(FixtureEntry {
            relative: child_relative.clone(),
            metadata,
            kind,
            payload,
        });
        if kind == 1 {
            collect_entries(root, &child_relative, entries)?;
        }
    }
    Ok(())
}

fn hex(bytes: &[u8]) -> String {
    const DIGITS: &[u8; 16] = b"0123456789abcdef";
    let mut encoded = String::with_capacity(bytes.len() * 2);
    for byte in bytes {
        encoded.push(char::from(DIGITS[usize::from(byte >> 4)]));
        encoded.push(char::from(DIGITS[usize::from(byte & 0x0f)]));
    }
    encoded
}

struct FixtureEntry {
    relative: PathBuf,
    metadata: fs::Metadata,
    kind: u8,
    payload: Vec<u8>,
}

struct MetadataResult {
    regular_files: u64,
    content_bytes: u64,
    fixture_identity: [u8; 32],
}