num-valid 0.3.3

#![deny(rustdoc::broken_intra_doc_links)]

//! Square root operations.
//!
//! This module provides the [`Sqrt`] trait for computing square roots of validated
//! real and complex numbers.

use crate::{
    core::{errors::capture_backtrace, policies::StrictFinitePolicy},
    functions::FunctionErrors,
    kernels::{RawComplexTrait, RawRealTrait, RawScalarTrait},
};
use num::Complex;
use std::{backtrace::Backtrace, fmt};
use thiserror::Error;
use try_create::ValidationPolicy;

//------------------------------------------------------------------------------------------------
/// Errors that can occur during the input validation phase when attempting to compute
/// the square root of a real number.
///
/// This enum is used as the source for the [`Input`](SqrtRealErrors::Input) variant of [`SqrtRealErrors`].
/// It is generic over `RawReal: RawRealTrait`, which defines the specific raw real number type and its associated
/// validation error type.
///
/// # Variants
///
/// - [`Self::NegativeValue`]: The input value is negative, which is not allowed for real square roots.
/// - [`Self::ValidationError`]: The input value failed basic validation checks (e.g., NaN, infinity, subnormal) according to the validation policy.
#[derive(Debug, Error)]
pub enum SqrtRealInputErrors<RawReal: RawRealTrait> {
    /// The input value for the square root computation is negative.
    ///
    /// The square root of a negative real number is not a real number.
    #[error("the input value ({value:?}) is negative!")]
    NegativeValue {
        /// The negative input value.
        value: RawReal,

        /// The backtrace of the error.
        backtrace: Backtrace,
    },

    /// The input value failed basic validation checks (e.g., it is NaN, infinite, or subnormal).
    ///
    /// This error occurs if the input value itself is considered invalid
    /// according to the validation policy (e.g., [`StrictFinitePolicy`]),
    /// before the domain-specific check (like negativity) for the square root
    /// operation is performed.
    #[error("the input value is invalid!")]
    ValidationError {
        /// The source error that occurred during validation.
        #[source]
        #[backtrace]
        source: <RawReal as RawScalarTrait>::ValidationErrors,
    },
}

/// Errors that can occur during the input validation phase when computing the square root of a complex number.
///
/// This enum is used as the source for the [`Input`](SqrtComplexErrors::Input) variant of [`SqrtComplexErrors`].
/// It is generic over `RawComplex: RawComplexTrait`.
///
/// # Variants
///
/// - [`Self::ValidationError`]: The input complex number failed basic validation checks (e.g., its components are NaN, infinite, or subnormal) according to the validation policy.
#[derive(Debug, Error)]
pub enum SqrtComplexInputErrors<RawComplex: RawComplexTrait> {
    /// The input complex number failed basic validation checks (e.g., its components are NaN, infinite, or subnormal).
    ///
    /// This error occurs if the input complex value itself is considered invalid
    /// according to the validation policy (e.g., [`StrictFinitePolicy`]),
    /// before any domain-specific checks for the complex square root operation are performed.
    /// It signifies that the real or imaginary part (or both) does not meet fundamental validity criteria.
    #[error("the input value is invalid!")]
    ValidationError {
        /// The detailed source error from the raw complex number's validation.
        ///
        /// This field encapsulates the specific error of type `<RawComplex as RawScalarTrait>::ValidationErrors`
        /// that was reported during the validation of the input complex number's components.
        #[source]
        #[backtrace]
        source: <RawComplex as RawScalarTrait>::ValidationErrors,
    },
}

/// A type alias for [`FunctionErrors`], specialized for errors that can occur during
/// the computation of the square root of a real number.
///
/// This type represents the possible failures when calling [`Sqrt::try_sqrt()`] on a real number.
///
/// # Generic Parameters
///
/// - `RawReal`: A type that implements [`RawRealTrait`]. This defines:
///   - The raw error type for the input real number via `SqrtRealInputErrors<RawReal>`.
///   - The raw error type for the output real number (the square root) also via `<RawReal as RawScalarTrait>::ValidationErrors`.
///
/// # Variants
///
/// This type alias wraps [`FunctionErrors`], which has the following variants in this context:
///
/// - `Input { source: SqrtRealInputErrors<RawReal> }`:
///   Indicates that the input real number was invalid for square root computation.
///   This could be due to the number being negative or failing general validation checks
///   (e.g., NaN, Infinity, subnormal). The `source` field provides more specific details
///   via [`SqrtRealInputErrors`].
///
/// - `Output { source: <RawReal as RawScalarTrait>::ValidationErrors }`:
///   Indicates that the computed square root (which should be a real number)
///   failed validation. This typically means the result of the `sqrt` operation yielded
///   a non-finite value (NaN or Infinity), which is unexpected if the input was valid
///   (non-negative and finite).
pub type SqrtRealErrors<RawReal> =
    FunctionErrors<SqrtRealInputErrors<RawReal>, <RawReal as RawScalarTrait>::ValidationErrors>;

/// A type alias for [`FunctionErrors`], specialized for errors that can occur during
/// the computation of the square root of a complex number.
///
/// This type represents the possible failures when calling [`Sqrt::try_sqrt()`] on a complex number.
///
/// # Generic Parameters
///
/// - `RawComplex`: A type that implements [`RawComplexTrait`]. This defines:
///   - The raw error type for the input complex number via `SqrtComplexInputErrors<RawComplex>`.
///   - The raw error type for the output complex number (the square root) also via `<RawComplex as RawScalarTrait>::ValidationErrors`.
///
/// # Variants
///
/// This type alias wraps [`FunctionErrors`], which has the following variants in this context:
///
/// - `Input { source: SqrtComplexInputErrors<RawComplex> }`:
///   Indicates that the input complex number was invalid for square root computation.
///   This is typically due to the complex number's components (real or imaginary parts)
///   failing general validation checks (e.g., NaN, Infinity, subnormal).
///   The `source` field provides more specific details via [`SqrtComplexInputErrors`].
///
/// - `Output { source: <RawComplex as RawScalarTrait>::ValidationErrors }`:
///   Indicates that the computed complex square root itself failed validation.
///   This typically means the result of the `sqrt` operation yielded a complex number
///   with non-finite components (NaN or Infinity), which is unexpected if the input was valid.
pub type SqrtComplexErrors<RawComplex> = FunctionErrors<
    SqrtComplexInputErrors<RawComplex>,
    <RawComplex as RawScalarTrait>::ValidationErrors,
>;
//--------------------------------------------------------------------------------------------

//--------------------------------------------------------------------------------------------
/// A trait for computing the principal square root of a number.
///
/// The principal square root of a non-negative real number `x` is the unique non-negative real number `y`
/// such that `y^2 = x`.
/// For a complex number `z`, its square root `w` satisfies `w^2 = z`. Complex numbers
/// (except 0) have two square roots; this trait computes the principal square root,
/// typically defined as `exp(0.5 * log(z))`, which usually has a non-negative real part.
///
/// This trait provides both a fallible version ([`try_sqrt`](Sqrt::try_sqrt)) that performs validation
/// and an infallible version ([`sqrt`](Sqrt::sqrt)) that may panic in debug builds if validation fails.
pub trait Sqrt: Sized {
    /// The error type that can be returned by the `try_sqrt` method.
    type Error: fmt::Debug;

    /// Attempts to compute the principal square root of `self`, returning a `Result`.
    ///
    /// Implementations should validate the input `self` according to the domain
    /// (e.g., non-negative for reals) and a general validation policy (e.g., [`StrictFinitePolicy`]).
    /// If the input is valid, the square root is computed, and then the result
    /// is also validated using the same policy.
    ///
    /// # Returns
    ///
    /// - `Ok(Self)`: If the input is valid for the square root operation and both the input
    ///   and the computed square root satisfy the validation policy.
    /// - `Err(SqrtRealErrors)`: If the input is invalid (e.g., negative for real sqrt, NaN, Infinity)
    ///   or if the computed square root is invalid (see below).
    ///
    /// # Errors
    ///
    /// - Returns [`SqrtRealErrors::Input`] (for reals) or [`SqrtComplexErrors::Input`] (for complex)
    ///   via [`FunctionErrors::Input`] if the input is invalid (e.g., negative real, NaN, Infinity, subnormal).
    /// - Returns [`SqrtRealErrors::Output`] or [`SqrtComplexErrors::Output`] via [`FunctionErrors::Output`]
    ///   if the result of the computation is not finite (e.g., NaN, Infinity) as per the validation policy.
    ///
    /// # Examples
    ///
    /// ```
    /// use num_valid::functions::Sqrt;
    /// use num::Complex;
    ///
    /// assert_eq!(4.0_f64.try_sqrt().unwrap(), 2.0);
    /// assert!((-1.0_f64).try_sqrt().is_err()); // Negative real
    /// assert!(f64::NAN.try_sqrt().is_err());
    ///
    /// let z = Complex::new(-4.0, 0.0); // sqrt(-4) = 2i
    /// let sqrt_z = z.try_sqrt().unwrap();
    /// assert!((sqrt_z.re).abs() < 1e-9 && (sqrt_z.im - 2.0).abs() < 1e-9);
    /// ```
    #[must_use = "this `Result` may contain an error that should be handled"]
    fn try_sqrt(self) -> Result<Self, <Self as Sqrt>::Error>;

    /// Computes the square principal square root of `self`.
    fn sqrt(self) -> Self;
}

impl Sqrt for f64 {
    type Error = SqrtRealErrors<f64>;

    /// Attempts to compute the principal square root of `self`, returning a `Result`.
    ///
    /// Implementations should validate the input `self` according to the domain
    /// (e.g., non-negative for reals) and a general validation policy (e.g., [`StrictFinitePolicy`]).
    /// If the input is valid, the square root is computed, and then the result
    /// is also validated using the same policy.
    ///
    /// # Returns
    ///
    /// - `Ok(Self)`: If the input is valid for the square root operation and both the input
    ///   and the computed square root satisfy the validation policy.
    /// - `Err(Self::Error)`: If the input is invalid (e.g., negative for real sqrt, NaN, Infinity)
    ///   or if the computed square root is invalid.
    ///
    /// # Examples
    ///
    /// ```
    /// use num_valid::functions::Sqrt;
    /// use num::Complex;
    ///
    /// assert_eq!(4.0_f64.try_sqrt().unwrap(), 2.0);
    /// assert!((-1.0_f64).try_sqrt().is_err()); // Negative real
    /// assert!(f64::NAN.try_sqrt().is_err());
    ///
    /// let z = Complex::new(-4.0, 0.0); // sqrt(-4) = 2i
    /// let sqrt_z = z.try_sqrt().unwrap();
    /// assert!((sqrt_z.re).abs() < 1e-9 && (sqrt_z.im - 2.0).abs() < 1e-9);
    /// ```
    #[inline(always)]
    fn try_sqrt(self) -> Result<f64, <f64 as Sqrt>::Error> {
        StrictFinitePolicy::<f64, 53>::validate(self)
            .map_err(|e| SqrtRealInputErrors::ValidationError { source: e }.into())
            .and_then(|validated_value| {
                if validated_value < 0.0 {
                    Err(SqrtRealInputErrors::NegativeValue {
                        value: validated_value,
                        backtrace: capture_backtrace(),
                    }
                    .into())
                } else {
                    StrictFinitePolicy::<f64, 53>::validate(f64::sqrt(validated_value))
                        .map_err(|e| SqrtRealErrors::Output { source: e })
                }
            })
    }

    /// Computes and returns the principal square root of `self`.
    ///
    /// # Behavior
    ///
    /// - **Debug Builds (`#[cfg(debug_assertions)]`)**: This method internally calls
    ///   [`try_sqrt().unwrap()`](Sqrt::try_sqrt). It will panic if `try_sqrt` returns an `Err`.
    /// - **Release Builds (`#[cfg(not(debug_assertions))]`)**: This method calls the underlying
    ///   square root function directly (e.g., `f64::sqrt`).
    ///   The behavior for invalid inputs (like `sqrt(-1.0)` for `f64` returning NaN)
    ///   depends on the underlying implementation.
    ///
    /// # Panics
    ///
    /// In debug builds, this method will panic if [`try_sqrt()`](Sqrt::try_sqrt) would return an `Err`.
    ///
    /// # Examples
    ///
    /// ```
    /// use num_valid::functions::Sqrt;
    /// use num::Complex;
    ///
    /// assert_eq!(9.0_f64.sqrt(), 3.0);
    ///
    /// // For f64, sqrt of negative is NaN in release, panics in debug with ftl's Sqrt
    /// #[cfg(not(debug_assertions))]
    /// assert!((-1.0_f64).sqrt().is_nan());
    ///
    /// let z = Complex::new(0.0, 4.0); // sqrt(4i) = sqrt(2) + i*sqrt(2)
    /// let sqrt_z = z.sqrt();
    /// let expected_val = std::f64::consts::SQRT_2;
    /// assert!((sqrt_z.re - expected_val).abs() < 1e-9 && (sqrt_z.im - expected_val).abs() < 1e-9);
    /// ```
    #[inline(always)]
    fn sqrt(self) -> Self {
        #[cfg(debug_assertions)]
        {
            self.try_sqrt().unwrap()
        }
        #[cfg(not(debug_assertions))]
        {
            f64::sqrt(self)
        }
    }
}

impl Sqrt for Complex<f64> {
    type Error = SqrtComplexErrors<Complex<f64>>;

    /// Attempts to compute the principal square root of `self` (a `Complex<f64>`).
    ///
    /// This method first validates `self` using [`StrictFinitePolicy`] (components must be finite and normal).
    /// If valid, it computes `Complex::sqrt` and validates the result using [`StrictFinitePolicy`].
    ///
    /// # Returns
    ///
    /// - `Ok(Complex<f64>)`: If `self` and the computed square root have finite and normal components.
    /// - `Err(SqrtComplexErrors<Complex<f64>>)`: If `self` or the result has invalid components.
    #[inline(always)]
    fn try_sqrt(self) -> Result<Self, <Self as Sqrt>::Error> {
        StrictFinitePolicy::<Complex<f64>, 53>::validate(self)
            .map_err(|e| SqrtComplexInputErrors::ValidationError { source: e }.into())
            .and_then(|validated_value| {
                StrictFinitePolicy::<Complex<f64>, 53>::validate(Complex::<f64>::sqrt(
                    validated_value,
                ))
                .map_err(|e| SqrtComplexErrors::Output { source: e })
            })
    }

    /// Computes and returns the principal square root of `self` (a `Complex<f64>`).
    ///
    /// # Behavior
    ///
    /// - **Debug Builds**: Calls `try_sqrt().unwrap()`. Panics on invalid input/output.
    /// - **Release Builds**: Calls `Complex::sqrt(self)` directly.
    ///
    /// # Panics
    ///
    /// In debug builds, if `try_sqrt()` would return an `Err`.
    #[inline(always)]
    fn sqrt(self) -> Self {
        #[cfg(debug_assertions)]
        {
            self.try_sqrt().unwrap()
        }
        #[cfg(not(debug_assertions))]
        {
            Complex::<f64>::sqrt(self)
        }
    }
}

//------------------------------------------------------------------------------------------------

//------------------------------------------------------------------------------------------------
#[cfg(test)]
mod tests {
    use super::*;
    use num::Complex;

    #[cfg(feature = "rug")]
    use try_create::TryNew;

    mod sqrt {
        use super::*;

        mod native64 {
            use super::*;

            mod real {
                use super::*;

                #[test]
                fn test_f64_sqrt_valid() {
                    let value = 4.0;

                    assert_eq!(value.try_sqrt().unwrap(), 2.0);
                    assert_eq!(<f64 as Sqrt>::sqrt(value), 2.0);
                }

                #[test]
                fn test_f64_sqrt_negative_value() {
                    let value = -4.0;
                    let result = value.try_sqrt();
                    assert!(matches!(result, Err(SqrtRealErrors::Input { .. })));
                }

                #[test]
                fn test_f64_sqrt_subnormal() {
                    let value = f64::MIN_POSITIVE / 2.0;
                    let result = value.try_sqrt();
                    assert!(matches!(result, Err(SqrtRealErrors::Input { .. })));
                }

                #[test]
                fn test_f64_sqrt_zero() {
                    let value = 0.0;
                    let result = value.try_sqrt();
                    assert!(matches!(result, Ok(0.0)));
                }

                #[test]
                fn test_f64_sqrt_nan() {
                    let value = f64::NAN;
                    let result = value.try_sqrt();
                    assert!(matches!(result, Err(SqrtRealErrors::Input { .. })));
                }

                #[test]
                fn test_f64_sqrt_infinite() {
                    let value = f64::INFINITY;
                    let result = value.try_sqrt();
                    println!("result: {result:?}");
                    assert!(matches!(result, Err(SqrtRealErrors::Input { .. })));
                }
            }

            mod complex {
                use super::*;

                #[test]
                fn test_complex_f64_sqrt_valid() {
                    let value = Complex::new(4.0, 0.0);

                    let expected_result = Complex::new(2.0, 0.0);

                    assert_eq!(value.try_sqrt().unwrap(), expected_result);
                    assert_eq!(<Complex<f64> as Sqrt>::sqrt(value), expected_result);
                }

                #[test]
                fn test_complex_f64_sqrt_invalid() {
                    let value = Complex::new(f64::NAN, 0.0);
                    let result = value.try_sqrt();
                    assert!(matches!(result, Err(SqrtComplexErrors::Input { .. })));

                    let value = Complex::new(0.0, f64::NAN);
                    let result = value.try_sqrt();
                    assert!(matches!(result, Err(SqrtComplexErrors::Input { .. })));

                    let value = Complex::new(f64::INFINITY, 0.0);
                    let result = value.try_sqrt();
                    assert!(matches!(result, Err(SqrtComplexErrors::Input { .. })));

                    let value = Complex::new(0.0, f64::INFINITY);
                    let result = value.try_sqrt();
                    assert!(matches!(result, Err(SqrtComplexErrors::Input { .. })));

                    let value = Complex::new(f64::MIN_POSITIVE / 2.0, 0.0);
                    let result = value.try_sqrt();
                    assert!(matches!(result, Err(SqrtComplexErrors::Input { .. })));

                    let value = Complex::new(0., f64::MIN_POSITIVE / 2.0);
                    let result = value.try_sqrt();
                    assert!(matches!(result, Err(SqrtComplexErrors::Input { .. })));
                }
            }
        }

        #[cfg(feature = "rug")]
        mod rug53 {
            use super::*;
            use crate::backends::rug::validated::{ComplexRugStrictFinite, RealRugStrictFinite};

            mod real {
                use super::*;

                #[test]
                fn test_rug_float_sqrt_valid() {
                    let value =
                        RealRugStrictFinite::<53>::try_new(rug::Float::with_val(53, 4.0)).unwrap();
                    let expected_result =
                        RealRugStrictFinite::<53>::try_new(rug::Float::with_val(53, 2.0)).unwrap();

                    assert_eq!(value.clone().try_sqrt().unwrap(), expected_result);
                    assert_eq!(value.sqrt(), expected_result);
                }

                #[test]
                fn test_rug_float_sqrt_negative_value() {
                    let value =
                        RealRugStrictFinite::<53>::try_new(rug::Float::with_val(53, -4.0)).unwrap();
                    let result = value.try_sqrt();
                    assert!(matches!(result, Err(SqrtRealErrors::Input { .. })));
                }
            }

            mod complex {
                use super::*;

                #[test]
                fn test_complex_rug_float_sqrt_valid() {
                    let value = ComplexRugStrictFinite::<53>::try_new(rug::Complex::with_val(
                        53,
                        (rug::Float::with_val(53, 4.0), rug::Float::with_val(53, 0.0)),
                    ))
                    .unwrap();

                    let expected_result =
                        ComplexRugStrictFinite::<53>::try_new(rug::Complex::with_val(
                            53,
                            (rug::Float::with_val(53, 2.0), rug::Float::with_val(53, 0.0)),
                        ))
                        .unwrap();

                    assert_eq!(value.clone().try_sqrt().unwrap(), expected_result);
                    assert_eq!(value.sqrt(), expected_result);
                }
            }
        }
    }
}