use serde::{Deserialize, Serialize};
use std::path::{Path, PathBuf};
use std::process::{Child, Command, Output, Stdio};
use std::thread;
use std::time::{Duration, Instant};
use thiserror::Error;
#[derive(Debug, Error)]
pub enum HeapParityRunnerError {
#[error("Lean heap parity runner binary not found at {0}")]
BinaryNotFound(PathBuf),
#[error("Lean heap parity runner failed with exit code {code}: {stderr}")]
ProcessFailed { code: i32, stderr: String },
#[error("Failed to parse Lean heap parity output: {0}")]
ParseError(String),
#[error("Lean heap parity runner '{operation}' timed out after {timeout_ms}ms")]
TimedOut { operation: String, timeout_ms: u64 },
#[error("IO error: {0}")]
Io(#[from] std::io::Error),
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct HeapParityResourceFixtureOutput {
pub canonical_bytes_hex: String,
pub resource_id_preimage_hex: String,
pub resource_leaf_preimage_hex: String,
pub nullifier_leaf_preimage_hex: String,
pub resource_id_digest_hex: String,
pub nullifier_leaf_hex: String,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct HeapParityProofStep {
pub direction: String,
pub sibling_hash_hex: String,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct HeapParityProof {
pub leaf_hash_hex: String,
pub path: Vec<HeapParityProofStep>,
pub root_hex: String,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct HeapParityResourceId {
pub digest_hex: String,
pub counter: u64,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct HeapParityStateSummary {
pub active_resource_ids: Vec<HeapParityResourceId>,
pub consumed_resource_ids: Vec<HeapParityResourceId>,
pub counter: u64,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct HeapParityCommitment {
pub resource_root_hex: String,
pub nullifier_root_hex: String,
pub counter: u64,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct HeapParityHeapFixtureOutput {
pub active_resource_ids_hex: Vec<String>,
pub consumed_resource_ids_hex: Vec<String>,
pub proof_index: u64,
pub resource_proof: Option<HeapParityProof>,
pub nullifier_proof: Option<HeapParityProof>,
pub merkle_node_preimage_hex: String,
pub heap_commitment: HeapParityCommitment,
pub heap_commitment_preimage_hex: String,
pub commitment_hash_hex: String,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct HeapParityReplayOutput {
pub first_run: HeapParityStateSummary,
pub second_run: HeapParityStateSummary,
pub stable: bool,
pub active_ids_match_fixture: bool,
pub consumed_ids_match_fixture: bool,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct HeapParityOutput {
pub schema_version: String,
pub heap_encoding_version: u16,
pub hasher: String,
pub resource_fixture: HeapParityResourceFixtureOutput,
pub heap_fixture: HeapParityHeapFixtureOutput,
pub replay: HeapParityReplayOutput,
}
pub struct HeapParityRunner {
binary_path: PathBuf,
}
impl HeapParityRunner {
pub const DEFAULT_BINARY_PATH: &'static str = "lean/.lake/build/bin/heap_parity_runner";
pub const DEFAULT_TIMEOUT_MS: u64 = 120_000;
fn process_timeout() -> Duration {
let ms = std::env::var("TELLTALE_LEAN_TIMEOUT_MS")
.ok()
.and_then(|raw| raw.parse::<u64>().ok())
.unwrap_or(Self::DEFAULT_TIMEOUT_MS);
Duration::from_millis(ms.max(1))
}
fn wait_with_timeout(
mut child: Child,
timeout: Duration,
operation: &str,
) -> Result<Output, HeapParityRunnerError> {
let start = Instant::now();
loop {
match child.try_wait()? {
Some(_) => {
return child
.wait_with_output()
.map_err(HeapParityRunnerError::from)
}
None => {
if start.elapsed() >= timeout {
if let Err(err) = child.kill() {
eprintln!(
"best-effort child.kill failed during timeout handling: {err}"
);
}
if let Err(err) = child.wait() {
eprintln!(
"best-effort child.wait failed during timeout handling: {err}"
);
}
return Err(HeapParityRunnerError::TimedOut {
operation: operation.to_string(),
timeout_ms: u64::try_from(timeout.as_millis()).unwrap_or(u64::MAX),
});
}
thread::sleep(Duration::from_millis(10));
}
}
}
}
fn run_command_with_timeout(
&self,
mut command: Command,
operation: &str,
) -> Result<Output, HeapParityRunnerError> {
command.stdout(Stdio::piped());
command.stderr(Stdio::piped());
let child = command.spawn()?;
Self::wait_with_timeout(child, Self::process_timeout(), operation)
}
fn find_workspace_root() -> Option<PathBuf> {
let manifest_dir = env!("CARGO_MANIFEST_DIR");
let mut path = PathBuf::from(manifest_dir);
for _ in 0..5 {
if path.join("lean/.lake").is_dir() {
return Some(path);
}
if !path.pop() {
break;
}
}
None
}
fn get_binary_path() -> Option<PathBuf> {
Self::find_workspace_root()
.map(|root| root.join(Self::DEFAULT_BINARY_PATH))
.filter(|path| path.exists())
}
pub fn new() -> Result<Self, HeapParityRunnerError> {
match Self::get_binary_path() {
Some(path) => Ok(Self { binary_path: path }),
None => Err(HeapParityRunnerError::BinaryNotFound(PathBuf::from(
Self::DEFAULT_BINARY_PATH,
))),
}
}
pub fn with_binary_path(path: impl AsRef<Path>) -> Result<Self, HeapParityRunnerError> {
let binary_path = PathBuf::from(path.as_ref());
if !binary_path.exists() || !binary_path.is_file() {
return Err(HeapParityRunnerError::BinaryNotFound(binary_path));
}
Ok(Self { binary_path })
}
#[must_use]
pub fn try_new() -> Option<Self> {
Self::new().ok()
}
#[must_use]
pub fn is_available() -> bool {
Self::get_binary_path().is_some()
}
pub fn require_available() {
if !Self::is_available() {
panic!(
"Lean heap parity runner required but unavailable. Run `cd lean && lake build heap_parity_runner`."
);
}
}
pub fn run_fixture(
&self,
fixture_path: impl AsRef<Path>,
) -> Result<HeapParityOutput, HeapParityRunnerError> {
let mut command = Command::new(&self.binary_path);
command.arg(fixture_path.as_ref());
let output = self.run_command_with_timeout(command, "run_fixture")?;
if !output.status.success() {
return Err(HeapParityRunnerError::ProcessFailed {
code: output.status.code().unwrap_or(-1),
stderr: String::from_utf8_lossy(&output.stderr).to_string(),
});
}
serde_json::from_slice(&output.stdout)
.map_err(|err| HeapParityRunnerError::ParseError(err.to_string()))
}
}