Accelerate odd-degree polynomial square-root recovery

libs/math/include/nil/crypto3/math/polynomial/quotient_ring/polynomial_frobenius.hpp

@@ -27,6 +27,8 @@
 
 #include <algorithm>
 #include <cstddef>
+#include <optional>
+#include <stdexcept>
 #include <utility>
 
 #include <nil/crypto3/algebra/fields/field_order.hpp>
@@ -147,6 +149,123 @@ namespace nil::crypto3::math {
         output = std::move(result);
     }
 
+    namespace detail {
+
+        // Shared chain for coefficient-fixed Frobenius products and the characteristic-power square-root seed.
+        // In the characteristic case the supplied image is X^p, and coefficients must support Frobenius_map.
+        template<bool CharacteristicFrobenius, SupportsDivrem Backend>
+        void frobenius_product_impl(typename Backend::polynomial_type &output,
+                                    const typename Backend::polynomial_type &input, std::size_t term_count,
+                                    const typename Backend::polynomial_type &x_to_frobenius_power,
+                                    const polynomial_divisor_context<Backend> &divisor_context,
+                                    polynomial_arithmetic::polynomial_context<Backend> &arithmetic_context) {
+            using polynomial_type = typename Backend::polynomial_type;
+            using value_type = typename polynomial_type::value_type;
+            using cache_type = polynomial_composition_precomputation<Backend>;
+            const std::size_t degree = divisor_context.degree();
+            const auto canonical = [](const polynomial_type &value) {
+                return !value.empty() && (value.size() == 1 || value.back() != value_type::zero());
+            };
+            if (!canonical(input) || !canonical(x_to_frobenius_power) ||
+                x_to_frobenius_power.size() > std::max<std::size_t>(1, degree)) {
+                throw std::invalid_argument(
+                    "Frobenius products require canonical inputs and a reduced Frobenius image");
+            }
+            if (degree == 0 || term_count == 0) {
+                output.assign(1, degree == 0 ? value_type::zero() : value_type::one());
+                return;
+            }
+            polynomial_type base;
+            remainder(base, input, divisor_context, arithmetic_context);
+            if (term_count == 1) {
+                output = std::move(base);
+                return;
+            }
+            if constexpr (!CharacteristicFrobenius) {
+                // Constants are fixed by Q-Frobenius, but generally not by p-Frobenius in an extension field.
+                if (base.size() == 1) {
+                    powmod(output, base, term_count, divisor_context, arithmetic_context);
+                    return;
+                }
+            }
+
+            const auto apply = [&](polynomial_type &result, const polynomial_type &operand, std::size_t power,
+                                   const cache_type &cache) {
+                if constexpr (CharacteristicFrobenius) {
+                    // Absolute Frobenius is semilinear: f^(p^power) = sigma^power(f)(X^(p^power)).
+                    // Copy before twisting so result may alias operand, including the cached-image input.
+                    polynomial_type twisted = operand;
+                    for (auto &coefficient : twisted) {
+                        coefficient = coefficient.Frobenius_map(power);
+                    }
+                    compose_mod(result, twisted, cache, divisor_context, arithmetic_context);
+                } else {
+                    compose_mod(result, operand, cache, divisor_context, arithmetic_context);
+                }
+            };
+
+            const cache_type first_cache(x_to_frobenius_power, degree, divisor_context, arithmetic_context);
+            const auto chain = [&](auto &&self, std::size_t count,
+                                   bool need_image) -> std::pair<polynomial_type, polynomial_type> {
+                if (count == 1) {
+                    return {base, x_to_frobenius_power};
+                }
+                auto [product, image] = self(self, count / 2, true);
+                std::optional<cache_type> doubling_cache;
+                if (count / 2 > 1) {
+                    doubling_cache.emplace(image, degree, divisor_context, arithmetic_context);
+                }
+                const auto &cache = doubling_cache ? *doubling_cache : first_cache;
+                // With P_k = product and H_k = image, P_2k = P_k * F^k(P_k) and H_2k = F^k(H_k).
+                polynomial_type shifted;
+                apply(shifted, product, count / 2, cache);
+                mulmod(product, product, shifted, divisor_context, arithmetic_context);
+                if (need_image) {
+                    apply(image, image, count / 2, cache);
+                }
+                if (count % 2 != 0) {
+                    // P_(2k+1) = base * F(P_2k); use the same initial cache at every odd step.
+                    apply(shifted, product, 1, first_cache);
+                    mulmod(product, base, shifted, divisor_context, arithmetic_context);
+                    if (need_image) {
+                        apply(image, image, 1, first_cache);
+                    }
+                }
+                return {std::move(product), std::move(image)};
+            };
+            // The final image is not used by the caller, so avoid its compositions.
+            auto result = chain(chain, term_count, false);
+            output = std::move(result.first);
+        }
+
+    }    // namespace detail
+
+    /**
+     * Compute input * F(input) * ... * F^(term_count - 1)(input) modulo B, where F is a power of the coefficient
+     * field's Frobenius map. The supplied x_to_frobenius_power is the canonical reduced image F(X); it can represent
+     * X^(Q^s) for any s >= 0, with Q the coefficient-field order. Its algebraic identity is a caller precondition,
+     * not checked by repeating the exponentiation. B need not be irreducible.
+     *
+     * Binary doubling uses O(log(term_count)) compositions and products. Each doubling shares one composition
+     * precomputation between the partial product and the next Frobenius image. Odd steps reuse the initial image's
+     * precomputation. Only the initial cache and one doubling cache are live at once.
+     *
+     * An empty product is the quotient-ring identity, and a singleton is the reduced input. The result is canonical
+     * and may alias either input polynomial. Reducing an unreduced input may require more inverse coefficients than
+     * products of reduced operands; otherwise degree(B) - 1 inverse coefficients suffice.
+     *
+     * @throws std::invalid_argument if an input is empty or noncanonical, the supplied image is not reduced, a
+     *         needed composition has a zero cached-power limit, or the divisor inverse has insufficient precision.
+     */
+    template<detail::SupportsDivrem Backend>
+    void frobenius_product(typename Backend::polynomial_type &output, const typename Backend::polynomial_type &input,
+                           std::size_t term_count, const typename Backend::polynomial_type &x_to_frobenius_power,
+                           const polynomial_divisor_context<Backend> &divisor_context,
+                           polynomial_arithmetic::polynomial_context<Backend> &arithmetic_context) {
+        detail::frobenius_product_impl<false>(output, input, term_count, x_to_frobenius_power, divisor_context,
+                                              arithmetic_context);
+    }
+
 }    // namespace nil::crypto3::math
 
 #endif    // CRYPTO3_MATH_POLYNOMIAL_FROBENIUS_HPP

libs/math/include/nil/crypto3/math/polynomial/quotient_ring/polynomial_square_root.hpp

@@ -32,10 +32,13 @@
 
 #include <boost/multiprecision/cpp_int.hpp>
 
+#include <nil/crypto3/algebra/fields/field_algorithms.hpp>
 #include <nil/crypto3/algebra/fields/field_order.hpp>
 #include <nil/crypto3/algebra/type_traits.hpp>
 
+#include <nil/crypto3/math/polynomial/operations/shift.hpp>
 #include <nil/crypto3/math/polynomial/quotient_ring/polynomial_exponentiation.hpp>
+#include <nil/crypto3/math/polynomial/quotient_ring/polynomial_frobenius.hpp>
 
 namespace nil::crypto3::math {
 
@@ -248,6 +251,154 @@ namespace nil::crypto3::math {
         return detail::is_square_mod_impl(input, divisor_context, arithmetic_context, nullptr);
     }
 
+    namespace detail {
+
+        // Return {X^((Q+1)/2), X^Q} modulo B. Called after the odd-degree root checks, with d > 1 and B(0) != 0.
+        template<SupportsDivrem Backend>
+        std::pair<typename Backend::polynomial_type, typename Backend::polynomial_type>
+            square_root_x_seed(const polynomial_divisor_context<Backend> &divisor_context,
+                               polynomial_arithmetic::polynomial_context<Backend> &arithmetic_context) {
+            using polynomial_type = typename Backend::polynomial_type;
+            using value_type = typename polynomial_type::value_type;
+            using field_type = typename value_type::field_type;
+            const auto &modulus = divisor_context.divisor();
+            const std::size_t degree = divisor_context.degree();
+            polynomial_type seed, x_to_q;
+            if constexpr (field_type::arity > 1 && requires(const value_type &value, std::size_t power) {
+                              { value.Frobenius_map(power) } -> std::convertible_to<value_type>;
+                          }) {
+                const auto multiply_by_x = [&](polynomial_type &value) {
+                    shift_left(value, value, 1);
+                    if (value.size() > degree) {
+                        const value_type leading = value[degree] * divisor_context.reversed_divisor_inverse()[0];
+                        value.resize(degree);
+                        for (std::size_t i = 0; i < degree; ++i) {
+                            value[i] -= leading * modulus[i];
+                        }
+                        condense(value);
+                    }
+                };
+                // For Q = p^e, a = x^((p-1)/2) gives x^p = x*a^2. Accumulate the e conjugates of a:
+                // T = a * a^p * ... * a^(p^(e-1)) = x^((Q-1)/2). Then b = x*T and x^Q = x*T^2.
+                // The chain shares its composition caches between products and images, avoiding a Q-sized power.
+                const boost::multiprecision::cpp_int exponent =
+                    (algebra::fields::field_characteristic<field_type>() - 1) >> 1;
+                polynomial_type characteristic_seed, x_to_p;
+                powmod_x(characteristic_seed, exponent, divisor_context, arithmetic_context);
+                squaremod(x_to_p, characteristic_seed, divisor_context, arithmetic_context);
+                multiply_by_x(x_to_p);
+                frobenius_product_impl<true>(seed, characteristic_seed, field_type::arity, x_to_p, divisor_context,
+                                             arithmetic_context);
+                squaremod(x_to_q, seed, divisor_context, arithmetic_context);
+                multiply_by_x(x_to_q);
+                multiply_by_x(seed);
+            } else {
+                // Prime fields and coefficient types without a Frobenius map retain the direct seed.
+                const boost::multiprecision::cpp_int exponent = (algebra::fields::field_order<field_type>() + 1) >> 1;
+                polynomial_type seed_squared;
+                powmod_x(seed, exponent, divisor_context, arithmetic_context);
+                squaremod(seed_squared, seed, divisor_context, arithmetic_context);
+                // For a reduced y, y / X = (y - c * B) / X with c = y[0] / B[0], also for nonmonic B.
+                const value_type c = seed_squared[0] * modulus[0].inversed();
+                x_to_q.assign(degree, value_type::zero());
+                for (std::size_t i = 0; i < degree; ++i) {
+                    x_to_q[i] =
+                        (i + 1 < seed_squared.size() ? seed_squared[i + 1] : value_type::zero()) - c * modulus[i + 1];
+                }
+                condense(x_to_q);
+            }
+            return {std::move(seed), std::move(x_to_q)};
+        }
+
+    }    // namespace detail
+
+    /**
+     * Compute a square root of X modulo an irreducible polynomial B of odd degree d over an odd-order field K.
+     * Write Q = |K|, x = X mod B, and N = Norm(x) = -B(0) / leading_coefficient(B). The element x is square exactly
+     * when N is square in K. For d = 1 the answer is simply sqrt(N), including zero when B is a multiple of X.
+     *
+     * For d > 1, let b = x^((Q + 1) / 2) and m = (d - 1) / 2. The product
+     *
+     *     D = b^Q * b^(Q^3) * ... * b^(Q^(2m - 1))
+     *
+     * satisfies D^2 = x^(Q + Q^2 + ... + Q^(d - 1)) = N / x. Therefore R = x * D / sqrt(N) satisfies R^2 = x.
+     * For K = F_(p^e) with coefficient Frobenius support, the seed uses one characteristic-sized exponentiation
+     * and O(log(e)) semilinear compositions, producing X^Q at the same time. Other coefficient types retain the
+     * direct Q-sized seed. The remaining Frobenius product uses O(log(d)) compositions and modular products,
+     * without constructing a quadratic nonresidue in K[X]/(B). Final multiplication by x is a linear pass.
+     *
+     * Return false and set output to zero if x is nonsquare. Otherwise output is a canonical reduced root.
+     * @throws std::invalid_argument if d is zero or even, K has characteristic two, B(0) is zero with d > 1,
+     *         a needed composition has a zero cached-power limit, or the divisor inverse has insufficient precision.
+     * @pre B is irreducible. For d > 1, degree(B) - 1 inverse coefficients suffice. K provides is_square() and
+     *      the coefficient-field sqrt_known_square operation. Irreducibility is not checked here.
+     */
+    template<detail::SupportsDivrem Backend>
+        requires algebra::FieldValue<typename Backend::polynomial_type::value_type> &&
+                 requires(const typename Backend::polynomial_type::value_type &value) {
+                     { value.is_square() } -> std::convertible_to<bool>;
+                 }
+    bool square_root_mod_x_odd_degree(typename Backend::polynomial_type &output,
+                                      const polynomial_divisor_context<Backend> &divisor_context,
+                                      polynomial_arithmetic::polynomial_context<Backend> &arithmetic_context) {
+        using polynomial_type = typename Backend::polynomial_type;
+        using value_type = typename polynomial_type::value_type;
+        using field_type = typename value_type::field_type;
+        const auto &modulus = divisor_context.divisor();
+        const std::size_t degree = divisor_context.degree();
+        if (degree % 2 == 0 || algebra::fields::field_characteristic<field_type>() == 2) {
+            throw std::invalid_argument("direct square roots of X require an odd degree and odd characteristic");
+        }
+        if (degree > 1 && modulus[0] == value_type::zero()) {
+            throw std::invalid_argument("an irreducible divisor of degree greater than one has nonzero constant term");
+        }
+
+        const value_type inverse_leading = divisor_context.reversed_divisor_inverse()[0];
+        const value_type norm = -modulus[0] * inverse_leading;
+        if (!norm.is_square()) {
+            output.assign(1, value_type::zero());
+            return false;
+        }
+        const value_type norm_root = algebra::fields::sqrt_known_square(norm);
+        if (degree == 1) {
+            output.assign(1, norm_root);
+            return true;
+        }
+        if (arithmetic_context.options().modular_composition_cached_power_limit == 0) {
+            throw std::invalid_argument("direct square roots of X require a positive composition cached-power limit");
+        }
+
+        auto [seed, x_to_q] = detail::square_root_x_seed(divisor_context, arithmetic_context);
+
+        const std::size_t term_count = (degree - 1) / 2;
+        polynomial_type product, x_to_q_squared;
+        {
+            // Reuse this cache for b^Q and X^(Q^2), then release it before allocating the product-chain caches.
+            const polynomial_composition_precomputation<Backend> cache(x_to_q, degree, divisor_context,
+                                                                       arithmetic_context);
+            compose_mod(product, seed, cache, divisor_context, arithmetic_context);
+            if (term_count > 1) {
+                compose_mod(x_to_q_squared, x_to_q, cache, divisor_context, arithmetic_context);
+            }
+        }
+        if (term_count > 1) {
+            frobenius_product(product, product, term_count, x_to_q_squared, divisor_context, arithmetic_context);
+        }
+
+        shift_left(product, product, 1);
+        if (product.size() > degree) {
+            // A shifted reduced polynomial exceeds degree(B) by at most one coefficient.
+            const value_type leading = product[degree] * inverse_leading;
+            product.resize(degree);
+            for (std::size_t i = 0; i < degree; ++i) {
+                product[i] -= leading * modulus[i];
+            }
+            condense(product);
+        }
+        scalar_multiplication(output, product, norm_root.inversed());
+        return true;
+    }
+
     /**
      * Compute a square root of input in K[X]/(B) with Tonelli-Shanks, using the multiplicative-group decomposition and
      * quadratic nonresidue cached in square_root_context. The irreducible polynomial B is stored in the divisor context