/*
* (C) 2018,2024,2025 Jack Lloyd
*
* Botan is released under the Simplified BSD License (see license.txt)
*/

#include <botan/internal/monty.h>

#include <botan/exceptn.h>
#include <botan/mem_ops.h>
#include <botan/internal/barrett.h>
#include <botan/internal/mp_core.h>
#include <array>

namespace Botan {

namespace {

// If the modulus is at most this many words, then use the stack instead
// of a heap variable for some temporary values
constexpr size_t MontgomeryUseStackLimit = 32;

}  // namespace

Montgomery_Params::Data::Data(const BigInt& p, const Barrett_Reduction& mod_p) {
   if(p.is_even() || p < 3) {
      throw Invalid_Argument("Montgomery_Params invalid modulus");
   }

   m_p = p;
   m_p_words = m_p.sig_words();
   m_p_dash = monty_inverse(m_p.word_at(0));

   const BigInt r = BigInt::power_of_2(m_p_words * WordInfo<word>::bits);

   m_r1 = mod_p.reduce(r);
   m_r2 = mod_p.square(m_r1);
   m_r3 = mod_p.multiply(m_r1, m_r2);

   // Barrett should be at least zero prefixing up to modulus size
   BOTAN_ASSERT_NOMSG(m_r1.size() >= m_p_words);
   BOTAN_ASSERT_NOMSG(m_r2.size() >= m_p_words);
   BOTAN_ASSERT_NOMSG(m_r3.size() >= m_p_words);
}

Montgomery_Params::Montgomery_Params(const BigInt& p, const Barrett_Reduction& mod_p) :
      m_data(std::make_shared<Data>(p, mod_p)) {}

Montgomery_Params::Montgomery_Params(const BigInt& p) :
      Montgomery_Params(p, Barrett_Reduction::for_secret_modulus(p)) {}

bool Montgomery_Params::operator==(const Montgomery_Params& other) const {
   if(this->m_data == other.m_data) {
      return true;
   }

   return (this->m_data->p() == other.m_data->p());
}

BigInt Montgomery_Params::redc(const BigInt& x, secure_vector<word>& ws) const {
   const size_t p_size = this->p_words();

   if(ws.size() < p_size) {
      ws.resize(p_size);
   }

   BigInt z = x;
   z.grow_to(2 * p_size);

   bigint_monty_redc_inplace(z.mutable_data(), this->p()._data(), p_size, this->p_dash(), ws.data(), ws.size());

   return z;
}

BigInt Montgomery_Params::mul(const BigInt& x, const BigInt& y, secure_vector<word>& ws) const {
   const size_t p_size = this->p_words();
   BigInt z = BigInt::with_capacity(2 * p_size);
   this->mul(z, x, y, ws);
   return z;
}

void Montgomery_Params::mul(BigInt& z, const BigInt& x, const BigInt& y, secure_vector<word>& ws) const {
   const size_t p_size = this->p_words();

   if(ws.size() < 2 * p_size) {
      ws.resize(2 * p_size);
   }

   BOTAN_DEBUG_ASSERT(x.sig_words() <= p_size);
   BOTAN_DEBUG_ASSERT(y.sig_words() <= p_size);

   if(z.size() < 2 * p_size) {
      z.grow_to(2 * p_size);
   }

   bigint_mul(z.mutable_data(),
              z.size(),
              x._data(),
              x.size(),
              std::min(p_size, x.size()),
              y._data(),
              y.size(),
              std::min(p_size, y.size()),
              ws.data(),
              ws.size());

   bigint_monty_redc_inplace(z.mutable_data(), this->p()._data(), p_size, this->p_dash(), ws.data(), ws.size());
}

void Montgomery_Params::mul(BigInt& z, const BigInt& x, std::span<const word> y, secure_vector<word>& ws) const {
   const size_t p_size = this->p_words();

   if(ws.size() < 2 * p_size) {
      ws.resize(2 * p_size);
   }
   if(z.size() < 2 * p_size) {
      z.grow_to(2 * p_size);
   }

   BOTAN_DEBUG_ASSERT(x.sig_words() <= p_size);

   bigint_mul(z.mutable_data(),
              z.size(),
              x._data(),
              x.size(),
              std::min(p_size, x.size()),
              y.data(),
              y.size(),
              std::min(p_size, y.size()),
              ws.data(),
              ws.size());

   bigint_monty_redc_inplace(z.mutable_data(), this->p()._data(), p_size, this->p_dash(), ws.data(), ws.size());
}

void Montgomery_Params::mul_by(BigInt& x, const BigInt& y, secure_vector<word>& ws) const {
   const size_t p_size = this->p_words();

   if(ws.size() < 4 * p_size) {
      ws.resize(4 * p_size);
   }

   word* z_data = ws.data();
   word* ws_data = &ws[2 * p_size];

   BOTAN_DEBUG_ASSERT(x.sig_words() <= p_size);

   bigint_mul(z_data,
              2 * p_size,
              x._data(),
              x.size(),
              std::min(p_size, x.size()),
              y._data(),
              y.size(),
              std::min(p_size, y.size()),
              ws_data,
              2 * p_size);

   bigint_monty_redc_inplace(z_data, this->p()._data(), p_size, this->p_dash(), ws_data, 2 * p_size);

   if(x.size() < 2 * p_size) {
      x.grow_to(2 * p_size);
   }
   copy_mem(x.mutable_data(), z_data, 2 * p_size);
}

BigInt Montgomery_Params::sqr(const BigInt& x, secure_vector<word>& ws) const {
   BOTAN_DEBUG_ASSERT(x.sig_words() <= this->p_words());
   return this->sqr(std::span{x._data(), x.size()}, ws);
}

BigInt Montgomery_Params::sqr(std::span<const word> x, secure_vector<word>& ws) const {
   const size_t p_size = this->p_words();
   BigInt z = BigInt::with_capacity(2 * p_size);
   this->sqr(z, x, ws);
   return z;
}

void Montgomery_Params::sqr(BigInt& z, const BigInt& x, secure_vector<word>& ws) const {
   this->sqr(z, std::span{x._data(), x.size()}, ws);
}

void Montgomery_Params::sqr(BigInt& z, std::span<const word> x, secure_vector<word>& ws) const {
   const size_t p_size = this->p_words();

   if(ws.size() < 2 * p_size) {
      ws.resize(2 * p_size);
   }

   if(z.size() < 2 * p_size) {
      z.grow_to(2 * p_size);
   }

   bigint_sqr(z.mutable_data(), z.size(), x.data(), x.size(), std::min(p_size, x.size()), ws.data(), ws.size());

   bigint_monty_redc_inplace(z.mutable_data(), this->p()._data(), p_size, this->p_dash(), ws.data(), ws.size());
}

Montgomery_Int::Montgomery_Int(const Montgomery_Params& params, secure_vector<word> words) :
      m_params(params), m_v(std::move(words)) {
   BOTAN_ASSERT_NOMSG(m_v.size() == m_params.p_words());
}

Montgomery_Int Montgomery_Int::one(const Montgomery_Params& params) {
   return Montgomery_Int(params, params.R1(), false);
}

Montgomery_Int Montgomery_Int::from_wide_int(const Montgomery_Params& params, const BigInt& x) {
   secure_vector<word> ws;
   auto redc_x = params.mul(params.redc(x, ws), params.R3(), ws);
   return Montgomery_Int(params, redc_x, false);
}

Montgomery_Int::Montgomery_Int(const Montgomery_Params& params, const BigInt& v, bool redc_needed) :
      m_params(params), m_v(m_params.p_words()) {
   BOTAN_ASSERT_NOMSG(v < m_params.p());

   const size_t p_size = m_params.p_words();

   auto v_span = v._as_span();

   if(v_span.size() > p_size) {
      // Safe to truncate the span since we already checked v < p
      v_span = v_span.first(p_size);
   }

   BOTAN_ASSERT_NOMSG(m_v.size() >= v_span.size());

   copy_mem(std::span{m_v}.first(v_span.size()), v_span);

   if(redc_needed) {
      secure_vector<word> ws;
      this->mul_by(m_params.R2()._as_span().first(p_size), ws);
   }
}

Montgomery_Int::Montgomery_Int(const Montgomery_Params& params, std::span<const word> words) :
      m_params(params), m_v(words.begin(), words.end()) {
   BOTAN_ARG_CHECK(m_v.size() == m_params.p_words(), "Invalid input span");
}

std::vector<uint8_t> Montgomery_Int::serialize() const {
   return value().serialize();
}

BigInt Montgomery_Int::value() const {
   secure_vector<word> ws(m_params.p_words());

   secure_vector<word> z = m_v;
   z.resize(2 * m_params.p_words());  // zero extend

   bigint_monty_redc_inplace(
      z.data(), m_params.p()._data(), m_params.p_words(), m_params.p_dash(), ws.data(), ws.size());

   return BigInt::_from_words(z);
}

Montgomery_Int Montgomery_Int::operator+(const Montgomery_Int& other) const {
   BOTAN_STATE_CHECK(other.m_params == m_params);

   const size_t p_size = m_params.p_words();
   BOTAN_ASSERT_NOMSG(m_v.size() == p_size && other.m_v.size() == p_size);

   secure_vector<word> z(2 * p_size);

   word* r = std::span{z}.first(p_size).data();
   word* t = std::span{z}.last(p_size).data();

   // t = this + other
   const word carry = bigint_add3(t, m_v.data(), p_size, other.m_v.data(), p_size);

   // Conditionally subtract r = t - p
   bigint_monty_maybe_sub(p_size, r, carry, t, m_params.p()._data());

   z.resize(p_size);  // truncate leaving only r
   return Montgomery_Int(m_params, std::move(z));
}

Montgomery_Int Montgomery_Int::operator-(const Montgomery_Int& other) const {
   BOTAN_STATE_CHECK(other.m_params == m_params);

   const size_t p_size = m_params.p_words();
   BOTAN_ASSERT_NOMSG(m_v.size() == p_size && other.m_v.size() == p_size);

   secure_vector<word> t(p_size);
   const word borrow = bigint_sub3(t.data(), m_v.data(), p_size, other.m_v.data(), p_size);

   bigint_cnd_add(borrow, t.data(), m_params.p()._data(), p_size);

   return Montgomery_Int(m_params, std::move(t));
}

Montgomery_Int Montgomery_Int::mul(const Montgomery_Int& other, secure_vector<word>& ws) const {
   BOTAN_STATE_CHECK(other.m_params == m_params);

   const size_t p_size = m_params.p_words();
   BOTAN_ASSERT_NOMSG(m_v.size() == p_size && other.m_v.size() == p_size);

   if(ws.size() < 2 * p_size) {
      ws.resize(2 * p_size);
   }

   secure_vector<word> z(2 * p_size);

   bigint_mul(z.data(), z.size(), m_v.data(), p_size, p_size, other.m_v.data(), p_size, p_size, ws.data(), ws.size());

   bigint_monty_redc_inplace(z.data(), m_params.p()._data(), p_size, m_params.p_dash(), ws.data(), ws.size());
   z.resize(p_size);  // truncate off high zero words

   return Montgomery_Int(m_params, std::move(z));
}

Montgomery_Int& Montgomery_Int::mul_by(const Montgomery_Int& other, secure_vector<word>& ws) {
   BOTAN_STATE_CHECK(other.m_params == m_params);
   return this->mul_by(std::span{other.m_v}, ws);
}

Montgomery_Int& Montgomery_Int::mul_by(std::span<const word> other, secure_vector<word>& ws) {
   const size_t p_size = m_params.p_words();
   BOTAN_ASSERT_NOMSG(m_v.size() == p_size && other.size() == p_size);

   if(ws.size() < 2 * p_size) {
      ws.resize(2 * p_size);
   }

   auto do_mul_by = [&](std::span<word> z) {
      bigint_mul(z.data(), z.size(), m_v.data(), p_size, p_size, other.data(), p_size, p_size, ws.data(), ws.size());

      bigint_monty_redc_inplace(z.data(), m_params.p()._data(), p_size, m_params.p_dash(), ws.data(), ws.size());

      copy_mem(m_v, z.first(p_size));
   };

   if(p_size <= MontgomeryUseStackLimit) {
      std::array<word, 2 * MontgomeryUseStackLimit> z{};
      do_mul_by(z);
   } else {
      secure_vector<word> z(2 * p_size);
      do_mul_by(z);
   }

   return (*this);
}

Montgomery_Int& Montgomery_Int::square_this_n_times(secure_vector<word>& ws, size_t n) {
   const size_t p_size = m_params.p_words();
   BOTAN_ASSERT_NOMSG(m_v.size() == p_size);

   if(ws.size() < 2 * p_size) {
      ws.resize(2 * p_size);
   }

   auto do_sqr_n = [&](std::span<word> z) {
      for(size_t i = 0; i != n; ++i) {
         bigint_sqr(z.data(), 2 * p_size, m_v.data(), p_size, p_size, ws.data(), ws.size());

         bigint_monty_redc_inplace(z.data(), m_params.p()._data(), p_size, m_params.p_dash(), ws.data(), ws.size());

         copy_mem(m_v, std::span{z}.first(p_size));
      }
   };

   if(p_size <= MontgomeryUseStackLimit) {
      std::array<word, 2 * MontgomeryUseStackLimit> z{};
      do_sqr_n(z);
   } else {
      secure_vector<word> z(2 * p_size);
      do_sqr_n(z);
   }

   return (*this);
}

Montgomery_Int Montgomery_Int::square(secure_vector<word>& ws) const {
   auto z = (*this);
   z.square_this_n_times(ws, 1);
   return z;
}

}  // namespace Botan
