qfall-math 0.2.1

// Copyright 2024 Marcel Luca Schmidt
//
// This file is part of qFALL-math.
//
// qFALL-math is free software: you can redistribute it and/or modify it under
// the terms of the Mozilla Public License Version 2.0 as published by the
// Mozilla Foundation. See <https://mozilla.org/en-US/MPL/2.0/>.

//! This module includes functionality for rounding instances of [`PolyOverQ`] coefficient-wise.

use super::PolyOverQ;
use crate::{
    error::MathError,
    integer::PolyOverZ,
    macros::seeded::seedable_function,
    rational::Q,
    traits::{GetCoefficient, SetCoefficient},
    utils::sample::uniform::SamplerRng,
};

impl PolyOverQ {
    /// Rounds all coefficients of the given rational polynomial [`PolyOverQ`] down to the next integer
    /// as a [`PolyOverZ`].
    ///
    /// # Examples
    /// ```
    /// use qfall_math::rational::PolyOverQ;
    /// use qfall_math::integer::PolyOverZ;
    /// use std::str::FromStr;
    ///
    /// let value = PolyOverQ::from_str("2  5/2 1").unwrap();
    /// assert_eq!(PolyOverZ::from_str("2  2 1").unwrap(), value.floor());
    ///
    /// let value = PolyOverQ::from_str("2  -5/2 1").unwrap();
    /// assert_eq!(PolyOverZ::from_str("2  -3 1").unwrap(), value.floor());
    /// ```
    pub fn floor(&self) -> PolyOverZ {
        let mut out = PolyOverZ::from(unsafe { self.get_coeff_unchecked(0).floor() });
        for i in 1..self.get_degree() + 1 {
            let coeff = unsafe { self.get_coeff_unchecked(i).floor() };
            unsafe { out.set_coeff_unchecked(i, coeff) };
        }

        out
    }

    /// Rounds all coefficients of the given rational polynomial [`PolyOverQ`] up to the next integer
    /// as a [`PolyOverZ`].
    ///
    /// # Examples
    /// ```
    /// use qfall_math::rational::PolyOverQ;
    /// use qfall_math::integer::PolyOverZ;
    /// use std::str::FromStr;
    ///
    /// let value = PolyOverQ::from_str("2  5/2 1").unwrap();
    /// assert_eq!(PolyOverZ::from_str("2  3 1").unwrap(), value.ceil());
    ///
    /// let value = PolyOverQ::from_str("2  -5/2 1").unwrap();
    /// assert_eq!(PolyOverZ::from_str("2  -2 1").unwrap(), value.ceil());
    /// ```
    pub fn ceil(&self) -> PolyOverZ {
        let mut out = PolyOverZ::from(unsafe { self.get_coeff_unchecked(0).ceil() });
        for i in 1..self.get_degree() + 1 {
            let coeff = unsafe { self.get_coeff_unchecked(i).ceil() };
            unsafe { out.set_coeff_unchecked(i, coeff) };
        }

        out
    }

    /// Rounds all coefficients of the given rational polynomial [`PolyOverQ`] to the closest integer
    /// as a [`PolyOverZ`].
    ///
    /// # Examples
    /// ```
    /// use qfall_math::rational::PolyOverQ;
    /// use qfall_math::integer::PolyOverZ;
    /// use std::str::FromStr;
    ///
    /// let value = PolyOverQ::from_str("2  5/2 1").unwrap();
    /// assert_eq!(PolyOverZ::from_str("2  3 1").unwrap(), value.round());
    ///
    /// let value = PolyOverQ::from_str("2  -5/2 1").unwrap();
    /// assert_eq!(PolyOverZ::from_str("2  -2 1").unwrap(), value.round());
    /// ```
    pub fn round(&self) -> PolyOverZ {
        let mut out = PolyOverZ::from(unsafe { self.get_coeff_unchecked(0).round() });

        for i in 1..self.get_degree() + 1 {
            let coeff = unsafe { self.get_coeff_unchecked(i).round() };
            unsafe { out.set_coeff_unchecked(i, coeff) };
        }

        out
    }

    seedable_function!(
        /// Performs the randomized rounding algorithm coefficient-wise
        /// by sampling from a discrete Gaussian over the integers shifted
        /// by `self` with gaussian parameter `r`.
        ///
        /// Parameters:
        /// - `r`: specifies the Gaussian parameter, which is proportional
        ///   to the standard deviation `sigma * sqrt(2 * pi) = r`
        #[seeded]
        /// - `seed`: specifies the 256-bit seed of the PRNG used for sampling
        #[optionally_seeded]
        /// - `seed`: specifies an optional 256-bit seed of the PRNG used for sampling.
        #[optionally_seeded]
        ///   If `None` is provided, a fresh [`ThreadRng`](rand::rngs::ThreadRng) is used instead.
        ///
        /// Returns the rounded polynomial as a [`PolyOverZ`] or an error if `r < 0`.
        ///
        /// # Examples
        /// ```
        /// use qfall_math::rational::PolyOverQ;
        /// use std::str::FromStr;
        ///
        /// let value = PolyOverQ::from_str("2  5/2 1").unwrap();
        #[unseeded]
        /// let rounded = value.randomized_rounding(3).unwrap();
        #[seeded]
        /// let rounded = value.randomized_rounding_seeded(3, [42; 32]).unwrap();
        /// ```
        ///
        /// # Errors and Failures
        /// - Returns a [`MathError`] of type [`InvalidIntegerInput`](MathError::InvalidIntegerInput)
        ///   if `r < 0`.
        ///
        /// This function implements randomized rounding according to:
        /// - \[1\] Peikert, C. (2010, August).
        ///   An efficient and parallel Gaussian sampler for lattices.
        ///   In: Annual Cryptology Conference (pp. 80-97).
        ///   <https://link.springer.com/chapter/10.1007/978-3-642-14623-7_5>
        pub(crate) fn randomized_rounding(
            &self,
            r: impl Into<Q>,
            seed: Option<[u8; 32]>,
        ) -> Result<PolyOverZ, MathError> {
            // derives an independent seed for every call of a seedable function below,
            // see the documentation of `seedable_function` in `src/macros/seeded.rs`
            let mut rng = SamplerRng::new(seed);

            let r = r.into();
            let mut out = PolyOverZ::from(unsafe {
                self.get_coeff_unchecked(0)
                    .randomized_rounding_optionally_seeded(&r, rng.derive_seed())?
            });
            for i in 1..self.get_degree() + 1 {
                let coeff = unsafe {
                    self.get_coeff_unchecked(i)
                        .randomized_rounding_optionally_seeded(&r, rng.derive_seed())?
                };
                unsafe { out.set_coeff_unchecked(i, coeff) };
            }

            Ok(out)
        }
    );
}

#[cfg(test)]
mod test_floor {
    use crate::{integer::PolyOverZ, rational::PolyOverQ};
    use std::str::FromStr;

    // Ensure that positive rationals are rounded correctly
    #[test]
    fn positive() {
        let value = PolyOverQ::from_str(&format!("2  1/{} {}/2", u64::MAX, i64::MAX)).unwrap();
        let cmp = PolyOverZ::from_str(&format!("2  0 {}", (i64::MAX - 1) / 2)).unwrap();

        assert_eq!(cmp, value.floor());
    }

    // Ensure that negative rationals are rounded correctly
    #[test]
    fn negative() {
        let value = PolyOverQ::from_str(&format!("2  -1/{} -{}/2", u64::MAX, i64::MAX)).unwrap();
        let cmp = PolyOverZ::from_str(&format!("2  -1 {}", (-i64::MAX - 1) / 2)).unwrap();

        assert_eq!(cmp, value.floor());
    }
}

#[cfg(test)]
mod test_ceil {
    use crate::{integer::PolyOverZ, rational::PolyOverQ};
    use std::str::FromStr;

    // Ensure that positive rationals are rounded correctly
    #[test]
    fn positive() {
        let value = PolyOverQ::from_str(&format!("2  1/{} {}/2", u64::MAX, i64::MAX)).unwrap();
        let cmp = PolyOverZ::from_str(&format!("2  1 {}", (i64::MAX - 1) / 2 + 1)).unwrap();

        assert_eq!(cmp, value.ceil());
    }

    // Ensure that negative rationals are rounded correctly
    #[test]
    fn negative() {
        let value = PolyOverQ::from_str(&format!("2  -1/{} -{}/2", u64::MAX, i64::MAX)).unwrap();
        let cmp = PolyOverZ::from_str(&format!("2  0 {}", (-i64::MAX - 1) / 2 + 1)).unwrap();

        assert_eq!(cmp, value.ceil());
    }
}

#[cfg(test)]
mod test_round {
    use crate::{integer::PolyOverZ, rational::PolyOverQ};
    use std::str::FromStr;

    // Ensure that positive rationals are rounded correctly
    #[test]
    fn positive() {
        let value = PolyOverQ::from_str(&format!("2  1/{} {}/2", u64::MAX, i64::MAX)).unwrap();
        let cmp = PolyOverZ::from_str(&format!("2  0 {}", (i64::MAX - 1) / 2 + 1)).unwrap();

        assert_eq!(cmp, value.round());
    }

    // Ensure that negative rationals are rounded correctly
    #[test]
    fn negative() {
        let value = PolyOverQ::from_str(&format!("2  -1/{} -{}/2", u64::MAX, i64::MAX)).unwrap();
        let cmp = PolyOverZ::from_str(&format!("2  0 {}", (-i64::MAX - 1) / 2 + 1)).unwrap();

        assert_eq!(cmp, value.round());
    }
}

#[cfg(test)]
mod test_randomized_rounding {
    use crate::rational::PolyOverQ;
    use crate::utils::sample::test_seed;
    use std::str::FromStr;

    /// Ensure that a `r < 0` throws an error
    #[test]
    fn negative_r() {
        for seeded in [false, true] {
            let value = PolyOverQ::from_str("2  5/2 1").unwrap();
            assert!(
                value
                    .randomized_rounding_optionally_seeded(-1, seeded.then(test_seed))
                    .is_err()
            );
        }
    }

    /// Checks whether the same seed results in the same sample.
    #[test]
    fn same_seed_same_sample() {
        use crate::rational::PolyOverQ;
        use std::str::FromStr;
        let value = PolyOverQ::from_str("2  5/2 1").unwrap();
        let sample_0 = value.randomized_rounding_seeded(3, [42; 32]).unwrap();
        let sample_1 = value.randomized_rounding_seeded(3, [42; 32]).unwrap();

        assert_eq!(sample_0, sample_1);
    }

    /// Checks whether the seeded variant outputs the known sample for a fixed seed.
    #[test]
    fn known_answer() {
        use crate::rational::PolyOverQ;
        use std::str::FromStr;
        let value = PolyOverQ::from_str("2  5/2 1").unwrap();
        let sample = value.randomized_rounding_seeded(3, [42; 32]).unwrap();

        assert_eq!("1  3", sample.to_string());
    }

    /// Checks whether the unseeded and the seeded function are available.
    #[test]
    fn availability() {
        use crate::rational::PolyOverQ;
        use std::str::FromStr;
        let value = PolyOverQ::from_str("2  5/2 1").unwrap();

        // explicitly checks the existence of the unseeded and the seeded function
        let _ = value.randomized_rounding(3);
        let _ = value.randomized_rounding_seeded(3, test_seed());
    }
}