feat(byte_array): add get(usize) and index(usize) to ByteSpan (#8427)

Co-authored-by: Gilad Chase <[email protected]>

Author: giladchase

corelib/src/byte_array.cairo

@@ -839,8 +839,7 @@ pub impl ByteSpanImpl of ByteSpanTrait {
     }
 
     /// Gets the element(s) at the given index.
-    /// Accepts ranges (returns Option<ByteSpan>), and (to-be-implemented) single indices (returns
-    /// Option<u8>).
+    /// Accepts ranges (returns Option<ByteSpan>), and single indices (returns Option<u8>).
     #[feature("corelib-get-trait")]
     fn get<I, impl TGet: crate::ops::Get<ByteSpan, I>, +Drop<I>>(
         self: @ByteSpan, index: I,
@@ -909,6 +908,16 @@ impl ByteSpanGetRangeInclusive of crate::ops::Get<ByteSpan, crate::ops::RangeInc
     }
 }
 
+impl ByteSpanGetUsize of crate::ops::Get<ByteSpan, usize> {
+    type Output = u8;
+
+    /// Returns the byte at the given index.
+    /// If out of bounds: returns `None`.
+    fn get(self: @ByteSpan, index: usize) -> Option<u8> {
+        helpers::byte_at(self, index)
+    }
+}
+
 #[feature("byte-span")]
 impl ByteSpanIndexViewRange of crate::ops::IndexView<ByteSpan, crate::ops::Range<usize>> {
     type Target = ByteSpan;
@@ -929,6 +938,14 @@ impl ByteSpanIndexViewRangeInclusive of crate::ops::IndexView<
     }
 }
 
+impl ByteSpanIndex of core::ops::index::IndexView<ByteSpan, usize> {
+    type Target = u8;
+
+    fn index(self: @ByteSpan, index: usize) -> u8 {
+        ByteSpanTrait::get(self, index).expect('Index out of bounds')
+    }
+}
+
 /// Trait for types that can be converted into a `ByteSpan`.
 #[unstable(feature: "byte-span")]
 pub trait ToByteSpanTrait<C> {
@@ -973,18 +990,21 @@ fn shift_right(word: felt252, word_len: usize, n_bytes: usize) -> felt252 {
 
 mod helpers {
     use core::num::traits::Bounded;
-    use crate::bytes_31::BYTES_IN_BYTES31;
+    use crate::bytes_31::{BYTES_IN_BYTES31, Bytes31Trait, u8_at_u256};
     #[feature("bounded-int-utils")]
     use crate::internal::bounded_int::{
-        self, AddHelper, BoundedInt, ConstrainHelper, MulHelper, SubHelper, UnitInt, downcast,
-        upcast,
+        self, AddHelper, BoundedInt, ConstrainHelper, DivRemHelper, MulHelper, SubHelper, UnitInt,
+        downcast, upcast,
     };
     use super::{BYTES_IN_BYTES31_MINUS_ONE, ByteSpan, Bytes31Index};
 
     type BytesInBytes31Typed = UnitInt<{ BYTES_IN_BYTES31.into() }>;
 
     const U32_MAX_TIMES_B31: felt252 = Bounded::<u32>::MAX.into() * BYTES_IN_BYTES31.into();
     const BYTES_IN_BYTES31_UNIT_INT: BytesInBytes31Typed = downcast(BYTES_IN_BYTES31).unwrap();
+    const NZ_BYTES_IN_BYTES31: NonZero<BytesInBytes31Typed> = 31;
+    const BYTES_IN_BYTES31_MINUS_ONE_TYPED: UnitInt<{ BYTES_IN_BYTES31_MINUS_ONE.into() }> = 30;
+    const ONE_TYPED: UnitInt<1> = 1;
 
     impl U32ByB31 of MulHelper<u32, BytesInBytes31Typed> {
         type Result = BoundedInt<0, U32_MAX_TIMES_B31>;
@@ -1002,6 +1022,50 @@ mod helpers {
             >;
     }
 
+    // For byte_at: usize + BoundedInt<0,30>
+    impl UsizeAddBytes31Index of AddHelper<usize, Bytes31Index> {
+        type Result =
+            BoundedInt<0, { Bounded::<usize>::MAX.into() + BYTES_IN_BYTES31_MINUS_ONE.into() }>;
+    }
+
+    // For byte_at: div_rem of (usize + BoundedInt<0,30>) by 31
+    const USIZE_PLUS_30_DIV_31: felt252 = (Bounded::<usize>::MAX / 31 + 1).into();
+    impl UsizePlusBytes31IndexDivRemB31 of DivRemHelper<
+        UsizeAddBytes31Index::Result, BytesInBytes31Typed,
+    > {
+        type DivT = BoundedInt<0, USIZE_PLUS_30_DIV_31>;
+        type RemT = Bytes31Index;
+    }
+
+    // For byte_at: 30 - BoundedInt<0,30>
+    impl B30SubBytes31Index of SubHelper<
+        UnitInt<{ BYTES_IN_BYTES31_MINUS_ONE.into() }>, Bytes31Index,
+    > {
+        type Result = Bytes31Index;
+    }
+
+    // For byte_at: BoundedInt<0,30> - 1
+    impl Bytes31IndexSub1 of SubHelper<Bytes31Index, UnitInt<1>> {
+        type Result = BoundedInt<-1, { BYTES_IN_BYTES31_MINUS_ONE.into() - 1 }>;
+    }
+
+    // For byte_at: (BoundedInt<0,30> - 1) - BoundedInt<0,30>
+    impl Bytes31IndexMinus1SubBytes31Index of SubHelper<Bytes31IndexSub1::Result, Bytes31Index> {
+        type Result =
+            BoundedInt<
+                { -BYTES_IN_BYTES31_MINUS_ONE.into() - 1 },
+                { BYTES_IN_BYTES31_MINUS_ONE.into() - 1 },
+            >;
+    }
+
+    // For byte_at: split BoundedInt<-31, 29> at 0.
+    impl ConstrainRemainderIndexAt0 of bounded_int::ConstrainHelper<
+        Bytes31IndexMinus1SubBytes31Index::Result, 0,
+    > {
+        type LowT = BoundedInt<{ -BYTES_IN_BYTES31_MINUS_ONE.into() - 1 }, -1>;
+        type HighT = BoundedInt<0, { BYTES_IN_BYTES31_MINUS_ONE.into() - 1 }>;
+    }
+
     /// Calculates the length of a `ByteSpan` in bytes.
     pub fn calc_bytespan_len(span: ByteSpan) -> usize {
         let data_bytes = bounded_int::mul(span.data.len(), BYTES_IN_BYTES31_UNIT_INT);
@@ -1091,5 +1155,39 @@ mod helpers {
     pub fn length_minus_one(len: BoundedInt<1, 31>) -> Bytes31Index {
         bounded_int::sub(len, 1)
     }
+    /// Returns the byte at the given index in the ByteSpan.
+    /// If out of bounds: returns `None`.
+    pub fn byte_at(self: @ByteSpan, index: usize) -> Option<u8> {
+        let absolute_index = bounded_int::add(index, *self.first_char_start_offset);
+        let (word_index_bounded, msb_index) = bounded_int::div_rem(
+            absolute_index, NZ_BYTES_IN_BYTES31,
+        );
+
+        let word_index = upcast(word_index_bounded);
+        match self.data.get(word_index) {
+            Some(word) => {
+                // Convert from MSB to LSB indexing.
+                let lsb_index = bounded_int::sub(BYTES_IN_BYTES31_MINUS_ONE_TYPED, msb_index);
+                Some(word.at(upcast(lsb_index)))
+            },
+            None => {
+                // Word index must equal data.len() for remainder word.
+                if word_index != self.data.len() {
+                    return None;
+                }
+
+                // Compute LSB index: remainder_len - 1 - msb_index.
+                let lsb_index_bounded = bounded_int::sub(
+                    bounded_int::sub(*self.remainder_len, ONE_TYPED), msb_index,
+                );
+
+                // Check if in bounds and extract non-negative index.
+                let Err(lsb_index) = bounded_int::constrain::<_, 0>(lsb_index_bounded) else {
+                    return None; // Out of bounds: index >= remainder_len.
+                };
+                Some(u8_at_u256((*self.remainder_word).into(), upcast(lsb_index)))
+            },
+        }
+    }
 }
 pub(crate) use helpers::len_parts;

corelib/src/ops/get.cairo

@@ -20,6 +20,10 @@
 /// let ba: ByteArray = "hello";
 /// let span = ba.span();
 ///
+/// // Using usize.
+/// let byte = span.get(1).unwrap();
+/// assert!(byte == 'e');
+///
 /// // Using Range<usize>.
 /// let slice = span.get(1..4).unwrap();
 /// assert_eq!(slice.to_byte_array(), "ell");
@@ -29,6 +33,7 @@
 /// assert_eq!(slice.to_byte_array(), "ell");
 ///
 /// // Out of bounds returns None.
+/// assert!(span.get(10).is_none());
 /// assert!(span.get(10..20).is_none());
 /// ```
 // TODO(giladchase): add examples for `usize` once supported.

corelib/src/test/byte_array_test.cairo

@@ -1,7 +1,7 @@
 #[feature("byte-span")]
-use crate::byte_array::{ByteSpan, ByteSpanTrait, ToByteSpanTrait};
-use crate::num::traits::Bounded;
-use crate::test::test_utils::{assert_eq, assert_ne};
+use core::byte_array::{ByteSpan, ByteSpanTrait, ToByteSpanTrait};
+use core::num::traits::Bounded;
+use core::test::test_utils::{assert_eq, assert_ne};
 
 #[test]
 fn test_append_byte() {
@@ -681,3 +681,60 @@ fn test_span_multiple_start_offset_slicing() {
     assert_eq!(slice2_inc.map(tba), Some("cdef"));
     assert_eq!(slice3_inc.map(tba), Some("def"));
 }
+
+#[test]
+fn test_span_at_and_index() {
+    // Test simple access.
+    let ba: ByteArray = "AB";
+    let span = ba.span();
+    assert_eq!(span[0_usize], 'A');
+    assert_eq!(span.get(1_usize), Some('B'));
+    assert_eq!(span.get(2_usize), None);
+
+    // Test with offset and two words.
+    let ba_33: ByteArray = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefg";
+    let mut span = ba_33.span();
+    span = span.get(1..33).unwrap();
+    assert_eq!(span.get(0_usize), Some('B'));
+    assert_eq!(span.get(30_usize), Some('f'));
+    assert_eq!(span[31_usize], 'g');
+    assert_eq!(span.get(32_usize), None);
+
+    // Test with offset and two words.
+    // 64 bytes: 31 + 31 + 2.
+    let ba_64: ByteArray = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789#$";
+    let mut span = ba_64.span();
+    span = span.get(1..64).unwrap();
+    assert_eq!(span.get(30_usize), Some('f'));
+    assert_eq!(span[31_usize], 'g');
+    assert_eq!(span.get(60_usize), Some('9'));
+    assert_eq!(span[61_usize], '#');
+    assert_eq!(span.get(62_usize), Some('$'));
+    assert_eq!(span.get(63_usize), None);
+
+    // Test empty span.
+    let empty: ByteArray = Default::default();
+    let empty_span = empty.span();
+    assert_eq!(empty_span.get(0_usize), None);
+}
+
+#[test]
+#[should_panic(expected: ('Index out of bounds',))]
+fn test_span_index_out_of_bounds() {
+    let ba: ByteArray = "AB";
+    let span = ba.span();
+    let _x = span[2_usize]; // Should panic
+}
+
+#[test]
+fn test_span_at_overflows() {
+    // Test overflow protection with large indices.
+    let ba: ByteArray = "test";
+    let span = ba.span();
+
+    assert_eq!(span.get(Bounded::<usize>::MAX), None);
+
+    let sliced = ba.span().get(1..3).unwrap();
+    assert_eq!(sliced.get(Bounded::<usize>::MAX - 1), None);
+    assert_eq!(sliced.get(Bounded::<usize>::MAX), None);
+}