From 1e2f9017a28a18105a266731d3fe0f4012c44dfa Mon Sep 17 00:00:00 2001
From: Preston Evans <32944016+preston-evans98@users.noreply.github.com>
Date: Fri, 16 Feb 2024 09:34:59 -0800
Subject: [PATCH] fix: debug deserialization for SP1ProofWithIO. Test. (#249)

---
 core/src/io.rs  | 79 +++++++++++++++++++++++++++++++++++++++++--------
 core/src/lib.rs | 10 +++----
 2 files changed, 72 insertions(+), 17 deletions(-)

diff --git a/core/src/io.rs b/core/src/io.rs
index 2ed448927..4dc4473de 100644
--- a/core/src/io.rs
+++ b/core/src/io.rs
@@ -1,8 +1,6 @@
 use serde::{de::DeserializeOwned, Deserialize, Serialize};
 
-use crate::stark::StarkGenericConfig;
 use crate::utils::Buffer;
-use crate::Proof;
 
 /// Standard input for the prover.
 #[derive(Serialize, Deserialize)]
@@ -88,14 +86,71 @@ impl SP1Stdout {
     }
 }
 
-pub fn serialize_proof<S, SC: StarkGenericConfig + Serialize>(
-    proof: &Proof<SC>,
-    serializer: S,
-) -> Result<S::Ok, S::Error>
-where
-    S: serde::Serializer,
-{
-    let bytes = bincode::serialize(proof).unwrap();
-    let hex_bytes = hex::encode(bytes);
-    serializer.serialize_str(&hex_bytes)
+pub mod proof_serde {
+    use serde::{de::DeserializeOwned, Deserialize, Deserializer, Serialize};
+
+    use crate::stark::{Proof, StarkGenericConfig};
+
+    pub fn serialize<S, SC: StarkGenericConfig + Serialize>(
+        proof: &Proof<SC>,
+        serializer: S,
+    ) -> Result<S::Ok, S::Error>
+    where
+        S: serde::Serializer,
+    {
+        if serializer.is_human_readable() {
+            let bytes = bincode::serialize(proof).unwrap();
+            let hex_bytes = hex::encode(bytes);
+            serializer.serialize_str(&hex_bytes)
+        } else {
+            proof.serialize(serializer)
+        }
+    }
+
+    pub fn deserialize<'de, D, SC: StarkGenericConfig + DeserializeOwned>(
+        deserializer: D,
+    ) -> Result<Proof<SC>, D::Error>
+    where
+        D: Deserializer<'de>,
+    {
+        if deserializer.is_human_readable() {
+            let hex_bytes = String::deserialize(deserializer).unwrap();
+            let bytes = hex::decode(hex_bytes).unwrap();
+            let proof = bincode::deserialize(&bytes).map_err(serde::de::Error::custom)?;
+            Ok(proof)
+        } else {
+            Proof::<SC>::deserialize(deserializer)
+        }
+    }
+
+    #[cfg(test)]
+    mod tests {
+        use crate::{
+            utils::{tests::FIBONACCI_IO_ELF, BabyBearBlake3},
+            SP1ProofWithIO, SP1Prover, SP1Stdin, SP1Verifier,
+        };
+
+        /// Tests serialization with a human-readable encoding
+        #[test]
+        fn test_json_roundtrip() {
+            let mut stdin = SP1Stdin::new();
+            stdin.write(&3u32);
+            let proof = SP1Prover::prove(FIBONACCI_IO_ELF, stdin).unwrap();
+            let json = serde_json::to_string(&proof).unwrap();
+            let output = serde_json::from_str::<SP1ProofWithIO<BabyBearBlake3>>(&json).unwrap();
+            SP1Verifier::verify(FIBONACCI_IO_ELF, &output).unwrap();
+        }
+
+        /// Tests serialization with a binary encoding
+        #[test]
+        fn test_bincode_roundtrip() {
+            let mut stdin = SP1Stdin::new();
+            stdin.write(&3u32);
+            let proof = SP1Prover::prove(FIBONACCI_IO_ELF, stdin).unwrap();
+            let serialized = bincode::serialize(&proof).unwrap();
+            let output =
+                bincode::deserialize::<SP1ProofWithIO<BabyBearBlake3>>(&serialized).unwrap();
+            SP1Verifier::verify(FIBONACCI_IO_ELF, &output).unwrap();
+        }
+    }
 }
diff --git a/core/src/lib.rs b/core/src/lib.rs
index 509c0d803..786c7980c 100644
--- a/core/src/lib.rs
+++ b/core/src/lib.rs
@@ -50,8 +50,8 @@ pub struct SP1Verifier;
 
 /// A proof of a RISCV ELF execution with given inputs and outputs.
 #[derive(Serialize, Deserialize)]
-pub struct SP1ProofWithIO<SC: StarkGenericConfig + Serialize> {
-    #[serde(serialize_with = "serialize_proof")]
+pub struct SP1ProofWithIO<SC: StarkGenericConfig + Serialize + DeserializeOwned> {
+    #[serde(with = "proof_serde")]
     pub proof: Proof<SC>,
     pub stdin: SP1Stdin,
     pub stdout: SP1Stdout,
@@ -91,7 +91,7 @@ impl SP1Prover {
         config: SC,
     ) -> Result<SP1ProofWithIO<SC>>
     where
-        SC: StarkUtils + Send + Sync + Serialize + Clone,
+        SC: StarkUtils + Send + Sync + Serialize + DeserializeOwned + Clone,
         SC::Challenger: Clone,
         OpeningProof<SC>: Send + Sync,
         <SC::Pcs as Pcs<SC::Val, RowMajorMatrix<SC::Val>>>::Commitment: Send + Sync,
@@ -134,7 +134,7 @@ impl SP1Verifier {
         config: SC,
     ) -> Result<(), ProgramVerificationError>
     where
-        SC: StarkUtils + Send + Sync + Serialize,
+        SC: StarkUtils + Send + Sync + Serialize + DeserializeOwned,
         SC::Challenger: Clone,
         OpeningProof<SC>: Send + Sync,
         <SC::Pcs as Pcs<SC::Val, RowMajorMatrix<SC::Val>>>::Commitment: Send + Sync,
@@ -150,7 +150,7 @@ impl SP1Verifier {
     }
 }
 
-impl<SC: StarkGenericConfig + Serialize> SP1ProofWithIO<SC> {
+impl<SC: StarkGenericConfig + Serialize + DeserializeOwned> SP1ProofWithIO<SC> {
     /// Saves the proof as a JSON to the given path.
     pub fn save(&self, path: &str) -> Result<()> {
         let data = serde_json::to_string(self).unwrap();