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