feat: add and sub traits for StorageKey

src/state.rs

@@ -7,6 +7,7 @@ use std::fmt::Debug;
 
 use indexmap::IndexMap;
 use serde::{Deserialize, Serialize};
+use starknet_crypto::FieldElement;
 
 use crate::block::{BlockHash, BlockNumber};
 use crate::core::{
@@ -147,6 +148,46 @@ impl From<u128> for StorageKey {
 
 impl_from_through_intermediate!(u128, StorageKey, u8, u16, u32, u64);
 
+impl std::ops::Add<i128> for StorageKey {
+    type Output = Self;
+
+    fn add(self, rhs: i128) -> Self::Output {
+        let offset = FieldElement::from(rhs.unsigned_abs());
+        let rhs = if rhs > 0 { offset } else { -offset };
+
+        let base_address = Into::<FieldElement>::into(*self.0.key()) + rhs;
+
+        StorageKey(
+            PatriciaKey::try_from(StarkFelt::from(base_address))
+                .expect("attempt to add to storage key with overflow"),
+        )
+    }
+}
+
+impl std::ops::Add<StorageKey> for i128 {
+    type Output = StorageKey;
+
+    fn add(self, rhs: StorageKey) -> Self::Output {
+        let offset = FieldElement::from(self.unsigned_abs());
+        let lhs = if self > 0 { offset } else { -offset };
+
+        let base_address = lhs + Into::<FieldElement>::into(*rhs.0.key());
+
+        StorageKey(
+            PatriciaKey::try_from(StarkFelt::from(base_address))
+                .expect("attempt to add to storage key with overflow"),
+        )
+    }
+}
+
+impl std::ops::Sub<i128> for StorageKey {
+    type Output = Self;
+
+    fn sub(self, rhs: i128) -> Self::Output {
+        self + (-rhs)
+    }
+}
+
 /// A contract class.
 #[derive(Debug, Clone, Default, Eq, PartialEq, Deserialize, Serialize)]
 pub struct ContractClass {

src/state_test.rs

@@ -3,6 +3,7 @@ use std::collections::HashMap;
 use serde_json::json;
 
 use crate::deprecated_contract_class::EntryPointOffset;
+use crate::state::StorageKey;
 
 #[test]
 fn entry_point_offset_from_json_str() {
@@ -22,3 +23,37 @@ fn entry_point_offset_into_json_str() {
     let offset = EntryPointOffset(123);
     assert_eq!(json!(offset), json!(format!("{:#x}", offset.0)));
 }
+
+#[test]
+fn offset_storage_key_add_rhs_ok() {
+    let key = StorageKey::from(123u128);
+    let offset = -23;
+    let expected = StorageKey::from(100u128);
+    let result: StorageKey = key + offset;
+    assert_eq!(expected, result);
+}
+
+#[test]
+#[should_panic(expected = "attempt to add to storage key with overflow")]
+fn offset_storage_key_add_rhs_err() {
+    let key = StorageKey::from(123u128);
+    let offset = -124;
+    let _: StorageKey = key + offset;
+}
+
+#[test]
+fn offset_storage_key_add_lhs_ok() {
+    let key = StorageKey::from(123u128);
+    let offset = -23;
+    let expected = StorageKey::from(100u128);
+    let result: StorageKey = offset + key;
+    assert_eq!(expected, result);
+}
+
+#[test]
+#[should_panic(expected = "attempt to add to storage key with overflow")]
+fn offset_storage_key_add_lhs_err() {
+    let key = StorageKey::from(123u128);
+    let offset = -124;
+    let _: StorageKey = offset + key;
+}