Fix tests

Sovereign-Labs · Dec 20, 2023 · 99b1ad3 · 99b1ad3
1 parent 46e8ddc
commit 99b1ad3
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 10 deletions.
diff --git a/full-node/sov-stf-runner/src/prover_service/parallel/prover.rs b/full-node/sov-stf-runner/src/prover_service/parallel/prover.rs
@@ -69,14 +69,20 @@ where
         let prover_manager_clone = self.prover_manager.clone();
         let mut prover_manager = self.prover_manager.write().expect("Lock was poisoned");
 
-        let (prover_status, state_transition_data) = prover_manager
+        let prover_status = prover_manager
             .remove(&block_header_hash)
             .ok_or_else(|| anyhow::anyhow!("Missing witness for block: {:?}", block_header_hash))?;
 
         match prover_status {
             ProverStatus::WitnessSubmitted => {
                 let start_prover = prover_manager.inc_task_count_if_not_busy(self.num_threads);
 
+                let state_transition_data = prover_manager
+                    .get_witness(&block_header_hash)
+                    .ok_or_else(|| {
+                        anyhow::anyhow!("Missing witness for block: {:?}", block_header_hash)
+                    })?;
+
                 // Initiate a new proving job only if the prover is not busy.
                 if start_prover {
                     vm.add_hint(state_transition_data);

diff --git a/full-node/sov-stf-runner/src/prover_service/parallel/prover_manager.rs b/full-node/sov-stf-runner/src/prover_service/parallel/prover_manager.rs
@@ -29,6 +29,13 @@ impl<StateRoot, Witness, Da: DaSpec> ProverState<StateRoot, Witness, Da> {
         self.witness.remove(hash)
     }
 
+    fn get_witness(
+        &mut self,
+        hash: &Da::SlotHash,
+    ) -> Option<&StateTransitionData<StateRoot, Witness, Da>> {
+        self.witness.get(hash)
+    }
+
     fn set_to_proving(&mut self, hash: Da::SlotHash) -> Option<ProverStatus> {
         self.prover_status
             .insert(hash, ProverStatus::ProvingInProgress)
@@ -117,8 +124,8 @@ impl<StateRoot, Witness, Da: DaSpec> ProverManager<StateRoot, Witness, Da> {
     pub(crate) fn get_witness(
         &mut self,
         hash: &Da::SlotHash,
-    ) -> &StateTransitionData<StateRoot, Witness, Da> {
-        self.prover_state.witness.get(hash).unwrap()
+    ) -> Option<&StateTransitionData<StateRoot, Witness, Da>> {
+        self.prover_state.witness.get(hash)
     }
 
     pub(crate) fn submit_witness(
@@ -149,13 +156,8 @@ impl<StateRoot, Witness, Da: DaSpec> ProverManager<StateRoot, Witness, Da> {
     }
 
     // TODO change name
-    pub(crate) fn remove(
-        &mut self,
-        hash: &Da::SlotHash,
-    ) -> Option<(ProverStatus, StateTransitionData<StateRoot, Witness, Da>)> {
-        let status = self.prover_state.remove(hash)?;
-        let witness = self.prover_state.remove_witness(hash)?;
-        Some((status, witness))
+    pub(crate) fn remove(&mut self, hash: &Da::SlotHash) -> Option<ProverStatus> {
+        self.prover_state.remove(hash)
     }
 
     pub(crate) fn get_prover_status(&mut self, hash: &Da::SlotHash) -> Option<&ProverStatus> {