feat(network): add macro for impl TryFrom<Vec<u8>> and Into<Vec<u8>> …

…and use it in query
starkware-libs · May 27, 2024 · 080f17b · 080f17b
1 parent 76dec5d
commit 080f17b
Show file tree

Hide file tree

Showing 4 changed files with 107 additions and 20 deletions.
diff --git a/crates/papyrus_protobuf/src/converters/header.rs b/crates/papyrus_protobuf/src/converters/header.rs
@@ -2,6 +2,7 @@
 #[path = "header_test.rs"]
 mod header_test;
 
+use prost::Message;
 use starknet_api::block::{
     BlockHash,
     BlockHeader,
@@ -21,8 +22,8 @@ use starknet_api::crypto::Signature;
 
 use super::common::{enum_int_to_l1_data_availability_mode, l1_data_availability_mode_to_enum_int};
 use super::ProtobufConversionError;
-use crate::protobuf;
-use crate::sync::{Query, SignedBlockHeader};
+use crate::sync::{HeaderQuery, Query, SignedBlockHeader};
+use crate::{auto_impl_into_and_try_from_vec_u8, protobuf};
 
 impl TryFrom<protobuf::BlockHeadersResponse> for Option<SignedBlockHeader> {
     type Error = ProtobufConversionError;
@@ -83,7 +84,7 @@ impl TryFrom<protobuf::SignedBlockHeader> for SignedBlockHeader {
         let transaction_commitment = value
             .transactions
             .map(|transactions| {
-                Ok(TransactionCommitment(
+                Ok::<_, ProtobufConversionError>(TransactionCommitment(
                     transactions
                         .root
                         .ok_or(ProtobufConversionError::MissingField {
@@ -102,7 +103,7 @@ impl TryFrom<protobuf::SignedBlockHeader> for SignedBlockHeader {
         let event_commitment = value
             .events
             .map(|events| {
-                Ok(EventCommitment(
+                Ok::<_, ProtobufConversionError>(EventCommitment(
                     events
                         .root
                         .ok_or(ProtobufConversionError::MissingField {
@@ -286,20 +287,39 @@ impl From<Option<SignedBlockHeader>> for protobuf::BlockHeadersResponse {
     }
 }
 
+// TODO(shahak): Erase this once network stops using it.
 impl TryFrom<protobuf::BlockHeadersRequest> for Query {
     type Error = ProtobufConversionError;
     fn try_from(value: protobuf::BlockHeadersRequest) -> Result<Self, Self::Error> {
-        value
-            .iteration
-            .ok_or(ProtobufConversionError::MissingField {
-                field_description: "BlockHeadersRequest::iteration",
-            })?
-            .try_into()
+        Ok(HeaderQuery::try_from(value)?.0)
     }
 }
 
+impl TryFrom<protobuf::BlockHeadersRequest> for HeaderQuery {
+    type Error = ProtobufConversionError;
+    fn try_from(value: protobuf::BlockHeadersRequest) -> Result<Self, Self::Error> {
+        Ok(HeaderQuery(
+            value
+                .iteration
+                .ok_or(ProtobufConversionError::MissingField {
+                    field_description: "BlockHeadersRequest::iteration",
+                })?
+                .try_into()?,
+        ))
+    }
+}
+
+// TODO(shahak): Erase this once network stops using it.
 impl From<Query> for protobuf::BlockHeadersRequest {
     fn from(value: Query) -> Self {
         protobuf::BlockHeadersRequest { iteration: Some(value.into()) }
     }
 }
+
+impl From<HeaderQuery> for protobuf::BlockHeadersRequest {
+    fn from(value: HeaderQuery) -> Self {
+        protobuf::BlockHeadersRequest { iteration: Some(value.0.into()) }
+    }
+}
+
+auto_impl_into_and_try_from_vec_u8!(HeaderQuery, protobuf::BlockHeadersRequest);
diff --git a/crates/papyrus_protobuf/src/converters/mod.rs b/crates/papyrus_protobuf/src/converters/mod.rs
@@ -6,6 +6,8 @@ mod receipt;
 pub mod state_diff;
 mod transaction;
 
+use prost::DecodeError;
+
 #[derive(thiserror::Error, Debug)]
 pub enum ProtobufConversionError {
     #[error("Type `{type_description}` got out of range value {value_as_str}")]
@@ -14,4 +16,33 @@ pub enum ProtobufConversionError {
     MissingField { field_description: &'static str },
     #[error("Type `{type_description}` should be {num_expected} bytes but it got {value:?}.")]
     BytesDataLengthMismatch { type_description: &'static str, num_expected: usize, value: Vec<u8> },
+    #[error(transparent)]
+    DecodeError(#[from] DecodeError),
+}
+
+#[macro_export]
+macro_rules! auto_impl_into_and_try_from_vec_u8 {
+    ($T:ty, $ProtobufT:ty) => {
+        impl From<$T> for Vec<u8> {
+            fn from(value: $T) -> Self {
+                let protobuf_value = <$ProtobufT>::from(value);
+                protobuf_value.encode_to_vec()
+            }
+        }
+        crate::auto_impl_try_from_vec_u8!($T, $ProtobufT);
+    };
+}
+
+// TODO(shahak): Remove this macro once all types implement both directions.
+#[macro_export]
+macro_rules! auto_impl_try_from_vec_u8 {
+    ($T:ty, $ProtobufT:ty) => {
+        impl TryFrom<Vec<u8>> for $T {
+            type Error = ProtobufConversionError;
+            fn try_from(value: Vec<u8>) -> Result<Self, Self::Error> {
+                let protobuf_value = <$ProtobufT>::decode(&value[..])?;
+                <$T>::try_from(protobuf_value)
+            }
+        }
+    };
 }
diff --git a/crates/papyrus_protobuf/src/converters/state_diff.rs b/crates/papyrus_protobuf/src/converters/state_diff.rs
@@ -1,12 +1,13 @@
 use indexmap::IndexMap;
+use prost::Message;
 use starknet_api::core::{ClassHash, CompiledClassHash, Nonce};
 use starknet_api::data_availability::DataAvailabilityMode;
 use starknet_api::hash::StarkFelt;
 use starknet_api::state::{StorageKey, ThinStateDiff};
 
 use super::ProtobufConversionError;
-use crate::protobuf;
-use crate::sync::Query;
+use crate::sync::{Query, StateDiffQuery};
+use crate::{auto_impl_into_and_try_from_vec_u8, protobuf};
 
 pub const DOMAIN: DataAvailabilityMode = DataAvailabilityMode::L1;
 
@@ -40,7 +41,12 @@ impl TryFrom<protobuf::ContractDiff> for ThinStateDiff {
 
         let deployed_contracts = value
             .class_hash
-            .map(|hash| Ok(IndexMap::from_iter([(contract_address, ClassHash(hash.try_into()?))])))
+            .map(|hash| {
+                Ok::<_, ProtobufConversionError>(IndexMap::from_iter([(
+                    contract_address,
+                    ClassHash(hash.try_into()?),
+                )]))
+            })
             .transpose()?
             .unwrap_or_default();
 
@@ -57,7 +63,12 @@ impl TryFrom<protobuf::ContractDiff> for ThinStateDiff {
 
         let nonces = value
             .nonce
-            .map(|nonce| Ok(IndexMap::from_iter([(contract_address, Nonce(nonce.try_into()?))])))
+            .map(|nonce| {
+                Ok::<_, ProtobufConversionError>(IndexMap::from_iter([(
+                    contract_address,
+                    Nonce(nonce.try_into()?),
+                )]))
+            })
             .transpose()?
             .unwrap_or_default();
 
@@ -132,20 +143,39 @@ impl TryFrom<protobuf::ContractStoredValue> for (StorageKey, StarkFelt) {
     }
 }
 
+// TODO(shahak): Erase this once network stops using it.
 impl TryFrom<protobuf::StateDiffsRequest> for Query {
     type Error = ProtobufConversionError;
     fn try_from(value: protobuf::StateDiffsRequest) -> Result<Self, Self::Error> {
-        value
-            .iteration
-            .ok_or(ProtobufConversionError::MissingField {
-                field_description: "StateDiffsRequest::iteration",
-            })?
-            .try_into()
+        Ok(StateDiffQuery::try_from(value)?.0)
     }
 }
 
+impl TryFrom<protobuf::StateDiffsRequest> for StateDiffQuery {
+    type Error = ProtobufConversionError;
+    fn try_from(value: protobuf::StateDiffsRequest) -> Result<Self, Self::Error> {
+        Ok(StateDiffQuery(
+            value
+                .iteration
+                .ok_or(ProtobufConversionError::MissingField {
+                    field_description: "StateDiffsRequest::iteration",
+                })?
+                .try_into()?,
+        ))
+    }
+}
+
+// TODO(shahak): Erase this once network stops using it.
 impl From<Query> for protobuf::StateDiffsRequest {
     fn from(value: Query) -> Self {
         protobuf::StateDiffsRequest { iteration: Some(value.into()) }
     }
 }
+
+impl From<StateDiffQuery> for protobuf::StateDiffsRequest {
+    fn from(value: StateDiffQuery) -> Self {
+        protobuf::StateDiffsRequest { iteration: Some(value.0.into()) }
+    }
+}
+
+auto_impl_into_and_try_from_vec_u8!(StateDiffQuery, protobuf::StateDiffsRequest);
diff --git a/crates/papyrus_protobuf/src/sync.rs b/crates/papyrus_protobuf/src/sync.rs
@@ -28,6 +28,12 @@ impl Default for BlockHashOrNumber {
     }
 }
 
+#[derive(Default, Clone, Debug, PartialEq, Eq, Hash)]
+pub struct HeaderQuery(pub Query);
+
+#[derive(Default, Clone, Debug, PartialEq, Eq, Hash)]
+pub struct StateDiffQuery(pub Query);
+
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct SignedBlockHeader {
     pub block_header: BlockHeader,