add assertion for delegated snapshot packing

EkuboProtocol · Apr 29, 2024 · 6b3c009 · 6b3c009
1 parent 7115778
commit 6b3c009
Show file tree

Hide file tree

Showing 3 changed files with 72 additions and 10 deletions.
diff --git a/src/execution_state_test.cairo b/src/execution_state_test.cairo
@@ -1,5 +1,5 @@
 use governance::utils::u64_tuple_storage::{ThreeU64TupleStorePacking, TwoU64TupleStorePacking};
-use governance::utils::u64_tuple_storage_test::{assert_pack_unpack};
+[use governance::utils::u64_tuple_storage_test::{assert_pack_unpack};
 use starknet::storage_access::{StorePacking};
 
 #[test]

diff --git a/src/staker.cairo b/src/staker.cairo
@@ -67,19 +67,23 @@ pub mod Staker {
         pub delegated_cumulative: u256,
     }
 
-    const TWO_POW_64: u128 = 0x10000000000000000_u128;
+    const TWO_POW_64: u128 = 0x10000000000000000;
+    const TWO_POW_192: u256 = 0x1000000000000000000000000000000000000000000000000;
+    const TWO_POW_192_DIVISOR: NonZero<u256> = 0x1000000000000000000000000000000000000000000000000;
 
     pub(crate) impl DelegatedSnapshotStorePacking of StorePacking<DelegatedSnapshot, felt252> {
         fn pack(value: DelegatedSnapshot) -> felt252 {
+            assert(value.delegated_cumulative < TWO_POW_192, 'MAX_DELEGATED_CUMULATIVE');
             (value.delegated_cumulative
                 + u256 { high: value.timestamp.into() * TWO_POW_64, low: 0 })
                 .try_into()
                 .unwrap()
         }
+
         fn unpack(value: felt252) -> DelegatedSnapshot {
-            let (timestamp, delegated_cumulative) = DivRem::<
-                u256
-            >::div_rem(value.into(), u256 { low: 0, high: TWO_POW_64 }.into().try_into().unwrap());
+            let (timestamp, delegated_cumulative) = DivRem::div_rem(
+                value.into(), TWO_POW_192_DIVISOR
+            );
             DelegatedSnapshot { timestamp: timestamp.low.try_into().unwrap(), delegated_cumulative }
         }
     }

diff --git a/src/staker_test.cairo b/src/staker_test.cairo
@@ -10,6 +10,7 @@ use governance::staker::{
     Staker::{DelegatedSnapshotStorePacking, DelegatedSnapshot},
 };
 use governance::test::test_token::{TestToken, deploy as deploy_token};
+use governance::utils::u64_tuple_storage_test::{assert_pack_unpack};
 use starknet::testing::{set_contract_address, set_block_timestamp, pop_log};
 use starknet::{
     get_contract_address, syscalls::deploy_syscall, ClassHash, contract_address_const,
@@ -90,7 +91,7 @@ mod stake_withdraw {
 }
 
 #[test]
-fn test_governance_token_delegated_snapshot_store_pack() {
+fn test_staker_delegated_snapshot_store_pack() {
     assert_eq!(
         DelegatedSnapshotStorePacking::pack(
             DelegatedSnapshot { timestamp: 0, delegated_cumulative: 0 }
@@ -119,16 +120,16 @@ fn test_governance_token_delegated_snapshot_store_pack() {
     assert_eq!(
         DelegatedSnapshotStorePacking::pack(
             DelegatedSnapshot {
-                timestamp: 576460752303423488, // this timestamp equal to 2**59 is so large it's invalid
-                delegated_cumulative: 6277101735386680763835789423207666416102355444464034512895 // max u192
+                timestamp: 576460752303423488,
+                delegated_cumulative: 6277101735386680763835789423207666416102355444464034512895
             }
         ),
         3618502788666131113263695016908177884250476444008934042335404944711319814143
     );
 }
 
 #[test]
-fn test_governance_token_delegated_snapshot_store_unpack() {
+fn test_staker_delegated_snapshot_store_unpack() {
     assert_eq!(
         DelegatedSnapshotStorePacking::unpack(0),
         DelegatedSnapshot { timestamp: 0, delegated_cumulative: 0 }
@@ -152,7 +153,64 @@ fn test_governance_token_delegated_snapshot_store_unpack() {
         ),
         DelegatedSnapshot {
             timestamp: 576460752303423488,
-            delegated_cumulative: 6277101735386680763835789423207666416102355444464034512895
+            delegated_cumulative: 6277101735386680763835789423207666416102355444464034512895 // 2**192 - 1
+        }
+    );
+
+    assert_eq!(
+        DelegatedSnapshotStorePacking::unpack(
+            // max felt252
+            3618502788666131213697322783095070105623107215331596699973092056135872020480
+        ),
+        DelegatedSnapshot { timestamp: 576460752303423505, delegated_cumulative: 0 }
+    );
+}
+
+#[test]
+fn test_staker_delegated_snapshot_store_pack_unpack() {
+    assert_pack_unpack(DelegatedSnapshot { timestamp: 0, delegated_cumulative: 0 });
+    assert_pack_unpack(DelegatedSnapshot { timestamp: 0, delegated_cumulative: 1 });
+    assert_pack_unpack(DelegatedSnapshot { timestamp: 1, delegated_cumulative: 0 });
+    assert_pack_unpack(DelegatedSnapshot { timestamp: 1, delegated_cumulative: 1 });
+    assert_pack_unpack(
+        DelegatedSnapshot {
+            timestamp: 0,
+            delegated_cumulative: 0x1000000000000000000000000000000000000000000000000 - 1
+        }
+    );
+    assert_pack_unpack(
+        DelegatedSnapshot { timestamp: 576460752303423505, delegated_cumulative: 0 }
+    );
+    assert_pack_unpack(
+        DelegatedSnapshot {
+            timestamp: 576460752303423504,
+            delegated_cumulative: 0x1000000000000000000000000000000000000000000000000 - 1
+        }
+    );
+}
+
+#[test]
+#[should_panic(expected: ('Option::unwrap failed.',))]
+fn test_staker_delegated_snapshot_pack_max_timestamp_and_delegated() {
+    DelegatedSnapshotStorePacking::pack(
+        DelegatedSnapshot { timestamp: 576460752303423505, delegated_cumulative: 1 }
+    );
+}
+
+#[test]
+#[should_panic(expected: ('Option::unwrap failed.',))]
+fn test_staker_delegated_snapshot_pack_max_timestamp_plus_one() {
+    DelegatedSnapshotStorePacking::pack(
+        DelegatedSnapshot { timestamp: 576460752303423506, delegated_cumulative: 0 }
+    );
+}
+
+#[test]
+#[should_panic(expected: ('MAX_DELEGATED_CUMULATIVE',))]
+fn test_staker_delegated_snapshot_pack_max_delegated_cumulative() {
+    DelegatedSnapshotStorePacking::pack(
+        DelegatedSnapshot {
+            timestamp: 0, delegated_cumulative: 0x1000000000000000000000000000000000000000000000000
         }
     );
 }