Improve the data-transformer::cairo-types::CairoU256 creation logic (#3622)

<!-- Reference any GitHub issues resolved by this PR -->

Closes
[#2576](#2576)

## Checklist

<!-- Make sure all of these are complete -->

- [x] Linked relevant issue
- [x] Performed self-review of the code

---------

Co-authored-by: Nnamdi Aninye <[email protected]>
Co-authored-by: Artur Michałek <[email protected]>
Co-authored-by: ksew1 <[email protected]>

Author: 4 people

crates/cheatnet/src/runtime_extensions/forge_runtime_extension/mod.rs

@@ -364,26 +364,46 @@ impl<'a> ExtensionLogic for ForgeExtension<'a> {
                 let curve: Felt = input_reader.read()?;
                 let curve = curve.to_short_string().ok();
 
-                let (signing_key_bytes, verifying_key_bytes) = {
+                let (signing_key_bytes, x_coordinate_bytes, y_coordinate_bytes) = {
+                    let extract_coordinates_from_verifying_key = |verifying_key: Box<[u8]>| {
+                        let verifying_key = verifying_key.iter().as_slice();
+                        (
+                            verifying_key[1..33].try_into().unwrap(),
+                            verifying_key[33..65].try_into().unwrap(),
+                        )
+                    };
+
                     match curve.as_deref() {
                         Some("Secp256k1") => {
                             let signing_key = k256::ecdsa::SigningKey::random(
                                 &mut k256::elliptic_curve::rand_core::OsRng,
                             );
-                            let verifying_key = signing_key.verifying_key();
+                            let verifying_key = signing_key
+                                .verifying_key()
+                                .to_encoded_point(false)
+                                .to_bytes();
+                            let (x_coordinate, y_coordinate) =
+                                extract_coordinates_from_verifying_key(verifying_key);
                             (
-                                signing_key.to_bytes(),
-                                verifying_key.to_encoded_point(false).to_bytes(),
+                                signing_key.to_bytes().as_slice()[0..32].try_into().unwrap(),
+                                x_coordinate,
+                                y_coordinate,
                             )
                         }
                         Some("Secp256r1") => {
                             let signing_key = p256::ecdsa::SigningKey::random(
                                 &mut p256::elliptic_curve::rand_core::OsRng,
                             );
-                            let verifying_key = signing_key.verifying_key();
+                            let verifying_key = signing_key
+                                .verifying_key()
+                                .to_encoded_point(false)
+                                .to_bytes();
+                            let (x_coordinate, y_coordinate) =
+                                extract_coordinates_from_verifying_key(verifying_key);
                             (
-                                signing_key.to_bytes(),
-                                verifying_key.to_encoded_point(false).to_bytes(),
+                                signing_key.to_bytes().as_slice()[0..32].try_into().unwrap(),
+                                x_coordinate,
+                                y_coordinate,
                             )
                         }
                         _ => return Ok(CheatcodeHandlingResult::Forwarded),
@@ -392,8 +412,8 @@ impl<'a> ExtensionLogic for ForgeExtension<'a> {
 
                 Ok(CheatcodeHandlingResult::from_serializable((
                     CairoU256::from_bytes(&signing_key_bytes),
-                    CairoU256::from_bytes(&verifying_key_bytes[1..]), // bytes of public_key's x-coordinate
-                    CairoU256::from_bytes(&verifying_key_bytes[33..]), // bytes of public_key's y-coordinate
+                    CairoU256::from_bytes(&x_coordinate_bytes), // bytes of public_key's x-coordinate
+                    CairoU256::from_bytes(&y_coordinate_bytes), // bytes of public_key's y-coordinate
                 )))
             }
             "ecdsa_sign_message" => {
@@ -441,6 +461,8 @@ impl<'a> ExtensionLogic for ForgeExtension<'a> {
                 };
 
                 let result = result.map(|(r_bytes, s_bytes)| {
+                    let r_bytes: [u8; 32] = r_bytes.as_slice()[0..32].try_into().unwrap();
+                    let s_bytes: [u8; 32] = s_bytes.as_slice()[0..32].try_into().unwrap();
                     (
                         CairoU256::from_bytes(&r_bytes),
                         CairoU256::from_bytes(&s_bytes),

crates/data-transformer/src/cairo_types/u256.rs

@@ -15,8 +15,7 @@ pub struct CairoU256 {
 
 impl CairoU256 {
     #[must_use]
-    pub fn from_bytes(bytes: &[u8]) -> Self {
-        // Takes slice without explicit size because of cheatnet's specific usages (See Issue #2575)
+    pub fn from_bytes(bytes: &[u8; 32]) -> Self {
         Self {
             low: u128::from_be_bytes(bytes[16..32].try_into().unwrap()),
             high: u128::from_be_bytes(bytes[0..16].try_into().unwrap()),
@@ -85,37 +84,19 @@ mod tests {
         244, 192, 163, 138, 253, 42, 253, 125, 53, 127, 44,
     ];
 
-    const TOO_BIG_NUMBER_BYTES: [u8; 48] = [
-        222, 9, 69, 194, 71, 77, 155, 51, 51, 49, 35, 229, 62, 55, 169, 63, 93, 228, 186, 10, 219,
-        244, 192, 163, 138, 253, 42, 253, 125, 53, 127, 44, 21, 37, 21, 37, 21, 37, 21, 37, 21, 37,
-        21, 37, 21, 37, 21, 37,
-    ];
-
     const BIG_NUMBER_LIMBS: [u128; 2] = [
         124_805_820_680_284_125_994_760_982_863_763_832_620,
         295_136_760_614_571_572_862_546_075_274_463_127_871,
     ];
 
     #[test_case(&[0; 32], [0, 0] ; "zeros")]
-    #[test_case(&BIG_NUMBER_BYTES[..], BIG_NUMBER_LIMBS; "big")]
-    fn test_happy_case_from_bytes(bytes: &[u8], expected_limbs: [u128; 2]) {
+    #[test_case(&BIG_NUMBER_BYTES, BIG_NUMBER_LIMBS; "big")]
+    fn test_happy_case_from_bytes(bytes: &[u8; 32], expected_limbs: [u128; 2]) {
         let result = CairoU256::from_bytes(bytes);
 
         assert_eq!([result.low, result.high], expected_limbs);
     }
 
-    #[should_panic(expected = "range end index 32 out of range for slice of length 4")]
-    #[test]
-    fn test_from_bytes_input_too_short() {
-        let _result = CairoU256::from_bytes(&[2, 1, 3, 7]);
-    }
-
-    #[test]
-    fn test_happy_case_from_bytes_longer_input() {
-        let result = CairoU256::from_bytes(&TOO_BIG_NUMBER_BYTES);
-        assert_eq!([result.low, result.high], BIG_NUMBER_LIMBS);
-    }
-
     #[test_case("0x0", [0, 0] ; "zero_hex")]
     #[test_case("0", [0, 0] ; "zero_dec")]
     #[test_case("0x237abc", [2_325_180, 0] ; "small_hex")]