runner & prover TESTED

crates/prover/Cargo.toml

@@ -13,7 +13,9 @@ cairo-proof-parser.workspace = true
 futures.workspace= true
 itertools.workspace = true
 serde_json.workspace = true
+serde.workspace = true
 sharp-p2p-common.workspace = true
+strum.workspace = true
 tempfile.workspace = true
 thiserror.workspace = true
 tokio.workspace = true

crates/prover/src/stone_prover/mod.rs

@@ -1,117 +1,113 @@
-use crate::{
-    errors::ProverControllerError,
-    traits::{Prover, ProverController},
-};
+use crate::{errors::ProverControllerError, traits::ProverController};
+use async_process::Stdio;
+use futures::Future;
 use itertools::{chain, Itertools};
-use sharp_p2p_common::{hash, job_trace::JobTrace, job_witness::JobWitness, vec252::VecFelt252};
+use sharp_p2p_common::{
+    hash, job_trace::JobTrace, job_witness::JobWitness, process::Process, vec252::VecFelt252,
+};
 use std::{
-    collections::HashMap,
     hash::{DefaultHasher, Hash, Hasher},
     io::Read,
+    pin::Pin,
 };
 use tempfile::NamedTempFile;
-use tokio::process::{Child, Command};
-use tracing::{debug, trace};
+use tokio::{process::Command, select, sync::mpsc};
+use tracing::debug;
+
+pub mod types;
+
+#[cfg(test)]
+pub mod tests;
 
 pub struct StoneProver {
-    tasks: HashMap<u64, Child>,
+    config: NamedTempFile,
+    params: NamedTempFile,
 }
 
-impl Prover for StoneProver {
-    fn init() -> impl ProverController {
-        Self { tasks: HashMap::new() }
+impl StoneProver {
+    pub fn new(config: NamedTempFile, params: NamedTempFile) -> Self {
+        Self { config, params }
     }
 }
 
 impl ProverController for StoneProver {
-    async fn prove(&mut self, job_trace: JobTrace) -> Result<JobWitness, ProverControllerError> {
-        let mut out_file = NamedTempFile::new()?;
-
-        let cpu_air_prover_config = NamedTempFile::new()?; // TODO implement default config and getting info from integrity verifier
-        let cpu_air_params = NamedTempFile::new()?; // TODO implement default config and getting info from integrity verifier
-
-        let task = Command::new("cpu_air_prover")
-            .arg("--out_file")
-            .arg(out_file.path())
-            .arg("--air_private_input")
-            .arg(job_trace.air_private_input.path())
-            .arg("--air_public_input")
-            .arg(job_trace.air_public_input.path())
-            .arg("--cpu_air_prover_config")
-            .arg(cpu_air_prover_config.path())
-            .arg("--cpu_air_params")
-            .arg(cpu_air_params.path())
-            .arg("--generate_annotations")
-            .spawn()?;
-
-        let job_trace_hash = hash!(job_trace);
-
-        debug!("task {} spawned", job_trace_hash);
-        self.tasks.insert(job_trace_hash.to_owned(), task);
-
-        let task_status = self
-            .tasks
-            .get_mut(&job_trace_hash)
-            .ok_or(ProverControllerError::TaskNotFound)?
-            .wait()
-            .await?;
-
-        trace!("task {} woke up", job_trace_hash);
-        if !task_status.success() {
-            debug!("task terminated {}", job_trace_hash);
-            return Err(ProverControllerError::TaskTerminated);
-        }
-
-        let task_output = self
-            .tasks
-            .remove(&job_trace_hash)
-            .ok_or(ProverControllerError::TaskNotFound)?
-            .wait_with_output()
-            .await?;
-        trace!("task {} output {:?}", job_trace_hash, task_output);
-
-        let mut raw_proof = String::new();
-        out_file.read_to_string(&mut raw_proof)?;
-
-        let parsed_proof = cairo_proof_parser::parse(raw_proof)
-            .map_err(|e| ProverControllerError::ProofParseError(e.to_string()))?;
-
-        let config: VecFelt252 = serde_json::from_str(&parsed_proof.config.to_string())?;
-        let public_input: VecFelt252 =
-            serde_json::from_str(&parsed_proof.public_input.to_string())?;
-        let unsent_commitment: VecFelt252 =
-            serde_json::from_str(&parsed_proof.unsent_commitment.to_string())?;
-        let witness: VecFelt252 = serde_json::from_str(&parsed_proof.witness.to_string())?;
-
-        let proof = chain!(
-            config.into_iter(),
-            public_input.into_iter(),
-            unsent_commitment.into_iter(),
-            witness.into_iter()
-        )
-        .collect_vec();
-
-        Ok(JobWitness { proof })
-    }
-
-    fn terminate(&mut self, job_trace_hash: u64) -> Result<(), ProverControllerError> {
-        self.tasks
-            .get_mut(&job_trace_hash)
-            .ok_or(ProverControllerError::TaskNotFound)?
-            .start_kill()?;
-        trace!("task scheduled for termination {}", job_trace_hash);
-        Ok(())
-    }
-
-    fn drop(mut self) -> Result<(), ProverControllerError> {
-        let keys: Vec<u64> = self.tasks.keys().cloned().collect();
-        for job_trace_hash in keys.iter() {
-            self.tasks
-                .get_mut(job_trace_hash)
-                .ok_or(ProverControllerError::TaskNotFound)?
-                .start_kill()?;
-            trace!("task scheduled for termination {}", job_trace_hash);
-        }
-        Ok(())
+    type ProcessResult = Result<JobWitness, ProverControllerError>;
+    fn run(
+        &self,
+        job_trace: JobTrace,
+    ) -> Result<Process<Self::ProcessResult>, ProverControllerError> {
+        let (terminate_tx, mut terminate_rx) = mpsc::channel::<()>(10);
+        let future: Pin<Box<dyn Future<Output = Self::ProcessResult> + '_>> =
+            Box::pin(async move {
+                let mut out_file = NamedTempFile::new()?;
+
+                // let mut cpu_air_prover_config = NamedTempFile::new()?;
+                // let mut cpu_air_params = NamedTempFile::new()?;
+
+                // cpu_air_prover_config
+                //     .write_all(&serde_json::to_string(&self.config)?.into_bytes())?;
+                // cpu_air_params.write_all(&serde_json::to_string(&self.params)?.into_bytes())?;
+
+                let mut task = Command::new("cpu_air_prover")
+                    .arg("--out_file")
+                    .arg(out_file.path())
+                    .arg("--private_input_file")
+                    .arg(job_trace.air_private_input.path())
+                    .arg("--public_input_file")
+                    .arg(job_trace.air_public_input.path())
+                    .arg("--prover_config_file")
+                    .arg(self.config.path())
+                    .arg("--parameter_file")
+                    .arg(self.params.path())
+                    .arg("--generate_annotations")
+                    .stdout(Stdio::null())
+                    .spawn()?;
+
+                let job_trace_hash = hash!(job_trace);
+
+                debug!("task {} spawned", job_trace_hash);
+
+                loop {
+                    select! {
+                        output = task.wait() => {
+                            debug!("{:?}", output);
+                            if !output?.success() {
+                                return Err(ProverControllerError::TaskTerminated);
+                            }
+                            let output = task.wait_with_output().await?;
+                            debug!("{:?}", output);
+                            break;
+                        }
+                        Some(()) = terminate_rx.recv() => {
+                            task.start_kill()?;
+                        }
+                    }
+                }
+
+                let mut raw_proof = String::new();
+                out_file.read_to_string(&mut raw_proof)?;
+
+                let parsed_proof = cairo_proof_parser::parse(raw_proof)
+                    .map_err(|e| ProverControllerError::ProofParseError(e.to_string()))?;
+
+                let config: VecFelt252 = serde_json::from_str(&parsed_proof.config.to_string())?;
+                let public_input: VecFelt252 =
+                    serde_json::from_str(&parsed_proof.public_input.to_string())?;
+                let unsent_commitment: VecFelt252 =
+                    serde_json::from_str(&parsed_proof.unsent_commitment.to_string())?;
+                let witness: VecFelt252 = serde_json::from_str(&parsed_proof.witness.to_string())?;
+
+                let proof = chain!(
+                    config.into_iter(),
+                    public_input.into_iter(),
+                    unsent_commitment.into_iter(),
+                    witness.into_iter()
+                )
+                .collect_vec();
+
+                Ok(JobWitness { proof })
+            });
+
+        Ok(Process::new(future, terminate_tx))
     }
 }

crates/prover/src/stone_prover/tests/cairo/air_private_input.json

@@ -0,0 +1,7 @@
+{
+    "trace_path": "./src/stone_prover/tests/cairo/trace",
+    "memory_path": "./src/stone_prover/tests/cairo/memory",
+    "pedersen": [],
+    "range_check": [],
+    "bitwise": []
+}