From ce75d154b6c98254816e92fff2242b282ad09eee Mon Sep 17 00:00:00 2001
From: Ahmad Kaouk <ahmadkaouk.93@gmail.com>
Date: Thu, 3 Aug 2023 09:20:17 +0300
Subject: [PATCH] Add ProofSizeMeter

---
 primitives/evm/src/metric.rs | 107 +++++++++++++++++++++++++----------
 1 file changed, 76 insertions(+), 31 deletions(-)

diff --git a/primitives/evm/src/metric.rs b/primitives/evm/src/metric.rs
index 838142abdd..15f8b43e03 100644
--- a/primitives/evm/src/metric.rs
+++ b/primitives/evm/src/metric.rs
@@ -1,14 +1,21 @@
-use evm::{gasometer::GasCost, Opcode};
-use sp_runtime::{
-	traits::{CheckedAdd, Zero},
-	Saturating,
-};
+use evm::{gasometer::GasCost, ExitError, Opcode};
+use sp_runtime::{traits::CheckedAdd, Saturating};
 
 #[derive(Debug, PartialEq)]
 /// Metric error.
-enum MetricError {
+pub enum MetricError {
 	/// The metric usage exceeds the limit.
 	LimitExceeded,
+	/// Invalid Base Cost.
+	InvalidBaseCost,
+}
+
+impl Into<ExitError> for MetricError {
+	fn into(self) -> ExitError {
+		match self {
+			MetricError::LimitExceeded | MetricError::InvalidBaseCost => ExitError::OutOfGas,
+		}
+	}
 }
 
 /// A struct that keeps track of metric usage and limit.
@@ -19,21 +26,23 @@ pub struct Metric<T> {
 
 impl<T> Metric<T>
 where
-	T: Zero,
+	T: CheckedAdd + Saturating + PartialOrd + Copy,
 {
-	/// Creates a new `Metric` instance with the given limit.
-	pub fn new(limit: T) -> Self {
-		Self {
-			limit,
-			usage: Zero::zero(),
+	/// Creates a new `Metric` instance with the given base cost and limit.
+	///
+	/// # Errors
+	///
+	/// Returns `MetricError::InvalidBaseCost` if the base cost is greater than the limit.
+	pub fn new(base_cost: T, limit: T) -> Result<Self, MetricError> {
+		if base_cost > limit {
+			return Err(MetricError::InvalidBaseCost);
 		}
+		Ok(Self {
+			limit,
+			usage: base_cost,
+		})
 	}
-}
 
-impl<T> Metric<T>
-where
-	T: CheckedAdd + Saturating + PartialOrd + Copy,
-{
 	/// Records the cost of an operation and updates the usage.
 	///
 	/// # Errors
@@ -44,6 +53,7 @@ where
 			.usage
 			.checked_add(&cost)
 			.ok_or(MetricError::LimitExceeded)?;
+
 		if usage > self.limit {
 			return Err(MetricError::LimitExceeded);
 		}
@@ -57,16 +67,35 @@ where
 	}
 }
 
-/// A struct that keeps track of new storage usage and limit.
-pub struct StorageMetric(Metric<u64>);
+/// A struct that keeps track of the size of the proof.
+pub struct ProofSizeMeter(Metric<u64>);
+
+impl ProofSizeMeter {
+	/// Creates a new `ProofSizeMetric` instance with the given limit.
+	pub fn new(base_cost: u64, limit: u64) -> Result<Self, MetricError> {
+		Ok(Self(Metric::new(base_cost, limit)?))
+	}
 
-impl StorageMetric {
+	/// Records the size of the proof and updates the usage.
+	///
+	/// # Errors
+	///
+	/// Returns `MetricError::LimitExceeded` if the proof size exceeds the limit.
+	fn record_proof_size(&mut self, size: u64) -> Result<(), MetricError> {
+		self.0.record_cost(size)
+	}
+}
+
+/// A struct that keeps track of storage usage (newly created storage) and limit.
+pub struct StorageMeter(Metric<u64>);
+
+impl StorageMeter {
 	/// Creates a new `StorageMetric` instance with the given limit.
-	pub fn new(limit: u64) -> Self {
-		Self(Metric::new(limit))
+	pub fn new(limit: u64) -> Result<Self, MetricError> {
+		Ok(Self(Metric::new(0, limit)?))
 	}
 
-	/// Refunds the given amount of storage gas.
+	/// Refunds the given amount of storage.
 	fn refund(&mut self, amount: u64) {
 		self.0.refund(amount)
 	}
@@ -75,15 +104,25 @@ impl StorageMetric {
 	///
 	/// # Errors
 	///
-	/// Returns `MetricError::LimitExceeded` if the storage gas usage exceeds the storage gas limit.
+	/// Returns `MetricError::LimitExceeded` if the storage usage exceeds the storage limit.
 	fn record_dynamic_opcode_cost(
 		&mut self,
 		_opcode: Opcode,
 		gas_cost: GasCost,
 	) -> Result<(), MetricError> {
 		let cost = match gas_cost {
-			GasCost::Create => 0,
-			GasCost::Create2 { len } => len.try_into().map_err(|_| MetricError::LimitExceeded)?,
+			GasCost::Create => {
+				// TODO record cost for create
+				0
+			}
+			GasCost::Create2 { len } => {
+				// len in bytes ??
+				len.try_into().map_err(|_| MetricError::LimitExceeded)?
+			}
+			GasCost::SStore { .. } => {
+				// TODO record cost for sstore
+				0
+			}
 			_ => return Ok(()),
 		};
 		self.0.record_cost(cost)
@@ -96,25 +135,31 @@ mod tests {
 
 	#[test]
 	fn test_init() {
-		let metric = Metric::<u64>::new(100);
+		let metric = Metric::<u64>::new(0, 100).unwrap();
 		assert_eq!(metric.limit, 100);
 		assert_eq!(metric.usage, 0);
+
+		// base cost > limit
+		let metric = Metric::<u64>::new(100, 0).err();
+		assert_eq!(metric, Some(MetricError::InvalidBaseCost));
 	}
 
 	#[test]
-	fn test_try_consume() {
-		let mut metric = Metric::<u64>::new(100);
+	fn test_record_cost() {
+		let mut metric = Metric::<u64>::new(0, 100).unwrap();
 		assert_eq!(metric.record_cost(10), Ok(()));
 		assert_eq!(metric.usage, 10);
 		assert_eq!(metric.record_cost(90), Ok(()));
 		assert_eq!(metric.usage, 100);
+
+		// exceed limit
 		assert_eq!(metric.record_cost(1), Err(MetricError::LimitExceeded));
 		assert_eq!(metric.usage, 100);
 	}
 
 	#[test]
 	fn test_refund() {
-		let mut metric = Metric::<u64>::new(100);
+		let mut metric = Metric::<u64>::new(0, 100).unwrap();
 		assert_eq!(metric.record_cost(10), Ok(()));
 		assert_eq!(metric.usage, 10);
 		metric.refund(10);
@@ -127,7 +172,7 @@ mod tests {
 
 	#[test]
 	fn test_storage_metric() {
-		let mut metric = StorageMetric::new(100);
+		let mut metric = StorageMeter::new(100).unwrap();
 		assert_eq!(metric.0.usage, 0);
 		assert_eq!(metric.0.limit, 100);
 		assert_eq!(metric.0.record_cost(10), Ok(()));