chore: clean up blst wrapper doc comment and types

crates/precompile/src/bls12_381/g1_add.rs

@@ -1,4 +1,5 @@
-use super::blst::{encode_g1_point, extract_g1_input_no_subgroup_check, p1_add_affine};
+use super::blst::{encode_g1_point, p1_add_affine, read_g1_no_subgroup_check};
+use super::utils::remove_g1_padding;
 use crate::bls12_381_const::{
     G1_ADD_ADDRESS, G1_ADD_BASE_GAS_FEE, G1_ADD_INPUT_LENGTH, PADDED_G1_LENGTH,
 };
@@ -25,12 +26,15 @@ pub(super) fn g1_add(input: &Bytes, gas_limit: u64) -> PrecompileResult {
         )));
     }
 
+    let [a_x, a_y] = remove_g1_padding(&input[..PADDED_G1_LENGTH])?;
+    let [b_x, b_y] = remove_g1_padding(&input[PADDED_G1_LENGTH..])?;
+
     // NB: There is no subgroup check for the G1 addition precompile because the time to do the subgroup
     // check would be more than the time it takes to to do the g1 addition.
     //
     // Users should be careful to note whether the points being added are indeed in the right subgroup.
-    let a_aff = &extract_g1_input_no_subgroup_check(&input[..PADDED_G1_LENGTH])?;
-    let b_aff = &extract_g1_input_no_subgroup_check(&input[PADDED_G1_LENGTH..])?;
+    let a_aff = &read_g1_no_subgroup_check(a_x, a_y)?;
+    let b_aff = &read_g1_no_subgroup_check(b_x, b_y)?;
     let p_aff = p1_add_affine(a_aff, b_aff);
 
     let out = encode_g1_point(&p_aff);

crates/precompile/src/bls12_381/g1_msm.rs

@@ -1,4 +1,5 @@
-use super::blst::{encode_g1_point, extract_g1_input, extract_scalar_input, p1_msm};
+use super::blst::{encode_g1_point, p1_msm, read_g1, read_scalar};
+use crate::bls12_381::utils::remove_g1_padding;
 use crate::bls12_381_const::{
     DISCOUNT_TABLE_G1_MSM, G1_MSM_ADDRESS, G1_MSM_BASE_GAS_FEE, G1_MSM_INPUT_LENGTH, NBITS,
     PADDED_G1_LENGTH, SCALAR_LENGTH,
@@ -34,7 +35,7 @@ pub(super) fn g1_msm(input: &Bytes, gas_limit: u64) -> PrecompileResult {
     }
 
     let mut g1_points: Vec<_> = Vec::with_capacity(k);
-    let mut scalars_bytes: Vec<u8> = Vec::with_capacity(k * SCALAR_LENGTH);
+    let mut scalars = Vec::with_capacity(k);
     for i in 0..k {
         let encoded_g1_element =
             &input[i * G1_MSM_INPUT_LENGTH..i * G1_MSM_INPUT_LENGTH + PADDED_G1_LENGTH];
@@ -48,8 +49,11 @@ pub(super) fn g1_msm(input: &Bytes, gas_limit: u64) -> PrecompileResult {
         if encoded_g1_element.iter().all(|i| *i == 0) {
             continue;
         }
+
+        let [a_x, a_y] = remove_g1_padding(encoded_g1_element)?;
+
         // NB: Scalar multiplications, MSMs and pairings MUST perform a subgroup check.
-        let p0_aff = extract_g1_input(encoded_g1_element)?;
+        let p0_aff = read_g1(a_x, a_y)?;
 
         // If the scalar is zero, then this is a no-op.
         //
@@ -61,8 +65,7 @@ pub(super) fn g1_msm(input: &Bytes, gas_limit: u64) -> PrecompileResult {
         }
 
         g1_points.push(p0_aff);
-
-        scalars_bytes.extend_from_slice(&extract_scalar_input(encoded_scalar)?.b);
+        scalars.push(read_scalar(encoded_scalar)?);
     }
 
     // Return the encoding for the point at the infinity according to EIP-2537
@@ -75,7 +78,7 @@ pub(super) fn g1_msm(input: &Bytes, gas_limit: u64) -> PrecompileResult {
         ));
     }
 
-    let multiexp_aff = p1_msm(g1_points, scalars_bytes, NBITS);
+    let multiexp_aff = p1_msm(g1_points, scalars, NBITS);
 
     let out = encode_g1_point(&multiexp_aff);
     Ok(PrecompileOutput::new(required_gas, out.into()))

crates/precompile/src/bls12_381/g2_add.rs

@@ -1,4 +1,5 @@
-use super::blst::{encode_g2_point, extract_g2_input_no_subgroup_check, p2_add_affine};
+use super::blst::{encode_g2_point, p2_add_affine, read_g2_no_subgroup_check};
+use super::utils::remove_g2_padding;
 use crate::bls12_381_const::{
     G2_ADD_ADDRESS, G2_ADD_BASE_GAS_FEE, G2_ADD_INPUT_LENGTH, PADDED_G2_LENGTH,
 };
@@ -26,12 +27,15 @@ pub(super) fn g2_add(input: &Bytes, gas_limit: u64) -> PrecompileResult {
         )));
     }
 
+    let [a_x_0, a_x_1, a_y_0, a_y_1] = remove_g2_padding(&input[..PADDED_G2_LENGTH])?;
+    let [b_x_0, b_x_1, b_y_0, b_y_1] = remove_g2_padding(&input[PADDED_G2_LENGTH..])?;
+
     // NB: There is no subgroup check for the G2 addition precompile because the time to do the subgroup
     // check would be more than the time it takes to to do the g1 addition.
     //
     // Users should be careful to note whether the points being added are indeed in the right subgroup.
-    let a_aff = &extract_g2_input_no_subgroup_check(&input[..PADDED_G2_LENGTH])?;
-    let b_aff = &extract_g2_input_no_subgroup_check(&input[PADDED_G2_LENGTH..])?;
+    let a_aff = &read_g2_no_subgroup_check(a_x_0, a_x_1, a_y_0, a_y_1)?;
+    let b_aff = &read_g2_no_subgroup_check(b_x_0, b_x_1, b_y_0, b_y_1)?;
 
     // Use the safe wrapper for G2 point addition
     let p_aff = p2_add_affine(a_aff, b_aff);

crates/precompile/src/bls12_381/g2_msm.rs

@@ -1,4 +1,5 @@
-use super::blst::{encode_g2_point, extract_g2_input, extract_scalar_input, p2_msm};
+use super::blst::{encode_g2_point, p2_msm, read_g2, read_scalar};
+use super::utils::remove_g2_padding;
 use crate::bls12_381_const::{
     DISCOUNT_TABLE_G2_MSM, G2_MSM_ADDRESS, G2_MSM_BASE_GAS_FEE, G2_MSM_INPUT_LENGTH, NBITS,
     PADDED_G2_LENGTH, SCALAR_LENGTH,
@@ -34,23 +35,26 @@ pub(super) fn g2_msm(input: &Bytes, gas_limit: u64) -> PrecompileResult {
     }
 
     let mut g2_points: Vec<_> = Vec::with_capacity(k);
-    let mut scalars_bytes: Vec<u8> = Vec::with_capacity(k * SCALAR_LENGTH);
+    let mut scalars = Vec::with_capacity(k);
     for i in 0..k {
-        let encoded_g2_elements =
+        let encoded_g2_element =
             &input[i * G2_MSM_INPUT_LENGTH..i * G2_MSM_INPUT_LENGTH + PADDED_G2_LENGTH];
         let encoded_scalar = &input[i * G2_MSM_INPUT_LENGTH + PADDED_G2_LENGTH
             ..i * G2_MSM_INPUT_LENGTH + PADDED_G2_LENGTH + SCALAR_LENGTH];
 
         // Filter out points infinity as an optimization, since it is a no-op.
         // Note: Previously, points were being batch converted from Jacobian to Affine. In `blst`, this would essentially,
         // zero out all of the points. Since all points are in affine, this bug is avoided.
-        if encoded_g2_elements.iter().all(|i| *i == 0) {
+        if encoded_g2_element.iter().all(|i| *i == 0) {
             continue;
         }
+
+        let [a_x_0, a_x_1, a_y_0, a_y_1] = remove_g2_padding(encoded_g2_element)?;
+
         // NB: Scalar multiplications, MSMs and pairings MUST perform a subgroup check.
         //
         // So we set the subgroup_check flag to `true`
-        let p0_aff = extract_g2_input(encoded_g2_elements)?;
+        let p0_aff = read_g2(a_x_0, a_x_1, a_y_0, a_y_1)?;
 
         // If the scalar is zero, then this is a no-op.
         //
@@ -63,17 +67,19 @@ pub(super) fn g2_msm(input: &Bytes, gas_limit: u64) -> PrecompileResult {
 
         // Convert affine point to Jacobian coordinates using our helper function
         g2_points.push(p0_aff);
-
-        scalars_bytes.extend_from_slice(&extract_scalar_input(encoded_scalar)?.b);
+        scalars.push(read_scalar(encoded_scalar)?);
     }
 
     // Return infinity point if all points are infinity
     if g2_points.is_empty() {
-        return Ok(PrecompileOutput::new(required_gas, [0; 256].into()));
+        return Ok(PrecompileOutput::new(
+            required_gas,
+            [0; PADDED_G2_LENGTH].into(),
+        ));
     }
 
     // Perform multi-scalar multiplication using the safe wrapper
-    let multiexp_aff = p2_msm(g2_points, scalars_bytes, NBITS);
+    let multiexp_aff = p2_msm(g2_points, scalars, NBITS);
 
     let out = encode_g2_point(&multiexp_aff);
     Ok(PrecompileOutput::new(required_gas, out.into()))

crates/precompile/src/bls12_381/map_fp2_to_g2.rs

@@ -1,6 +1,6 @@
 use super::{
-    blst::{check_canonical_fp2, encode_g2_point, map_fp2_to_g2 as blst_map_fp2_to_g2},
-    utils::remove_padding,
+    blst::{encode_g2_point, map_fp2_to_g2 as blst_map_fp2_to_g2, read_fp2},
+    utils::remove_fp_padding,
 };
 use crate::bls12_381_const::{
     MAP_FP2_TO_G2_ADDRESS, MAP_FP2_TO_G2_BASE_GAS_FEE, PADDED_FP2_LENGTH, PADDED_FP_LENGTH,
@@ -29,9 +29,9 @@ pub(super) fn map_fp2_to_g2(input: &Bytes, gas_limit: u64) -> PrecompileResult {
         )));
     }
 
-    let input_p0_x = remove_padding(&input[..PADDED_FP_LENGTH])?;
-    let input_p0_y = remove_padding(&input[PADDED_FP_LENGTH..PADDED_FP2_LENGTH])?;
-    let fp2 = check_canonical_fp2(input_p0_x, input_p0_y)?;
+    let input_p0_x = remove_fp_padding(&input[..PADDED_FP_LENGTH])?;
+    let input_p0_y = remove_fp_padding(&input[PADDED_FP_LENGTH..PADDED_FP2_LENGTH])?;
+    let fp2 = read_fp2(input_p0_x, input_p0_y)?;
     let p_aff = blst_map_fp2_to_g2(&fp2);
 
     let out = encode_g2_point(&p_aff);

crates/precompile/src/bls12_381/map_fp_to_g1.rs

@@ -1,6 +1,6 @@
 use super::{
-    blst::{encode_g1_point, fp_from_bendian, map_fp_to_g1 as blst_map_fp_to_g1},
-    utils::remove_padding,
+    blst::{encode_g1_point, map_fp_to_g1 as blst_map_fp_to_g1, read_fp},
+    utils::remove_fp_padding,
 };
 use crate::bls12_381_const::{MAP_FP_TO_G1_ADDRESS, MAP_FP_TO_G1_BASE_GAS_FEE, PADDED_FP_LENGTH};
 use crate::{PrecompileError, PrecompileOutput, PrecompileResult, PrecompileWithAddress};
@@ -25,8 +25,8 @@ pub(super) fn map_fp_to_g1(input: &Bytes, gas_limit: u64) -> PrecompileResult {
         )));
     }
 
-    let input_p0 = remove_padding(input)?;
-    let fp = fp_from_bendian(input_p0)?;
+    let input_p0 = remove_fp_padding(input)?;
+    let fp = read_fp(input_p0)?;
     let p_aff = blst_map_fp_to_g1(&fp);
 
     let out = encode_g1_point(&p_aff);

crates/precompile/src/bls12_381/pairing.rs

@@ -1,4 +1,5 @@
-use super::blst::{extract_g1_input, extract_g2_input, pairing_check};
+use super::blst::{pairing_check, read_g1, read_g2};
+use super::utils::{remove_g1_padding, remove_g2_padding};
 use crate::bls12_381_const::{
     PADDED_G1_LENGTH, PADDED_G2_LENGTH, PAIRING_ADDRESS, PAIRING_INPUT_LENGTH,
     PAIRING_MULTIPLIER_BASE, PAIRING_OFFSET_BASE,
@@ -55,10 +56,13 @@ pub(super) fn pairing(input: &Bytes, gas_limit: u64) -> PrecompileResult {
         let g1_is_zero = encoded_g1_element.iter().all(|i| *i == 0);
         let g2_is_zero = encoded_g2_element.iter().all(|i| *i == 0);
 
+        let [a_x, a_y] = remove_g1_padding(encoded_g1_element)?;
+        let [b_x_0, b_x_1, b_y_0, b_y_1] = remove_g2_padding(encoded_g2_element)?;
+
         // NB: Scalar multiplications, MSMs and pairings MUST perform a subgroup check.
         // extract_g1_input and extract_g2_input perform the necessary checks
-        let p1_aff = extract_g1_input(encoded_g1_element)?;
-        let p2_aff = extract_g2_input(encoded_g2_element)?;
+        let p1_aff = read_g1(a_x, a_y)?;
+        let p2_aff = read_g2(b_x_0, b_x_1, b_y_0, b_y_1)?;
 
         if !g1_is_zero & !g2_is_zero {
             pairs.push((p1_aff, p2_aff));

crates/precompile/src/bls12_381/utils.rs

@@ -1,8 +1,10 @@
-use crate::bls12_381_const::{FP_LENGTH, FP_PAD_BY, MODULUS_REPR, PADDED_FP_LENGTH};
+use crate::bls12_381_const::{
+    FP_LENGTH, FP_PAD_BY, PADDED_FP_LENGTH, PADDED_G1_LENGTH, PADDED_G2_LENGTH,
+};
 use crate::PrecompileError;
 
 /// Removes zeros with which the precompile inputs are left padded to 64 bytes.
-pub(super) fn remove_padding(input: &[u8]) -> Result<&[u8; FP_LENGTH], PrecompileError> {
+pub(super) fn remove_fp_padding(input: &[u8]) -> Result<&[u8; FP_LENGTH], PrecompileError> {
     if input.len() != PADDED_FP_LENGTH {
         return Err(PrecompileError::Other(format!(
             "Padded input should be {PADDED_FP_LENGTH} bytes, was {}",
@@ -17,8 +19,34 @@ pub(super) fn remove_padding(input: &[u8]) -> Result<&[u8; FP_LENGTH], Precompil
     }
     Ok(unpadded.try_into().unwrap())
 }
+/// remove_g1_padding removes the padding applied to the Fp elements that constitute the
+/// encoded G1 element.
+pub(super) fn remove_g1_padding(input: &[u8]) -> Result<[&[u8; FP_LENGTH]; 2], PrecompileError> {
+    if input.len() != PADDED_G1_LENGTH {
+        return Err(PrecompileError::Other(format!(
+            "Input should be {PADDED_G1_LENGTH} bytes, was {}",
+            input.len()
+        )));
+    }
+
+    let x = remove_fp_padding(&input[..PADDED_FP_LENGTH])?;
+    let y = remove_fp_padding(&input[PADDED_FP_LENGTH..PADDED_G1_LENGTH])?;
+    Ok([x, y])
+}
 
-/// Checks if the input is a valid big-endian representation of a field element.
-pub(super) fn is_valid_be(input: &[u8; 48]) -> bool {
-    *input < MODULUS_REPR
+/// remove_g2_padding removes the padding applied to the Fp elements that constitute the
+/// encoded G2 element.
+pub(super) fn remove_g2_padding(input: &[u8]) -> Result<[&[u8; FP_LENGTH]; 4], PrecompileError> {
+    if input.len() != PADDED_G2_LENGTH {
+        return Err(PrecompileError::Other(format!(
+            "Input should be {PADDED_G2_LENGTH} bytes, was {}",
+            input.len()
+        )));
+    }
+
+    let mut input_fps = [&[0; FP_LENGTH]; 4];
+    for i in 0..4 {
+        input_fps[i] = remove_fp_padding(&input[i * PADDED_FP_LENGTH..(i + 1) * PADDED_FP_LENGTH])?;
+    }
+    Ok(input_fps)
 }