crypto: fast fe_batch_invert using Montgomery's trick

diff --git a/src/crypto/crypto-ops.c b/src/crypto/crypto-ops.c
index eca5c47..210bb80 100644
--- a/src/crypto/crypto-ops.c
+++ b/src/crypto/crypto-ops.c
@@ -29,6 +29,7 @@
 // Parts of this file are originally copyright (c) 2012-2013 The Cryptonote developers
 
 #include <assert.h>
+#include <sodium/crypto_verify_32.h>
 #include <stdint.h>
 
 #include "warnings.h"
@@ -313,6 +314,61 @@ void fe_invert(fe out, const fe z) {
   return;
 }
 
+// return 1 if a == b, else 0
+int fe_equals(const fe a, const fe b) {
+  unsigned char a_bytes[32];
+  unsigned char b_bytes[32];
+  fe_tobytes(a_bytes, a);
+  fe_tobytes(b_bytes, b);
+  return crypto_verify_32(a_bytes, b_bytes) == 0;
+}
+
+/* From fe_isnonzero.c, modified */
+
+static int fe_isnonzero(const fe f) {
+  unsigned char s[32];
+  fe_tobytes(s, f);
+  return (((int) (s[0] | s[1] | s[2] | s[3] | s[4] | s[5] | s[6] | s[7] | s[8] |
+    s[9] | s[10] | s[11] | s[12] | s[13] | s[14] | s[15] | s[16] | s[17] |
+    s[18] | s[19] | s[20] | s[21] | s[22] | s[23] | s[24] | s[25] | s[26] |
+    s[27] | s[28] | s[29] | s[30] | s[31]) - 1) >> 8) + 1;
+}
+
+// Montgomery's trick
+// https://iacr.org/archive/pkc2004/29470042/29470042.pdf 2.2
+int fe_batch_invert(fe* __restrict out, const fe* __restrict in, const unsigned int n) {
+  if (n == 0) {
+    return 0;
+  }
+
+  assert(out);
+  assert(in);
+  assert(in != out); // also should not overlap
+
+  // Step 1: collect initial muls
+  fe_copy(out[0], in[0]);
+  for (unsigned int i = 1; i < n; ++i) {
+    fe_mul(out[i], out[i-1], in[i]);
+  }
+
+  // Step 2: get the inverse of all elems multiplied together
+  if (!fe_isnonzero(out[n-1])) {
+    // Don't divide by 0
+    return -1;
+  }
+  fe a;
+  fe_invert(a, out[n-1]);
+
+  // Step 3: get each inverse
+  for (unsigned int i = n; i > 1; --i) {
+    fe_mul(out[i-1], a, out[i-2]);
+    fe_mul(a, a, in[i-1]);
+  }
+  fe_copy(out[0], a);
+
+  return 0;
+}
+
 /* From fe_isnegative.c */
 
 /*
@@ -329,17 +385,6 @@ static int fe_isnegative(const fe f) {
   return s[0] & 1;
 }
 
-/* From fe_isnonzero.c, modified */
-
-static int fe_isnonzero(const fe f) {
-  unsigned char s[32];
-  fe_tobytes(s, f);
-  return (((int) (s[0] | s[1] | s[2] | s[3] | s[4] | s[5] | s[6] | s[7] | s[8] |
-    s[9] | s[10] | s[11] | s[12] | s[13] | s[14] | s[15] | s[16] | s[17] |
-    s[18] | s[19] | s[20] | s[21] | s[22] | s[23] | s[24] | s[25] | s[26] |
-    s[27] | s[28] | s[29] | s[30] | s[31]) - 1) >> 8) + 1;
-}
-
 /* From fe_mul.c */
 
 /*
@@ -1328,16 +1373,9 @@ void ge_double_scalarmult_base_vartime_p3(ge_p3 *r3, const unsigned char *a, con
   }
 }
 
-/* From ge_frombytes.c, modified */
-
-int ge_frombytes_vartime(ge_p3 *h, const unsigned char *s) {
-  fe u;
-  fe v;
-  fe vxx;
-  fe check;
-
-  /* From fe_frombytes.c */
+/* From fe_frombytes.c */
 
+int fe_frombytes_vartime(fe y, const unsigned char *s) {
   int64_t h0 = load_4(s);
   int64_t h1 = load_3(s + 4) << 6;
   int64_t h2 = load_3(s + 7) << 5;
@@ -1378,18 +1416,31 @@ int ge_frombytes_vartime(ge_p3 *h, const unsigned char *s) {
   carry6 = (h6 + (int64_t) (1<<25)) >> 26; h7 += carry6; h6 -= carry6 << 26;
   carry8 = (h8 + (int64_t) (1<<25)) >> 26; h9 += carry8; h8 -= carry8 << 26;
 
-  h->Y[0] = h0;
-  h->Y[1] = h1;
-  h->Y[2] = h2;
-  h->Y[3] = h3;
-  h->Y[4] = h4;
-  h->Y[5] = h5;
-  h->Y[6] = h6;
-  h->Y[7] = h7;
-  h->Y[8] = h8;
-  h->Y[9] = h9;
+  y[0] = h0;
+  y[1] = h1;
+  y[2] = h2;
+  y[3] = h3;
+  y[4] = h4;
+  y[5] = h5;
+  y[6] = h6;
+  y[7] = h7;
+  y[8] = h8;
+  y[9] = h9;
 
-  /* End fe_frombytes.c */
+  return 0;
+}
+
+/* From ge_frombytes.c, modified */
+
+int ge_frombytes_vartime(ge_p3 *h, const unsigned char *s) {
+  fe u;
+  fe v;
+  fe vxx;
+  fe check;
+
+  if (fe_frombytes_vartime(h->Y, s) != 0) {
+    return -1;
+  }
 
   fe_1(h->Z);
   fe_sq(u, h->Y);
diff --git a/src/crypto/crypto-ops.h b/src/crypto/crypto-ops.h
index dcb335d..cf5a684 100644
--- a/src/crypto/crypto-ops.h
+++ b/src/crypto/crypto-ops.h
@@ -88,6 +88,7 @@ void ge_double_scalarmult_base_vartime_p3(ge_p3 *, const unsigned char *, const
 
 extern const fe fe_sqrtm1;
 extern const fe fe_d;
+int fe_frombytes_vartime(fe, const unsigned char *);
 int ge_frombytes_vartime(ge_p3 *, const unsigned char *);
 
 /* From ge_p1p1_to_p2.c */
@@ -164,6 +165,17 @@ void ge_sub(ge_p1p1 *r, const ge_p3 *p, const ge_cached *q);
 void fe_add(fe h, const fe f, const fe g);
 void fe_tobytes(unsigned char *, const fe);
 void fe_invert(fe out, const fe z);
+int fe_equals(const fe a, const fe b);
+/**
+@brief: out[i] = 1/in[i] for i in [0, n). Uses Montgomery's trick
+@return: 0 on success, some other value otherwise
+
+Unlike other crypto functions, `out` and `in` memory sections CANNOT be aliased.
+If `out` and `in` overlap, it will cause undefined output.
+
+No 0 fe's are expected for `in`, otherwise fails.
+**/
+int fe_batch_invert(fe* __restrict out, const fe* __restrict in, const unsigned int n);
 void fe_mul(fe out, const fe, const fe);
 void fe_0(fe h);
 
diff --git a/tests/performance_tests/CMakeLists.txt b/tests/performance_tests/CMakeLists.txt
index 4c32e5e..8b5bf03 100644
--- a/tests/performance_tests/CMakeLists.txt
+++ b/tests/performance_tests/CMakeLists.txt
@@ -36,6 +36,7 @@ set(performance_tests_headers
   construct_tx.h
   derive_public_key.h
   derive_secret_key.h
+  fe_batch_invert.h
   ge_frombytes_vartime.h
   generate_key_derivation.h
   generate_key_image.h
diff --git a/tests/performance_tests/fe_batch_invert.h b/tests/performance_tests/fe_batch_invert.h
new file mode 100644
index 0000000..447abf2
--- /dev/null
+++ b/tests/performance_tests/fe_batch_invert.h
@@ -0,0 +1,79 @@
+// Copyright (c) 2025, The Monero Project
+// 
+// All rights reserved.
+// 
+// Redistribution and use in source and binary forms, with or without modification, are
+// permitted provided that the following conditions are met:
+// 
+// 1. Redistributions of source code must retain the above copyright notice, this list of
+//    conditions and the following disclaimer.
+// 
+// 2. Redistributions in binary form must reproduce the above copyright notice, this list
+//    of conditions and the following disclaimer in the documentation and/or other
+//    materials provided with the distribution.
+// 
+// 3. Neither the name of the copyright holder nor the names of its contributors may be
+//    used to endorse or promote products derived from this software without specific
+//    prior written permission.
+// 
+// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY
+// EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
+// MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL
+// THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
+// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
+// STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF
+// THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+// 
+// Parts of this file are originally copyright (c) 2012-2013 The Cryptonote developers
+
+#pragma once
+
+#include <memory>
+
+#include "crypto/crypto.h"
+
+template<bool batched>
+class test_fe_batch_invert
+{
+public:
+  static const size_t loop_count = 50;
+  static const size_t n_elems = 1000;
+
+  bool init()
+  {
+    m_fes = std::make_unique<fe[]>(n_elems);
+
+    for (std::size_t i = 0; i < n_elems; ++i)
+    {
+      crypto::secret_key r;
+      crypto::random32_unbiased((unsigned char*)r.data);
+
+      ge_p3 point;
+      ge_scalarmult_base(&point, (unsigned char*)r.data);
+
+      memcpy(m_fes[i], &point.Y, sizeof(fe));
+    }
+
+    return true;
+  }
+
+  bool test()
+  {
+    std::unique_ptr<fe[]> inv_fes = std::make_unique<fe[]>(n_elems);
+
+    if constexpr (batched)
+      fe_batch_invert(inv_fes.get(), m_fes.get(), n_elems);
+    else
+    {
+      for (std::size_t i = 0; i < n_elems; ++i)
+        fe_invert(inv_fes[i], m_fes[i]);
+    }
+
+    return true;
+  }
+
+private:
+  std::unique_ptr<fe[]> m_fes;
+};
diff --git a/tests/performance_tests/main.cpp b/tests/performance_tests/main.cpp
index c40656e..99fe1bb 100644
--- a/tests/performance_tests/main.cpp
+++ b/tests/performance_tests/main.cpp
@@ -43,6 +43,7 @@
 #include "derive_public_key.h"
 #include "derive_secret_key.h"
 #include "derive_view_tag.h"
+#include "fe_batch_invert.h"
 #include "ge_frombytes_vartime.h"
 #include "ge_tobytes.h"
 #include "generate_key_derivation.h"
@@ -198,6 +199,8 @@ int main(int argc, char** argv)
   TEST_PERFORMANCE0(filter, p, test_generate_key_image);
   TEST_PERFORMANCE0(filter, p, test_derive_public_key);
   TEST_PERFORMANCE0(filter, p, test_derive_secret_key);
+  TEST_PERFORMANCE1(filter, p, test_fe_batch_invert, true); // batched
+  TEST_PERFORMANCE1(filter, p, test_fe_batch_invert, false); // individual inversions
   TEST_PERFORMANCE0(filter, p, test_ge_frombytes_vartime);
   TEST_PERFORMANCE0(filter, p, test_ge_tobytes);
   TEST_PERFORMANCE0(filter, p, test_generate_keypair);
diff --git a/tests/unit_tests/crypto.cpp b/tests/unit_tests/crypto.cpp
index 42500dd..8583185 100644
--- a/tests/unit_tests/crypto.cpp
+++ b/tests/unit_tests/crypto.cpp
@@ -70,6 +70,14 @@ namespace
     out << "BEGIN" << value << "END";  
     return out.str() == "BEGIN<" + std::string{expected, sizeof(T) * 2} + ">END";
   }
+
+  void random_fe(fe rand_fe)
+  {
+    unsigned char s[32];
+    crypto::random32_unbiased(s);
+    if (fe_frombytes_vartime(rand_fe, s) != 0)
+      throw std::runtime_error("invalid random fe");
+  }
 }
 
 TEST(Crypto, Ostream)
@@ -343,3 +351,81 @@ TEST(Crypto, generator_consistency)
   // ringct/rctTypes.h
   ASSERT_TRUE(memcmp(H.data, rct::H.bytes, 32) == 0);
 }
+
+TEST(Crypto, batch_inversion)
+{
+  const std::size_t MAX_TEST_ELEMS = 1000;
+
+  std::unique_ptr<fe[]> init_elems = std::make_unique<fe[]>(MAX_TEST_ELEMS);
+  std::unique_ptr<fe[]> norm_inverted = std::make_unique<fe[]>(MAX_TEST_ELEMS);
+
+  // Init test elems and individual inversions
+  for (std::size_t i = 0; i < MAX_TEST_ELEMS; ++i)
+  {
+    random_fe(init_elems[i]);
+    fe_invert(norm_inverted[i], init_elems[i]);
+  }
+
+  // Do batch inversions and compare to individual inversions
+  for (std::size_t n_elems = 1; n_elems <= MAX_TEST_ELEMS; ++n_elems)
+  {
+    std::unique_ptr<fe[]> batch_inverted = std::make_unique<fe[]>(n_elems);
+    ASSERT_EQ(fe_batch_invert(batch_inverted.get(), init_elems.get(), n_elems), 0);
+    for (std::size_t i = 0; i < n_elems; ++i)
+      ASSERT_EQ(fe_equals(batch_inverted[i], norm_inverted[i]), 1);
+  }
+}
+
+TEST(Crypto, batch_inversion_touching)
+{
+  // vals[0] and vals[1] are input, vals[2] and vals[3] are output
+  fe vals[4];
+
+  // Init input elems
+  for (std::size_t i = 0; i < 2; ++i)
+    random_fe(vals[i]);
+
+  fe_batch_invert(vals + 2, vals, 2);
+
+  // Do batch inversions and compare to individual inversions
+  for (std::size_t i = 0; i < 2; ++i)
+  {
+    fe inv_simple;
+    fe_invert(inv_simple, vals[i]);
+    ASSERT_EQ(1, fe_equals(inv_simple, vals[2 + i]));
+  }
+}
+
+TEST(Crypto, batch_invert_zero)
+{
+  const std::size_t TEST_ELEMS = 2;
+  std::unique_ptr<fe[]> init_elems = std::make_unique<fe[]>(TEST_ELEMS);
+
+  // Init test elems
+  random_fe(init_elems[0]);
+  fe_0(init_elems[1]);
+
+  // Do the batch inversion, should fail
+  std::unique_ptr<fe[]> batch_inverted = std::make_unique<fe[]>(TEST_ELEMS);
+  ASSERT_EQ(fe_batch_invert(batch_inverted.get(), init_elems.get(), TEST_ELEMS), -1);
+}
+
+TEST(Crypto, fe_equals)
+{
+  // Test equality
+  ASSERT_EQ(fe_equals(fe_d, fe_d), 1);
+
+  // Test inequality
+  fe fe_d2;
+  fe_add(fe_d2, fe_d, fe_d);
+  ASSERT_EQ(fe_equals(fe_d2, fe_d), 0);
+
+  // Test different fe reprs that are actually equal
+  unsigned char fe_d2_bytes[32];
+  fe fe_d2_reduced;
+  fe_tobytes(fe_d2_bytes, fe_d2);
+  fe_frombytes_vartime(fe_d2_reduced, fe_d2_bytes);
+  // We expect distinct fe_d2_reduced and fe_d2 reprs, since fe_add produces fe repr elems in a larger subdomain.
+  ASSERT_NE(memcmp(fe_d2_reduced, fe_d2, sizeof(fe)), 0);
+  ASSERT_EQ(fe_equals(fe_d2_reduced, fe_d2), 1);
+}