feat(byte_array): add at and index to ByteSpan

Author: Gilad Chase

corelib/src/byte_array.cairo

@@ -42,6 +42,7 @@
 //! assert!(first_byte == 0x41);
 //! ```
 
+use crate::IndexView;
 use crate::array::{ArrayTrait, Span, SpanTrait};
 use crate::bytes_31::split_bytes31;
 #[allow(unused_imports)]
@@ -653,6 +654,27 @@ pub impl ByteArraySpanImpl of ByteSpanTrait {
         self.data.len() == 0 && *self.remainder_len == 0
     }
 
+    fn at(self: @ByteSpan, index: usize) -> Option<u8> {
+        let actual_index = index.checked_add(upcast(self.first_char_start_offset))?;
+        let (word_index, index_in_word) = DivRem::div_rem(actual_index, BYTES_IN_BYTES31_NONZERO);
+
+        match self.data.get(word_index) {
+            Some(word) => {
+                // index_in_word is from MSB, we need index from LSB.
+                Some(word.at(BYTES_IN_BYTES31 - 1 - index_in_word))
+            },
+            None => {
+                if word_index == self.data.len() && index_in_word < upcast(self.remainder_len) {
+                    // index_in_word is from MSB, we need index from LSB.
+                    let index_in_remainder = upcast(self.remainder_len) - 1 - index_in_word;
+                    Some(u8_at_u256((*self.remainder_word).into(), index_in_remainder))
+                } else {
+                    None // Out of bounds.
+                }
+            },
+        }
+    }
+
     /// Returns a slice of the ByteSpan from the given start position with the given length.
     fn slice(self: @ByteSpan, start: usize, len: usize) -> Option<ByteSpan> {
         if len == 0 {
@@ -694,6 +716,12 @@ impl ByteSpanDefault of Default<ByteSpan> {
     }
 }
 
+impl ByteSpanIndex of IndexView<ByteSpan, usize, u8> {
+    fn index(self: @ByteSpan, index: usize) -> u8 {
+        self.at(index).expect('Index out of bounds')
+    }
+}
+
 pub trait ToByteSpanTrait<C> {
     #[must_use]
     fn span(self: @C) -> ByteSpan;

corelib/src/test/byte_array_test.cairo

@@ -726,3 +726,63 @@ fn test_span_multiple_start_offset_slicing() {
     assert_eq!(result3, "def", "third slice");
 }
 
+#[test]
+fn test_span_at_and_index() {
+    // Test simple access.
+    let ba: ByteArray = "AB";
+    let span = ba.span();
+    assert_eq!(span.at(0), Some('A'));
+    assert_eq!(span.at(1), Some('B'));
+    assert_eq!(span.at(2), None);
+
+    // Test index operator on same span.
+    assert_eq!(span[0], 'A');
+    assert_eq!(span[1], 'B');
+
+    // Test with offset and two words.
+    let ba_33: ByteArray = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefg";
+    let mut span = ba_33.span();
+    span = span.slice(1, 32).unwrap();
+    assert_eq!(span.at(0), Option::Some('B'));
+    assert_eq!(span.at(30), Option::Some('f'));
+    assert_eq!(span.at(31), Option::Some('g'));
+    assert_eq!(span.at(32), Option::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.slice(1, 63).unwrap();
+    assert_eq!(span.at(30), Some('f'), "byte 30 - last of 1nd word");
+    assert_eq!(span.at(31), Some('g'), "byte 31 - first of 2nd word");
+    assert_eq!(span.at(60), Some('9'), "byte 60 - last of 2nd word");
+    assert_eq!(span.at(61), Some('#'), "byte 61 - first in last_word");
+    assert_eq!(span.at(62), Some('$'), "byte 62 - last in last_word");
+    assert_eq!(span.at(63), None);
+
+    // Test empty span.
+    let empty: ByteArray = Default::default();
+    let empty_span = empty.span();
+    assert_eq!(empty_span.at(0), 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]; // 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.at(Bounded::<usize>::MAX), None);
+
+    let sliced = ba.span().slice(1, 3).unwrap();
+    assert_eq!(sliced.at(Bounded::<usize>::MAX - 1), None);
+    assert_eq!(sliced.at(Bounded::<usize>::MAX), None);
+}