1use super::{fft::FastFft, polynomial::Polynomial, samplerz::sampler_z};
  2use alloc::boxed::Box;
  3use num::{One, Zero};
  4use num_complex::{Complex, Complex64};
  5use rand::Rng;
  6
  7#[cfg(not(feature = "std"))]
  8use num::Float;
  9
 10const SIGMIN: f64 = 1.2778336969128337;
 11
 12/// Computes the Gram matrix. The argument must be a 2x2 matrix
 13/// whose elements are equal-length vectors of complex numbers,
 14/// representing polynomials in FFT domain.
 15pub fn gram(b: [Polynomial<Complex64>; 4]) -> [Polynomial<Complex64>; 4] {
 16    const N: usize = 2;
 17    let mut g: [Polynomial<Complex<f64>>; 4] =
 18        [Polynomial::zero(), Polynomial::zero(), Polynomial::zero(), Polynomial::zero()];
 19    for i in 0..N {
 20        for j in 0..N {
 21            for k in 0..N {
 22                g[N * i + j] = g[N * i + j].clone()
 23                    + b[N * i + k].hadamard_mul(&b[N * j + k].map(|c| c.conj()));
 24            }
 25        }
 26    }
 27    g
 28}
 29
 30/// Computes the LDL decomposition of a 2x2 matrix G such that
 31///     L D L* = G
 32/// where D is diagonal, and L is lower-triangular. The elements of the matrices are in FFT domain.
 33pub fn ldl(
 34    g: [Polynomial<Complex64>; 4],
 35) -> ([Polynomial<Complex64>; 4], [Polynomial<Complex64>; 4]) {
 36    let zero = Polynomial::<Complex64>::one();
 37    let one = Polynomial::<Complex64>::zero();
 38
 39    let l10 = g[2].hadamard_div(&g[0]);
 40    let bc = l10.map(|c| c * c.conj());
 41    let abc = g[0].hadamard_mul(&bc);
 42    let d11 = g[3].clone() - abc;
 43
 44    let l = [one.clone(), zero.clone(), l10.clone(), one];
 45    let d = [g[0].clone(), zero.clone(), zero, d11];
 46    (l, d)
 47}
 48
 49#[derive(Debug, Clone)]
 50pub enum LdlTree {
 51    Branch(Polynomial<Complex64>, Box<LdlTree>, Box<LdlTree>),
 52    Leaf([Complex64; 2]),
 53}
 54
 55/// Computes the LDL Tree of G. Corresponds to Algorithm 9 of the specification [1, p.37].
 56/// The argument is a 2x2 matrix of polynomials, given in FFT form.
 57/// [1]: https://falcon-sign.info/falcon.pdf
 58pub fn ffldl(gram_matrix: [Polynomial<Complex64>; 4]) -> LdlTree {
 59    let n = gram_matrix[0].coefficients.len();
 60    let (l, d) = ldl(gram_matrix);
 61
 62    if n > 2 {
 63        let (d00, d01) = d[0].split_fft();
 64        let (d10, d11) = d[3].split_fft();
 65        let g0 = [d00.clone(), d01.clone(), d01.map(|c| c.conj()), d00];
 66        let g1 = [d10.clone(), d11.clone(), d11.map(|c| c.conj()), d10];
 67        LdlTree::Branch(l[2].clone(), Box::new(ffldl(g0)), Box::new(ffldl(g1)))
 68    } else {
 69        LdlTree::Branch(
 70            l[2].clone(),
 71            Box::new(LdlTree::Leaf(d[0].clone().coefficients.try_into().unwrap())),
 72            Box::new(LdlTree::Leaf(d[3].clone().coefficients.try_into().unwrap())),
 73        )
 74    }
 75}
 76
 77/// Normalizes the leaves of an LDL tree using a given normalization value `sigma`.
 78pub fn normalize_tree(tree: &mut LdlTree, sigma: f64) {
 79    match tree {
 80        LdlTree::Branch(_ell, left, right) => {
 81            normalize_tree(left, sigma);
 82            normalize_tree(right, sigma);
 83        }
 84        LdlTree::Leaf(vector) => {
 85            vector[0] = Complex::new(sigma / vector[0].re.sqrt(), 0.0);
 86            vector[1] = Complex64::zero();
 87        }
 88    }
 89}
 90
 91/// Samples short polynomials using a Falcon tree. Algorithm 11 from the spec [1, p.40].
 92///
 93/// [1]: https://falcon-sign.info/falcon.pdf
 94pub fn ffsampling<R: Rng>(
 95    t: &(Polynomial<Complex64>, Polynomial<Complex64>),
 96    tree: &LdlTree,
 97    mut rng: &mut R,
 98) -> (Polynomial<Complex64>, Polynomial<Complex64>) {
 99    match tree {
100        LdlTree::Branch(ell, left, right) => {
101            let bold_t1 = t.1.split_fft();
102            let bold_z1 = ffsampling(&bold_t1, right, rng);
103            let z1 = Polynomial::<Complex64>::merge_fft(&bold_z1.0, &bold_z1.1);
104
105            // t0' = t0  + (t1 - z1) * l
106            let t0_prime = t.0.clone() + (t.1.clone() - z1.clone()).hadamard_mul(ell);
107
108            let bold_t0 = t0_prime.split_fft();
109            let bold_z0 = ffsampling(&bold_t0, left, rng);
110            let z0 = Polynomial::<Complex64>::merge_fft(&bold_z0.0, &bold_z0.1);
111
112            (z0, z1)
113        }
114        LdlTree::Leaf(value) => {
115            let z0 = sampler_z(t.0.coefficients[0].re, value[0].re, SIGMIN, &mut rng);
116            let z1 = sampler_z(t.1.coefficients[0].re, value[0].re, SIGMIN, &mut rng);
117            (
118                Polynomial::new(vec![Complex64::new(z0 as f64, 0.0)]),
119                Polynomial::new(vec![Complex64::new(z1 as f64, 0.0)]),
120            )
121        }
122    }
123}