From a05f753d2695e047c3eff080fde80a1a9a30455d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Maciej=20Zwoli=C5=84ski?= <mac.zwolinski@gmail.com>
Date: Mon, 24 Oct 2022 15:52:54 +0200
Subject: [PATCH] perf: optimize Record::is_owner

Signed-off-by: Maciej Zwolinski <mac.zwolinski@gmail.com>
---
 console/program/src/data/record/is_owner.rs   |  8 +++++---
 .../short_weierstrass_jacobian/affine.rs      | 17 ++++++++++++++++
 .../twisted_edwards_extended/affine.rs        | 20 +++++++++++++++++++
 curves/src/traits/group.rs                    |  6 ++++++
 4 files changed, 48 insertions(+), 3 deletions(-)

diff --git a/console/program/src/data/record/is_owner.rs b/console/program/src/data/record/is_owner.rs
index 4e10c547a4..fc3136e774 100644
--- a/console/program/src/data/record/is_owner.rs
+++ b/console/program/src/data/record/is_owner.rs
@@ -29,9 +29,11 @@ impl<N: Network> Record<N, Ciphertext<N>> {
                 // Compute the 0th randomizer.
                 let randomizer = N::hash_many_psd8(&[N::encryption_domain(), record_view_key], 1);
                 // Decrypt the owner.
-                match Address::from_field(&(ciphertext[0] - randomizer[0])) {
-                    Ok(owner) => &owner == address,
-                    Err(_) => false,
+                let x_coordinate = ciphertext[0] - randomizer[0];
+                if let Some((point1, point2)) = N::Affine::from_x_coordinate_variants(*x_coordinate) {
+                    point1 == ***address || point2 == ***address
+                } else {
+                    false
                 }
             }
         }
diff --git a/curves/src/templates/short_weierstrass_jacobian/affine.rs b/curves/src/templates/short_weierstrass_jacobian/affine.rs
index 6bab36023f..99586f47a9 100644
--- a/curves/src/templates/short_weierstrass_jacobian/affine.rs
+++ b/curves/src/templates/short_weierstrass_jacobian/affine.rs
@@ -142,6 +142,23 @@ impl<P: Parameters> AffineCurve for Affine<P> {
         })
     }
 
+    /// Attempts to construct an affine point given an x-coordinate. The
+    /// point is not guaranteed to be in the prime order subgroup.
+    /// Returns variants with and without the lexicographically largest
+    /// y-coordinate selected.
+    fn from_x_coordinate_variants(x: Self::BaseField) -> Option<(Self, Self)> {
+        // Compute x^3 + ax + b
+        let x3b = P::add_b(&((x.square() * x) + P::mul_by_a(&x)));
+
+        x3b.sqrt().map(|y| {
+            let negy = -y;
+
+            let y1 = if (y < negy) ^ false { y } else { negy };
+            let y2 = if (y < negy) ^ true { y } else { negy };
+            (Self::new(x, y1, false), Self::new(x, y2, false))
+        })
+    }
+
     /// Attempts to construct an affine point given a y-coordinate. The
     /// point is not guaranteed to be in the prime order subgroup.
     ///
diff --git a/curves/src/templates/twisted_edwards_extended/affine.rs b/curves/src/templates/twisted_edwards_extended/affine.rs
index db4b9e2c79..8376fdd9a5 100644
--- a/curves/src/templates/twisted_edwards_extended/affine.rs
+++ b/curves/src/templates/twisted_edwards_extended/affine.rs
@@ -142,6 +142,26 @@ impl<P: Parameters> AffineCurve for Affine<P> {
         })
     }
 
+    /// Attempts to construct an affine point given an x-coordinate. The
+    /// point is not guaranteed to be in the prime order subgroup.
+    /// Returns variants with and without the lexicographically largest
+    /// y-coordinate selected.
+    #[inline]
+    fn from_x_coordinate_variants(x: Self::BaseField) -> Option<(Self, Self)> {
+        // y = sqrt( (a * x^2 - 1)  / (d * x^2 - 1) )
+        let x2 = x.square();
+        let one = Self::BaseField::one();
+        let numerator = P::mul_by_a(&x2) - one;
+        let denominator = P::EDWARDS_D * x2 - one;
+        let y2 = denominator.inverse().map(|denom| denom * numerator);
+        y2.and_then(|y2| y2.sqrt()).map(|y| {
+            let negy = -y;
+            let y1 = if (y < negy) ^ false { y } else { negy };
+            let y2 = if (y < negy) ^ true { y } else { negy };
+            (Self::new(x, y1, x * y1), Self::new(x, y2, x * y2))
+        })
+    }
+
     /// Attempts to construct an affine point given a y-coordinate. The
     /// point is not guaranteed to be in the prime order subgroup.
     ///
diff --git a/curves/src/traits/group.rs b/curves/src/traits/group.rs
index 1a936271e5..c4863578ef 100644
--- a/curves/src/traits/group.rs
+++ b/curves/src/traits/group.rs
@@ -163,6 +163,12 @@ pub trait AffineCurve:
     /// largest y-coordinate be selected.
     fn from_x_coordinate(x: Self::BaseField, greatest: bool) -> Option<Self>;
 
+    /// Attempts to construct an affine point given an x-coordinate. The
+    /// point is not guaranteed to be in the prime order subgroup.
+    /// Returns variants with and without the lexicographically largest
+    /// y-coordinate selected.
+    fn from_x_coordinate_variants(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.
     ///