Improved concat code

aurora-is-near · Aug 9, 2022 · 54faa98 · 54faa98
1 parent 55bdbe6
commit 54faa98
Showing 1 changed file with 27 additions and 54 deletions.
diff --git a/engine-precompiles/src/alt_bn256.rs b/engine-precompiles/src/alt_bn256.rs
@@ -69,10 +69,10 @@ trait HostFnEncode {
 }
 
 #[cfg(feature = "contract")]
-fn concat_high_low(low: [u8; 8], high: [u8; 8]) -> [u8; consts::SCALAR_LEN] {
-    let mut bytes = [0u8; consts::SCALAR_LEN];
-    bytes[0..consts::SCALAR_PART_LEN].copy_from_slice(&low);
-    bytes[consts::SCALAR_PART_LEN..consts::SCALAR_LEN * 2].copy_from_slice(&high);
+fn concat_low_high<const P: usize, const S: usize>(low: [u8; P], high: [u8; P]) -> [u8; S] {
+    let mut bytes = [0u8; S];
+    bytes[0..P].copy_from_slice(&low);
+    bytes[P..P * 2].copy_from_slice(&high);
     bytes
 }
 
@@ -81,8 +81,8 @@ impl HostFnEncode for bn::Fr {
     type Encoded = [u8; consts::SCALAR_LEN];
 
     fn sdk_encode(self) -> Self::Encoded {
-        let [low_int, high_int] = self.into_u256().0;
-        concat_high_low(low_int.to_le_bytes(), high_int.to_le_bytes())
+        let [low, high] = self.into_u256().0;
+        concat_low_high(low.to_le_bytes(), high.to_le_bytes())
     }
 }
 
@@ -91,8 +91,23 @@ impl HostFnEncode for bn::Fq {
     type Encoded = [u8; consts::SCALAR_LEN];
 
     fn sdk_encode(self) -> Self::Encoded {
-        let [low_int, high_int] = self.into_u256().0;
-        concat_high_low(low_int.to_le_bytes(), high_int.to_le_bytes())
+        let [low, high] = self.into_u256().0;
+        concat_low_high(low.to_le_bytes(), high.to_le_bytes())
+    }
+}
+
+impl HostFnEncode for bn::Fq2 {
+    type Encoded = [u8; consts::SCALAR_LEN * 2];
+
+    fn sdk_encode(self) -> Self::Encoded {
+        let [real_low, real_high] = self.real().into_u256().0;
+        let real: [u8; consts::SCALAR_LEN] =
+            concat_low_high(real_low.to_le_bytes(), real_high.to_le_bytes());
+
+        let [imaginary_low, imaginary_high] = self.imaginary().into_u256().0;
+        let imaginary: [u8; consts::SCALAR_LEN] =
+            concat_low_high(imaginary_low.to_le_bytes(), imaginary_high.to_le_bytes());
+        concat_low_high(real, imaginary)
     }
 }
 
@@ -104,10 +119,7 @@ impl HostFnEncode for bn::G1 {
         bn::AffineG1::from_jacobian(self)
             .map(|p| {
                 let (px, py) = (p.x().sdk_encode(), p.y().sdk_encode());
-                let mut bytes = [0u8; consts::POINT_LEN];
-                bytes[0..consts::SCALAR_LEN * 2].copy_from_slice(&px);
-                bytes[consts::SCALAR_LEN * 2..consts::SCALAR_LEN * 4].copy_from_slice(&py);
-                bytes
+                concat_low_high(px, py)
             })
             .unwrap_or_else(|| [0u8; consts::POINT_LEN])
     }
@@ -120,48 +132,9 @@ impl HostFnEncode for bn::G2 {
     fn sdk_encode(self) -> Self::Encoded {
         bn::AffineG2::from_jacobian(self)
             .map(|g2| {
-                let (x_real_low, x_real_high, x_imaginary_low, x_imaginary_high) = {
-                    let x = g2.x();
-                    let (x_real, x_imaginary) = (x.real(), x.imaginary());
-                    let [x_real_low, x_real_high] = x_real.into_u256().0;
-                    let [x_imaginary_low, x_imaginary_high] = x_imaginary.into_u256().0;
-                    (
-                        x_real_low.to_le_bytes(),
-                        x_real_high.to_le_bytes(),
-                        x_imaginary_low.to_le_bytes(),
-                        x_imaginary_high.to_le_bytes(),
-                    )
-                };
-                let (y_real_low, y_real_high, y_imaginary_low, y_imaginary_high) = {
-                    let y = g2.y();
-                    let (y_real, y_imaginary) = (y.real(), y.imaginary());
-                    let [y_real_low, y_real_high] = y_real.into_u256().0;
-                    let [y_imaginary_low, y_imaginary_high] = y_imaginary.into_u256().0;
-                    (
-                        y_real_low.to_le_bytes(),
-                        y_real_high.to_le_bytes(),
-                        y_imaginary_low.to_le_bytes(),
-                        y_imaginary_high.to_le_bytes(),
-                    )
-                };
-
-                let mut bytes = [0u8; consts::POINT_PAIR_LEN];
-                bytes[0..consts::SCALAR_PART_LEN].copy_from_slice(&x_real_low);
-                bytes[consts::SCALAR_PART_LEN..consts::SCALAR_LEN * 2]
-                    .copy_from_slice(&x_real_high);
-                bytes[consts::SCALAR_LEN * 2..consts::SCALAR_PART_LEN * 3]
-                    .copy_from_slice(&x_imaginary_low);
-                bytes[consts::SCALAR_PART_LEN * 3..consts::SCALAR_PART_LEN * 4]
-                    .copy_from_slice(&x_imaginary_high);
-                bytes[consts::SCALAR_PART_LEN * 4..consts::SCALAR_PART_LEN * 5]
-                    .copy_from_slice(&y_real_low);
-                bytes[consts::SCALAR_PART_LEN * 5..consts::SCALAR_PART_LEN * 6]
-                    .copy_from_slice(&y_real_high);
-                bytes[consts::SCALAR_PART_LEN * 6..consts::SCALAR_PART_LEN * 7]
-                    .copy_from_slice(&y_imaginary_low);
-                bytes[consts::SCALAR_PART_LEN * 7..consts::SCALAR_PART_LEN * 8]
-                    .copy_from_slice(&y_imaginary_high);
-                bytes
+                let x = g2.x().sdk_encode();
+                let y = g2.y().sdk_encode();
+                concat_low_high(x, y)
             })
             .unwrap_or_else(|| [0u8; consts::POINT_PAIR_LEN])
     }