modular_arithmetic.H

Go to the documentation of this file.

/*
                          Aleph_w
 
  Data structures & Algorithms
  version 2.0.0b
  https://github.com/lrleon/Aleph-w
 
  This file is part of Aleph-w library
 
  Copyright (c) 2002-2026 Leandro Rabindranath Leon
 
  Permission is hereby granted, free of charge, to any person obtaining a copy
  of this software and associated documentation files (the "Software"), to deal
  in the Software without restriction, including without limitation the rights
  to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
  copies of the Software, and to permit persons to whom the Software is
  furnished to do so, subject to the following conditions:
 
  The above copyright notice and this permission notice shall be included in all
  copies or substantial portions of the Software.
 
  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
  IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
  AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
  SOFTWARE.
*/
 
# ifndef MODULAR_ARITHMETIC_H
# define MODULAR_ARITHMETIC_H
 
# include <cstdint>
# include <type_traits>
# ifdef _WIN32
#   include <intrin.h>
# endif
 
# include <ah-errors.H>
# include <tpl_array.H>
 
namespace Aleph
{
  [[nodiscard]] inline uint64_t mod_mul(uint64_t a, uint64_t b, uint64_t m)
  {
    ah_invalid_argument_if(m == 0) << "mod_mul: modulus must be > 0";
 
# if defined(__SIZEOF_INT128__) && !defined(_WIN32)
    return static_cast<uint64_t>((static_cast<__uint128_t>(a) * b) % m);
# elif defined(_MSC_VER) && !defined(__clang__) && (defined(_M_X64) || defined(_M_ARM64))
    // MSVC cl (64-bit): no __int128, but _umul128 + _udiv128 are available
    unsigned long long hi;
    unsigned long long lo = _umul128(a, b, &hi);
    unsigned long long rem;
    _udiv128(hi, lo, m, &rem);
    return static_cast<uint64_t>(rem);
# else
    // clang-cl: has __uint128_t but no __umodti3 runtime and no _udiv128.
    // MSVC 32-bit: no _umul128/_udiv128. Fall through to binary method.
    uint64_t res = 0;
    a %= m;
    while (b)
      {
        if (b & 1)
          {
            if (m - res <= a) res = a - (m - res);
            else res += a;
          }
        if (m - a <= a) a = a - (m - a);
        else a <<= 1;
        b >>= 1;
      }
    return res;
# endif
  }
 
  [[nodiscard]] inline uint64_t mod_exp(uint64_t base, uint64_t exp, const uint64_t m)
  {
    ah_invalid_argument_if(m == 0) << "mod_exp: modulus must be > 0";
    if (m == 1) return 0;
    uint64_t res = 1;
    base %= m;
    while (exp > 0)
      {
        if (exp & 1)
          res = mod_mul(res, base, m);
        base = mod_mul(base, base, m);
        exp >>= 1;
      }
    return res;
  }
 
  template <typename T>
    requires (std::is_integral_v<T> and std::is_signed_v<T>)
  [[nodiscard]] T ext_gcd(T a, T b, T & x, T & y) noexcept
  {
    if (b == 0)
      {
        x = 1;
        y = 0;
        return a;
      }
    T x1, y1;
    T d = ext_gcd(b, a % b, x1, y1);
    x = y1;
    y = x1 - static_cast<T>(a / b) * y1;
    return d;
  }
 
  [[nodiscard]] inline uint64_t mod_inv(const uint64_t a, const uint64_t m)
  {
    ah_domain_error_if(m == 0) << "mod_inv: modulus cannot be 0";
 
    if (m == 1)
      return 0;
 
    const uint64_t a_mod = a % m;
    ah_invalid_argument_if(a_mod == 0)
      << "Modular inverse does not exist (" << a << " is 0 mod " << m << ")";
 
    // Iterative extended Euclidean algorithm using unsigned arithmetic.
    // We track t0 such that a * t0 ≡ r0 (mod m) at each step,
    // keeping t0 in [0, m) to avoid signed overflow.
    uint64_t r0 = a_mod, r1 = m;
    uint64_t t0 = 1, t1 = 0;
 
    while (r1 != 0)
      {
        const uint64_t q = r0 / r1;
 
        const uint64_t tmp_r = r0 - q * r1;
        r0 = r1;
        r1 = tmp_r;
 
        // t_new = t0 - q * t1 (mod m), computed in unsigned
        const uint64_t qt1 = mod_mul(q, t1, m);
        const uint64_t tmp_t = (t0 >= qt1) ? (t0 - qt1) : (m - (qt1 - t0));
        t0 = t1;
        t1 = tmp_t;
      }
 
    ah_invalid_argument_if(r0 != 1)
      << "Modular inverse does not exist (numbers " << a
      << " and " << m << " are not coprime)";
 
    return t0;
  }
 
# if defined(__SIZEOF_INT128__) && !defined(_WIN32)
  struct MontgomeryCtx;
 
  namespace detail
  {
    [[nodiscard]] constexpr MontgomeryCtx
    montgomery_ctx_unchecked(const uint64_t mod) noexcept;
  }
 
  struct MontgomeryCtx
  {
  public:
    MontgomeryCtx() = delete;
 
    [[nodiscard]] constexpr uint64_t mod() const noexcept { return mod_; }
 
    [[nodiscard]] constexpr uint64_t mod2() const noexcept { return mod2_; }
 
    [[nodiscard]] constexpr uint64_t r() const noexcept { return r_; }
 
    [[nodiscard]] constexpr uint64_t r2() const noexcept { return r2_; }
 
    [[nodiscard]] constexpr uint64_t mod_inv_neg() const noexcept
    {
      return mod_inv_neg_;
    }
 
  private:
    friend constexpr MontgomeryCtx
    detail::montgomery_ctx_unchecked(const uint64_t mod) noexcept;
 
    constexpr MontgomeryCtx(const uint64_t mod,
                            const uint64_t mod2,
                            const uint64_t r,
                            const uint64_t r2,
                            const uint64_t mod_inv_neg) noexcept
      : mod_(mod), mod2_(mod2), r_(r), r2_(r2), mod_inv_neg_(mod_inv_neg)
    {
      /* empty */
    }
 
    uint64_t mod_; 
    uint64_t mod2_; 
    uint64_t r_; 
    uint64_t r2_; 
    uint64_t mod_inv_neg_; 
  };
 
  namespace detail
  {
    [[nodiscard]] constexpr uint64_t
    montgomery_neg_inverse(const uint64_t mod) noexcept
    {
      uint64_t inv = 1;
      for (size_t i = 0; i < 6; ++i)
        inv *= 2 - mod * inv;
      return ~inv + 1;
    }
 
    [[nodiscard]] constexpr MontgomeryCtx
    montgomery_ctx_unchecked(const uint64_t mod) noexcept
    {
      const __uint128_t r128 = static_cast<__uint128_t>(1) << 64;
      const auto r = static_cast<uint64_t>(r128 % mod);
      const auto r2 = static_cast<uint64_t>((static_cast<__uint128_t>(r) * r) % mod);
 
      return MontgomeryCtx(mod,
                           mod <= UINT64_MAX - mod ? mod + mod : 0,
                           r,
                           r2,
                           montgomery_neg_inverse(mod));
    }
  }
 
  [[nodiscard]] inline MontgomeryCtx
  montgomery_ctx(const uint64_t mod)
  {
    ah_invalid_argument_if(mod <= 1)
      << "montgomery_ctx: modulus must be > 1";
    ah_invalid_argument_if((mod & 1ULL) == 0)
      << "montgomery_ctx: modulus " << mod << " must be odd";
    return detail::montgomery_ctx_unchecked(mod);
  }
 
  template <uint64_t Mod>
  [[nodiscard]] consteval MontgomeryCtx
  montgomery_ctx_for_mod()
  {
    static_assert(Mod > 1, "montgomery_ctx_for_mod: modulus must be > 1");
    static_assert((Mod & 1ULL) == 1ULL,
                  "montgomery_ctx_for_mod: modulus must be odd");
    return detail::montgomery_ctx_unchecked(Mod);
  }
 
  [[nodiscard]] constexpr uint64_t
  mont_reduce(const __uint128_t x,
              const MontgomeryCtx & ctx) noexcept
  {
    const uint64_t q = static_cast<uint64_t>(x) * ctx.mod_inv_neg();
    const auto x_lo = static_cast<uint64_t>(x);
    const auto x_hi = static_cast<uint64_t>(x >> 64);
    const __uint128_t qmod = static_cast<__uint128_t>(q) * ctx.mod();
    const auto qmod_lo = static_cast<uint64_t>(qmod);
    const auto qmod_hi = static_cast<uint64_t>(qmod >> 64);
    const uint64_t carry = qmod_lo > UINT64_MAX - x_lo ? 1ULL : 0ULL;
 
    // t = (x + q·p) / R; REDC guarantees 0 ≤ t < 2p.
    const __uint128_t t = static_cast<__uint128_t>(x_hi) + qmod_hi + carry;
 
    // Fast path: 2p < 2^64 so t fits in uint64_t; one conditional subtraction
    // reduces to [0, p) using only additions and comparisons.
    if (ctx.mod2() != 0)
      {
        const uint64_t t64 = static_cast<uint64_t>(t);
        return t64 >= ctx.mod() ? t64 - ctx.mod() : t64;
      }
 
    // Fallback for large primes (p ≥ 2^63): t may exceed UINT64_MAX, so keep
    // it in __uint128_t and use a single 128-bit division to finish.
    return static_cast<uint64_t>(t % ctx.mod());
  }
 
  [[nodiscard]] constexpr uint64_t
  mont_mul(const uint64_t a,
           const uint64_t b,
           const MontgomeryCtx & ctx) noexcept
  {
    return mont_reduce(static_cast<__uint128_t>(a) * b, ctx);
  }
 
  [[nodiscard]] constexpr uint64_t
  to_mont(const uint64_t a,
          const MontgomeryCtx & ctx) noexcept
  {
    return mont_mul(a % ctx.mod(), ctx.r2(), ctx);
  }
 
  [[nodiscard]] constexpr uint64_t
  from_mont(const uint64_t a,
            const MontgomeryCtx & ctx) noexcept
  {
    return mont_reduce(a, ctx);
  }
 
  [[nodiscard]] constexpr uint64_t
  mont_exp(uint64_t base,
           uint64_t exp,
           const MontgomeryCtx & ctx) noexcept
  {
    uint64_t result = to_mont(1, ctx);
    while (exp > 0)
      {
        if (exp & 1ULL)
          result = mont_mul(result, base, ctx);
        base = mont_mul(base, base, ctx);
        exp >>= 1;
      }
    return result;
  }
# endif
 
  [[nodiscard]] inline uint64_t crt(const Array<uint64_t> & rem,
                                    const Array<uint64_t> & mod)
  {
    ah_invalid_argument_if(rem.size() != mod.size())
      << "crt: arrays must have the same size (got " << rem.size()
      << " vs " << mod.size() << ")";
 
    const size_t n = rem.size();
    if (n == 0)
      return 0;
 
    // Compute product of all moduli with overflow detection
    uint64_t prod = 1;
    for (size_t i = 0; i < n; ++i)
      {
        ah_invalid_argument_if(mod[i] <= 1)
          << "crt: all moduli must be > 1 (got " << mod[i] << " at index " << i << ")";
 
        ah_overflow_error_if(prod > UINT64_MAX / mod[i])
          << "crt: product of moduli overflows uint64_t at index " << i;
        prod *= mod[i];
      }
 
    uint64_t result = 0;
    for (size_t i = 0; i < n; ++i)
      {
        const uint64_t p = prod / mod[i];
        const uint64_t inv = mod_inv(p, mod[i]);
        uint64_t term = mod_mul(rem[i], p, prod);
        term = mod_mul(term, inv, prod);
        if (result >= prod - term)
          result -= (prod - term);
        else
          result += term;
      }
 
    return result;
  }
} // namespace Aleph
 
# endif // MODULAR_ARITHMETIC_H