refactor(sha256): Compute ByteArray sha256 via iterator

Author: Gilad Chase

corelib/src/sha256.cairo

@@ -13,6 +13,10 @@
 //! assert!(hash == [0x185f8db3, 0x2271fe25, 0xf561a6fc, 0x938b2e26, 0x4306ec30, 0x4eda5180,
 //! 0x7d17648, 0x26381969]);
 //! ```
+#[feature("byte-span")]
+use core::byte_array::ToByteSpanTrait;
+#[feature("bounded-int-utils")]
+use core::internal::bounded_int::upcast;
 use starknet::SyscallResultTrait;
 
 /// A handle to the state of a SHA-256 hash.
@@ -82,29 +86,30 @@ pub fn compute_sha256_u32_array(
 /// ```
 pub fn compute_sha256_byte_array(arr: @ByteArray) -> [u32; 8] {
     let mut word_arr = array![];
-    let len = arr.len();
-    let rem = len % 4;
-    let mut index = 0;
-    let rounded_len = len - rem;
-    while index != rounded_len {
-        let word = arr.at(index + 3).unwrap().into()
-            + arr.at(index + 2).unwrap().into() * 0x100
-            + arr.at(index + 1).unwrap().into() * 0x10000
-            + arr.at(index).unwrap().into() * 0x1000000;
-        word_arr.append(word);
-        index = index + 4;
-    }
+    let mut iter = arr.span().into_iter();
 
-    let last = match rem {
-        0 => 0,
-        1 => arr.at(len - 1).unwrap().into(),
-        2 => arr.at(len - 1).unwrap().into() + arr.at(len - 2).unwrap().into() * 0x100,
-        _ => arr.at(len - 1).unwrap().into()
-            + arr.at(len - 2).unwrap().into() * 0x100
-            + arr.at(len - 3).unwrap().into() * 0x10000,
+    // Process bytes from the iterator in chunks of 4.
+    let (last_word, last_word_len): (u32, u32) = loop {
+        let Some(b0) = iter.next() else {
+            break (0, 0);
+        };
+        let Some(b1) = iter.next() else {
+            break (upcast(b0), 1);
+        };
+        let b0_b1 = conversions::u8_shift_append_byte(b0, b1);
+        let Some(b2) = iter.next() else {
+            break (upcast(b0_b1), 2);
+        };
+        let b0_b1_b2 = conversions::u16_shift_append_byte(b0_b1, b2);
+        let Some(b3) = iter.next() else {
+            break (upcast(b0_b1_b2), 3);
+        };
+
+        let word = upcast(conversions::u24_shift_append_byte(b0_b1_b2, b3));
+        word_arr.append(word);
     };
 
-    compute_sha256_u32_array(word_arr, last, rem.into())
+    compute_sha256_u32_array(word_arr, last_word, last_word_len)
 }
 
 /// Adds padding to the input array according to the SHA-256 specification.
@@ -208,3 +213,44 @@ fn append_zeros(ref arr: Array<u32>, count: felt252) {
     }
     arr.append(0);
 }
+
+mod conversions {
+    #[feature("bounded-int-utils")]
+    use core::internal::bounded_int::{self, AddHelper, BoundedInt, MulHelper, UnitInt};
+
+    const POW_2_8_TYPED: UnitInt<0x100> = 0x100;
+
+    impl U8Shift of MulHelper<u8, UnitInt<0x100>> {
+        type Result = BoundedInt<0, 0xFF00>;
+    }
+    impl U8ShiftAddU8 of AddHelper<U8Shift::Result, u8> {
+        type Result = BoundedInt<0, 0xFFFF>;
+    }
+    impl U16Shift of MulHelper<U8ShiftAddU8::Result, UnitInt<0x100>> {
+        type Result = BoundedInt<0, 0xFFFF00>;
+    }
+    impl U16ShiftAddU8 of AddHelper<U16Shift::Result, u8> {
+        type Result = BoundedInt<0, 0xFFFFFF>;
+    }
+    impl U24Shift of MulHelper<U16ShiftAddU8::Result, UnitInt<0x100>> {
+        type Result = BoundedInt<0, 0xFFFFFF00>;
+    }
+    impl U24ShiftAddU8 of AddHelper<U24Shift::Result, u8> {
+        type Result = BoundedInt<0, 0xFFFFFFFF>;
+    }
+
+    /// Shifts a u8 left by 8 bits and appends a byte.
+    pub fn u8_shift_append_byte(word: u8, byte: u8) -> U8ShiftAddU8::Result {
+        bounded_int::add(bounded_int::mul(word, POW_2_8_TYPED), byte)
+    }
+
+    /// Shifts a 16-bit value left by 8 bits and appends a byte.
+    pub fn u16_shift_append_byte(word: U8ShiftAddU8::Result, byte: u8) -> U16ShiftAddU8::Result {
+        bounded_int::add(bounded_int::mul(word, POW_2_8_TYPED), byte)
+    }
+
+    /// Shifts a 24-bit value left by 8 bits and appends a byte.
+    pub fn u24_shift_append_byte(word: U16ShiftAddU8::Result, byte: u8) -> U24ShiftAddU8::Result {
+        bounded_int::add(bounded_int::mul(word, POW_2_8_TYPED), byte)
+    }
+}