Fix: replace ProverInput generic for an enum (#57)

* replaced ProverInput generic for an enum * fix serialization * cargo-fmt
cartridge-gg · Sep 12, 2024 · 63509b4 · 63509b4
1 parent 20f712b
commit 63509b4
Show file tree

Hide file tree

Showing 12 changed files with 41 additions and 50 deletions.
diff --git a/bin/cairo-prove/src/prove.rs b/bin/cairo-prove/src/prove.rs
@@ -2,10 +2,7 @@ use crate::errors::ProveErrors;
 use crate::validate_input;
 use crate::Args;
 use crate::CairoVersion;
-use common::{
-    cairo0_prover_input::{Cairo0CompiledProgram, Cairo0ProverInput},
-    cairo_prover_input::{CairoCompiledProgram, CairoProverInput},
-};
+use common::prover_input::*;
 use prover_sdk::sdk::ProverSDK;
 use serde_json::Value;
 

diff --git a/common/src/lib.rs b/common/src/lib.rs
@@ -1,7 +1,3 @@
-pub mod cairo0_prover_input;
-pub mod cairo_prover_input;
 pub mod models;
+pub mod prover_input;
 pub mod requests;
-pub trait ProverInput {
-    fn serialize(self) -> serde_json::Value;
-}
diff --git a/common/src/cairo_prover_input.rs → common/src/prover_input/cairo.rs b/common/src/cairo_prover_input.rs → common/src/prover_input/cairo.rs
@@ -1,19 +1,13 @@
 use serde::{Deserialize, Serialize};
 use starknet_types_core::felt::Felt;
 
-use crate::ProverInput;
-
 #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
 pub struct CairoProverInput {
     pub program: CairoCompiledProgram,
     pub program_input: Vec<Felt>,
     pub layout: String,
 }
-impl ProverInput for CairoProverInput {
-    fn serialize(self) -> serde_json::Value {
-        serde_json::to_value(self).unwrap()
-    }
-}
+
 #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
 pub struct CairoCompiledProgram {
     //pub version: u64,

diff --git a/common/src/cairo0_prover_input.rs → common/src/prover_input/cairo0.rs b/common/src/cairo0_prover_input.rs → common/src/prover_input/cairo0.rs
@@ -1,20 +1,12 @@
 use serde::{Deserialize, Serialize};
 
-use crate::ProverInput;
-
 #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
 pub struct Cairo0ProverInput {
     pub program: Cairo0CompiledProgram,
     pub program_input: serde_json::Value,
     pub layout: String,
 }
 
-impl ProverInput for Cairo0ProverInput {
-    fn serialize(self) -> serde_json::Value {
-        serde_json::to_value(self).unwrap()
-    }
-}
-
 #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
 pub struct Cairo0CompiledProgram {
     pub attributes: Vec<String>,

diff --git a/common/src/prover_input/mod.rs b/common/src/prover_input/mod.rs
@@ -0,0 +1,20 @@
+mod cairo;
+mod cairo0;
+
+pub use cairo::{CairoCompiledProgram, CairoProverInput};
+pub use cairo0::{Cairo0CompiledProgram, Cairo0ProverInput};
+
+#[derive(Debug)]
+pub enum ProverInput {
+    Cairo0(Cairo0ProverInput),
+    Cairo(CairoProverInput),
+}
+
+impl ProverInput {
+    pub fn to_json_value(&self) -> serde_json::Value {
+        match self {
+            ProverInput::Cairo0(input) => serde_json::to_value(input).unwrap(),
+            ProverInput::Cairo(input) => serde_json::to_value(input).unwrap(),
+        }
+    }
+}
diff --git a/prover-sdk/src/lib.rs b/prover-sdk/src/lib.rs
@@ -3,5 +3,4 @@ pub mod errors;
 pub mod sdk;
 pub mod sdk_builder;
 
-pub use common::cairo0_prover_input::Cairo0ProverInput;
-pub use common::cairo_prover_input::CairoProverInput;
+pub use common::prover_input::ProverInput;
diff --git a/prover-sdk/src/sdk.rs b/prover-sdk/src/sdk.rs
@@ -1,5 +1,8 @@
 use crate::{access_key::ProverAccessKey, errors::SdkErrors, sdk_builder::ProverSDKBuilder};
-use common::{requests::AddKeyRequest, ProverInput};
+use common::{
+    prover_input::{Cairo0ProverInput, CairoProverInput, ProverInput},
+    requests::AddKeyRequest,
+};
 use ed25519_dalek::{ed25519::signature::SignerMut, VerifyingKey};
 use futures::StreamExt;
 use reqwest::{Client, Response};
@@ -39,28 +42,21 @@ impl ProverSDK {
             .build()
     }
 
-    pub async fn prove_cairo0<T>(&self, data: T) -> Result<u64, SdkErrors>
-    where
-        T: ProverInput + Send + 'static,
-    {
-        self.prove(data, self.prover_cairo0.clone()).await
+    pub async fn prove_cairo0(&self, data: Cairo0ProverInput) -> Result<u64, SdkErrors> {
+        self.prove(ProverInput::Cairo0(data), self.prover_cairo0.clone())
+            .await
     }
 
-    pub async fn prove_cairo<T>(&self, data: T) -> Result<u64, SdkErrors>
-    where
-        T: ProverInput + Send + 'static,
-    {
-        self.prove(data, self.prover_cairo.clone()).await
+    pub async fn prove_cairo(&self, data: CairoProverInput) -> Result<u64, SdkErrors> {
+        self.prove(ProverInput::Cairo(data), self.prover_cairo.clone())
+            .await
     }
 
-    async fn prove<T>(&self, data: T, url: Url) -> Result<u64, SdkErrors>
-    where
-        T: ProverInput + Send + 'static,
-    {
+    async fn prove(&self, data: ProverInput, url: Url) -> Result<u64, SdkErrors> {
         let response = self
             .client
             .post(url.clone())
-            .json(&data.serialize())
+            .json(&data.to_json_value())
             .send()
             .await?;
 

diff --git a/prover-sdk/tests/prove_test.rs b/prover-sdk/tests/prove_test.rs
@@ -1,7 +1,4 @@
-use common::{
-    cairo0_prover_input::{Cairo0CompiledProgram, Cairo0ProverInput},
-    cairo_prover_input::{CairoCompiledProgram, CairoProverInput},
-};
+use common::prover_input::*;
 use helpers::fetch_job;
 use prover_sdk::{access_key::ProverAccessKey, sdk::ProverSDK};
 use serde_json::Value;

diff --git a/prover-sdk/tests/verify_test.rs b/prover-sdk/tests/verify_test.rs
@@ -1,4 +1,4 @@
-use common::cairo_prover_input::{CairoCompiledProgram, CairoProverInput};
+use common::prover_input::*;
 use helpers::fetch_job;
 use prover_sdk::{access_key::ProverAccessKey, sdk::ProverSDK};
 use starknet_types_core::felt::Felt;

diff --git a/prover/src/prove/cairo.rs b/prover/src/prove/cairo.rs
@@ -4,7 +4,7 @@ use crate::server::AppState;
 use crate::threadpool::CairoVersionedInput;
 use axum::Json;
 use axum::{extract::State, http::StatusCode, response::IntoResponse};
-use common::cairo_prover_input::CairoProverInput;
+use common::prover_input::CairoProverInput;
 use serde_json::json;
 
 pub async fn root(

diff --git a/prover/src/prove/cairo0.rs b/prover/src/prove/cairo0.rs
@@ -4,7 +4,7 @@ use crate::server::AppState;
 use crate::threadpool::CairoVersionedInput;
 use axum::Json;
 use axum::{extract::State, http::StatusCode, response::IntoResponse};
-use common::cairo0_prover_input::Cairo0ProverInput;
+use common::prover_input::Cairo0ProverInput;
 use serde_json::json;
 
 pub async fn root(

diff --git a/prover/src/threadpool/run.rs b/prover/src/threadpool/run.rs
@@ -1,6 +1,6 @@
 use std::{fs, path::PathBuf};
 
-use common::{cairo0_prover_input::Cairo0ProverInput, cairo_prover_input::CairoProverInput};
+use common::prover_input::{Cairo0ProverInput, CairoProverInput};
 use starknet_types_core::felt::Felt;
 use tokio::process::Command;
 use tracing::trace;