simd-0284: alt_bn128 little endian compatibility

Author: 0x0ece

src/ballet/bn254/fd_bn254.c

@@ -10,13 +10,14 @@
 
 uchar *
 fd_bn254_g1_compress( uchar       out[32],
-                      uchar const in [64] ) {
+                      uchar const in [64],
+                      int         big_endian ) {
   fd_bn254_g1_t p[1] = { 0 };
-  if( FD_UNLIKELY( !fd_bn254_g1_frombytes_internal( p, in ) ) ) {
+  if( FD_UNLIKELY( !fd_bn254_g1_frombytes_internal( p, in, big_endian ) ) ) {
     return NULL;
   }
   int is_inf   = fd_bn254_g1_is_zero( p );
-  int flag_inf = in[32] & FLAG_INF;
+  int flag_inf = in[ big_endian ? 32 : 63 ] & FLAG_INF;
 
   /* Serialize compressed point:
      https://github.com/arkworks-rs/algebra/blob/v0.4.2/ec/src/models/short_weierstrass/mod.rs#L122
@@ -32,7 +33,7 @@ fd_bn254_g1_compress( uchar       out[32],
   }
 
   int is_neg = fd_bn254_fp_is_neg_nm( &p->Y );
-  memmove( out, in, 32 );
+  fd_bn254_fp_tobytes_nm( out, &p->X, 1 /* compressed form is always big endian */ );
   if( is_neg ) {
     out[0] = (uchar)( out[0] | FLAG_NEG );
   }
@@ -41,16 +42,17 @@ fd_bn254_g1_compress( uchar       out[32],
 
 uchar *
 fd_bn254_g1_decompress( uchar       out[64],
-                        uchar const in [32] ) {
+                        uchar const in [32],
+                        int         big_endian ) {
   /* Special case: all zeros in => all zeros out, no flags */
   const uchar zero[32] = { 0 };
   if( fd_memeq( in, zero, 32 ) ) {
     return fd_memset( out, 0, 64UL );
   }
 
-  fd_bn254_fp_t x[1], x2[1], x3_plus_b[1], y[1];
+  fd_bn254_fp_t x_nm[1], x[1], x2[1], x3_plus_b[1], y[1];
   int is_inf, is_neg;
-  if( FD_UNLIKELY( !fd_bn254_fp_frombytes_be_nm( x, in, &is_inf, &is_neg ) ) ) {
+  if( FD_UNLIKELY( !fd_bn254_fp_frombytes_nm( x_nm, in, 1 /* always BE */, &is_inf, &is_neg ) ) ) {
     return NULL;
   }
 
@@ -65,7 +67,7 @@ fd_bn254_g1_decompress( uchar       out[64],
     return out;
   }
 
-  fd_bn254_fp_to_mont( x, x );
+  fd_bn254_fp_to_mont( x, x_nm );
   fd_bn254_fp_sqr( x2, x );
   fd_bn254_fp_mul( x3_plus_b, x2, x );
   fd_bn254_fp_add( x3_plus_b, x3_plus_b, fd_bn254_const_b_mont );
@@ -78,17 +80,18 @@ fd_bn254_g1_decompress( uchar       out[64],
     fd_bn254_fp_neg_nm( y, y );
   }
 
-  memmove( out, in, 32 ); out[0] &= FLAG_MASK;
-  fd_bn254_fp_tobytes_be_nm( &out[32], y );
+  fd_bn254_fp_tobytes_nm(  out,     x_nm, big_endian );
+  fd_bn254_fp_tobytes_nm( &out[32], y,    big_endian );
   /* no flags */
   return out;
 }
 
 uchar *
 fd_bn254_g2_compress( uchar       out[64],
-                      uchar const in[128] ) {
+                      uchar const in[128],
+                      int         big_endian ) {
   fd_bn254_g2_t p[1] = { 0 };
-  if( FD_UNLIKELY( !fd_bn254_g2_frombytes_internal( p, in ) ) ) {
+  if( FD_UNLIKELY( !fd_bn254_g2_frombytes_internal( p, in, big_endian ) ) ) {
     return NULL;
   }
   int is_inf   = fd_bn254_g2_is_zero( p );
@@ -106,7 +109,7 @@ fd_bn254_g2_compress( uchar       out[64],
   /* Serialize x coordinate. The flags are on the 2nd element.
      https://github.com/arkworks-rs/algebra/blob/v0.4.2/ff/src/fields/models/quadratic_extension.rs#L700-L702 */
   int is_neg = fd_bn254_fp2_is_neg_nm( &p->Y );
-  memmove( out, in, 64 );
+  fd_bn254_fp2_tobytes_nm( out, &p->X, 1 /* compressed form is always big endian */ );
   if( is_neg ) {
     out[0] = (uchar)( out[0] | FLAG_NEG );
   }
@@ -115,16 +118,17 @@ fd_bn254_g2_compress( uchar       out[64],
 
 uchar *
 fd_bn254_g2_decompress( uchar       out[128],
-                        uchar const in  [64] ) {
+                        uchar const in  [64],
+                        int         big_endian ) {
   /* Special case: all zeros in => all zeros out, no flags */
   const uchar zero[64] = { 0 };
   if( fd_memeq( in, zero, 64 ) ) {
     return fd_memset( out, 0, 128UL );
   }
 
-  fd_bn254_fp2_t x[1], x2[1], x3_plus_b[1], y[1];
+  fd_bn254_fp2_t x_nm[1], x[1], x2[1], x3_plus_b[1], y[1];
   int is_inf, is_neg;
-  if( FD_UNLIKELY( !fd_bn254_fp2_frombytes_be_nm( x, in, &is_inf, &is_neg ) ) ) {
+  if( FD_UNLIKELY( !fd_bn254_fp2_frombytes_nm( x_nm, in, 1, &is_inf, &is_neg ) ) ) {
     return NULL;
   }
 
@@ -138,7 +142,7 @@ fd_bn254_g2_decompress( uchar       out[128],
     return out;
   }
 
-  fd_bn254_fp2_to_mont( x, x );
+  fd_bn254_fp2_to_mont( x, x_nm );
   fd_bn254_fp2_sqr( x2, x );
   fd_bn254_fp2_mul( x3_plus_b, x2, x );
   fd_bn254_fp2_add( x3_plus_b, x3_plus_b, fd_bn254_const_twist_b_mont );
@@ -151,8 +155,8 @@ fd_bn254_g2_decompress( uchar       out[128],
     fd_bn254_fp2_neg_nm( y, y );
   }
 
-  memmove( out, in, 64 ); out[0] &= FLAG_MASK;
-  fd_bn254_fp2_tobytes_be_nm( &out[64], y );
+  fd_bn254_fp2_tobytes_nm(  out,     x_nm, big_endian );
+  fd_bn254_fp2_tobytes_nm( &out[64], y,    big_endian );
   /* no flags */
   return out;
 }
@@ -162,40 +166,43 @@ fd_bn254_g2_decompress( uchar       out[128],
 int
 fd_bn254_g1_add_syscall( uchar       out[64],
                          uchar const in[],
-                         ulong       in_sz ) {
-  /* Expected 128-byte input (2 points). Pad input with 0s. */
+                         ulong       in_sz,
+                         int         big_endian ) {
+  /* Expected 128-byte input (2 points). Pad input with 0s (only big endian). */
   if( FD_UNLIKELY( in_sz > 128UL ) ) {
     return -1;
   }
+  if( FD_UNLIKELY( !big_endian && in_sz != 128UL ) ) {
+    return -1;
+  }
   uchar FD_ALIGNED buf[128] = { 0 };
   fd_memcpy( buf, in, in_sz );
 
   /* Validate inputs */
   fd_bn254_g1_t r[1], a[1], b[1];
-  if( FD_UNLIKELY( !fd_bn254_g1_frombytes_check_subgroup( a, &buf[ 0] ) ) ) {
+  if( FD_UNLIKELY( !fd_bn254_g1_frombytes_check_subgroup( a, &buf[ 0], big_endian ) ) ) {
     return -1;
   }
-  if( FD_UNLIKELY( !fd_bn254_g1_frombytes_check_subgroup( b, &buf[64] ) ) ) {
+  if( FD_UNLIKELY( !fd_bn254_g1_frombytes_check_subgroup( b, &buf[64], big_endian ) ) ) {
     return -1;
   }
 
   /* Compute point add and serialize result */
   fd_bn254_g1_affine_add( r, a, b );
-  fd_bn254_g1_tobytes( out, r );
+  fd_bn254_g1_tobytes( out, r, big_endian );
   return 0;
 }
 
 int
 fd_bn254_g1_scalar_mul_syscall( uchar       out[64],
                                 uchar const in[],
                                 ulong       in_sz,
-                                int         check_correct_sz ) {
-  /* Expected 96-byte input (1 point + 1 scalar). Pad input with 0s.
-     Note: Agave checks for 128 bytes instead of 96. We have to do the same check.
-     https://github.com/anza-xyz/agave/blob/v1.18.6/sdk/program/src/alt_bn128/mod.rs#L17
-     Update: https://github.com/anza-xyz/agave/blob/d2df66d3/programs/bpf_loader/src/syscalls/mod.rs#L1654-L1658 */
-  ulong check_sz = check_correct_sz ? 96UL : 128UL;
-  if( FD_UNLIKELY( in_sz > check_sz ) ) {
+                                int         big_endian ) {
+  /* Expected 96-byte input (1 point + 1 scalar). Pad input with 0s (only big endian). */
+  if( FD_UNLIKELY( in_sz > 96UL ) ) {
+    return -1;
+  }
+  if( FD_UNLIKELY( !big_endian && in_sz != 96UL ) ) {
     return -1;
   }
   uchar FD_ALIGNED buf[96] = { 0 };
@@ -204,25 +211,30 @@ fd_bn254_g1_scalar_mul_syscall( uchar       out[64],
   /* Validate inputs */
   fd_bn254_g1_t r[1], a[1];
   fd_bn254_scalar_t s[1];
-  if( FD_UNLIKELY( !fd_bn254_g1_frombytes_check_subgroup( a, &buf[ 0] ) ) ) {
+  if( FD_UNLIKELY( !fd_bn254_g1_frombytes_check_subgroup( a, &buf[ 0], big_endian ) ) ) {
     return -1;
   }
 
   /* Scalar is big endian and NOT validated
      https://github.com/anza-xyz/agave/blob/v1.18.6/sdk/program/src/alt_bn128/mod.rs#L211-L214 */
-  fd_uint256_bswap( s, fd_type_pun_const( &buf[64] ) ); /* &buf[64] is always FD_ALIGNED */
+  if( FD_BIG_ENDIAN_LIKELY( big_endian ) ) {
+    fd_uint256_bswap( s, fd_type_pun_const( &buf[64] ) ); /* &buf[64] is always FD_ALIGNED */
+  } else {
+    memcpy( s, &buf[64], 32 );
+  }
   // no: if( FD_UNLIKELY( !fd_bn254_scalar_validate( s ) ) ) return -1;
 
   /* Compute scalar mul and serialize result */
   fd_bn254_g1_scalar_mul( r, a, s );
-  fd_bn254_g1_tobytes( out, r );
+  fd_bn254_g1_tobytes( out, r, big_endian );
   return 0;
 }
 
 int
 fd_bn254_pairing_is_one_syscall( uchar       out[32],
                                  uchar const in[],
                                  ulong       in_sz,
+                                 int         big_endian,
                                  int         check_len ) {
   /* https://github.com/anza-xyz/agave/blob/v1.18.6/sdk/program/src/alt_bn128/mod.rs#L244
      Note: Solana had a bug where it checked if input.len().checked_rem(192).is_none(),
@@ -245,11 +257,11 @@ fd_bn254_pairing_is_one_syscall( uchar       out[32],
   ulong sz=0;
   for( ulong i=0; i<elements_len; i++ ) {
     /* G1: deserialize and check subgroup membership */
-    if( FD_UNLIKELY( !fd_bn254_g1_frombytes_check_subgroup( &p[sz], &in[i*192   ] ) ) ) {
+    if( FD_UNLIKELY( !fd_bn254_g1_frombytes_check_subgroup( &p[sz], &in[i*192   ], big_endian ) ) ) {
       return -1;
     }
     /* G2: deserialize and check subgroup membership */
-    if( FD_UNLIKELY( !fd_bn254_g2_frombytes_check_subgroup( &q[sz], &in[i*192+64] ) ) ) {
+    if( FD_UNLIKELY( !fd_bn254_g2_frombytes_check_subgroup( &q[sz], &in[i*192+64], big_endian ) ) ) {
       return -1;
     }
     /* Skip any pair where either P or Q is the point at infinity */
@@ -278,7 +290,7 @@ fd_bn254_pairing_is_one_syscall( uchar       out[32],
   /* Output is 0 or 1, serialized as big endian uint256. */
   fd_memset( out, 0, 32 );
   if( FD_LIKELY( fd_bn254_fp12_is_one( r ) ) ) {
-    out[31] = 1;
+    out[ big_endian ? 31 : 0 ] = 1;
   }
   return 0;
 }

src/ballet/bn254/fd_bn254.h

@@ -8,24 +8,27 @@
 #include "./fd_bn254_scalar.h"
 
 #define FD_BN254_PAIRING_BATCH_MAX 16UL
+#define FD_BIG_ENDIAN_LIKELY FD_LIKELY
 
 FD_PROTOTYPES_BEGIN
 
 int
 fd_bn254_g1_add_syscall( uchar       out[64],
                          uchar const in[],
-                         ulong       in_sz );
+                         ulong       in_sz,
+                         int         big_endian );
 
 int
 fd_bn254_g1_scalar_mul_syscall( uchar       out[64],
                                 uchar const in[],
                                 ulong       in_sz,
-                                int         check_correct_sz );
+                                int         big_endian );
 
 int
 fd_bn254_pairing_is_one_syscall( uchar       out[32],
                                  uchar const in[],
                                  ulong       in_sz,
+                                 int         big_endian,
                                  int         check_len );
 
 /* fd_bn254_g1_compress compresses a point in G1.
@@ -37,7 +40,8 @@ fd_bn254_pairing_is_one_syscall( uchar       out[32],
    Note: this function does NOT check that (x, y) is in G1. */
 uchar *
 fd_bn254_g1_compress( uchar       out[32],
-                      uchar const in [64] );
+                      uchar const in [64],
+                      int         big_endian );
 
 /* fd_bn254_g1_decompress decompresses a point in G1.
    Input in is a 32-byte big endian buffer representing the x coord of a point,
@@ -48,7 +52,8 @@ fd_bn254_g1_compress( uchar       out[32],
    (Success implies that (x, y) is in G1.) */
 uchar *
 fd_bn254_g1_decompress( uchar       out[64],
-                        uchar const in [32] );
+                        uchar const in [32],
+                        int         big_endian );
 
 /* fd_bn254_g2_compress compresses a point in G2.
    Same as fd_bn254_g1_compress, but x, y are in Fp2, so twice as long.
@@ -60,7 +65,8 @@ fd_bn254_g1_decompress( uchar       out[64],
    Note: this function does NOT check that (x, y) is in G2. */
 uchar *
 fd_bn254_g2_compress( uchar       out[64],
-                      uchar const in[128] );
+                      uchar const in[128],
+                      int         big_endian );
 
 /* fd_bn254_g2_decompress decompresses a point in G2.
    Same as fd_bn254_g1_decompress, but x, y are in Fp2, so twice as long.
@@ -73,7 +79,8 @@ fd_bn254_g2_compress( uchar       out[64],
    imply that). */
 uchar *
 fd_bn254_g2_decompress( uchar       out[128],
-                        uchar const in  [64] );
+                        uchar const in  [64],
+                        int         big_endian );
 
 FD_PROTOTYPES_END
 

src/ballet/bn254/fd_bn254_field.c

@@ -65,14 +65,15 @@ fd_bn254_fp_is_neg_nm( fd_bn254_fp_t * x ) {
 }
 
 static inline fd_bn254_fp_t *
-fd_bn254_fp_frombytes_be_nm( fd_bn254_fp_t * r,
-                             uchar const     buf[32],
-                             int *           is_inf,
-                             int *           is_neg ) {
+fd_bn254_fp_frombytes_nm( fd_bn254_fp_t * r,
+                          uchar const     buf[32],
+                          int             big_endian,
+                          int *           is_inf,
+                          int *           is_neg ) {
   /* Flags (optional) */
   if( is_inf != NULL /* && is_neg != NULL */ ) {
-    *is_inf = !!(buf[0] & FLAG_INF);
-    *is_neg = !!(buf[0] & FLAG_NEG);
+    *is_inf = !!(buf[ big_endian ? 0 : 31 ] & FLAG_INF);
+    *is_neg = !!(buf[ big_endian ? 0 : 31 ] & FLAG_NEG);
     /* If both flags are set (bit 6, 7), return error.
        https://github.com/arkworks-rs/algebra/blob/v0.4.2/ec/src/models/short_weierstrass/serialization_flags.rs#L75 */
     if( FD_UNLIKELY( *is_inf && *is_neg ) ) {
@@ -81,9 +82,11 @@ fd_bn254_fp_frombytes_be_nm( fd_bn254_fp_t * r,
   }
 
   fd_memcpy( r, buf, 32 );
-  fd_uint256_bswap( r, r );
+  if( FD_BIG_ENDIAN_LIKELY( big_endian ) ) {
+    fd_uint256_bswap( r, r );
+  }
   if( is_inf != NULL ) {
-    r->buf[31] &= FLAG_MASK;
+    r->buf[ 31 ] &= FLAG_MASK;
   }
 
   /* Field element */
@@ -94,9 +97,12 @@ fd_bn254_fp_frombytes_be_nm( fd_bn254_fp_t * r,
 }
 
 static inline uchar *
-fd_bn254_fp_tobytes_be_nm( uchar           buf[32],
-                           fd_bn254_fp_t * a ) {
-  fd_uint256_bswap( a, a );
+fd_bn254_fp_tobytes_nm( uchar           buf[32],
+                        fd_bn254_fp_t * a,
+                        int             big_endian ) {
+  if( FD_BIG_ENDIAN_LIKELY( big_endian ) ) {
+    fd_uint256_bswap( a, a );
+  }
   fd_memcpy( buf, a, 32 );
   return buf;
 }