plat-core 0.1.1

//! RLWE (Ring Learning With Errors) encryption primitives.
//!
//! Single-key encrypt/decrypt and homomorphic addition in Z_q[X]/(X^N + 1).
//! This forms the foundation for both single-key FHE and Multi-Key FHE.
//!
//! Encryption scheme:
//!   sk = s (ternary polynomial)
//!   pk = (a, b = -a*s + e) where a ← uniform, e ← Gaussian
//!   Encrypt(pk, m): ct = (a*u + e1 + ⌊q/t⌋*m, b*u + e2)  → (c0, c1)
//!   Decrypt(sk, ct): m = ⌊t/q * (c0 + c1*s)⌉ mod t

use crate::modular::center;
use crate::ntt::NttTables;
use crate::params::Params;
use crate::poly::Poly;
use crate::sampling::{sample_gaussian, sample_ternary, sample_uniform};
use rand::Rng;

/// RLWE secret key: a ternary polynomial.
#[derive(Clone, Debug)]
pub struct SecretKey {
    pub poly: Poly,
}

/// RLWE public key: (a, b) where b = -(a*s) + e.
#[derive(Clone, Debug)]
pub struct PublicKey {
    pub a: Poly,
    pub b: Poly,
}

/// RLWE ciphertext: (c0, c1) such that c0 + c1*s ≈ ⌊q/t⌋*m.
#[derive(Clone, Debug)]
pub struct Ciphertext {
    pub c0: Poly,
    pub c1: Poly,
}

/// Generate a secret/public key pair.
pub fn keygen<R: Rng>(params: &Params, ntt: &NttTables, rng: &mut R) -> (SecretKey, PublicKey) {
    let s = sample_ternary(params, rng);
    let a = sample_uniform(params, rng);
    let e = sample_gaussian(params, rng);

    // b = -(a*s) + e = e - a*s
    let as_prod = ntt.mul(&a, &s);
    let b = e.sub(&as_prod, params);

    (SecretKey { poly: s }, PublicKey { a, b })
}

/// Encrypt a plaintext polynomial (coefficients in [0, t)).
pub fn encrypt<R: Rng>(
    params: &Params,
    ntt: &NttTables,
    pk: &PublicKey,
    plaintext: &Poly,
    rng: &mut R,
) -> Ciphertext {
    let u = sample_ternary(params, rng);
    let e1 = sample_gaussian(params, rng);
    let e2 = sample_gaussian(params, rng);

    // delta = ⌊q/t⌋
    let delta = params.q / params.t;

    // Scale message: m_scaled = delta * m
    let m_scaled = plaintext.scalar_mul(delta, params);

    // c0 = b*u + e1 + delta*m
    let bu = ntt.mul(&pk.b, &u);
    let c0 = bu.add(&e1, params).add(&m_scaled, params);

    // c1 = a*u + e2
    let au = ntt.mul(&pk.a, &u);
    let c1 = au.add(&e2, params);

    Ciphertext { c0, c1 }
}

/// Decrypt a ciphertext to recover the plaintext polynomial.
pub fn decrypt(params: &Params, ntt: &NttTables, sk: &SecretKey, ct: &Ciphertext) -> Poly {
    // phase = c0 + c1 * s ≈ delta * m + small_noise
    let c1s = ntt.mul(&ct.c1, &sk.poly);
    let phase = ct.c0.add(&c1s, params);

    // Round: m_i = ⌊t * phase_i / q⌉ mod t
    let coeffs = phase
        .coeffs
        .iter()
        .map(|&c| {
            // centered representative
            let centered = center(c, params.q);
            // t * centered / q, rounded
            let scaled = (centered as f64 * params.t as f64) / params.q as f64;
            let rounded = scaled.round() as i64;
            let result = rounded.rem_euclid(params.t as i64) as u64;
            result
        })
        .collect();

    Poly { coeffs }
}

/// Homomorphic addition of two ciphertexts.
pub fn add(params: &Params, ct1: &Ciphertext, ct2: &Ciphertext) -> Ciphertext {
    Ciphertext {
        c0: ct1.c0.add(&ct2.c0, params),
        c1: ct1.c1.add(&ct2.c1, params),
    }
}

/// Homomorphic subtraction: ct1 - ct2.
pub fn sub(params: &Params, ct1: &Ciphertext, ct2: &Ciphertext) -> Ciphertext {
    Ciphertext {
        c0: ct1.c0.sub(&ct2.c0, params),
        c1: ct1.c1.sub(&ct2.c1, params),
    }
}

/// Multiply ciphertext by a plaintext scalar (no noise growth from second operand).
pub fn scalar_mul(params: &Params, ct: &Ciphertext, scalar: u64) -> Ciphertext {
    Ciphertext {
        c0: ct.c0.scalar_mul(scalar, params),
        c1: ct.c1.scalar_mul(scalar, params),
    }
}

/// Encrypt a single u64 value (placed in the constant coefficient).
pub fn encrypt_u64<R: Rng>(
    params: &Params,
    ntt: &NttTables,
    pk: &PublicKey,
    value: u64,
    rng: &mut R,
) -> Ciphertext {
    let m = Poly::constant(value % params.t, params.n);
    encrypt(params, ntt, pk, &m, rng)
}

/// Decrypt to a single u64 value (from the constant coefficient).
pub fn decrypt_u64(params: &Params, ntt: &NttTables, sk: &SecretKey, ct: &Ciphertext) -> u64 {
    let m = decrypt(params, ntt, sk, ct);
    m.coeffs[0]
}

#[cfg(test)]
mod tests {
    use super::*;
    use rand::SeedableRng;
    use rand::rngs::StdRng;

    fn setup() -> (Params, NttTables, StdRng) {
        // Use test_small (N=256, q=12289, t=257) for adequate noise budget.
        // test_tiny (q=769, t=17) has too little headroom for scalar multiplication.
        let p = Params::test_small();
        let ntt = NttTables::new(&p);
        let rng = StdRng::seed_from_u64(42);
        (p, ntt, rng)
    }

    #[test]
    fn test_encrypt_decrypt_roundtrip() {
        let (p, ntt, mut rng) = setup();
        let (sk, pk) = keygen(&p, &ntt, &mut rng);

        // Values must be < p.t
        let m = Poly::from_coeffs(&[5 % p.t, 3 % p.t, 7 % p.t, 1], p.n);
        let ct = encrypt(&p, &ntt, &pk, &m, &mut rng);
        let m_dec = decrypt(&p, &ntt, &sk, &ct);

        for i in 0..4 {
            assert_eq!(m.coeffs[i], m_dec.coeffs[i], "mismatch at {i}");
        }
    }

    #[test]
    fn test_encrypt_decrypt_u64() {
        let (p, ntt, mut rng) = setup();
        let (sk, pk) = keygen(&p, &ntt, &mut rng);

        for val in 0..p.t {
            let ct = encrypt_u64(&p, &ntt, &pk, val, &mut rng);
            let dec = decrypt_u64(&p, &ntt, &sk, &ct);
            assert_eq!(val, dec, "roundtrip failed for {val}");
        }
    }

    #[test]
    fn test_homomorphic_add() {
        let (p, ntt, mut rng) = setup();
        let (sk, pk) = keygen(&p, &ntt, &mut rng);

        let a = 3u64;
        let b = 5u64;
        let ct_a = encrypt_u64(&p, &ntt, &pk, a, &mut rng);
        let ct_b = encrypt_u64(&p, &ntt, &pk, b, &mut rng);
        let ct_sum = add(&p, &ct_a, &ct_b);
        let dec = decrypt_u64(&p, &ntt, &sk, &ct_sum);
        assert_eq!((a + b) % p.t, dec);
    }

    #[test]
    fn test_scalar_mul_ciphertext() {
        let (p, ntt, mut rng) = setup();
        let (sk, pk) = keygen(&p, &ntt, &mut rng);

        let val = 2u64;
        let scalar = 3u64;
        let ct = encrypt_u64(&p, &ntt, &pk, val, &mut rng);
        let ct_scaled = scalar_mul(&p, &ct, scalar);
        let dec = decrypt_u64(&p, &ntt, &sk, &ct_scaled);
        assert_eq!((val * scalar) % p.t, dec);
    }

    #[test]
    fn test_homomorphic_weighted_sum() {
        let (p, ntt, mut rng) = setup();
        let (sk, pk) = keygen(&p, &ntt, &mut rng);

        let a_val = 2u64;
        let b_val = 3u64;
        let c_val = 1u64;

        let ct_a = encrypt_u64(&p, &ntt, &pk, a_val, &mut rng);
        let ct_b = encrypt_u64(&p, &ntt, &pk, b_val, &mut rng);
        let ct_c = encrypt_u64(&p, &ntt, &pk, c_val, &mut rng);

        // 3*a + 2*b + 1*c = 6 + 6 + 1 = 13
        let s1 = scalar_mul(&p, &ct_a, 3);
        let s2 = scalar_mul(&p, &ct_b, 2);
        let s3 = scalar_mul(&p, &ct_c, 1);

        let sum = add(&p, &add(&p, &s1, &s2), &s3);
        let dec = decrypt_u64(&p, &ntt, &sk, &sum);
        let expected = (3 * a_val + 2 * b_val + 1 * c_val) % p.t;
        assert_eq!(expected, dec);
    }

    #[test]
    fn test_multiple_seeds() {
        let (p, ntt, _) = setup();

        for seed in 0..20u64 {
            let mut rng = StdRng::seed_from_u64(seed);
            let (sk, pk) = keygen(&p, &ntt, &mut rng);
            let val = seed % p.t;
            let ct = encrypt_u64(&p, &ntt, &pk, val, &mut rng);
            let dec = decrypt_u64(&p, &ntt, &sk, &ct);
            assert_eq!(val, dec, "failed for seed {seed}");
        }
    }
}