use crate::div::{WIDE_LIMBS, div_rem_512, low_limbs};
use crate::{LIMBS, U256};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
pub enum Rounding {
#[default]
Floor,
Ceil,
Trunc,
Expand,
}
impl Rounding {
#[inline]
pub const fn unsigned_rounds_up(self) -> bool {
matches!(self, Self::Ceil | Self::Expand)
}
#[inline]
pub const fn rounds_away_from_zero(self, negative: bool) -> bool {
match self {
Self::Floor => negative,
Self::Ceil => !negative,
Self::Trunc => false,
Self::Expand => true,
}
}
}
impl U256 {
#[inline]
pub const fn min(self, other: Self) -> Self {
if self.const_lt(other) { self } else { other }
}
#[inline]
pub const fn max(self, other: Self) -> Self {
if self.const_gt(other) { self } else { other }
}
#[inline]
pub const fn clamp(self, min: Self, max: Self) -> Self {
if min.const_gt(max) {
panic!("min must be less than or equal to max");
}
self.max(min).min(max)
}
#[inline]
pub const fn average(self, other: Self) -> Self {
self.const_bitand(other)
.wrapping_add(self.const_bitxor(other).wrapping_shr(1))
}
#[inline]
pub const fn checked_ceil_div(self, rhs: Self) -> Option<Self> {
if self.is_zero() {
if rhs.is_zero() {
return None;
}
return Some(Self::ZERO);
}
match self.wrapping_sub(Self::ONE).checked_div(rhs) {
Some(q) => Some(q.wrapping_add(Self::ONE)),
None => None,
}
}
#[inline]
pub const fn ceil_div(self, rhs: Self) -> Self {
match self.checked_ceil_div(rhs) {
Some(v) => v,
None => panic!("attempt to divide by zero"),
}
}
#[inline]
pub const fn checked_div_rounding(
self,
rhs: Self,
rounding: Rounding,
) -> Option<Self> {
match self.checked_div_rem(rhs) {
Some((q, r)) => {
if rounding.unsigned_rounds_up() && !r.is_zero() {
Some(q.wrapping_add(Self::ONE))
} else {
Some(q)
}
}
None => None,
}
}
#[inline]
pub const fn div_rounding(self, rhs: Self, rounding: Rounding) -> Self {
match self.checked_div_rounding(rhs, rounding) {
Some(v) => v,
None => panic!("attempt to divide by zero"),
}
}
#[inline]
pub const fn checked_mul_div(x: Self, y: Self, d: Self) -> Option<Self> {
match Self::mul_div_parts(x, y, d) {
Some((q, _)) => Some(q),
None => None,
}
}
#[inline]
pub const fn mul_div(x: Self, y: Self, d: Self) -> Self {
match Self::checked_mul_div(x, y, d) {
Some(v) => v,
None => {
if d.is_zero() {
panic!("attempt to divide by zero");
}
panic!("mul_div quotient exceeds U256");
}
}
}
#[inline]
pub const fn checked_mul_div_rounding(
x: Self,
y: Self,
d: Self,
rounding: Rounding,
) -> Option<Self> {
match Self::mul_div_parts(x, y, d) {
Some((q, r)) => {
if rounding.unsigned_rounds_up() && !r.is_zero() {
q.checked_add(Self::ONE)
} else {
Some(q)
}
}
None => None,
}
}
#[inline]
pub const fn mul_div_rounding(
x: Self,
y: Self,
d: Self,
rounding: Rounding,
) -> Self {
match Self::checked_mul_div_rounding(x, y, d, rounding) {
Some(v) => v,
None => {
if d.is_zero() {
panic!("attempt to divide by zero");
}
panic!("mul_div quotient exceeds U256");
}
}
}
#[inline]
pub(crate) const fn mul_div_parts(
x: Self,
y: Self,
d: Self,
) -> Option<(Self, Self)> {
let (lo, hi) = x.widening_mul(y);
let (q, r) = match div_rem_512(hi, lo, d) {
Some(v) => v,
None => return None,
};
let mut i = LIMBS;
while i < WIDE_LIMBS {
if q[i] != 0 {
return None;
}
i += 1;
}
Some((Self::from_limbs(low_limbs(&q)), r))
}
#[inline]
pub const fn checked_mul_shr(x: Self, y: Self, n: u32) -> Option<Self> {
match Self::mul_shr_parts(x, y, n) {
Some((v, _)) => Some(v),
None => None,
}
}
#[inline]
const fn mul_shr_parts(x: Self, y: Self, n: u32) -> Option<(Self, bool)> {
if n >= Self::BITS {
return None;
}
let (lo, hi) = x.widening_mul(y);
if hi.bit_len() > n {
return None;
}
if n == 0 {
return Some((lo, false));
}
let dropped = !lo
.const_bitand(Self::MAX.wrapping_shr(Self::BITS - n))
.is_zero();
Some((
hi.wrapping_shl(Self::BITS - n)
.const_bitor(lo.wrapping_shr(n)),
dropped,
))
}
#[inline]
pub const fn mul_shr(x: Self, y: Self, n: u32) -> Self {
match Self::checked_mul_shr(x, y, n) {
Some(v) => v,
None => {
if n >= Self::BITS {
panic!("attempt to shift right with overflow");
}
panic!("mul_shr result exceeds U256");
}
}
}
#[inline]
pub const fn checked_mul_shr_rounding(
x: Self,
y: Self,
n: u32,
rounding: Rounding,
) -> Option<Self> {
let (v, dropped) = match Self::mul_shr_parts(x, y, n) {
Some(pair) => pair,
None => return None,
};
if dropped && rounding.unsigned_rounds_up() {
v.checked_add(Self::ONE)
} else {
Some(v)
}
}
#[inline]
pub const fn mul_shr_rounding(
x: Self,
y: Self,
n: u32,
rounding: Rounding,
) -> Self {
match Self::checked_mul_shr_rounding(x, y, n, rounding) {
Some(v) => v,
None => {
if n >= Self::BITS {
panic!("attempt to shift right with overflow");
}
panic!("mul_shr result exceeds U256");
}
}
}
#[inline]
pub const fn checked_add_mod(x: Self, y: Self, m: Self) -> Option<Self> {
let (x, y) = match (x.checked_rem(m), y.checked_rem(m)) {
(Some(x), Some(y)) => (x, y),
_ => return None,
};
let (sum, carry) = x.overflowing_add(y);
if carry || sum.const_ge(m) {
Some(sum.wrapping_sub(m))
} else {
Some(sum)
}
}
#[inline]
pub const fn add_mod(x: Self, y: Self, m: Self) -> Self {
match Self::checked_add_mod(x, y, m) {
Some(v) => v,
None => panic!("attempt to divide by zero"),
}
}
#[inline]
pub const fn checked_mul_mod(x: Self, y: Self, m: Self) -> Option<Self> {
let (lo, hi) = x.widening_mul(y);
match div_rem_512(hi, lo, m) {
Some((_, r)) => Some(r),
None => None,
}
}
#[inline]
pub const fn mul_mod(x: Self, y: Self, m: Self) -> Self {
match Self::checked_mul_mod(x, y, m) {
Some(v) => v,
None => panic!("attempt to divide by zero"),
}
}
pub const fn checked_mod_exp(b: Self, e: Self, m: Self) -> Option<Self> {
if m.is_zero() {
return None;
}
if m.const_eq(Self::ONE) {
return Some(Self::ZERO);
}
let mut result = Self::ONE;
let mut base = match b.checked_rem(m) {
Some(v) => v,
None => return None,
};
let bits = e.bit_len();
let mut i = 0;
while i < bits {
if e.bit(i) {
result = match Self::checked_mul_mod(result, base, m) {
Some(v) => v,
None => return None,
};
}
i += 1;
if i < bits {
base = match Self::checked_mul_mod(base, base, m) {
Some(v) => v,
None => return None,
};
}
}
Some(result)
}
pub const fn mod_exp(b: Self, e: Self, m: Self) -> Self {
match Self::checked_mod_exp(b, e, m) {
Some(v) => v,
None => panic!("attempt to divide by zero"),
}
}
pub const fn inv_mod(a: Self, n: Self) -> Self {
if n.is_zero() {
return Self::ZERO;
}
let mut r0 = n;
let mut r1 = match a.checked_rem(n) {
Some(v) => v,
None => return Self::ZERO,
};
let mut t0 = Self::ZERO;
let mut t1 = Self::ONE;
let mut k_odd = true;
while !r1.is_zero() {
let (q, r) = match r0.checked_div_rem(r1) {
Some(v) => v,
None => return Self::ZERO,
};
r0 = r1;
r1 = r;
let t2 = t0.wrapping_add(q.wrapping_mul(t1));
t0 = t1;
t1 = t2;
k_odd = !k_odd;
}
if !r0.const_eq(Self::ONE) {
return Self::ZERO;
}
if !k_odd || t0.is_zero() {
t0
} else {
n.wrapping_sub(t0)
}
}
pub const fn inv_mod_prime(a: Self, p: Self) -> Self {
Self::mod_exp(a, p.wrapping_sub(Self::from_u32(2)), p)
}
pub const fn overflowing_pow(self, exp: u32) -> (Self, bool) {
let mut base = self;
let mut exp = exp;
let mut acc = Self::ONE;
let mut overflow = false;
while exp > 0 {
if exp & 1 == 1 {
let (v, o) = acc.overflowing_mul(base);
acc = v;
overflow |= o;
}
exp >>= 1;
if exp > 0 {
let (v, o) = base.overflowing_mul(base);
base = v;
overflow |= o;
}
}
(acc, overflow)
}
#[inline]
pub const fn checked_pow(self, exp: u32) -> Option<Self> {
match self.overflowing_pow(exp) {
(v, false) => Some(v),
(_, true) => None,
}
}
#[inline]
pub const fn wrapping_pow(self, exp: u32) -> Self {
self.overflowing_pow(exp).0
}
#[inline]
pub const fn saturating_pow(self, exp: u32) -> Self {
match self.overflowing_pow(exp) {
(v, false) => v,
(_, true) => Self::MAX,
}
}
#[inline]
pub const fn pow(self, exp: u32) -> Self {
let (v, overflow) = self.overflowing_pow(exp);
if crate::arith::OVERFLOW_CHECKS && overflow {
panic!("attempt to multiply with overflow");
}
v
}
#[inline]
pub const fn sqrt_rounding(self, rounding: Rounding) -> Self {
let r = self.integer_sqrt();
if rounding.unsigned_rounds_up() && !r.wrapping_mul(r).const_eq(self) {
r.wrapping_add(Self::ONE)
} else {
r
}
}
#[inline]
pub const fn log2(self) -> u32 {
self.bit_len().saturating_sub(1)
}
#[inline]
pub const fn log2_rounding(self, rounding: Rounding) -> u32 {
let r = self.log2();
if rounding.unsigned_rounds_up()
&& Self::ONE.wrapping_shl(r).const_lt(self)
{
r + 1
} else {
r
}
}
pub const fn log10(self) -> u32 {
self.log10_parts().0
}
const fn log10_parts(self) -> (u32, bool) {
const P8: u32 = 100_000_000;
if self.is_zero() {
return (0, false);
}
let mut v = self;
let mut r = 0u32;
let mut exact = true;
while v.limb_len() > 1 || v.to_limbs()[0] >= P8 {
let (q, rem) = v.div_rem_small(P8);
if rem != 0 {
exact = false;
}
v = q;
r += 8;
}
let (tail, tail_exact) = log10_u32(v.to_limbs()[0]);
(r + tail, !(exact && tail_exact))
}
#[inline]
pub const fn log10_rounding(self, rounding: Rounding) -> u32 {
let (r, inexact) = self.log10_parts();
if rounding.unsigned_rounds_up() && inexact {
r + 1
} else {
r
}
}
#[inline]
pub const fn log256(self) -> u32 {
self.bit_len().saturating_sub(1) / 8
}
#[inline]
pub const fn log256_rounding(self, rounding: Rounding) -> u32 {
let r = self.log256();
if rounding.unsigned_rounds_up()
&& Self::ONE.wrapping_shl(r * 8).const_lt(self)
{
r + 1
} else {
r
}
}
}
#[inline]
const fn log10_u32(x: u32) -> (u32, bool) {
let mut x = x;
let mut r = 0;
let mut exact = true;
if x >= 10_000 {
exact &= x.is_multiple_of(10_000);
x /= 10_000;
r += 4;
}
if x >= 100 {
exact &= x.is_multiple_of(100);
x /= 100;
r += 2;
}
if x >= 10 {
exact &= x.is_multiple_of(10);
x /= 10;
r += 1;
}
(r, exact && x == 1)
}