vaea-ntt 0.1.1

// Copyright (C) 2024-2026 Vaea SAS
// SPDX-License-Identifier: AGPL-3.0-or-later
//
// This file is part of VaeaNTT.
//
// VaeaNTT is free software: you can redistribute it and/or modify it under
// the terms of the GNU Affero General Public License as published by the
// Free Software Foundation, either version 3 of the License, or (at your
// option) any later version.
//
// VaeaNTT is distributed in the hope that it will be useful, but WITHOUT
// ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
// FITNESS FOR A PARTICULAR PURPOSE. See the GNU Affero General Public
// License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with VaeaNTT. If not, see <https://www.gnu.org/licenses/>.

//! # VaeaNTT Foreign Function Interface (Diplomat)
//!
//! This module exposes VaeaNTT's NTT32 pipeline to C, C++, JavaScript/WASM,
//! Python, Dart, and Kotlin via the [Diplomat](https://github.com/rust-diplomat/diplomat)
//! binding generator.
//!
//! ## Generating bindings
//!
//! ```bash
//! # Install diplomat-tool
//! cargo install diplomat-tool
//!
//! # Generate C headers
//! diplomat-tool c bindings/c/ -l src/ffi.rs
//!
//! # Generate C++ headers
//! diplomat-tool cpp bindings/cpp/ -l src/ffi.rs
//!
//! # Generate JavaScript/TypeScript (WASM)
//! diplomat-tool js bindings/js/ -l src/ffi.rs
//!
//! # Generate Python
//! diplomat-tool python bindings/python/ -l src/ffi.rs
//! ```

#[cfg(feature = "ffi")]
#[diplomat::bridge]
/// FFI bindings generated by Diplomat for C, C++, JS/WASM, Python, Dart.
#[allow(clippy::module_inception)]
pub mod ffi {
    use alloc::boxed::Box;
    use alloc::vec::Vec;

    // =========================================================================
    // Error type
    // =========================================================================

    /// Error codes returned by VaeaNTT operations.
    pub enum VaeaNttError {
        /// N must be a power of 2 >= 2.
        InvalidSize,
        /// q must be prime.
        NotPrime,
        /// q must satisfy q ≡ 1 (mod 2N).
        NotNttFriendly,
        /// q must be < 2^28.
        PrimeTooLarge,
        /// Buffer length does not match the context's N.
        InvalidBufferLength,
    }

    // =========================================================================
    // Ntt32Context — opaque handle
    // =========================================================================

    /// Opaque handle to a pre-computed NTT context for 28-bit primes.
    ///
    /// Create with `VaeaNtt32::try_new()`, then call `forward()`, `inverse()`,
    /// or `negacyclic_mul()` on polynomial buffers.
    #[diplomat::opaque]
    pub struct VaeaNtt32(crate::ntt32::Ntt32Context);

    impl VaeaNtt32 {
        /// Creates a new NTT context.
        ///
        /// # Arguments
        /// - `n` — polynomial size, must be a power of 2 ≥ 2
        /// - `q` — prime < 2^28, must satisfy q ≡ 1 (mod 2N)
        ///
        /// Returns an error if parameters are invalid.
        pub fn try_new(n: usize, q: u32) -> Result<Box<VaeaNtt32>, VaeaNttError> {
            match crate::ntt32::Ntt32Context::try_new(n, q) {
                Ok(ctx) => Ok(Box::new(VaeaNtt32(ctx))),
                Err(crate::NttError::InvalidSize(_)) => Err(VaeaNttError::InvalidSize),
                Err(crate::NttError::NotPrime(_)) => Err(VaeaNttError::NotPrime),
                Err(crate::NttError::NotNttFriendly { .. }) => Err(VaeaNttError::NotNttFriendly),
                Err(crate::NttError::PrimeTooLarge(_)) => Err(VaeaNttError::PrimeTooLarge),
            }
        }

        /// Returns the polynomial size N.
        pub fn get_n(&self) -> usize {
            self.0.n
        }

        /// Returns the prime modulus q.
        pub fn get_q(&self) -> u32 {
            self.0.q
        }

        /// Forward NTT transform (in-place).
        ///
        /// `data` must have exactly N elements. All elements must be < q.
        pub fn forward(&self, data: &mut [u32]) -> Result<(), VaeaNttError> {
            if data.len() != self.0.n {
                return Err(VaeaNttError::InvalidBufferLength);
            }
            self.0.forward(data);
            Ok(())
        }

        /// Inverse NTT transform (in-place, with N⁻¹ normalization).
        ///
        /// `data` must have exactly N elements.
        pub fn inverse(&self, data: &mut [u32]) -> Result<(), VaeaNttError> {
            if data.len() != self.0.n {
                return Err(VaeaNttError::InvalidBufferLength);
            }
            self.0.inverse(data);
            Ok(())
        }

        /// Inverse NTT without N⁻¹ normalization (matches concrete-ntt behavior).
        ///
        /// `data` must have exactly N elements.
        pub fn inverse_lazy(&self, data: &mut [u32]) -> Result<(), VaeaNttError> {
            if data.len() != self.0.n {
                return Err(VaeaNttError::InvalidBufferLength);
            }
            self.0.inverse_lazy(data);
            Ok(())
        }

        /// Negacyclic polynomial multiplication in Z_q\[X\]/(X^N + 1).
        ///
        /// Computes `result = a × b mod (X^N + 1)` using NTT.
        /// All three buffers must have exactly N elements.
        /// `a` and `b` are modified (transformed to NTT domain).
        pub fn negacyclic_mul(
            &self,
            a: &mut [u32],
            b: &mut [u32],
            result: &mut [u32],
        ) -> Result<(), VaeaNttError> {
            if a.len() != self.0.n || b.len() != self.0.n || result.len() != self.0.n {
                return Err(VaeaNttError::InvalidBufferLength);
            }
            self.0.negacyclic_mul_into(a, b, result);
            Ok(())
        }

        /// Pointwise multiplication of two NTT-domain polynomials.
        ///
        /// All three buffers must have exactly N elements.
        /// Inputs must already be in NTT domain.
        pub fn pointwise_mul(
            &self,
            a: &[u32],
            b: &[u32],
            result: &mut [u32],
        ) -> Result<(), VaeaNttError> {
            if a.len() != self.0.n || b.len() != self.0.n || result.len() != self.0.n {
                return Err(VaeaNttError::InvalidBufferLength);
            }
            self.0.pointwise_mul(a, b, result);
            Ok(())
        }
    }

    // =========================================================================
    // Convenience: generate NTT-friendly 28-bit primes
    // =========================================================================

    /// Result of prime generation: a list of primes.
    #[diplomat::opaque]
    pub struct VaeaPrimeList(Vec<u32>);

    impl VaeaPrimeList {
        /// Generates `count` NTT-friendly 28-bit primes for polynomial size `n`.
        pub fn generate(n: usize, count: usize) -> Box<VaeaPrimeList> {
            let primes = crate::ntt32::generate_primes_28(n, count);
            Box::new(VaeaPrimeList(primes))
        }

        /// Returns the number of primes in the list.
        pub fn len(&self) -> usize {
            self.0.len()
        }

        /// Returns true if the list is empty.
        pub fn is_empty(&self) -> bool {
            self.0.is_empty()
        }

        /// Returns the prime at the given index (0-based).
        pub fn get(&self, index: usize) -> u32 {
            self.0[index]
        }
    }
}