refactor(katana-provider): move compiled_class_hash_of_class_hash f…

…n to `StateProvider` trait
dojoengine · Dec 4, 2023 · 7eae957 · 7eae957
1 parent 31d1f2b
commit 7eae957
Show file tree

Hide file tree

Showing 5 changed files with 61 additions and 61 deletions.
diff --git a/crates/katana/storage/provider/src/lib.rs b/crates/katana/storage/provider/src/lib.rs
@@ -150,6 +150,13 @@ where
     ) -> Result<Option<StorageValue>> {
         self.provider.storage(address, storage_key)
     }
+
+    fn compiled_class_hash_of_class_hash(
+        &self,
+        hash: ClassHash,
+    ) -> Result<Option<CompiledClassHash>> {
+        self.provider.compiled_class_hash_of_class_hash(hash)
+    }
 }
 
 impl<Db> StateProviderExt for BlockchainProvider<Db>
@@ -159,13 +166,6 @@ where
     fn sierra_class(&self, hash: ClassHash) -> Result<Option<SierraClass>> {
         self.provider.sierra_class(hash)
     }
-
-    fn compiled_class_hash_of_class_hash(
-        &self,
-        hash: ClassHash,
-    ) -> Result<Option<CompiledClassHash>> {
-        self.provider.compiled_class_hash_of_class_hash(hash)
-    }
 }
 
 impl<Db> StateFactoryProvider for BlockchainProvider<Db>

diff --git a/crates/katana/storage/provider/src/providers/fork/backend.rs b/crates/katana/storage/provider/src/providers/fork/backend.rs
@@ -393,6 +393,20 @@ impl StateProvider for SharedStateProvider {
 
         Ok(Some(class_hash))
     }
+
+    fn compiled_class_hash_of_class_hash(
+        &self,
+        hash: ClassHash,
+    ) -> Result<Option<CompiledClassHash>> {
+        if let hash @ Some(_) = self.0.compiled_class_hashes.read().get(&hash) {
+            return Ok(hash.cloned());
+        }
+
+        let compiled_hash = self.0.do_get_compiled_class_hash(hash).unwrap();
+        self.0.compiled_class_hashes.write().insert(hash, compiled_hash);
+
+        Ok(Some(hash))
+    }
 }
 
 impl StateProviderExt for SharedStateProvider {
@@ -415,20 +429,6 @@ impl StateProviderExt for SharedStateProvider {
             }
         }
     }
-
-    fn compiled_class_hash_of_class_hash(
-        &self,
-        hash: ClassHash,
-    ) -> Result<Option<CompiledClassHash>> {
-        if let hash @ Some(_) = self.0.compiled_class_hashes.read().get(&hash) {
-            return Ok(hash.cloned());
-        }
-
-        let compiled_hash = self.0.do_get_compiled_class_hash(hash).unwrap();
-        self.0.compiled_class_hashes.write().insert(hash, compiled_hash);
-
-        Ok(Some(hash))
-    }
 }
 
 #[cfg(test)]

diff --git a/crates/katana/storage/provider/src/providers/fork/state.rs b/crates/katana/storage/provider/src/providers/fork/state.rs
@@ -55,19 +55,19 @@ impl StateProvider for ForkedStateDb {
         }
         StateProvider::storage(&self.db, address, storage_key)
     }
-}
 
-impl StateProviderExt for CacheStateDb<SharedStateProvider> {
     fn compiled_class_hash_of_class_hash(
         &self,
         hash: ClassHash,
     ) -> Result<Option<CompiledClassHash>> {
         if let hash @ Some(_) = self.compiled_class_hashes.read().get(&hash) {
             return Ok(hash.cloned());
         }
-        StateProviderExt::compiled_class_hash_of_class_hash(&self.db, hash)
+        StateProvider::compiled_class_hash_of_class_hash(&self.db, hash)
     }
+}
 
+impl StateProviderExt for CacheStateDb<SharedStateProvider> {
     fn sierra_class(&self, hash: ClassHash) -> Result<Option<SierraClass>> {
         if let class @ Some(_) = self.shared_contract_classes.sierra_classes.read().get(&hash) {
             return Ok(class.cloned());
@@ -98,18 +98,18 @@ impl StateProvider for LatestStateProvider {
     fn class_hash_of_contract(&self, address: ContractAddress) -> Result<Option<ClassHash>> {
         StateProvider::class_hash_of_contract(&self.0, address)
     }
-}
-
-impl StateProviderExt for LatestStateProvider {
-    fn sierra_class(&self, hash: ClassHash) -> Result<Option<SierraClass>> {
-        StateProviderExt::sierra_class(&self.0, hash)
-    }
 
     fn compiled_class_hash_of_class_hash(
         &self,
         hash: ClassHash,
     ) -> Result<Option<CompiledClassHash>> {
-        StateProviderExt::compiled_class_hash_of_class_hash(&self.0, hash)
+        StateProvider::compiled_class_hash_of_class_hash(&self.0, hash)
+    }
+}
+
+impl StateProviderExt for LatestStateProvider {
+    fn sierra_class(&self, hash: ClassHash) -> Result<Option<SierraClass>> {
+        StateProviderExt::sierra_class(&self.0, hash)
     }
 }
 
@@ -147,15 +147,6 @@ impl StateProvider for ForkedSnapshot {
         }
         StateProvider::class_hash_of_contract(&self.inner.db, address)
     }
-}
-
-impl StateProviderExt for ForkedSnapshot {
-    fn sierra_class(&self, hash: ClassHash) -> Result<Option<SierraClass>> {
-        if let class @ Some(_) = self.classes.sierra_classes.read().get(&hash).cloned() {
-            return Ok(class);
-        }
-        StateProviderExt::sierra_class(&self.inner.db, hash)
-    }
 
     fn compiled_class_hash_of_class_hash(
         &self,
@@ -164,6 +155,15 @@ impl StateProviderExt for ForkedSnapshot {
         if let hash @ Some(_) = self.inner.compiled_class_hashes.get(&hash).cloned() {
             return Ok(hash);
         }
-        StateProviderExt::compiled_class_hash_of_class_hash(&self.inner.db, hash)
+        StateProvider::compiled_class_hash_of_class_hash(&self.inner.db, hash)
+    }
+}
+
+impl StateProviderExt for ForkedSnapshot {
+    fn sierra_class(&self, hash: ClassHash) -> Result<Option<SierraClass>> {
+        if let class @ Some(_) = self.classes.sierra_classes.read().get(&hash).cloned() {
+            return Ok(class);
+        }
+        StateProviderExt::sierra_class(&self.inner.db, hash)
     }
 }
diff --git a/crates/katana/storage/provider/src/providers/in_memory/state.rs b/crates/katana/storage/provider/src/providers/in_memory/state.rs
@@ -138,13 +138,6 @@ impl StateProvider for InMemorySnapshot {
         let class_hash = self.inner.contract_state.get(&address).map(|info| info.class_hash);
         Ok(class_hash)
     }
-}
-
-impl StateProviderExt for InMemorySnapshot {
-    fn sierra_class(&self, hash: ClassHash) -> Result<Option<SierraClass>> {
-        let class = self.classes.sierra_classes.read().get(&hash).cloned();
-        Ok(class)
-    }
 
     fn compiled_class_hash_of_class_hash(
         &self,
@@ -155,6 +148,13 @@ impl StateProviderExt for InMemorySnapshot {
     }
 }
 
+impl StateProviderExt for InMemorySnapshot {
+    fn sierra_class(&self, hash: ClassHash) -> Result<Option<SierraClass>> {
+        let class = self.classes.sierra_classes.read().get(&hash).cloned();
+        Ok(class)
+    }
+}
+
 pub(super) struct LatestStateProvider(pub(super) Arc<InMemoryStateDb>);
 
 impl StateProvider for LatestStateProvider {
@@ -181,13 +181,6 @@ impl StateProvider for LatestStateProvider {
         let class_hash = self.0.contract_state.read().get(&address).map(|info| info.class_hash);
         Ok(class_hash)
     }
-}
-
-impl StateProviderExt for LatestStateProvider {
-    fn sierra_class(&self, hash: ClassHash) -> Result<Option<SierraClass>> {
-        let class = self.0.shared_contract_classes.sierra_classes.read().get(&hash).cloned();
-        Ok(class)
-    }
 
     fn compiled_class_hash_of_class_hash(
         &self,
@@ -198,6 +191,13 @@ impl StateProviderExt for LatestStateProvider {
     }
 }
 
+impl StateProviderExt for LatestStateProvider {
+    fn sierra_class(&self, hash: ClassHash) -> Result<Option<SierraClass>> {
+        let class = self.0.shared_contract_classes.sierra_classes.read().get(&hash).cloned();
+        Ok(class)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use katana_primitives::block::BlockHashOrNumber;

diff --git a/crates/katana/storage/provider/src/traits/state.rs b/crates/katana/storage/provider/src/traits/state.rs
@@ -22,19 +22,19 @@ pub trait StateProvider: Send + Sync {
 
     /// Returns the class hash of a contract.
     fn class_hash_of_contract(&self, address: ContractAddress) -> Result<Option<ClassHash>>;
+
+    /// Returns the compiled class hash for the given class hash.
+    fn compiled_class_hash_of_class_hash(
+        &self,
+        hash: ClassHash,
+    ) -> Result<Option<CompiledClassHash>>;
 }
 
 /// An extension of the `StateProvider` trait which provides additional methods.
 #[auto_impl::auto_impl(&, Box, Arc)]
 pub trait StateProviderExt: StateProvider + Send + Sync {
     /// Retrieves the Sierra class definition of a contract class given its class hash.
     fn sierra_class(&self, hash: ClassHash) -> Result<Option<SierraClass>>;
-
-    /// Returns the compiled class hash for the given class hash.
-    fn compiled_class_hash_of_class_hash(
-        &self,
-        hash: ClassHash,
-    ) -> Result<Option<CompiledClassHash>>;
 }
 
 /// A state factory provider is a provider which can create state providers for