From 300d28c081e5de0b7cf2773e463adff97cca3925 Mon Sep 17 00:00:00 2001
From: Gilad Chase <gilad@starkware.com>
Date: Tue, 16 Sep 2025 16:14:07 +0300
Subject: [PATCH] feat(byte_array): add `at` and `index` to `ByteSpan`

---
 corelib/src/byte_array.cairo           | 28 ++++++++++++
 corelib/src/test/byte_array_test.cairo | 61 ++++++++++++++++++++++++++
 2 files changed, 89 insertions(+)

diff --git a/corelib/src/byte_array.cairo b/corelib/src/byte_array.cairo
index 9fd795378b3..0ea59f4137d 100644
--- a/corelib/src/byte_array.cairo
+++ b/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)]
@@ -696,6 +697,27 @@ pub impl ByteSpanImpl of ByteSpanTrait {
         ba
     }
 
+    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 for the given range.
     fn slice(self: @ByteSpan, range: crate::ops::Range<usize>) -> Option<ByteSpan> {
         let len = range.end.checked_sub(range.start)?;
@@ -738,6 +760,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')
+    }
+}
+
 /// Trait for types that can be converted into a `ByteSpan`.
 pub trait ToByteSpanTrait<C> {
     #[must_use]
diff --git a/corelib/src/test/byte_array_test.cairo b/corelib/src/test/byte_array_test.cairo
index 6eeb8383086..82ab2c9526e 100644
--- a/corelib/src/test/byte_array_test.cairo
+++ b/corelib/src/test/byte_array_test.cairo
@@ -730,3 +730,64 @@ fn test_span_multiple_start_offset_slicing() {
     let result3: ByteArray = slice3.to_byte_array();
     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);
+}