use crate::InnerFloat::{Infinity, NaN, Zero};
use crate::{Float, emulate_float_to_float_fn, emulate_rational_to_float_fn, float_nan};
use core::cmp::Ordering::{self, *};
use malachite_base::fail_on_untested_path;
use malachite_base::num::arithmetic::traits::{
Abs, CeilingLogBase2, CheckedRoot, DivRound, DivisibleBy, IsPowerOf2, Parity, Reciprocal, Root,
RootAssign, RootRem, Sign,
};
use malachite_base::num::basic::floats::PrimitiveFloat;
use malachite_base::num::basic::integers::PrimitiveInt;
use malachite_base::num::basic::traits::One;
use malachite_base::num::comparison::traits::PartialOrdAbs;
use malachite_base::num::conversion::traits::{ExactFrom, RoundingFrom};
use malachite_base::num::logic::traits::SignificantBits;
use malachite_base::rounding_modes::RoundingMode::{self, *};
use malachite_nz::natural::Natural;
use malachite_nz::natural::arithmetic::float_extras::float_can_round;
use malachite_nz::platform::Limb;
use malachite_q::Rational;
const HIGH_K_THRESHOLD: u64 = 100;
fn root_u_exact(x: &Float, k: u64, prec: u64) -> (Float, Ordering) {
let m = x.significand_ref().unwrap();
let nu = m.trailing_zeros().unwrap();
let m_odd = m >> nu;
let e =
i128::from(x.get_exponent().unwrap()) - i128::from(m.significant_bits()) + i128::from(nu);
let root = Float::from_natural_prec_round(
(&m_odd).checked_root(k).expect("Inexact root"),
prec,
Exact,
)
.0 << e.div_round(i128::from(k), Exact).0;
(if x.is_sign_negative() { -root } else { root }, Equal)
}
fn root_u_integer(x: &Float, k: u64, prec: u64, rm: RoundingMode) -> (Float, Ordering) {
let negative = x.is_sign_negative();
let mut m = x.significand_ref().unwrap().clone();
let mut e = i128::from(x.get_exponent().unwrap()) - i128::from(m.significant_bits());
let ki = i128::from(k);
let r = e.rem_euclid(ki); let n = i128::from(prec) + i128::from(rm == Nearest);
let size_m = i128::from(m.significant_bits());
let t = (ki * n - size_m - r).div_euclid(ki) * ki + r;
let mut inexact = false;
if t >= 0 {
m <<= u64::exact_from(t);
} else {
let cut = u64::exact_from(-t);
inexact = m.trailing_zeros().unwrap() < cut;
m >>= cut;
}
e -= t;
let (mut root, rem) = m.root_rem(k);
inexact = inexact || rem != 0u32;
let size_root = root.significant_bits();
let n = u64::exact_from(n);
if size_root > n {
fail_on_untested_path(
"root_u_integer, root wider than n: the truncation above sizes m so that its root has \
exactly n bits",
);
let sh2 = size_root - n;
inexact = inexact || root.trailing_zeros().unwrap() < sh2;
root >>= sh2;
e += i128::from(sh2) * ki;
}
let mut o = Equal;
if inexact {
assert_ne!(rm, Exact, "Inexact root");
let rm_abs = if negative { -rm } else { rm };
o = if rm_abs == Ceiling || rm_abs == Up || (rm_abs == Nearest && root.odd()) {
root += Natural::ONE;
Greater
} else {
Less
};
}
let (y, o2) = Float::from_natural_prec(root, prec);
debug_assert!(o == Equal || o2 == Equal);
if o == Equal {
o = o2;
}
let y = y << e.div_round(ki, Exact).0;
if negative { (-y, o.reverse()) } else { (y, o) }
}
fn root_u_aux(x: &Float, k: u64, prec: u64, rm: RoundingMode) -> (Float, Ordering) {
let negative = x.is_sign_negative();
let rm_abs = if negative { -rm } else { rm };
let abs_x = x.abs();
let ex = i64::from(x.get_exponent().unwrap());
let mut working_prec = prec
+ 10
+ if ex == 0 {
0
} else {
ex.unsigned_abs().ceiling_log_base_2()
};
let kf = Float::exact_from(k);
let mut increment = Limb::WIDTH;
loop {
let mut t = abs_x.ln_prec_ref(working_prec).0;
t.div_prec_assign_ref(&kf, working_prec);
let exp_t = i64::from(t.get_exponent().unwrap());
t.exp_prec_assign(working_prec);
let err = match (exp_t + 2).sign() {
Greater => u64::exact_from(exp_t + 3),
Equal => 2,
Less => 1,
};
if working_prec > err
&& float_can_round(
t.significand_ref().unwrap(),
working_prec - err,
prec,
rm_abs,
)
{
let (root, o) = Float::from_float_prec_round(t, prec, rm_abs);
return if negative {
(-root, o.reverse())
} else {
(root, o)
};
}
let z = Float::from_float_prec_ref(&t, prec + u64::from(rm == Nearest)).0;
let (zk, o_pow) = z.pow_u_prec_ref(k, x.significant_bits());
if o_pow == Equal && zk == abs_x {
let (root, o) = Float::from_float_prec_round(z, prec, rm_abs);
return if negative {
(-root, o.reverse())
} else {
(root, o)
};
}
working_prec += increment;
increment = working_prec >> 1;
}
}
fn root_u_rational_generic(x: &Rational, k: u64, prec: u64, rm: RoundingMode) -> (Float, Ordering) {
let mut working_prec = prec + 10;
let mut increment = Limb::WIDTH;
let e = x.floor_log_base_2_abs();
if e.gt_abs(&0x3fff_0000) && k > 0x3fff_0000 {
let xm = x >> e; let p2_exp = Rational::from_signeds(e, i64::exact_from(k));
loop {
let mut root = Float::from_rational_prec_round_ref(&xm, working_prec, Floor)
.0
.root_u_prec(k, working_prec)
.0;
let p2 = Float::power_of_2_rational_prec_ref(&p2_exp, working_prec).0;
root.mul_prec_assign(p2, working_prec);
if float_can_round(root.significand_ref().unwrap(), working_prec - 3, prec, rm) {
return Float::from_float_prec_round(root, prec, rm);
}
working_prec += increment;
increment = working_prec >> 1;
}
}
let ki = i64::exact_from(k);
let mut end_shift = e;
let x2;
let reduced_x: &Rational;
if end_shift.gt_abs(&0x3fff_0000) {
end_shift -= end_shift.rem_euclid(ki);
x2 = x >> end_shift;
reduced_x = &x2;
} else {
end_shift = 0;
reduced_x = x;
}
loop {
let root = Float::from_rational_prec_round_ref(reduced_x, working_prec, Floor)
.0
.root_u_prec(k, working_prec)
.0;
if float_can_round(root.significand_ref().unwrap(), working_prec - 2, prec, rm) {
let (mut root, mut o) = Float::from_float_prec_round(root, prec, rm);
if end_shift != 0 {
o = root.shl_prec_round_assign_helper(i128::from(end_shift / ki), prec, rm, o);
}
return (root, o);
}
working_prec += increment;
increment = working_prec >> 1;
}
}
impl Float {
#[inline]
pub fn root_u_prec_round(self, k: u64, prec: u64, rm: RoundingMode) -> (Self, Ordering) {
self.root_u_prec_round_ref(k, prec, rm)
}
pub fn root_u_prec_round_ref(&self, k: u64, prec: u64, rm: RoundingMode) -> (Self, Ordering) {
assert_ne!(prec, 0);
if k == 0 {
return (float_nan!(), Equal);
} else if k == 1 {
return Self::from_float_prec_round_ref(self, prec, rm);
}
match self {
Self(NaN) => (float_nan!(), Equal),
Self(Infinity { sign }) => {
if *sign {
(Self(Infinity { sign: true }), Equal)
} else if k.odd() {
(Self(Infinity { sign: false }), Equal)
} else {
(float_nan!(), Equal)
}
}
Self(Zero { sign }) => (
Self(Zero {
sign: *sign || k.even(),
}),
Equal,
),
_ => {
if self.is_sign_negative() && k.even() {
return (float_nan!(), Equal);
}
if *self == 1u32 || *self == -1i32 {
return Self::from_float_prec_round_ref(self, prec, rm);
}
if k > HIGH_K_THRESHOLD {
if rm == Exact {
root_u_exact(self, k, prec)
} else {
root_u_aux(self, k, prec, rm)
}
} else {
root_u_integer(self, k, prec, rm)
}
}
}
}
#[inline]
pub fn root_u_prec(self, k: u64, prec: u64) -> (Self, Ordering) {
self.root_u_prec_round(k, prec, Nearest)
}
#[inline]
pub fn root_u_prec_ref(&self, k: u64, prec: u64) -> (Self, Ordering) {
self.root_u_prec_round_ref(k, prec, Nearest)
}
#[inline]
pub fn root_u_round(self, k: u64, rm: RoundingMode) -> (Self, Ordering) {
let prec = self.significant_bits();
self.root_u_prec_round(k, prec, rm)
}
#[inline]
pub fn root_u_round_ref(&self, k: u64, rm: RoundingMode) -> (Self, Ordering) {
self.root_u_prec_round_ref(k, self.significant_bits(), rm)
}
#[inline]
pub fn root_u_prec_round_assign(&mut self, k: u64, prec: u64, rm: RoundingMode) -> Ordering {
let (root, o) = core::mem::take(self).root_u_prec_round(k, prec, rm);
*self = root;
o
}
#[inline]
pub fn root_u_prec_assign(&mut self, k: u64, prec: u64) -> Ordering {
self.root_u_prec_round_assign(k, prec, Nearest)
}
#[inline]
pub fn root_u_round_assign(&mut self, k: u64, rm: RoundingMode) -> Ordering {
let prec = self.significant_bits();
self.root_u_prec_round_assign(k, prec, rm)
}
#[inline]
pub fn root_s_prec_round(self, k: i64, prec: u64, rm: RoundingMode) -> (Self, Ordering) {
self.root_s_prec_round_ref(k, prec, rm)
}
pub fn root_s_prec_round_ref(&self, k: i64, prec: u64, rm: RoundingMode) -> (Self, Ordering) {
if k >= 0 {
return self.root_u_prec_round_ref(k.unsigned_abs(), prec, rm);
}
assert_ne!(prec, 0);
match self {
Self(NaN) => (float_nan!(), Equal),
Self(Infinity { sign }) => {
if *sign || k.odd() {
(Self(Zero { sign: *sign }), Equal)
} else {
(float_nan!(), Equal)
}
}
Self(Zero { sign }) => {
(
Self(Infinity {
sign: *sign || k.even(),
}),
Equal,
)
}
_ => {
if self.is_sign_negative() && k.even() {
return (float_nan!(), Equal);
}
if *self == 1u32 || *self == -1i32 {
return Self::from_float_prec_round_ref(self, prec, rm);
}
if k == -1 {
return self.reciprocal_prec_round_ref(prec, rm);
} else if k == -2 {
return self.reciprocal_sqrt_prec_round_ref(prec, rm);
}
let ku = k.unsigned_abs();
if rm == Exact {
let e = i64::from(self.get_exponent().unwrap()) - 1;
if self.significand_ref().unwrap().is_power_of_2() && e.divisible_by(k) {
let (root, o) =
Self::from_float_prec_round_ref(&(Self::ONE << (e / k)), prec, Exact);
debug_assert_eq!(o, Equal);
return if self.is_sign_negative() {
(-root, Equal)
} else {
(root, Equal)
};
}
panic!("Inexact root");
}
let exact_representable = self.significand_ref().unwrap().is_power_of_2()
&& (i64::from(self.get_exponent().unwrap()) - 1).divisible_by(k);
let mut working_prec = prec + 10;
let mut increment = Limb::WIDTH;
loop {
let mut t = self.root_u_prec_ref(ku, working_prec).0;
let o = t.reciprocal_prec_round_assign(working_prec, rm);
if float_can_round(t.significand_ref().unwrap(), working_prec - 3, prec, rm)
|| (o == Equal && exact_representable)
{
return Self::from_float_prec_round(t, prec, rm);
}
working_prec += increment;
increment = working_prec >> 1;
}
}
}
}
#[inline]
pub fn root_s_prec(self, k: i64, prec: u64) -> (Self, Ordering) {
self.root_s_prec_round(k, prec, Nearest)
}
#[inline]
pub fn root_s_prec_ref(&self, k: i64, prec: u64) -> (Self, Ordering) {
self.root_s_prec_round_ref(k, prec, Nearest)
}
#[inline]
pub fn root_s_round(self, k: i64, rm: RoundingMode) -> (Self, Ordering) {
let prec = self.significant_bits();
self.root_s_prec_round(k, prec, rm)
}
#[inline]
pub fn root_s_round_ref(&self, k: i64, rm: RoundingMode) -> (Self, Ordering) {
self.root_s_prec_round_ref(k, self.significant_bits(), rm)
}
#[inline]
pub fn root_s_prec_round_assign(&mut self, k: i64, prec: u64, rm: RoundingMode) -> Ordering {
let (root, o) = core::mem::take(self).root_s_prec_round(k, prec, rm);
*self = root;
o
}
#[inline]
pub fn root_s_prec_assign(&mut self, k: i64, prec: u64) -> Ordering {
self.root_s_prec_round_assign(k, prec, Nearest)
}
#[inline]
pub fn root_s_round_assign(&mut self, k: i64, rm: RoundingMode) -> Ordering {
let prec = self.significant_bits();
self.root_s_prec_round_assign(k, prec, rm)
}
#[allow(clippy::needless_pass_by_value)]
#[inline]
pub fn root_u_rational_prec_round(
x: Rational,
k: u64,
prec: u64,
rm: RoundingMode,
) -> (Self, Ordering) {
Self::root_u_rational_prec_round_ref(&x, k, prec, rm)
}
pub fn root_u_rational_prec_round_ref(
x: &Rational,
k: u64,
prec: u64,
rm: RoundingMode,
) -> (Self, Ordering) {
assert_ne!(prec, 0);
if k == 0 {
return (float_nan!(), Equal);
} else if k == 1 {
return Self::from_rational_prec_round_ref(x, prec, rm);
}
if *x == 0u32 {
return (Self(Zero { sign: true }), Equal);
}
if *x < 0u32 && k.even() {
return (float_nan!(), Equal);
}
if let Some(root) = x.checked_root(k) {
return Self::from_rational_prec_round(root, prec, rm);
}
assert_ne!(rm, Exact, "Inexact root");
root_u_rational_generic(x, k, prec, rm)
}
#[inline]
pub fn root_u_rational_prec(x: Rational, k: u64, prec: u64) -> (Self, Ordering) {
Self::root_u_rational_prec_round(x, k, prec, Nearest)
}
#[inline]
pub fn root_u_rational_prec_ref(x: &Rational, k: u64, prec: u64) -> (Self, Ordering) {
Self::root_u_rational_prec_round_ref(x, k, prec, Nearest)
}
#[allow(clippy::needless_pass_by_value)]
#[inline]
pub fn root_s_rational_prec_round(
x: Rational,
k: i64,
prec: u64,
rm: RoundingMode,
) -> (Self, Ordering) {
Self::root_s_rational_prec_round_ref(&x, k, prec, rm)
}
pub fn root_s_rational_prec_round_ref(
x: &Rational,
k: i64,
prec: u64,
rm: RoundingMode,
) -> (Self, Ordering) {
if k >= 0 {
return Self::root_u_rational_prec_round_ref(x, k.unsigned_abs(), prec, rm);
}
assert_ne!(prec, 0);
if *x == 0u32 {
return (Self(Infinity { sign: true }), Equal);
}
if *x < 0u32 && k.even() {
return (float_nan!(), Equal);
}
Self::root_u_rational_prec_round(x.reciprocal(), k.unsigned_abs(), prec, rm)
}
#[inline]
pub fn root_s_rational_prec(x: Rational, k: i64, prec: u64) -> (Self, Ordering) {
Self::root_s_rational_prec_round(x, k, prec, Nearest)
}
#[inline]
pub fn root_s_rational_prec_ref(x: &Rational, k: i64, prec: u64) -> (Self, Ordering) {
Self::root_s_rational_prec_round_ref(x, k, prec, Nearest)
}
}
impl Root<u64> for Float {
type Output = Self;
#[inline]
fn root(self, pow: u64) -> Self {
let prec = self.significant_bits();
self.root_u_prec(pow, prec).0
}
}
impl Root<u64> for &Float {
type Output = Float;
#[inline]
fn root(self, pow: u64) -> Float {
self.root_u_prec_ref(pow, self.significant_bits()).0
}
}
impl RootAssign<u64> for Float {
#[inline]
fn root_assign(&mut self, pow: u64) {
let prec = self.significant_bits();
self.root_u_prec_round_assign(pow, prec, Nearest);
}
}
impl Root<i64> for Float {
type Output = Self;
#[inline]
fn root(self, pow: i64) -> Self {
let prec = self.significant_bits();
self.root_s_prec(pow, prec).0
}
}
impl Root<i64> for &Float {
type Output = Float;
#[inline]
fn root(self, pow: i64) -> Float {
self.root_s_prec_ref(pow, self.significant_bits()).0
}
}
impl RootAssign<i64> for Float {
#[inline]
fn root_assign(&mut self, pow: i64) {
let prec = self.significant_bits();
self.root_s_prec_round_assign(pow, prec, Nearest);
}
}
#[inline]
#[allow(clippy::type_repetition_in_bounds)]
pub fn primitive_float_root_u<T: PrimitiveFloat>(x: T, k: u64) -> T
where
Float: From<T> + PartialOrd<T>,
for<'a> T: ExactFrom<&'a Float> + RoundingFrom<&'a Float>,
{
emulate_float_to_float_fn(|f, prec| Float::root_u_prec(f, k, prec), x)
}
#[inline]
#[allow(clippy::type_repetition_in_bounds)]
pub fn primitive_float_root_s<T: PrimitiveFloat>(x: T, k: i64) -> T
where
Float: From<T> + PartialOrd<T>,
for<'a> T: ExactFrom<&'a Float> + RoundingFrom<&'a Float>,
{
emulate_float_to_float_fn(|f, prec| Float::root_s_prec(f, k, prec), x)
}
#[inline]
#[allow(clippy::type_repetition_in_bounds)]
pub fn primitive_float_root_u_rational<T: PrimitiveFloat>(x: &Rational, k: u64) -> T
where
Float: PartialOrd<T>,
for<'a> T: ExactFrom<&'a Float> + RoundingFrom<&'a Float>,
{
emulate_rational_to_float_fn(|q, prec| Float::root_u_rational_prec_ref(q, k, prec), x)
}
#[inline]
#[allow(clippy::type_repetition_in_bounds)]
pub fn primitive_float_root_s_rational<T: PrimitiveFloat>(x: &Rational, k: i64) -> T
where
Float: PartialOrd<T>,
for<'a> T: ExactFrom<&'a Float> + RoundingFrom<&'a Float>,
{
emulate_rational_to_float_fn(|q, prec| Float::root_s_rational_prec_ref(q, k, prec), x)
}