From 14f56a6fe22d3dc6b70723576339d1f490253374 Mon Sep 17 00:00:00 2001
From: "keroroxx520@gmail.com" <keroroxx520@gmail.com>
Date: Mon, 13 Jan 2025 00:06:08 +0800
Subject: [PATCH] feat(lib): adjust IdStore

---
 lib/src/input.rs     | 2 +-
 lib/src/prover.rs    | 6 ++++--
 taskdb/src/lib.rs    | 2 +-
 taskdb/src/mem_db.rs | 2 +-
 4 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/lib/src/input.rs b/lib/src/input.rs
index 67bc5b11..c97b849e 100644
--- a/lib/src/input.rs
+++ b/lib/src/input.rs
@@ -174,7 +174,7 @@ impl TryFrom<Vec<TransactionSigned>> for TaikoGuestInput {
     }
 }
 
-#[derive(Clone, Debug, Serialize, Deserialize, Default)]
+#[derive(Clone, Debug, Serialize, Deserialize, Default, Eq, PartialEq, Ord, PartialOrd, Hash)]
 #[serde(rename_all = "snake_case")]
 pub enum BlobProofType {
     /// Guest runs through the entire computation from blob to Kzg commitment
diff --git a/lib/src/prover.rs b/lib/src/prover.rs
index 0b1bf349..043c5b3f 100644
--- a/lib/src/prover.rs
+++ b/lib/src/prover.rs
@@ -26,7 +26,9 @@ pub type ProverResult<T, E = ProverError> = core::result::Result<T, E>;
 pub type ProverConfig = serde_json::Value;
 pub type ProofKey = (ChainId, u64, B256, u8);
 
-#[derive(Clone, Debug, Serialize, ToSchema, Deserialize, Default, PartialEq, Eq, Hash)]
+#[derive(
+    Clone, Debug, Serialize, ToSchema, Deserialize, Default, PartialEq, Eq, PartialOrd, Ord, Hash,
+)]
 /// The response body of a proof request.
 pub struct Proof {
     /// The proof either TEE or ZK.
@@ -50,7 +52,7 @@ pub trait IdWrite: Send {
 
 #[async_trait::async_trait]
 pub trait IdStore: IdWrite {
-    async fn read_id(&self, key: ProofKey) -> ProverResult<String>;
+    async fn read_id(&mut self, key: ProofKey) -> ProverResult<String>;
 }
 
 #[allow(async_fn_in_trait)]
diff --git a/taskdb/src/lib.rs b/taskdb/src/lib.rs
index 513843a1..246cb179 100644
--- a/taskdb/src/lib.rs
+++ b/taskdb/src/lib.rs
@@ -321,7 +321,7 @@ impl<T: TaskManager> IdWrite for TaskManagerWrapper<T> {
 
 #[async_trait::async_trait]
 impl<T: TaskManager> IdStore for TaskManagerWrapper<T> {
-    async fn read_id(&self, key: ProofKey) -> ProverResult<String> {
+    async fn read_id(&mut self, key: ProofKey) -> ProverResult<String> {
         self.manager.read_id(key).await
     }
 }
diff --git a/taskdb/src/mem_db.rs b/taskdb/src/mem_db.rs
index 8ebab1b7..aa4bd76b 100644
--- a/taskdb/src/mem_db.rs
+++ b/taskdb/src/mem_db.rs
@@ -291,7 +291,7 @@ impl IdWrite for InMemoryTaskManager {
 
 #[async_trait::async_trait]
 impl IdStore for InMemoryTaskManager {
-    async fn read_id(&self, key: ProofKey) -> ProverResult<String> {
+    async fn read_id(&mut self, key: ProofKey) -> ProverResult<String> {
         let mut db = self.db.lock().await;
         db.read_id(key)
             .map_err(|e| ProverError::StoreError(e.to_string()))