Cherry-pick fixes. (#6859)

Author: orizi

Cargo.lock

@@ -45,8 +45,8 @@
 use crate::array::{ArrayTrait, SpanTrait};
 #[allow(unused_imports)]
 use crate::bytes_31::{
-    BYTES_IN_BYTES31, Bytes31Trait, one_shift_left_bytes_felt252, one_shift_left_bytes_u128,
-    POW_2_128, POW_2_8, U128IntoBytes31, U8IntoBytes31,
+    BYTES_IN_BYTES31, Bytes31Trait, POW_2_128, POW_2_8, U128IntoBytes31, U8IntoBytes31,
+    one_shift_left_bytes_felt252, one_shift_left_bytes_u128, split_u128, u8_at_u256,
 };
 use crate::clone::Clone;
 use crate::cmp::min;
@@ -277,29 +277,21 @@ pub impl ByteArrayImpl of ByteArrayTrait {
     /// assert!(byte == 98);
     /// ```
     fn at(self: @ByteArray, index: usize) -> Option<u8> {
-        let (word_index, index_in_word) = DivRem::div_rem(
-            index, BYTES_IN_BYTES31.try_into().unwrap(),
-        );
-
-        let data_len = self.data.len();
-        if word_index == data_len {
+        let (word_index, index_in_word) = DivRem::div_rem(index, 31);
+        if word_index == self.data.len() {
             // Index is in pending word.
             if index_in_word >= *self.pending_word_len {
                 return Option::None;
             }
             // index_in_word is from MSB, we need index from LSB.
-            let index_from_lsb = *self.pending_word_len - 1 - index_in_word;
-            let pending_bytes31: bytes31 = (*self.pending_word).try_into().unwrap();
-            return Option::Some(pending_bytes31.at(index_from_lsb));
-        }
-
-        if word_index > data_len {
-            return Option::None;
+            return Option::Some(
+                u8_at_u256((*self.pending_word).into(), *self.pending_word_len - 1 - index_in_word),
+            );
         }
 
+        let data_word: bytes31 = *self.data.get(word_index)?.deref();
         // index_in_word is from MSB, we need index from LSB.
-        let index_from_lsb = BYTES_IN_BYTES31 - 1 - index_in_word;
-        Option::Some(self.data.at(word_index).at(index_from_lsb))
+        Option::Some(data_word.at(BYTES_IN_BYTES31 - 1 - index_in_word))
     }
 
     /// Returns a `ByteArray` with the reverse order of `self`.
@@ -350,9 +342,7 @@ pub impl ByteArrayImpl of ByteArrayTrait {
             if index == low_part_limit {
                 break;
             }
-            let curr_byte_as_u128 = (low / one_shift_left_bytes_u128(index)) % POW_2_8;
-
-            self.append_byte(curr_byte_as_u128.try_into().unwrap());
+            self.append_byte(core::bytes_31::get_lsb(split_u128(low, index).high));
             index += 1;
         };
         if low_part_limit == BYTES_IN_U128 {
@@ -362,10 +352,10 @@ pub impl ByteArrayImpl of ByteArrayTrait {
                 if index_in_high_part == high_part_len {
                     break;
                 }
-                let curr_byte_as_u128 = (high
-                    / one_shift_left_bytes_u128(index_in_high_part)) % POW_2_8;
-
-                self.append_byte(curr_byte_as_u128.try_into().unwrap());
+                self
+                    .append_byte(
+                        core::bytes_31::get_lsb(split_u128(high, index_in_high_part).high),
+                    );
                 index_in_high_part += 1;
             }
         }
@@ -402,13 +392,11 @@ pub impl ByteArrayImpl of ByteArrayTrait {
     fn append_split_index_lt_16(ref self: ByteArray, word: felt252, split_index: usize) {
         let u256 { low, high } = word.into();
 
-        let (low_quotient, low_remainder) = u128_safe_divmod(
-            low, one_shift_left_bytes_u128(split_index).try_into().unwrap(),
-        );
+        let low_result = split_u128(low, split_index);
         let left = high.into() * one_shift_left_bytes_u128(BYTES_IN_U128 - split_index).into()
-            + low_quotient.into();
+            + low_result.high.into();
 
-        self.append_split(left, low_remainder.into());
+        self.append_split(left, low_result.low.into());
     }
 
     /// Appends a single word to the end of `self`, given that the index of splitting `word` is
@@ -437,12 +425,10 @@ pub impl ByteArrayImpl of ByteArrayTrait {
     fn append_split_index_gt_16(ref self: ByteArray, word: felt252, split_index: usize) {
         let u256 { low, high } = word.into();
 
-        let (high_quotient, high_remainder) = u128_safe_divmod(
-            high, one_shift_left_bytes_u128(split_index - BYTES_IN_U128).try_into().unwrap(),
-        );
-        let right = high_remainder.into() * POW_2_128 + low.into();
+        let high_result = split_u128(high, split_index - BYTES_IN_U128);
+        let right = high_result.low.into() * POW_2_128 + low.into();
 
-        self.append_split(high_quotient.into(), right);
+        self.append_split(high_result.high.into(), right);
     }
 
     /// A helper function to append a remainder to `self`, by:

corelib/src/byte_array.cairo

@@ -1,33 +1,55 @@
-use crate::traits::{Into, TryInto};
-use crate::option::OptionTrait;
+//! Definitions and utilities for the `bytes31` type.
+//!
+//! The `bytes31` type is a compact, indexable 31-byte type.
+//!
+//! # Examples
+//!
+//! Creating a `bytes31` from a `felt252`:
+//! ```
+//! let value: bytes31 = 0xaabb.try_into().unwrap();
+//! ```
+//!
+//! Accessing a byte by index:
+//! ```
+//! assert!(value[0] == 0xbb);
+//! ```
+
 #[allow(unused_imports)]
 use crate::integer::{u128_safe_divmod, u128_to_felt252};
+#[allow(unused_imports)]
+use crate::option::OptionTrait;
 use crate::RangeCheck;
+use crate::traits::{Into, TryInto};
 
 pub(crate) const BYTES_IN_BYTES31: usize = 31;
 const BYTES_IN_U128: usize = 16;
 pub(crate) const POW_2_128: felt252 = 0x100000000000000000000000000000000;
 pub(crate) const POW_2_8: u128 = 0x100;
 
+/// Represents a 31-byte fixed-size byte type.
 #[derive(Copy, Drop)]
 pub extern type bytes31;
 
 pub(crate) extern fn bytes31_const<const value: felt252>() -> bytes31 nopanic;
 extern fn bytes31_try_from_felt252(value: felt252) -> Option<bytes31> implicits(RangeCheck) nopanic;
 extern fn bytes31_to_felt252(value: bytes31) -> felt252 nopanic;
 
+/// A trait for accessing a specific byte of a `bytes31` type.
 #[generate_trait]
 pub impl Bytes31Impl of Bytes31Trait {
-    /// Gets the byte at the given index (LSB's index is 0), assuming that
-    /// `index < BYTES_IN_BYTES31`. If the assumption is not met, the behavior is undefined.
+    /// Returns the byte at the given index (LSB's index is 0).
+    ///
+    /// Assumes that `index < BYTES_IN_BYTES31`. If the assumption is not met, the behavior is
+    /// undefined.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// let bytes: bytes31 = 1_u8.into();
+    /// assert!(bytes.at(0) == 1);
+    /// ```
     fn at(self: @bytes31, index: usize) -> u8 {
-        let u256 { low, high } = (*self).into();
-        let res_u128 = if index < BYTES_IN_U128 {
-            (low / one_shift_left_bytes_u128(index)) % POW_2_8
-        } else {
-            (high / one_shift_left_bytes_u128(index - BYTES_IN_U128)) % POW_2_8
-        };
-        res_u128.try_into().unwrap()
+        u8_at_u256((*self).into(), index)
     }
 }
 
@@ -70,34 +92,40 @@ pub(crate) impl U8IntoBytes31 of Into<u8, bytes31> {
         crate::integer::upcast(self)
     }
 }
+
 impl U16IntoBytes31 of Into<u16, bytes31> {
     fn into(self: u16) -> bytes31 {
         crate::integer::upcast(self)
     }
 }
+
 impl U32IntoBytes31 of Into<u32, bytes31> {
     fn into(self: u32) -> bytes31 {
         crate::integer::upcast(self)
     }
 }
+
 impl U64IntoBytes31 of Into<u64, bytes31> {
     fn into(self: u64) -> bytes31 {
         crate::integer::upcast(self)
     }
 }
+
 pub(crate) impl U128IntoBytes31 of Into<u128, bytes31> {
     fn into(self: u128) -> bytes31 {
         crate::integer::upcast(self)
     }
 }
 
-/// Splits a bytes31 into two bytes31s at the given index (LSB's index is 0).
-/// The bytes31s are represented using felt252s to improve performance.
+/// Splits a `bytes31` into two `bytes31`s at the given index (LSB's index is 0).
+/// The input `bytes31` and the output `bytes31`s are represented using `felt252`s to improve
+/// performance.
+///
 /// Note: this function assumes that:
-/// 1. `word` is validly convertible to a bytes31 which has no more than `len` bytes of data.
-/// 2. index <= len.
-/// 3. len <= BYTES_IN_BYTES31.
-/// If these assumptions are not met, it can corrupt the ByteArray. Thus, this should be a
+/// 1. `word` is validly convertible to a `bytes31`` which has no more than `len` bytes of data.
+/// 2. `index <= len`.
+/// 3. `len <= BYTES_IN_BYTES31`.
+/// If these assumptions are not met, it can corrupt the `byte31`s. Thus, this should be a
 /// private function. We could add masking/assertions but it would be more expansive.
 pub(crate) fn split_bytes31(word: felt252, len: usize, index: usize) -> (felt252, felt252) {
     if index == 0 {
@@ -114,32 +142,27 @@ pub(crate) fn split_bytes31(word: felt252, len: usize, index: usize) -> (felt252
     }
 
     if len <= BYTES_IN_U128 {
-        let (quotient, remainder) = u128_safe_divmod(
-            low, one_shift_left_bytes_u128(index).try_into().unwrap(),
-        );
-        return (remainder.into(), quotient.into());
+        let result = split_u128(low, index);
+        return (result.low.into(), result.high.into());
     }
 
     // len > BYTES_IN_U128
     if index < BYTES_IN_U128 {
-        let (low_quotient, low_remainder) = u128_safe_divmod(
-            low, one_shift_left_bytes_u128(index).try_into().unwrap(),
-        );
+        let low_result = split_u128(low, index);
         let right = high.into() * one_shift_left_bytes_u128(BYTES_IN_U128 - index).into()
-            + low_quotient.into();
-        return (low_remainder.into(), right);
+            + low_result.high.into();
+        return (low_result.low.into(), right);
     }
 
     // len > BYTES_IN_U128 && index > BYTES_IN_U128
-    let (high_quotient, high_remainder) = u128_safe_divmod(
-        high, one_shift_left_bytes_u128(index - BYTES_IN_U128).try_into().unwrap(),
-    );
-    let left = high_remainder.into() * POW_2_128 + low.into();
-    return (left, high_quotient.into());
+
+    let high_result = split_u128(high, index - BYTES_IN_U128);
+    let left = high_result.low.into() * POW_2_128 + low.into();
+    return (left, high_result.high.into());
 }
 
 
-/// Returns 1 << (8 * `n_bytes`) as felt252, assuming that `n_bytes < BYTES_IN_BYTES31`.
+/// Returns `1 << (8 * n_bytes)` as `felt252`, assuming that `n_bytes < BYTES_IN_BYTES31`.
 ///
 /// Note: if `n_bytes >= BYTES_IN_BYTES31`, the behavior is undefined. If one wants to
 /// assert that in the callsite, it's sufficient to assert that `n_bytes != BYTES_IN_BYTES31`
@@ -152,10 +175,33 @@ pub(crate) fn one_shift_left_bytes_felt252(n_bytes: usize) -> felt252 {
     }
 }
 
-/// Returns 1 << (8 * `n_bytes`) as u128, where `n_bytes` must be < BYTES_IN_U128.
+/// Returns `1 << (8 * n_bytes)` as `u128`, where `n_bytes` must be < `BYTES_IN_U128`.
 ///
 /// Panics if `n_bytes >= BYTES_IN_U128`.
 pub(crate) fn one_shift_left_bytes_u128(n_bytes: usize) -> u128 {
+    one_shift_left_bytes_u128_nz(n_bytes).into()
+}
+
+/// Splits a u128 into two words as a `u256` - the `low` is the first `n_bytes` the `high` is the
+/// rest.
+pub(crate) fn split_u128(value: u128, n_bytes: usize) -> u256 {
+    let (high, low) = DivRem::div_rem(value, one_shift_left_bytes_u128_nz(n_bytes));
+    u256 { low, high }
+}
+
+/// Returns the `u8` at `index` if you look at `value` as an array of 32 `u8`s.
+pub(crate) fn u8_at_u256(value: u256, index: usize) -> u8 {
+    get_lsb(
+        if index < BYTES_IN_U128 {
+            split_u128(value.low, index).high
+        } else {
+            split_u128(value.high, index - BYTES_IN_U128).high
+        },
+    )
+}
+
+/// Same as `one_shift_left_bytes_u128` but returns `NonZero` value.
+fn one_shift_left_bytes_u128_nz(n_bytes: usize) -> NonZero<u128> {
     match n_bytes {
         0 => 0x1,
         1 => 0x100,
@@ -184,3 +230,20 @@ impl Bytes31PartialEq of PartialEq<bytes31> {
         lhs_as_felt252 == rhs_as_felt252
     }
 }
+
+mod helpers {
+    use core::internal::bounded_int::{DivRemHelper, BoundedInt, div_rem};
+
+    impl DivRemU128By256 of DivRemHelper<u128, BoundedInt<256, 256>> {
+        type DivT = BoundedInt<0, 0xffffffffffffffffffffffffffffff>;
+        type RemT = BoundedInt<0, 0xff>;
+    }
+
+    /// Returns the least significant byte of the given u128.
+    pub fn get_lsb(value: u128) -> u8 {
+        let (_, res) = div_rem::<_, BoundedInt<256, 256>>(value, 256);
+        core::integer::upcast(res)
+    }
+}
+
+pub(crate) use helpers::get_lsb;