From 557337afe1a699db909e7369b92f06fdf78cbd82 Mon Sep 17 00:00:00 2001
From: Gregory Edison <gregory.edison1993@gmail.com>
Date: Wed, 13 Dec 2023 13:55:19 +0100
Subject: [PATCH 1/3] implement Add and Sub for StorageKey and i128

---
 src/state.rs | 39 +++++++++++++++++++++++++++++++++++++++
 1 file changed, 39 insertions(+)

diff --git a/src/state.rs b/src/state.rs
index 133b5807..84489023 100644
--- a/src/state.rs
+++ b/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,44 @@ 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("storage key out of range"),
+        )
+    }
+}
+
+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("storage key out of range"),
+        )
+    }
+}
+
+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 {

From e8b5c34a3da1c71f3d0bd798c117a76747b6c869 Mon Sep 17 00:00:00 2001
From: Gregory Edison <gregory.edison1993@gmail.com>
Date: Wed, 13 Dec 2023 13:55:21 +0100
Subject: [PATCH 2/3] test

---
 src/state_test.rs | 35 +++++++++++++++++++++++++++++++++++
 1 file changed, 35 insertions(+)

diff --git a/src/state_test.rs b/src/state_test.rs
index 3e49b9de..3e6fb1c2 100644
--- a/src/state_test.rs
+++ b/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 = "storage key out of range")]
+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 = "storage key out of range")]
+fn offset_storage_key_add_lhs_err() {
+    let key = StorageKey::from(123u128);
+    let offset = -124;
+    let _: StorageKey = offset + key;
+}

From 77344ec16de90d626aa65520096085d9c7386e90 Mon Sep 17 00:00:00 2001
From: Gregory Edison <gregory.edison1993@gmail.com>
Date: Wed, 13 Dec 2023 14:00:02 +0100
Subject: [PATCH 3/3] renaming

---
 src/state.rs      | 6 ++++--
 src/state_test.rs | 4 ++--
 2 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/src/state.rs b/src/state.rs
index 84489023..5ef75b5d 100644
--- a/src/state.rs
+++ b/src/state.rs
@@ -158,7 +158,8 @@ impl std::ops::Add<i128> for StorageKey {
         let base_address = Into::<FieldElement>::into(*self.0.key()) + rhs;
 
         StorageKey(
-            PatriciaKey::try_from(StarkFelt::from(base_address)).expect("storage key out of range"),
+            PatriciaKey::try_from(StarkFelt::from(base_address))
+                .expect("attempt to add to storage key with overflow"),
         )
     }
 }
@@ -173,7 +174,8 @@ impl std::ops::Add<StorageKey> for i128 {
         let base_address = lhs + Into::<FieldElement>::into(*rhs.0.key());
 
         StorageKey(
-            PatriciaKey::try_from(StarkFelt::from(base_address)).expect("storage key out of range"),
+            PatriciaKey::try_from(StarkFelt::from(base_address))
+                .expect("attempt to add to storage key with overflow"),
         )
     }
 }
diff --git a/src/state_test.rs b/src/state_test.rs
index 3e6fb1c2..a6ba6ed3 100644
--- a/src/state_test.rs
+++ b/src/state_test.rs
@@ -34,7 +34,7 @@ fn offset_storage_key_add_rhs_ok() {
 }
 
 #[test]
-#[should_panic(expected = "storage key out of range")]
+#[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;
@@ -51,7 +51,7 @@ fn offset_storage_key_add_lhs_ok() {
 }
 
 #[test]
-#[should_panic(expected = "storage key out of range")]
+#[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;