feat(byte_array): add ByteSpan iterator

Also extract msb->lsb indexing into helper.

Author: Gilad Chase

corelib/src/byte_array.cairo

@@ -42,12 +42,11 @@
 //! assert!(first_byte == 0x41);
 //! ```
 
-use crate::array::{ArrayTrait, Span, SpanTrait};
-use crate::bytes_31::split_bytes31;
+use crate::array::{ArrayTrait, Span, SpanIter, SpanTrait};
 #[allow(unused_imports)]
 use crate::bytes_31::{
     BYTES_IN_BYTES31, Bytes31Trait, POW_2_128, POW_2_8, U128IntoBytes31, U8IntoBytes31,
-    one_shift_left_bytes_u128, split_u128, u8_at_u256,
+    one_shift_left_bytes_u128, split_bytes31, split_u128, u8_at_u256,
 };
 use crate::clone::Clone;
 use crate::cmp::min;
@@ -238,8 +237,8 @@ pub impl ByteArrayImpl of ByteArrayTrait {
             }
             // index_in_word is from MSB, we need index from LSB.
             return Some(
-                u8_at_u256(
-                    (*self.pending_word).into(), upcast(*self.pending_word_len) - 1 - index_in_word,
+                helpers::byte_at_msb(
+                    *self.pending_word, upcast(*self.pending_word_len), index_in_word,
                 ),
             );
         }
@@ -707,7 +706,7 @@ pub(crate) impl ByteArrayIndexView of crate::traits::IndexView<ByteArray, usize,
     }
 }
 
-// TODO: Implement a more efficient version of this iterator.
+// TODO(giladchase): Delegate to byte span iterator instead of current at-based implementation.
 /// An iterator struct over a ByteArray.
 #[derive(Drop, Clone)]
 pub struct ByteArrayIter {
@@ -957,6 +956,106 @@ impl ByteSpanToByteSpan of ToByteSpanTrait<ByteSpan> {
     }
 }
 
+/// An iterator struct over a ByteSpan.
+#[derive(Drop, Clone)]
+pub struct ByteSpanIter {
+    data_iter: SpanIter<bytes31>,
+    current_word: felt252,
+    current_word_len: usize,
+    byte_in_word_index: usize,
+    remainder_word: felt252,
+    remainder_len: usize,
+}
+
+impl ByteSpanIterDefault of Default<ByteSpanIter> {
+    fn default() -> ByteSpanIter {
+        ByteSpanIter {
+            data_iter: [].span().into_iter(),
+            current_word: Default::default(),
+            current_word_len: Default::default(),
+            byte_in_word_index: Default::default(),
+            remainder_word: Default::default(),
+            remainder_len: Default::default(),
+        }
+    }
+}
+
+
+#[generate_trait]
+impl ByteSpanIterImpl of ByteSpanIterTrait {
+    /// Loads a new word as the the current word and returns its first byte if it's not empty,
+    /// otherwise returns None.
+    #[inline]
+    fn load_word_next(ref self: ByteSpanIter, word: felt252, len: usize) -> Option<u8> {
+        if len == 0 {
+            return None;
+        }
+
+        self.current_word = word;
+        self.current_word_len = len;
+        self.byte_in_word_index = 1;
+        Some(helpers::byte_at_msb(word, len, 0))
+    }
+}
+
+impl ByteSpanIterator of crate::iter::Iterator<ByteSpanIter> {
+    type Item = u8;
+
+    #[inline]
+    fn next(ref self: ByteSpanIter) -> Option<u8> {
+        // Return the next byte from the current word if it has bytes left.
+        if self.byte_in_word_index < self.current_word_len {
+            let byte = helpers::byte_at_msb(
+                self.current_word, self.current_word_len, self.byte_in_word_index,
+            );
+            self.byte_in_word_index += 1;
+            return Some(byte);
+        }
+
+        // Current word exhausted, try advancing to the next word from the data iterator.
+        if let Some(word) = self.data_iter.next() {
+            self.load_word_next((*word).into(), BYTES_IN_BYTES31)
+        } else {
+            // No more data words, advance into the remainder word if it exists.
+            let len = self.remainder_len;
+            self.remainder_len = 0; // Mark remainder as consumed.
+            self.load_word_next(self.remainder_word, len)
+        }
+    }
+}
+
+impl ByteSpanIntoIterator of crate::iter::IntoIterator<ByteSpan> {
+    type IntoIter = ByteSpanIter;
+
+    /// Creates an iterator over the bytes in the `ByteSpan`.
+    #[inline]
+    fn into_iter(self: ByteSpan) -> Self::IntoIter {
+        let mut data_iter = self.data.into_iter();
+
+        let Some(first_word) = data_iter.next() else {
+            // No data words, start iterating from the remainder word.
+            return ByteSpanIter {
+                data_iter,
+                current_word: self.remainder_word,
+                current_word_len: upcast(self.remainder_len),
+                byte_in_word_index: upcast(self.first_char_start_offset),
+                remainder_word: self.remainder_word,
+                remainder_len: 0 // Mark as consumed since we're starting with it.
+            };
+        };
+
+        ByteSpanIter {
+            data_iter,
+            current_word: (*first_word).into(),
+            current_word_len: BYTES_IN_BYTES31,
+            byte_in_word_index: upcast(self.first_char_start_offset),
+            remainder_word: self.remainder_word,
+            remainder_len: upcast(self.remainder_len),
+        }
+    }
+}
+
+
 /// Shifts a word right by `n_bytes`.
 /// The input `bytes31` and the output `bytes31`s are represented using `felt252`s to improve
 /// performance.
@@ -1173,5 +1272,14 @@ mod helpers {
             },
         }
     }
+
+    /// Extracts a byte from a word at the given MSB index.
+    /// Converts from MSB indexing (0 = first/leftmost byte) to LSB indexing.
+    #[inline]
+    pub fn byte_at_msb(word: felt252, word_len: usize, msb_index: usize) -> u8 {
+        // Convert from MSB to LSB indexing: index 0 MSB = index (len-1) LSB.
+        let lsb_index = word_len - 1 - msb_index;
+        u8_at_u256(word.into(), lsb_index)
+    }
 }
 pub(crate) use helpers::len_parts;

corelib/src/test/byte_array_test.cairo

@@ -1,7 +1,7 @@
 #[feature("byte-span")]
-use crate::byte_array::{ByteSpanTrait, ToByteSpanTrait};
-use crate::num::traits::Bounded;
-use crate::test::test_utils::{assert_eq, assert_ne};
+use core::byte_array::{ByteSpanTrait, ToByteSpanTrait};
+use core::num::traits::Bounded;
+use core::test::test_utils::{assert_eq, assert_ne};
 
 #[test]
 fn test_append_byte() {
@@ -843,10 +843,221 @@ 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);
 }
+
+#[test]
+fn test_byte_span_simple() {
+    let empty: ByteArray = "";
+    let mut iter = empty.span().into_iter();
+    assert_eq!(iter.next(), None);
+
+    let ba: ByteArray = "A";
+    let mut iter = ba.span().into_iter();
+    assert_eq!(iter.next(), Some('A'));
+    assert_eq!(iter.next(), None);
+
+    let ba: ByteArray = "ABC";
+    let mut iter = ba.span().into_iter();
+    assert_eq!(iter.next(), Some('A'));
+    assert_eq!(iter.next(), Some('B'));
+    assert_eq!(iter.next(), Some('C'));
+    assert_eq!(iter.next(), None);
+    assert_eq!(iter.next(), None, "Idempotent empty");
+}
+
+#[test]
+fn test_byte_span_iterator_word_boundaries() {
+    // Test 30, 31, 32 bytes (1 word boundary).
+    let ba_30: ByteArray = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcd";
+    let mut iter = ba_30.span().into_iter();
+    for _ in 0_usize..29 {
+        let _ = iter.next();
+    }
+    assert_eq!(iter.next(), Some('d'), "30 bytes - last byte");
+    assert_eq!(iter.next(), None);
+
+    let ba_31: ByteArray = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcde";
+    let mut iter = ba_31.span().into_iter();
+    assert_eq!(iter.next(), Some('A'));
+    for _ in 1_usize..30 {
+        let _ = iter.next();
+    }
+    assert_eq!(iter.next(), Some('e'), "31 bytes - last byte");
+    assert_eq!(iter.next(), None);
+
+    let ba_32: ByteArray = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdef";
+    let mut iter = ba_32.span().into_iter();
+    for _ in 0_usize..30 {
+        let _ = iter.next();
+    }
+    assert_eq!(iter.next(), Some('e'), "32 bytes - byte 30");
+    assert_eq!(iter.next(), Some('f'), "32 bytes - byte 31");
+    assert_eq!(iter.next(), None);
+
+    // Test 62, 63, 64 bytes (2 word boundaries).
+    let ba_62: ByteArray = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789";
+    let mut iter = ba_62.span().into_iter();
+    for _ in 0_usize..61 {
+        let _ = iter.next();
+    }
+    assert_eq!(iter.next(), Some('9'), "62 bytes - last byte");
+    assert_eq!(iter.next(), None);
+
+    let ba_63: ByteArray = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789!";
+    let mut iter = ba_63.span().into_iter();
+    for _ in 0_usize..62 {
+        let _ = iter.next();
+    }
+    assert_eq!(iter.next(), Some('!'), "63 bytes - last byte");
+    assert_eq!(iter.next(), None);
+
+    let ba_64: ByteArray = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789!@";
+    let mut iter = ba_64.span().into_iter();
+    for _ in 0_usize..62 {
+        let _ = iter.next();
+    }
+    assert_eq!(iter.next(), Some('!'), "64 bytes - byte 62");
+    assert_eq!(iter.next(), Some('@'), "64 bytes - byte 63");
+    assert_eq!(iter.next(), None);
+}
+
+#[test]
+fn test_byte_span_iterator_multiple_words() {
+    // Test with 3+ words to verify iteration works across multiple word boundaries.
+    // 92 bytes: 31 + 31 + 30.
+    let ba_92: ByteArray =
+        "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789!@#$%^&*()_+-=[]{}|;':,.<>?/~`";
+    let span = ba_92.span();
+    let mut iter = span.into_iter();
+    // Verify we can iterate through all bytes.
+    let mut count = 0;
+    while let Some(_) = iter.next() {
+        count += 1;
+    }
+    assert_eq!(count, 92, "should iterate all 92 bytes");
+
+    // Verify correctness at specific positions.
+    let mut iter = span.into_iter();
+    assert_eq!(iter.next(), Some('A'));
+
+    // Skip to last byte.
+    for _ in 1_usize..91 {
+        let _ = iter.next();
+    }
+    assert_eq!(iter.next(), Some('`'));
+    assert_eq!(iter.next(), None);
+}
+
+#[test]
+fn test_byte_span_iterator_for_loop_collect() {
+    let small_ba: ByteArray = "Hello";
+    let span = small_ba.span();
+
+    let mut collected = Default::default();
+    let mut count = 0;
+    for byte in span {
+        collected.append_byte(byte);
+        count += 1;
+    }
+    assert_eq!(collected, small_ba);
+    assert_eq!(count, 5);
+    assert_eq!(span.into_iter().collect(), small_ba);
+
+    // Test with 2 words.
+    let ba_40: ByteArray = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmn";
+    collected = Default::default();
+    count = 0;
+    for byte in ba_40.span() {
+        collected.append_byte(byte);
+        count += 1;
+    }
+    assert_eq!(collected, ba_40);
+    assert_eq!(count, 40);
+    assert_eq!(ba_40.span().into_iter().collect(), ba_40);
+
+    // Test with 3 words.
+    let ba_70: ByteArray = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789!@#$%^&*";
+    collected = Default::default();
+    count = 0;
+    for byte in ba_70.span() {
+        collected.append_byte(byte);
+        count += 1;
+    }
+    assert_eq!(collected, ba_70);
+    assert_eq!(count, 70);
+    assert_eq!(ba_70.span().into_iter().collect(), ba_70);
+}
+
+#[test]
+fn test_byte_span_iterator_slices() {
+    // Slice within remainder word (< 31 bytes).
+    let ba_13: ByteArray = "Hello Shmello";
+    let span = ba_13.span().get(2..7).unwrap();
+
+    let mut iter = span.into_iter();
+    assert_eq!(iter.next(), Some('l'));
+    assert_eq!(iter.next(), Some('l'));
+    assert_eq!(iter.next(), Some('o'));
+    assert_eq!(iter.next(), Some(' '));
+    assert_eq!(iter.next(), Some('S'));
+    assert_eq!(iter.next(), None);
+
+    // Iterate slice across 2 words (1 data + remainder).
+    let ba_33: ByteArray = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefg";
+    let span = ba_33.span().get(27..32).unwrap();
+
+    let mut iter = span.into_iter();
+    assert_eq!(iter.next(), Some('b'));
+    assert_eq!(iter.next(), Some('c'));
+    assert_eq!(iter.next(), Some('d'));
+    assert_eq!(iter.next(), Some('e'));
+    assert_eq!(iter.next(), Some('f'));
+    assert_eq!(iter.next(), None);
+
+    // Iterate slice across 3 words.
+    let ba_66: ByteArray = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789!@#$";
+    let span = ba_66.span().get(29..64).unwrap();
+
+    let mut iter = span.into_iter();
+    assert_eq!(iter.next(), Some('d'), "First word");
+    assert_eq!(iter.next(), Some('e'));
+    assert_eq!(iter.next(), Some('f'), "Second word");
+    assert_eq!(iter.next(), Some('g'));
+    assert_eq!(iter.next(), Some('h'));
+    assert_eq!(iter.next(), Some('i'));
+    assert_eq!(iter.next(), Some('j'));
+    assert_eq!(iter.next(), Some('k'));
+    assert_eq!(iter.next(), Some('l'));
+    assert_eq!(iter.next(), Some('m'));
+    assert_eq!(iter.next(), Some('n'));
+    assert_eq!(iter.next(), Some('o'));
+    assert_eq!(iter.next(), Some('p'));
+    assert_eq!(iter.next(), Some('q'));
+    assert_eq!(iter.next(), Some('r'));
+    assert_eq!(iter.next(), Some('s'));
+    assert_eq!(iter.next(), Some('t'));
+    assert_eq!(iter.next(), Some('u'));
+    assert_eq!(iter.next(), Some('v'));
+    assert_eq!(iter.next(), Some('w'));
+    assert_eq!(iter.next(), Some('x'));
+    assert_eq!(iter.next(), Some('y'));
+    assert_eq!(iter.next(), Some('z'));
+    assert_eq!(iter.next(), Some('0'));
+    assert_eq!(iter.next(), Some('1'));
+    assert_eq!(iter.next(), Some('2'));
+    assert_eq!(iter.next(), Some('3'));
+    assert_eq!(iter.next(), Some('4'));
+    assert_eq!(iter.next(), Some('5'));
+    assert_eq!(iter.next(), Some('6'));
+    assert_eq!(iter.next(), Some('7'));
+    assert_eq!(iter.next(), Some('8'));
+    assert_eq!(iter.next(), Some('9'));
+    assert_eq!(iter.next(), Some('!'), "Remainder word");
+    assert_eq!(iter.next(), Some('@'));
+    assert_eq!(iter.next(), None);
+}