zisk-rom-setup 1.1.0-alpha

use anyhow::{Context, Result};
use proofman_common::{write_custom_commit_trace, ProofCtx, ProofmanError, ProofmanResult};
use proofman_fields::{Goldilocks, PrimeField64};
use std::fs;
use std::fs::File;
use std::io::Read;
use std::path::{Path, PathBuf};
use zisk_common::{ZiskPaths, PROGRAM_VK_LEN};
use zisk_pil::{RomRomTrace, PILOUT_HASH};
use zisk_sm_rom::CustomRom;

pub use zisk_common::HashMode;

/// Resolve the artifact output directory, defaulting to the ZisK cache when
/// `output_dir` is `None`. Creates the directory if needed and returns its
/// canonical absolute path.
pub fn get_output_path(output_dir: &Option<PathBuf>) -> Result<PathBuf> {
    let output_path = if output_dir.is_none() {
        let cache_path = ZiskPaths::global().cache.clone();
        ensure_dir_exists(&cache_path);
        cache_path
    } else {
        ensure_dir_exists(output_dir.as_ref().unwrap());
        output_dir.clone().unwrap()
    };

    let output_path = fs::canonicalize(&output_path)
        .with_context(|| format!("Failed to get absolute path for {output_path:?}"))?;

    Ok(output_path)
}

/// Build the ROM custom-commit trace for `elf`, write it to `rom_buffer_path`,
/// and return its Merkle root (the raw material for the verkey).
pub fn gen_elf_hash<F: PrimeField64>(
    pctx: &ProofCtx<F>,
    elf: &[u8],
    rom_buffer_path: &Path,
    hash_mode: HashMode,
) -> ProofmanResult<Vec<F>> {
    let mut custom_rom_trace =
        CustomRom::build::<F>(elf).map_err(|e| ProofmanError::InvalidParameters(e.to_string()))?;

    write_custom_commit_trace(
        pctx,
        &mut custom_rom_trace,
        hash_mode.blowup_factor(),
        hash_mode.merkle_tree_arity(),
        rom_buffer_path,
    )
}

/// Read a verification key (a `PROGRAM_VK_LEN` array of little-endian `u64`s)
/// from `verkey_path`. Returns `None` if the file does not exist.
pub fn get_elf_vk(verkey_path: &Path) -> Result<Option<Vec<u64>>> {
    if !verkey_path.exists() {
        return Ok(None);
    }

    let mut file = File::open(verkey_path)?;
    let mut vk = vec![0u64; PROGRAM_VK_LEN];
    for word in vk.iter_mut() {
        let mut buf = [0u8; 8];
        file.read_exact(&mut buf)?;
        *word = u64::from_le_bytes(buf);
    }
    Ok(Some(vk))
}

/// Read the ELF at `elf_path` and return the hex-encoded blake3 hash of its
/// bytes. See [`get_elf_data_hash`] for the in-memory variant.
pub fn get_elf_data_hash_from_path(elf_path: &Path) -> Result<String> {
    let elf_data =
        fs::read(elf_path).with_context(|| format!("Error reading ELF file: {elf_path:?}"))?;

    let hash = blake3::hash(&elf_data).to_hex().to_string();

    Ok(hash)
}

/// The hex-encoded blake3 hash of `elf`. This is the content address used to
/// name and locate all of a program's cached artifacts.
pub fn get_elf_data_hash(elf: &[u8]) -> String {
    blake3::hash(elf).to_hex().to_string()
}

/// Build the cache path of the ROM custom-commit binary for `hash`.
///
/// The filename encodes the ELF hash, PILOUT hash, row count, and hash-mode
/// parameters (plus a `_gpu` suffix when `gpu` is set) so incompatible setups
/// never collide.
pub fn get_elf_bin_file_path_with_hash(
    hash: &str,
    default_cache_path: &Path,
    gpu: bool,
    hash_mode: HashMode,
) -> Result<PathBuf> {
    let pilout_hash = PILOUT_HASH;

    let n = RomRomTrace::<Goldilocks>::NUM_ROWS;

    let gpu_suffix = if gpu { "_gpu" } else { "" };
    let rom_cache_file_name = format!(
        "{}_{}_{}_{}_{}_{}{}.bin",
        hash,
        pilout_hash,
        n,
        hash_mode.file_tag(),
        hash_mode.blowup_factor(),
        hash_mode.merkle_tree_arity(),
        gpu_suffix
    );

    Ok(default_cache_path.join(rom_cache_file_name))
}

/// Build the cache path of the verkey (`.verkey.bin`) file for `hash`,
/// mirroring the naming scheme of [`get_elf_bin_file_path_with_hash`].
pub fn get_elf_bin_verkey_file_path_with_hash(
    hash: &str,
    default_cache_path: &Path,
    hash_mode: HashMode,
) -> Result<PathBuf> {
    let pilout_hash = PILOUT_HASH;

    let n = RomRomTrace::<Goldilocks>::NUM_ROWS;

    let rom_cache_file_name = format!(
        "{}_{}_{}_{}_{}_{}.verkey.bin",
        hash,
        pilout_hash,
        n,
        hash_mode.file_tag(),
        hash_mode.blowup_factor(),
        hash_mode.merkle_tree_arity(),
    );

    Ok(default_cache_path.join(rom_cache_file_name))
}

/// Create `path` (and any missing parents) if it does not already exist.
///
/// # Panics
/// Panics if the directory cannot be created for a reason other than it
/// already existing.
pub fn ensure_dir_exists(path: &Path) {
    if let Err(e) = std::fs::create_dir_all(path) {
        if e.kind() != std::io::ErrorKind::AlreadyExists {
            panic!("Failed to create cache directory {path:?}: {e}");
        }
    }
}