Merge pull request #2277 from bishopcheckmate/perf/fast-fail-group-from-x-coord

perf: optimize console's group from_x_coordinate

Author: howardwu

.github/workflows/benchmarks.yml

@@ -62,6 +62,12 @@ jobs:
           cargo bench --bench merkle_tree -- --output-format bencher | tee -a ../../output.txt
           cd ../..
           
+      - name: Benchmark console/types
+        run: |
+          cd console/types
+          cargo bench --bench group -- --output-format bencher | tee -a ../../output.txt
+          cd ../..
+
       - name: Benchmark curves
         run: |
           cd curves

Cargo.lock

@@ -6,6 +6,11 @@ description = "Console types for a decentralized virtual machine"
 license = "Apache-2.0"
 edition = "2021"
 
+[[bench]]
+name = "group"
+path = "benches/group.rs"
+harness = false
+
 [dependencies.snarkvm-console-network-environment]
 path = "../network/environment"
 version = "=0.16.16"
@@ -45,6 +50,13 @@ path = "./string"
 version = "=0.16.16"
 optional = true
 
+[dev-dependencies.criterion]
+version = "0.5.1"
+
+[dev-dependencies.snarkvm-console-network]
+path = "../network"
+version = "=0.16.16"
+
 [features]
 default = [
   "address",

console/types/Cargo.toml

@@ -0,0 +1,84 @@
+// Copyright (C) 2019-2023 Aleo Systems Inc.
+// This file is part of the snarkVM library.
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at:
+// http://www.apache.org/licenses/LICENSE-2.0
+
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#[macro_use]
+extern crate criterion;
+
+use criterion::Criterion;
+use snarkvm_console_network::{environment::prelude::*, Testnet3};
+use snarkvm_console_types::{Field, Group};
+
+type CurrentNetwork = Testnet3;
+
+fn group_from_field(c: &mut Criterion) {
+    let rng = &mut TestRng::default();
+
+    c.bench_function("group_from_field", move |b| {
+        let fields: Vec<Field<CurrentNetwork>> = (0..1000).map(|_| rng.gen()).collect();
+
+        b.iter(|| {
+            for field in &fields {
+                let group = Group::from_field(field);
+                let _ = std::hint::black_box(group);
+            }
+        })
+    });
+}
+
+fn group_from_field_on_curve(c: &mut Criterion) {
+    let rng = &mut TestRng::default();
+
+    c.bench_function("group_from_field_on_curve", move |b| {
+        type Projective = <CurrentNetwork as Environment>::Projective;
+
+        let fields: Vec<Field<CurrentNetwork>> =
+            (0..1000).map(|_| rng.gen::<Projective>().to_affine().to_x_coordinate()).map(Field::new).collect();
+
+        b.iter(|| {
+            for field in &fields {
+                let group = Group::from_field(field);
+                let _ = std::hint::black_box(group);
+            }
+        })
+    });
+}
+
+fn group_from_field_off_curve(c: &mut Criterion) {
+    let rng = &mut TestRng::default();
+
+    c.bench_function("group_from_field_off_curve", move |b| {
+        type Affine = <CurrentNetwork as Environment>::Affine;
+
+        let fields: Vec<_> = std::iter::repeat(())
+            .map(|_| rng.gen::<Field<CurrentNetwork>>())
+            .filter(|&field| Affine::from_x_coordinate(*field, true).is_none())
+            .take(1000)
+            .collect();
+
+        b.iter(|| {
+            for field in &fields {
+                let group = Group::from_field(field);
+                let _ = std::hint::black_box(group);
+            }
+        })
+    });
+}
+
+criterion_group! {
+    name = group;
+    config = Criterion::default().sample_size(20);
+    targets = group_from_field, group_from_field_on_curve, group_from_field_off_curve
+}
+
+criterion_main!(group);

console/types/benches/group.rs

@@ -18,14 +18,11 @@ impl<E: Environment> Group<E> {
     /// Attempts to recover an affine group element from a given x-coordinate field element.
     /// For safety, the resulting point is always enforced to be on the curve and in the correct subgroup.
     pub fn from_x_coordinate(x_coordinate: Field<E>) -> Result<Self> {
-        if let Some(point) = E::Affine::from_x_coordinate(*x_coordinate, true) {
-            if point.is_in_correct_subgroup_assuming_on_curve() {
-                return Ok(Self::new(point));
-            }
-        }
-        if let Some(point) = E::Affine::from_x_coordinate(*x_coordinate, false) {
-            if point.is_in_correct_subgroup_assuming_on_curve() {
-                return Ok(Self::new(point));
+        if let Some((p1, p2)) = E::Affine::pair_from_x_coordinate(*x_coordinate) {
+            for point in [p2, p1] {
+                if point.is_in_correct_subgroup_assuming_on_curve() {
+                    return Ok(Self::new(point));
+                }
             }
         }
         bail!("Failed to recover an affine group from an x-coordinate of {x_coordinate}")

console/types/group/src/from_x_coordinate.rs

@@ -148,6 +148,18 @@ impl<P: Parameters> AffineCurve for Affine<P> {
         })
     }
 
+    /// Attempts to construct both possible affine points given an x-coordinate.
+    /// Points are not guaranteed to be in the prime order subgroup.
+    ///
+    /// The affine points returned should be in lexicographically growing order.
+    ///
+    /// Calling this should be equivalent (but likely more performant) to
+    /// `(AffineCurve::from_x_coordinate(x, false), AffineCurve::from_x_coordinate(x, true))`.
+    #[inline]
+    fn pair_from_x_coordinate(x: Self::BaseField) -> Option<(Self, Self)> {
+        Self::from_x_coordinate(x, false).map(|p1| (p1, Self::new(p1.x, -p1.y, false)))
+    }
+
     /// Attempts to construct an affine point given a y-coordinate. The
     /// point is not guaranteed to be in the prime order subgroup.
     ///

curves/src/templates/short_weierstrass_jacobian/affine.rs

@@ -29,6 +29,7 @@ pub const ITERATIONS: usize = 10;
 pub fn sw_tests<P: ShortWeierstrassParameters>(rng: &mut TestRng) {
     sw_curve_serialization_test::<P>(rng);
     sw_from_random_bytes::<P>(rng);
+    sw_from_x_coordinate::<P>(rng);
 }
 
 pub fn sw_curve_serialization_test<P: ShortWeierstrassParameters>(rng: &mut TestRng) {
@@ -138,3 +139,23 @@ pub fn sw_from_random_bytes<P: ShortWeierstrassParameters>(rng: &mut TestRng) {
         }
     }
 }
+
+pub fn sw_from_x_coordinate<P: ShortWeierstrassParameters>(rng: &mut TestRng) {
+    for _ in 0..ITERATIONS {
+        let a = Projective::<P>::rand(rng);
+        let a = a.to_affine();
+        {
+            let x = a.x;
+
+            let a1 = Affine::<P>::from_x_coordinate(x, true).unwrap();
+            let a2 = Affine::<P>::from_x_coordinate(x, false).unwrap();
+
+            assert!(a == a1 || a == a2);
+
+            let (b2, b1) = Affine::<P>::pair_from_x_coordinate(x).unwrap();
+
+            assert_eq!(a1, b1);
+            assert_eq!(a2, b2);
+        }
+    }
+}

curves/src/templates/short_weierstrass_jacobian/tests.rs

@@ -148,6 +148,18 @@ impl<P: Parameters> AffineCurve for Affine<P> {
         })
     }
 
+    /// Attempts to construct both possible affine points given an x-coordinate.
+    /// Points are not guaranteed to be in the prime order subgroup.
+    ///
+    /// The affine points returned should be in lexicographically growing order.
+    ///
+    /// Calling this should be equivalent (but likely more performant) to
+    /// `(AffineCurve::from_x_coordinate(x, false), AffineCurve::from_x_coordinate(x, true))`.
+    #[inline]
+    fn pair_from_x_coordinate(x: Self::BaseField) -> Option<(Self, Self)> {
+        Self::from_x_coordinate(x, false).map(|p1| (p1, Self::new(p1.x, -p1.y, -p1.t)))
+    }
+
     /// Attempts to construct an affine point given a y-coordinate. The
     /// point is not guaranteed to be in the prime order subgroup.
     ///

curves/src/templates/twisted_edwards_extended/affine.rs

@@ -201,6 +201,11 @@ where
             let a2 = Affine::<P>::from_x_coordinate(x, false).unwrap();
 
             assert!(a == a1 || a == a2);
+
+            let (b2, b1) = Affine::<P>::pair_from_x_coordinate(x).unwrap();
+
+            assert_eq!(a1, b1);
+            assert_eq!(a2, b2);
         }
         {
             let y = a.y;

curves/src/templates/twisted_edwards_extended/tests.rs

@@ -164,6 +164,15 @@ pub trait AffineCurve:
     /// largest y-coordinate be selected.
     fn from_x_coordinate(x: Self::BaseField, greatest: bool) -> Option<Self>;
 
+    /// Attempts to construct both possible affine points given an x-coordinate.
+    /// Points are not guaranteed to be in the prime order subgroup.
+    ///
+    /// The affine points returned should be in lexicographically growing order.
+    ///
+    /// Calling this should be equivalent (but likely more performant) to
+    /// `(AffineCurve::from_x_coordinate(x, false), AffineCurve::from_x_coordinate(x, true))`.
+    fn pair_from_x_coordinate(x: Self::BaseField) -> Option<(Self, Self)>;
+
     /// Attempts to construct an affine point given a y-coordinate. The
     /// point is not guaranteed to be in the prime order subgroup.
     ///