fips204 0.5.0

// This file implements functionality from FIPS 204 section 7.4 High-Order and Low-Order Bits and Hints

use crate::helpers::full_reduce32;
use crate::types::{Zq, R};
use crate::{D, Q};

// Some arith routines leverage dilithium https://github.com/PQClean/PQClean/tree/master/crypto_sign


/// # Algorithm 35: `Power2Round(r)` on page 40, but applied over all coefficients of K polys.
/// Decomposes `r` into `(r1, r0)` such that `r ≡ r1·2^d + r0 mod q`.
///
/// **Input**: `r ∈ Z_q`. <br>
/// **Output**: Integers `(r1, r0)`.
pub(crate) fn power2round<const K: usize>(r: &[R; K]) -> ([R; K], [R; K]) {
    // 1: r+ ← r mod q
    // 2: r0 ← r+ mod±2^d
    // 3: return ((r+ − r0)/2^d, r0)

    // Check input ranges
    debug_assert!(
        r.iter().flat_map(|row| row.0).all(|element| (0..Q).contains(&element)),
        "power2round input"
    );

    let r_1: [R; K] = core::array::from_fn(|k| {
        R(core::array::from_fn(|n| (r[k].0[n] + (1 << (D - 1)) - 1) >> D))
    });

    let r_0: [R; K] =
        core::array::from_fn(|k| R(core::array::from_fn(|n| r[k].0[n] - (r_1[k].0[n] << D))));

    // Check output reconstruction
    debug_assert!(
        {
            let mut result = true;
            for k in 0..K {
                for n in 0..256 {
                    result &= r[k].0[n] == ((r_1[k].0[n] << D) + r_0[k].0[n]);
                }
            }
            result
        },
        "Alg 35: fails"
    );

    (r_1, r_0)
}


/// # Algorithm 36: `Decompose(r)` on page 40.
/// Decomposes `r` into `(r1, r0)` such that `r ≡ r1·(2·γ_2) + r0 mod q`.
///
/// If using the 'centered' `r0` yields `r1 = (q - 1)/(2*gamma2)` then the result should
/// be adjusted by setting `r1` to `0`, and subtracting `1` from `r0`. In practice, this
/// happens only with one input value of `rp` for each possible `gamma2`: <br>
/// `        gamma2        rp      r0'   r1'        r0     r1` <br>
/// `         95232   8285185   -95231    44    -95232      0` <br>
/// `        261888   8118529  -261887    16   -261888      0` <br>
/// For each `gamma2` and the specific `r mod q` value, `r0'` is `rp mod± 2*gamma2`, and
/// `r1'` is `rp - 2*gamma2*r0`. Values `(r0,r1)` are the value that just be returned in
/// that case.
///
/// **Input**: `r ∈ Z_q` <br>
/// **Output**: Integers `(r1, r0)`.
pub(crate) fn decompose(gamma2: i32, r: Zq) -> (Zq, Zq) {
    // 1: r+ ← r mod q
    // 2: r0 ← r+ mod±(2γ_2)
    // 3: if r+ − r0 = q − 1 then
    // 4: r1 ← 0
    // 5: r0 ← r0 − 1
    // 6: else r_1 ← (r+ − r0)/(2γ2)
    // 7: end if

    let rp = full_reduce32(r);
    let mut xr1;
    if gamma2 & (1 << 17) == 0 {
        // ml-dsa-44
        xr1 = (rp + 127) >> 7;
        xr1 = (xr1 * 11275 + (1 << 23)) >> 24;
        xr1 ^= ((43 - xr1) >> 31) & xr1;
    } else {
        // ml-dsa-65 and ml-dsa-87
        xr1 = (rp + 127) >> 7;
        xr1 = (xr1 * 1025 + (1 << 21)) >> 22;
        xr1 &= 15;
    }

    let xr0 = rp - xr1 * 2 * gamma2;
    let xr0 = xr0 - ((((Q - 1) / 2 - xr0) >> 31) & Q);

    // Reconstruct/validate outputs
    debug_assert_eq!(r.rem_euclid(Q), (xr1 * 2 * gamma2 + xr0).rem_euclid(Q), "Alg 36: fails");

    (xr1, xr0)
}


/// # Algorithm 37: `HighBits(r)` on page 40.
/// Returns `r1` from the output of `Decompose(r)`.
///
/// **Input**: `r ∈ Z_q` <br>
/// **Output**: Integer `r_1`.
pub(crate) fn high_bits(gamma2: i32, r: Zq) -> Zq {
    //
    // 1: (r1, r0) ← Decompose(r)
    let (r1, _r0) = decompose(gamma2, r);

    // 2: return r1
    r1
}


/// # Algorithm 38: `LowBits(r)` on page 41.
/// Returns `r_0` from the output of Decompose (r).
///
/// **Input**: `r ∈ Z_q` <br>
/// **Output**: Integer `r_0`.
pub(crate) fn low_bits(gamma2: i32, r: Zq) -> Zq {
    //
    // 1: (r1, r0) ← Decompose(r)
    let (_r1, r0) = decompose(gamma2, r);

    // 2: return r0
    r0
}


/// # Algorithm 39: `MakeHint(z,r)` on page 41.
/// Compute hint bit indicating whether adding `z` to `r` alters the high bits of `r`.
///
/// Input: `z`, `r` ∈ `Z_q` <br>
/// Output: Boolean
pub(crate) fn make_hint(gamma2: i32, z: Zq, r: Zq) -> bool {
    //
    // 1: r1 ← HighBits(r)
    let r1 = high_bits(gamma2, r);

    // 2: v1 ← HighBits(r + z)
    let v1 = high_bits(gamma2, r + z);

    // 3: return [[r1 != v1]]
    r1 != v1
}


/// # Algorithm 40: `UseHint(h,r)` on page 41.
/// Returns the high bits of `r` adjusted according to hint `h`.
///
/// This non-CT function uses public data from the signature; thus
/// does not need to be constant time
///
/// **Input**: Boolean `h` (cast in `Z_q`), `r` ∈ `Z_q` <br>
/// **Output**: `r_1 ∈ Z` with `0 ≤ r_1 ≤ (q − 1)/(2·γ_2)`
pub(crate) fn use_hint(gamma2: i32, h: Zq, r: Zq) -> Zq {
    //
    // 1: m ← (q− 1)/(2*γ_2)
    // dissolved into steps 3 and 4 below

    // 2: (r1, r0) ← Decompose(r)
    let (r1, r0) = decompose(gamma2, r);

    // Step 5 here, to simplify later logic
    if h == 0 {
        return r1;
    }

    // 3: if h = 1 and r0 > 0 return (r1 + 1) mod m
    // 4: if h = 1 and r0 ≤ 0 return (r1 − 1) mod m
    if gamma2 & (1 << 17) == 0 {
        // ml-dsa-44; explicit r1 + 1 mod m(44)
        if r0 > 0 {
            if r1 == 43 {
                return 0;
            }
            return r1 + 1;
        } // explicit r1 - 1 mod m(44)
        if r1 == 0 {
            return 43;
        }
        r1 - 1
    } else {
        // ml-dsa-65 and ml-dsa-87; explicit r1 + 1 mod m(16)
        if r0 > 0 {
            return (r1 + 1) & 15;
        } // explicit r1 - 1 mod m(16)
        (r1 - 1) & 15
    }

    // 5: return r1
    // r1 see first 'if' above
}


#[cfg(test)]
mod tests {
    use super::*;

    // gamma2 for ML-DSA-44, and for ML-DSA-65 and ML-DSA-87 (Table 1 on page 15)
    const GAMMA2S: [i32; 2] = [(Q - 1) / 88, (Q - 1) / 32];

    // m mod± α for even α: the m' ≡ m (mod α) with −α/2 < m' ≤ α/2 (section 2.3)
    fn mod_pm(m: i32, alpha: i32) -> i32 {
        let r = m.rem_euclid(alpha);
        if r > alpha / 2 {
            r - alpha
        } else {
            r
        }
    }

    // Algorithm 36, step by step
    fn spec_decompose(gamma2: i32, r: i32) -> (i32, i32) {
        let rp = r.rem_euclid(Q);
        let r0 = mod_pm(rp, 2 * gamma2);
        if rp - r0 == Q - 1 {
            (0, r0 - 1)
        } else {
            ((rp - r0) / (2 * gamma2), r0)
        }
    }

    // Algorithm 40, step by step
    fn spec_use_hint(gamma2: i32, h: i32, r: i32) -> i32 {
        let m = (Q - 1) / (2 * gamma2);
        let (r1, r0) = spec_decompose(gamma2, r);
        if h == 1 && r0 > 0 {
            return (r1 + 1).rem_euclid(m);
        }
        if h == 1 && r0 <= 0 {
            return (r1 - 1).rem_euclid(m);
        }
        r1
    }

    // Every r in [0, q), and r − q as callers pass it after `partial_reduce32()`
    #[test]
    fn decompose_and_use_hint_exhaustive() {
        for gamma2 in GAMMA2S {
            for r in 0..Q {
                let (r1, r0) = spec_decompose(gamma2, r);
                assert_eq!(decompose(gamma2, r), (r1, r0), "Decompose({r}), gamma2 {gamma2}");
                assert_eq!(
                    decompose(gamma2, r - Q),
                    (r1, r0),
                    "Decompose({r} - q), gamma2 {gamma2}"
                );
                assert_eq!(high_bits(gamma2, r), r1, "HighBits({r}), gamma2 {gamma2}");
                assert_eq!(low_bits(gamma2, r), r0, "LowBits({r}), gamma2 {gamma2}");
                for h in 0..=1 {
                    let expected = spec_use_hint(gamma2, h, r);
                    assert_eq!(
                        use_hint(gamma2, h, r),
                        expected,
                        "UseHint({h}, {r}), gamma2 {gamma2}"
                    );
                }
            }
        }
        // Hand-computed values: either side of the step-3 edge q − γ2 for ML-DSA-44, and
        // UseHint(1, r) with r0 = 0, which must step down (r1 − 1) mod m
        assert_eq!(decompose(GAMMA2S[0], 8_285_184), (43, 95_232));
        assert_eq!(decompose(GAMMA2S[0], 8_285_185), (0, -95_232));
        assert_eq!(use_hint(GAMMA2S[0], 1, 0), 43);
        assert_eq!(use_hint(GAMMA2S[1], 1, 0), 15);
    }

    // Algorithm 35 over every r in [0, q): r0 = r mod± 2^d and r1 = (r − r0)/2^d
    #[test]
    fn power2round_exhaustive() {
        for start in (0..Q).step_by(256) {
            let r = [R(core::array::from_fn(|n| (start + i32::try_from(n).unwrap()).min(Q - 1)))];
            let (r_1, r_0) = power2round(&r);
            for ((&x, &x1), &x0) in r[0].0.iter().zip(&r_1[0].0).zip(&r_0[0].0) {
                let e0 = mod_pm(x, 1 << D);
                assert_eq!((x1, x0), ((x - e0) / (1 << D), e0), "Power2Round({x})");
            }
        }
    }

    // Algorithm 39 around every HighBits bin edge, with z as given and as Q + z (signing
    // passes Q − c·t0)
    #[test]
    fn make_hint_bin_edges() {
        for gamma2 in GAMMA2S {
            for k in 0..(Q - 1) / (2 * gamma2) {
                let edge = k * 2 * gamma2 + gamma2; // r0 = γ2, the top of bin k
                for r in edge - 2..=edge + 2 {
                    for z in -2..=2 {
                        let expected =
                            spec_decompose(gamma2, r).0 != spec_decompose(gamma2, r + z).0;
                        assert_eq!(make_hint(gamma2, z, r), expected, "MakeHint({z}, {r})");
                        assert_eq!(make_hint(gamma2, Q + z, r), expected, "MakeHint(q + {z}, {r})");
                    }
                }
            }
        }
    }
}