From 0467b892d9a4784a166f407a52ce8c24d5adcb5c Mon Sep 17 00:00:00 2001
From: Nimrod Weiss <nimrod@starkware.co>
Date: Sun, 18 Aug 2024 11:38:42 +0300
Subject: [PATCH] build(fee): define valid resources bounds enum

---
 .../src/execution/syscalls/hint_processor.rs  |  1 +
 .../src/stateless_transaction_validator.rs    |  1 +
 .../src/serialization/serializers.rs          |  1 +
 crates/starknet_api/src/transaction.rs        | 56 ++++++++++++++++++-
 4 files changed, 56 insertions(+), 3 deletions(-)

diff --git a/crates/blockifier/src/execution/syscalls/hint_processor.rs b/crates/blockifier/src/execution/syscalls/hint_processor.rs
index 1c1a9cb9c9b..df847f4c59a 100644
--- a/crates/blockifier/src/execution/syscalls/hint_processor.rs
+++ b/crates/blockifier/src/execution/syscalls/hint_processor.rs
@@ -472,6 +472,7 @@ impl<'a> SyscallHintProcessor<'a> {
                 let resource = match resource {
                     Resource::L1Gas => l1_gas,
                     Resource::L2Gas => l2_gas,
+                    Resource::L1DataGas => todo!(),
                 };
 
                 vec![
diff --git a/crates/gateway/src/stateless_transaction_validator.rs b/crates/gateway/src/stateless_transaction_validator.rs
index 266946ed4cb..393ce2693c7 100644
--- a/crates/gateway/src/stateless_transaction_validator.rs
+++ b/crates/gateway/src/stateless_transaction_validator.rs
@@ -181,6 +181,7 @@ fn validate_resource_is_non_zero(
     let resource_bounds = match resource {
         Resource::L1Gas => resource_bounds_mapping.l1_gas,
         Resource::L2Gas => resource_bounds_mapping.l2_gas,
+        Resource::L1DataGas => todo!(),
     };
     if resource_bounds.max_amount == 0 || resource_bounds.max_price_per_unit == 0 {
         return Err(StatelessTransactionValidatorError::ZeroResourceBounds {
diff --git a/crates/papyrus_storage/src/serialization/serializers.rs b/crates/papyrus_storage/src/serialization/serializers.rs
index 051389bdc3a..95923339d98 100644
--- a/crates/papyrus_storage/src/serialization/serializers.rs
+++ b/crates/papyrus_storage/src/serialization/serializers.rs
@@ -355,6 +355,7 @@ auto_storage_serde! {
     pub enum Resource {
         L1Gas = 0,
         L2Gas = 1,
+        L1DataGas = 2,
     }
     pub struct ResourceBounds {
         pub max_amount: u64,
diff --git a/crates/starknet_api/src/transaction.rs b/crates/starknet_api/src/transaction.rs
index 2c812ad9033..b14ce1ab04e 100644
--- a/crates/starknet_api/src/transaction.rs
+++ b/crates/starknet_api/src/transaction.rs
@@ -1,4 +1,4 @@
-use std::collections::{BTreeMap, HashSet};
+use std::collections::{BTreeMap, HashMap, HashSet};
 use std::fmt::Display;
 use std::sync::Arc;
 
@@ -867,6 +867,7 @@ pub enum Resource {
     L1Gas,
     #[serde(rename = "L2_GAS")]
     L2Gas,
+    L1DataGas,
 }
 
 /// Fee bounds for an execution resource.
@@ -884,6 +885,12 @@ pub struct ResourceBounds {
     pub max_price_per_unit: u128,
 }
 
+impl ResourceBounds {
+    pub fn is_zero(&self) -> bool {
+        self.max_amount == 0 && self.max_price_per_unit == 0
+    }
+}
+
 fn u64_to_hex<S>(value: &u64, serializer: S) -> Result<S::Ok, S::Error>
 where
     S: Serializer,
@@ -916,6 +923,7 @@ where
 
 /// A mapping from execution resources to their corresponding fee bounds..
 #[derive(Clone, Debug, Default, Deserialize, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize)]
+// TODO(Nimrod): Remove this struct definition.
 pub struct ResourceBoundsMapping(pub BTreeMap<Resource, ResourceBounds>);
 
 impl TryFrom<Vec<(Resource, ResourceBounds)>> for ResourceBoundsMapping {
@@ -924,11 +932,13 @@ impl TryFrom<Vec<(Resource, ResourceBounds)>> for ResourceBoundsMapping {
         resource_resource_bounds_pairs: Vec<(Resource, ResourceBounds)>,
     ) -> Result<Self, Self::Error> {
         let n_variants = Resource::iter().count();
+        let allowed_singed_variants = [n_variants, n_variants - 1];
         let unique_resources: HashSet<Resource> =
             HashSet::from_iter(resource_resource_bounds_pairs.iter().map(|(k, _)| *k));
-        if unique_resources.len() != n_variants
-            || resource_resource_bounds_pairs.len() != n_variants
+        if !allowed_singed_variants.contains(&unique_resources.len())
+            || !allowed_singed_variants.contains(&resource_resource_bounds_pairs.len())
         {
+            // TODO(Nimrod): Consider making this check more strict.
             Err(StarknetApiError::InvalidResourceMappingInitializer(format!(
                 "{:?}",
                 resource_resource_bounds_pairs
@@ -939,6 +949,46 @@ impl TryFrom<Vec<(Resource, ResourceBounds)>> for ResourceBoundsMapping {
     }
 }
 
+pub enum ValidResourceBounds {
+    L1Gas(ResourceBounds), // Pre 0.13.3. L2 bounds are signed but never used.
+    AllResources { l1_gas: ResourceBounds, l2_gas: ResourceBounds, l1_data_gas: ResourceBounds },
+}
+
+impl TryFrom<HashMap<Resource, ResourceBounds>> for ValidResourceBounds {
+    type Error = StarknetApiError;
+    fn try_from(
+        raw_resource_bounds: HashMap<Resource, ResourceBounds>,
+    ) -> Result<Self, Self::Error> {
+        if let (Some(l1_bounds), Some(l2_bounds)) =
+            (raw_resource_bounds.get(&Resource::L1Gas), raw_resource_bounds.get(&Resource::L2Gas))
+        {
+            match raw_resource_bounds.get(&Resource::L1DataGas) {
+                Some(data_bounds) => Ok(Self::AllResources {
+                    l1_gas: *l1_bounds,
+                    l1_data_gas: *data_bounds,
+                    l2_gas: *l2_bounds,
+                }),
+                None => {
+                    if l2_bounds.is_zero() {
+                        Ok(Self::L1Gas(*l1_bounds))
+                    } else {
+                        Ok(Self::AllResources {
+                            l1_gas: *l1_bounds,
+                            l2_gas: *l2_bounds,
+                            l1_data_gas: ResourceBounds::default(),
+                        })
+                    }
+                }
+            }
+        } else {
+            Err(StarknetApiError::InvalidResourceMappingInitializer(format!(
+                "{:?}",
+                raw_resource_bounds
+            )))
+        }
+    }
+}
+
 /// Paymaster-related data.
 #[derive(Clone, Debug, Default, Deserialize, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize)]
 pub struct PaymasterData(pub Vec<Felt>);