From 43a15766902aa14a81bbf0bbe189d1326466e837 Mon Sep 17 00:00:00 2001
From: Eagle941 <8973725+Eagle941@users.noreply.github.com>
Date: Mon, 12 Aug 2024 23:48:44 +0100
Subject: [PATCH] Changed  to match previous format

---
 .../src/blockifier/transaction_executor.rs    |  6 ++---
 crates/blockifier/src/state/cached_state.rs   | 22 ++++++++++++++++++-
 2 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/crates/blockifier/src/blockifier/transaction_executor.rs b/crates/blockifier/src/blockifier/transaction_executor.rs
index 1bf6b62c9a..204d139934 100644
--- a/crates/blockifier/src/blockifier/transaction_executor.rs
+++ b/crates/blockifier/src/blockifier/transaction_executor.rs
@@ -13,9 +13,9 @@ use crate::bouncer::{Bouncer, BouncerWeights};
 #[cfg(feature = "concurrency")]
 use crate::concurrency::worker_logic::WorkerExecutor;
 use crate::context::BlockContext;
-use crate::state::cached_state::{
-    CachedState, CommitmentStateDiff, TransactionalState, VisitedPcs,
-};
+#[cfg(feature = "concurrency")]
+use crate::state::cached_state::VisitedPcs;
+use crate::state::cached_state::{CachedState, CommitmentStateDiff, TransactionalState};
 use crate::state::errors::StateError;
 use crate::state::state_api::StateReader;
 use crate::transaction::errors::TransactionExecutionError;
diff --git a/crates/blockifier/src/state/cached_state.rs b/crates/blockifier/src/state/cached_state.rs
index e68ff1038c..fcfdefef2f 100644
--- a/crates/blockifier/src/state/cached_state.rs
+++ b/crates/blockifier/src/state/cached_state.rs
@@ -105,8 +105,28 @@ impl<S: StateReader> CachedState<S> {
     }
 
     pub fn update_visited_pcs_cache(&mut self, visited_pcs: &VisitedPcs) {
+        #[cfg(not(feature = "vector_of_visited_program_counters"))]
+        fn from_set<S: StateReader>(
+            cached_state: &mut CachedState<S>,
+            class_hash: &ClassHash,
+            visited_pcs: &Pcs,
+        ) {
+            cached_state.add_visited_pcs(*class_hash, visited_pcs);
+        }
+
+        #[cfg(feature = "vector_of_visited_program_counters")]
+        fn from_set<S: StateReader>(
+            cached_state: &mut CachedState<S>,
+            class_hash: &ClassHash,
+            visited_pcs: &Vec<Vec<usize>>,
+        ) {
+            for pcs in visited_pcs {
+                cached_state.add_visited_pcs(*class_hash, pcs);
+            }
+        }
+
         for (class_hash, class_visited_pcs) in visited_pcs {
-            self.visited_pcs.entry(*class_hash).or_default().extend(class_visited_pcs.clone());
+            from_set(self, class_hash, class_visited_pcs);
         }
     }