test(target_chains/starknet): add byte array tests

pyth-network · Apr 29, 2024 · f212907 · f212907
1 parent ef92222
commit f212907
Show file tree

Hide file tree

Showing 2 changed files with 148 additions and 2 deletions.
diff --git a/target_chains/starknet/contracts/src/byte_array.cairo b/target_chains/starknet/contracts/src/byte_array.cairo
@@ -1,4 +1,6 @@
+use core::array::ArrayTrait;
 use core::fmt::{Debug, Formatter};
+use super::util::one_shift_left_bytes_u256;
 
 /// A byte array with storage format similar to `core::ByteArray`, but
 /// suitable for reading data from it.
@@ -35,8 +37,13 @@ pub impl ByteArrayImpl of ByteArrayTrait {
         if data.len() == 0 {
             assert!(num_last_bytes == 0);
         } else {
+            assert!(num_last_bytes > 0);
             assert!(num_last_bytes <= 31);
-        // TODO: check that unused bytes are zeroed.
+            let last: u256 = (*data.at(data.len() - 1)).into();
+            assert!(
+                last / one_shift_left_bytes_u256(num_last_bytes) == 0,
+                "ByteArrayImpl::new: last value is too large"
+            );
         }
         ByteArray { num_last_bytes, data }
     }
@@ -62,3 +69,100 @@ pub impl ByteArrayImpl of ByteArrayTrait {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::{ByteArray, ByteArrayImpl};
+    use pyth::util::array_felt252_to_bytes31;
+
+    #[test]
+    fn empty_byte_array() {
+        let mut array = ByteArrayImpl::new(array![], 0);
+        assert!(array.len() == 0);
+        assert!(array.pop_front() == Option::None);
+    }
+
+    #[test]
+    fn byte_array_3_zeros() {
+        let mut array = ByteArrayImpl::new(array_felt252_to_bytes31(array![0]), 3);
+        assert!(array.len() == 3);
+        assert!(array.pop_front() == Option::Some((0.try_into().unwrap(), 3)));
+        assert!(array.len() == 0);
+        assert!(array.pop_front() == Option::None);
+    }
+
+    #[test]
+    fn byte_array_3_bytes() {
+        let mut array = ByteArrayImpl::new(array_felt252_to_bytes31(array![0x010203]), 3);
+        assert!(array.len() == 3);
+        assert!(array.pop_front() == Option::Some((0x010203.try_into().unwrap(), 3)));
+        assert!(array.len() == 0);
+        assert!(array.pop_front() == Option::None);
+    }
+
+    #[test]
+    fn byte_array_single_full() {
+        let value_31_bytes = 0x0102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f;
+        let mut array = ByteArrayImpl::new(array_felt252_to_bytes31(array![value_31_bytes]), 31);
+        assert!(array.len() == 31);
+        assert!(array.pop_front() == Option::Some((value_31_bytes.try_into().unwrap(), 31)));
+        assert!(array.len() == 0);
+        assert!(array.pop_front() == Option::None);
+    }
+
+    #[test]
+    fn byte_array_two_full() {
+        let value_31_bytes = 0x0102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f;
+        let value2_31_bytes = 0x2122232425262728292a2b2c2d2e2f303132333435363738393a3b3c3d3e3f;
+        let mut array = ByteArrayImpl::new(
+            array_felt252_to_bytes31(array![value_31_bytes, value2_31_bytes]), 31
+        );
+        assert!(array.len() == 62);
+        assert!(array.pop_front() == Option::Some((value_31_bytes.try_into().unwrap(), 31)));
+        assert!(array.len() == 31);
+        assert!(array.pop_front() == Option::Some((value2_31_bytes.try_into().unwrap(), 31)));
+        assert!(array.len() == 0);
+        assert!(array.pop_front() == Option::None);
+    }
+
+    #[test]
+    fn byte_array_3_items() {
+        let value_31_bytes = 0x0102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f;
+        let value2_31_bytes = 0x2122232425262728292a2b2c2d2e2f303132333435363738393a3b3c3d3e3f;
+        let value3_5_bytes = 0x4142434445;
+        let mut array = ByteArrayImpl::new(
+            array_felt252_to_bytes31(array![value_31_bytes, value2_31_bytes, value3_5_bytes]), 5
+        );
+        assert!(array.len() == 67);
+        assert!(array.pop_front() == Option::Some((value_31_bytes.try_into().unwrap(), 31)));
+        assert!(array.len() == 36);
+        assert!(array.pop_front() == Option::Some((value2_31_bytes.try_into().unwrap(), 31)));
+        assert!(array.len() == 5);
+        assert!(array.pop_front() == Option::Some((value3_5_bytes.try_into().unwrap(), 5)));
+        assert!(array.pop_front() == Option::None);
+    }
+
+    #[test]
+    #[should_panic]
+    fn byte_array_empty_invalid() {
+        ByteArrayImpl::new(array![], 5);
+    }
+
+    #[test]
+    #[should_panic]
+    fn byte_array_last_too_large() {
+        ByteArrayImpl::new(array_felt252_to_bytes31(array![1, 2, 3]), 35);
+    }
+
+    #[test]
+    #[should_panic]
+    fn byte_array_last_zero_invalid() {
+        ByteArrayImpl::new(array_felt252_to_bytes31(array![1, 2, 0]), 0);
+    }
+
+    #[test]
+    #[should_panic]
+    fn byte_array_last_too_many_bytes() {
+        ByteArrayImpl::new(array_felt252_to_bytes31(array![1, 2, 0x010203]), 2);
+    }
+}
diff --git a/target_chains/starknet/contracts/src/util.cairo b/target_chains/starknet/contracts/src/util.cairo
@@ -8,7 +8,49 @@ pub const ONE_SHIFT_64: u128 = 0x10000000000000000;
 pub const UNEXPECTED_OVERFLOW: felt252 = 'unexpected overflow';
 pub const UNEXPECTED_ZERO: felt252 = 'unexpected zero';
 
-// Returns 1 << (8 * `n_bytes`) as u128, where `n_bytes` must be < BYTES_IN_U128.
+// Returns 1 << (8 * `n_bytes`) as u256.
+//
+// Panics if `n_bytes >= 32`.
+pub fn one_shift_left_bytes_u256(n_bytes: u8) -> u256 {
+    match n_bytes {
+        0 => 0x1,
+        1 => 0x100,
+        2 => 0x10000,
+        3 => 0x1000000,
+        4 => 0x100000000,
+        5 => 0x10000000000,
+        6 => 0x1000000000000,
+        7 => 0x100000000000000,
+        8 => 0x10000000000000000,
+        9 => 0x1000000000000000000,
+        10 => 0x100000000000000000000,
+        11 => 0x10000000000000000000000,
+        12 => 0x1000000000000000000000000,
+        13 => 0x100000000000000000000000000,
+        14 => 0x10000000000000000000000000000,
+        15 => 0x1000000000000000000000000000000,
+        16 => 0x100000000000000000000000000000000,
+        17 => 0x10000000000000000000000000000000000,
+        18 => 0x1000000000000000000000000000000000000,
+        19 => 0x100000000000000000000000000000000000000,
+        20 => 0x10000000000000000000000000000000000000000,
+        21 => 0x1000000000000000000000000000000000000000000,
+        22 => 0x100000000000000000000000000000000000000000000,
+        23 => 0x10000000000000000000000000000000000000000000000,
+        24 => 0x1000000000000000000000000000000000000000000000000,
+        25 => 0x100000000000000000000000000000000000000000000000000,
+        26 => 0x10000000000000000000000000000000000000000000000000000,
+        27 => 0x1000000000000000000000000000000000000000000000000000000,
+        28 => 0x100000000000000000000000000000000000000000000000000000000,
+        29 => 0x10000000000000000000000000000000000000000000000000000000000,
+        30 => 0x1000000000000000000000000000000000000000000000000000000000000,
+        31 => 0x100000000000000000000000000000000000000000000000000000000000000,
+        _ => core::panic_with_felt252('n_bytes too big'),
+    }
+}
+
+
+// Returns 1 << (8 * `n_bytes`) as u128.
 //
 // Panics if `n_bytes >= 16`.
 pub fn one_shift_left_bytes_u128(n_bytes: u8) -> u128 {