poulpy-hal 0.8.1

//! Data layout types and trait definitions for the hardware abstraction layer.
//!
//! This module aggregates all layout-related types and re-exports them from
//! their respective sub-modules, including convolution kernels, matrix and
//! vector representations over polynomial rings, serialization support,
//! statistical utilities, and scratch-space management.
//!
//! It also defines the shared storage trait aliases used throughout the crate.
//! `Data` models backend-owned storage in the abstract, while
//! `HostDataRef`/`HostDataMut` capture host-byte-readable buffers for the
//! portions of the API that still require direct byte access.

mod convolution;
mod crt;
mod encoding;
mod layout_compat;
mod mat_znx;
mod module;
mod plan_cache;
mod scalar_znx;
mod scratch;
mod scratch_views;
mod serialization;
mod stats;
mod svp_ppol;
mod vec_znx;
mod vec_znx_big;
mod vec_znx_dft;
mod vmp_pmat;
mod word;
mod znx_base;

pub use convolution::*;
pub use crt::*;
pub use layout_compat::*;
pub use mat_znx::*;
pub use module::*;
pub use plan_cache::*;
pub use scalar_znx::*;
pub use scratch::*;
pub use scratch_views::*;
pub use serialization::*;
pub use stats::*;
pub use svp_ppol::*;
pub use vec_znx::*;
pub use vec_znx_big::*;
pub use vec_znx_dft::*;
pub use vmp_pmat::*;
pub use word::*;
pub use znx_base::*;

use anyhow::Result;
use std::ptr::NonNull;

use crate::oep::HalModuleImpl;

/// Base trait alias for all data containers.
///
/// Requires equality comparison ([`PartialEq`], [`Eq`]), a known size at
/// compile time ([`Sized`]), and a default value ([`Default`]). Every
/// layout type that holds raw data must satisfy at least this bound.
pub trait Data = PartialEq + Eq + Sized + Default;

/// Trait alias for read-only host-byte-accessible containers.
///
/// Extends [`Data`] with byte-level shared access via [`AsRef<[u8]>`] and
/// thread-safe sharing via [`Sync`]. Types satisfying this bound can be
/// borrowed immutably and read across threads.
pub trait HostDataRef = Data + AsRef<[u8]> + Sync;

/// Trait alias for mutable host-byte-accessible containers.
///
/// Extends [`HostDataRef`] with byte-level mutable access via [`AsMut<[u8]>`]
/// and cross-thread transfer via [`Send`]. Types satisfying this bound
/// support in-place modification and can be moved between threads.
pub trait HostDataMut = HostDataRef + AsMut<[u8]> + Send;

#[inline]
pub(crate) fn checked_product(factors: &[usize], context: &str) -> usize {
    factors
        .iter()
        .copied()
        .try_fold(1usize, usize::checked_mul)
        .unwrap_or_else(|| {
            panic!("{context} overflows usize");
        })
}

mod private {
    pub trait Sealed {}
}

/// Sealed trait identifying the residency of a [`Backend`]'s buffers.
///
/// Implemented only by [`Host`] and [`Device`]. Each [`Backend`] declares
/// its residency via its [`Backend::Location`] associated type, which lets
/// generic code discriminate host- and device-resident backends at the
/// type level.
pub trait Location: private::Sealed {}

/// Marker type for host-resident buffers.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct Host;

/// Marker type for device-resident buffers.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct Device;

impl private::Sealed for Host {}
impl private::Sealed for Device {}
impl Location for Host {}
impl Location for Device {}

/// Convenience marker for host-resident backends.
pub trait HostBackend: Backend<Location = Host> {}
impl<BE: Backend<Location = Host>> HostBackend for BE {}

/// Convenience marker for host-resident backends whose borrowed views are directly readable and writable as host bytes.
pub trait HostVisibleBackend: HostBackend
where
    for<'a> Self::BufRef<'a>: AsRef<[u8]>,
    for<'a> Self::BufMut<'a>: AsRef<[u8]> + AsMut<[u8]>,
{
}

impl<BE> HostVisibleBackend for BE
where
    BE: HostBackend,
    for<'a> BE::BufRef<'a>: AsRef<[u8]>,
    for<'a> BE::BufMut<'a>: AsRef<[u8]> + AsMut<[u8]>,
{
}

/// Minimal host-resident backend used as the default backend adapter for
/// host-visible byte-slice views in generic helper code.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Hash)]
pub struct HostBytesBackend;

impl Backend for HostBytesBackend {
    type ZnxWord = i64;
    type BigWord = i128;
    type DftWord = i64;
    type OwnedBuf = Vec<u8>;
    type BufRef<'a> = &'a [u8];
    type BufMut<'a> = &'a mut [u8];
    type Handle = ();
    type Location = Host;

    fn alloc_bytes(len: usize) -> Self::OwnedBuf {
        crate::alloc_aligned::<u8>(len)
    }

    fn alloc_zeroed_bytes(len: usize) -> Self::OwnedBuf {
        crate::alloc_aligned::<u8>(len)
    }

    fn from_host_bytes(bytes: &[u8]) -> Self::OwnedBuf {
        let mut out = crate::alloc_aligned::<u8>(bytes.len());
        out.copy_from_slice(bytes);
        out
    }

    fn from_bytes(bytes: Vec<u8>) -> Self::OwnedBuf {
        if crate::is_aligned(bytes.as_ptr()) {
            bytes
        } else {
            let mut out = crate::alloc_aligned::<u8>(bytes.len());
            out.copy_from_slice(&bytes);
            out
        }
    }

    fn to_host_bytes(buf: &Self::OwnedBuf) -> Vec<u8> {
        buf.clone()
    }

    fn copy_to_host(buf: &Self::OwnedBuf, dst: &mut [u8]) {
        assert!(
            buf.len() >= dst.len(),
            "backend buffer length {} is smaller than destination host slice length {}",
            buf.len(),
            dst.len()
        );
        dst.copy_from_slice(&buf[..dst.len()]);
    }

    fn copy_from_host(buf: &mut Self::OwnedBuf, src: &[u8]) {
        assert!(
            buf.len() >= src.len(),
            "backend buffer length {} is smaller than source host slice length {}",
            buf.len(),
            src.len()
        );
        let src_len = src.len();
        buf[..src_len].copy_from_slice(src);
        buf[src_len..].fill(0);
    }

    fn copy_view_to_host(buf: &Self::BufRef<'_>, dst: &mut [u8]) {
        assert_eq!(buf.len(), dst.len());
        dst.copy_from_slice(buf);
    }

    fn copy_host_to_view(buf: &mut Self::BufMut<'_>, src: &[u8]) {
        assert_eq!(buf.len(), src.len());
        buf.copy_from_slice(src);
    }

    fn len_bytes(buf: &Self::OwnedBuf) -> usize {
        buf.len()
    }

    fn len_bytes_ref(buf: &Self::BufRef<'_>) -> usize {
        buf.len()
    }

    fn len_bytes_mut(buf: &Self::BufMut<'_>) -> usize {
        buf.len()
    }

    fn view(buf: &Self::OwnedBuf) -> Self::BufRef<'_> {
        buf.as_slice()
    }

    fn view_ref<'a, 'b>(buf: &'a Self::BufRef<'b>) -> Self::BufRef<'a>
    where
        Self: 'b,
    {
        buf
    }

    fn view_ref_mut<'a, 'b>(buf: &'a Self::BufMut<'b>) -> Self::BufRef<'a>
    where
        Self: 'b,
    {
        buf
    }

    fn view_mut_ref<'a, 'b>(buf: &'a mut Self::BufMut<'b>) -> Self::BufMut<'a>
    where
        Self: 'b,
    {
        buf
    }

    fn view_mut(buf: &mut Self::OwnedBuf) -> Self::BufMut<'_> {
        buf.as_mut_slice()
    }

    fn region(buf: &Self::OwnedBuf, offset: usize, len: usize) -> Self::BufRef<'_> {
        &buf[offset..offset + len]
    }

    fn region_mut(buf: &mut Self::OwnedBuf, offset: usize, len: usize) -> Self::BufMut<'_> {
        &mut buf[offset..offset + len]
    }

    fn region_ref<'a, 'b>(buf: &'a Self::BufRef<'b>, offset: usize, len: usize) -> Self::BufRef<'a>
    where
        Self: 'b,
    {
        &buf[offset..offset + len]
    }

    fn region_ref_mut<'a, 'b>(buf: &'a Self::BufMut<'b>, offset: usize, len: usize) -> Self::BufRef<'a>
    where
        Self: 'b,
    {
        &buf[offset..offset + len]
    }

    fn region_mut_ref<'a, 'b>(buf: &'a mut Self::BufMut<'b>, offset: usize, len: usize) -> Self::BufMut<'a>
    where
        Self: 'b,
    {
        &mut buf[offset..offset + len]
    }

    unsafe fn destroy(_handle: NonNull<Self::Handle>) {}
}

unsafe impl HalModuleImpl<HostBytesBackend> for HostBytesBackend {
    fn new(n: u64) -> crate::layouts::Module<Self> {
        assert!(n.is_power_of_two(), "n must be a power of two, got {n}");
        unsafe { crate::layouts::Module::from_nonnull(NonNull::dangling(), n) }
    }
}

/// Convenience marker for device-resident backends.
pub trait DeviceBackend: Backend<Location = Device> {}
impl<BE: Backend<Location = Device>> DeviceBackend for BE {}

/// Deep-clone a borrowed layout into a fully owned variant.
///
/// Unlike the standard [`Clone`] trait, `ToOwnedDeep` is intended for
/// types that may borrow their underlying storage. Calling
/// [`to_owned_deep`](ToOwnedDeep::to_owned_deep) produces an independent
/// copy whose lifetime is not tied to the original.
pub trait ToOwnedDeep {
    type Owned;
    fn to_owned_deep(&self) -> Self::Owned;
}

/// Compute a `u64` hash digest of a layout's contents.
///
/// Provides a lightweight fingerprint suitable for fast equality checks
/// and debugging. This is **not** cryptographically secure; it is a
/// convenience mechanism for detecting whether two values hold identical
/// data without performing a full byte-by-byte comparison.
pub trait DigestU64 {
    fn digest_u64(&self) -> u64;
}

/// Backend-owned byte buffer type alias.
pub type OwnedBuf<BE> = <BE as Backend>::OwnedBuf;

/// A buffer whose bytes can be read out to the host.
///
/// Implemented by the buffer rather than the backend: a transfer needs to know
/// how to read its donor and write its receiver, not which backend each came
/// from. That is what lets the layout-level move infer everything from its two
/// operands, with no backend named anywhere.
pub trait CopyToHost {
    /// Bytes spanned by this buffer.
    fn len_bytes(&self) -> usize;

    /// Reads the whole buffer into `dst`, which must be [`Self::len_bytes`] long.
    fn copy_to_host(&self, dst: &mut [u8]);

    /// Host buffers lend their bytes directly; device buffers cannot, and
    /// return `None` so the move stages through [`Self::copy_to_host`].
    fn as_host_bytes(&self) -> Option<&[u8]> {
        None
    }
}

/// A buffer whose bytes can be written from the host. Counterpart of [`CopyToHost`].
pub trait CopyFromHost {
    /// Bytes spanned by this buffer.
    fn len_bytes(&self) -> usize;

    /// Overwrites the whole buffer from `src`, which must be [`Self::len_bytes`] long.
    fn copy_from_host(&mut self, src: &[u8]);

    /// Host buffers lend their bytes directly; device buffers return `None`.
    fn as_host_bytes_mut(&mut self) -> Option<&mut [u8]> {
        None
    }
}

impl CopyToHost for Vec<u8> {
    fn len_bytes(&self) -> usize {
        self.len()
    }
    fn copy_to_host(&self, dst: &mut [u8]) {
        dst.copy_from_slice(self);
    }
    fn as_host_bytes(&self) -> Option<&[u8]> {
        Some(self)
    }
}

impl CopyFromHost for Vec<u8> {
    fn len_bytes(&self) -> usize {
        self.len()
    }
    fn copy_from_host(&mut self, src: &[u8]) {
        self.copy_from_slice(src);
    }
    fn as_host_bytes_mut(&mut self) -> Option<&mut [u8]> {
        Some(self)
    }
}

impl CopyToHost for &[u8] {
    fn len_bytes(&self) -> usize {
        <[u8]>::len(self)
    }
    fn copy_to_host(&self, dst: &mut [u8]) {
        dst.copy_from_slice(self);
    }
    fn as_host_bytes(&self) -> Option<&[u8]> {
        Some(self)
    }
}

impl CopyToHost for &mut [u8] {
    fn len_bytes(&self) -> usize {
        <[u8]>::len(self)
    }
    fn copy_to_host(&self, dst: &mut [u8]) {
        dst.copy_from_slice(self);
    }
    fn as_host_bytes(&self) -> Option<&[u8]> {
        Some(self)
    }
}

impl CopyFromHost for &mut [u8] {
    fn len_bytes(&self) -> usize {
        <[u8]>::len(self)
    }
    fn copy_from_host(&mut self, src: &[u8]) {
        self.copy_from_slice(src);
    }
    fn as_host_bytes_mut(&mut self) -> Option<&mut [u8]> {
        Some(self)
    }
}

/// Moves `src`'s bytes into `dst`, which the caller has already allocated.
///
/// One copy whenever either side is host-visible, which covers host to host,
/// host to device and device to host. Only device to device stages, and a
/// device backend that can move directly should offer its own path.
///
/// # Panics
///
/// If the two buffers do not span the same number of bytes.
pub fn transfer_buf_into<S: CopyToHost + ?Sized, D: CopyFromHost + ?Sized>(src: &S, dst: &mut D) {
    let len: usize = src.len_bytes();
    assert_eq!(
        len,
        dst.len_bytes(),
        "transfer_buf_into: source is {} bytes, destination is {}",
        len,
        dst.len_bytes()
    );
    if let Some(bytes) = src.as_host_bytes() {
        dst.copy_from_host(bytes);
        return;
    }
    if let Some(bytes) = dst.as_host_bytes_mut() {
        src.copy_to_host(bytes);
        return;
    }
    let mut staging: Vec<u8> = vec![0u8; len];
    src.copy_to_host(&mut staging);
    dst.copy_from_host(&staging);
}

/// A backend that can exchange coefficient data with [`HostBytesBackend`].
///
/// Bundles the two requirements of a host-staged transfer: buffers that can be
/// read out to and written from host bytes, and agreement on the coefficient
/// word, since the move is a byte copy and cannot re-decompose limbs.
pub trait HostStaged: Backend<ZnxWord = i64, OwnedBuf: CopyToHost + CopyFromHost> {}

impl<BE> HostStaged for BE where BE: Backend<ZnxWord = i64, OwnedBuf: CopyToHost + CopyFromHost> {}

/// Implement a backend marker by forwarding all storage- and handle-level
/// behavior to an existing backend.
///
/// This is useful for proof or delegating backends that want to remain a
/// distinct backend type while reusing the same owned buffer, borrowed views,
/// scalar types, and handle representation as a source backend.
#[macro_export]
macro_rules! impl_backend_from {
    ($be:ty, $from:ty) => {
        impl poulpy_hal::layouts::Backend for $be {
            const DFT_IS_EXACT: bool = <$from as poulpy_hal::layouts::Backend>::DFT_IS_EXACT;

            type ZnxWord = <$from as poulpy_hal::layouts::Backend>::ZnxWord;
            type BigWord = <$from as poulpy_hal::layouts::Backend>::BigWord;
            type DftWord = <$from as poulpy_hal::layouts::Backend>::DftWord;
            type OwnedBuf = <$from as poulpy_hal::layouts::Backend>::OwnedBuf;
            type BufRef<'a> = <$from as poulpy_hal::layouts::Backend>::BufRef<'a>;
            type BufMut<'a> = <$from as poulpy_hal::layouts::Backend>::BufMut<'a>;
            type Handle = <$from as poulpy_hal::layouts::Backend>::Handle;
            type Location = <$from as poulpy_hal::layouts::Backend>::Location;

            fn alloc_bytes(len: usize) -> Self::OwnedBuf {
                <$from as poulpy_hal::layouts::Backend>::alloc_bytes(len)
            }

            fn alloc_zeroed_bytes(len: usize) -> Self::OwnedBuf {
                <$from as poulpy_hal::layouts::Backend>::alloc_zeroed_bytes(len)
            }

            fn from_host_bytes(bytes: &[u8]) -> Self::OwnedBuf {
                <$from as poulpy_hal::layouts::Backend>::from_host_bytes(bytes)
            }

            fn from_bytes(bytes: Vec<u8>) -> Self::OwnedBuf {
                <$from as poulpy_hal::layouts::Backend>::from_bytes(bytes)
            }

            fn to_host_bytes(buf: &Self::OwnedBuf) -> Vec<u8> {
                <$from as poulpy_hal::layouts::Backend>::to_host_bytes(buf)
            }

            fn copy_to_host(buf: &Self::OwnedBuf, dst: &mut [u8]) {
                <$from as poulpy_hal::layouts::Backend>::copy_to_host(buf, dst)
            }

            fn copy_from_host(buf: &mut Self::OwnedBuf, src: &[u8]) {
                <$from as poulpy_hal::layouts::Backend>::copy_from_host(buf, src)
            }

            fn copy_view_to_host(buf: &Self::BufRef<'_>, dst: &mut [u8]) {
                <$from as poulpy_hal::layouts::Backend>::copy_view_to_host(buf, dst)
            }

            fn copy_host_to_view(buf: &mut Self::BufMut<'_>, src: &[u8]) {
                <$from as poulpy_hal::layouts::Backend>::copy_host_to_view(buf, src)
            }

            fn len_bytes(buf: &Self::OwnedBuf) -> usize {
                <$from as poulpy_hal::layouts::Backend>::len_bytes(buf)
            }

            fn len_bytes_ref(buf: &Self::BufRef<'_>) -> usize {
                <$from as poulpy_hal::layouts::Backend>::len_bytes_ref(buf)
            }

            fn len_bytes_mut(buf: &Self::BufMut<'_>) -> usize {
                <$from as poulpy_hal::layouts::Backend>::len_bytes_mut(buf)
            }

            fn view(buf: &Self::OwnedBuf) -> Self::BufRef<'_> {
                <$from as poulpy_hal::layouts::Backend>::view(buf)
            }

            fn view_ref<'a, 'b>(buf: &'a Self::BufRef<'b>) -> Self::BufRef<'a>
            where
                Self: 'b,
            {
                <$from as poulpy_hal::layouts::Backend>::view_ref(buf)
            }

            fn view_ref_mut<'a, 'b>(buf: &'a Self::BufMut<'b>) -> Self::BufRef<'a>
            where
                Self: 'b,
            {
                <$from as poulpy_hal::layouts::Backend>::view_ref_mut(buf)
            }

            fn view_mut_ref<'a, 'b>(buf: &'a mut Self::BufMut<'b>) -> Self::BufMut<'a>
            where
                Self: 'b,
            {
                <$from as poulpy_hal::layouts::Backend>::view_mut_ref(buf)
            }

            fn view_mut(buf: &mut Self::OwnedBuf) -> Self::BufMut<'_> {
                <$from as poulpy_hal::layouts::Backend>::view_mut(buf)
            }

            fn region(buf: &Self::OwnedBuf, offset: usize, len: usize) -> Self::BufRef<'_> {
                <$from as poulpy_hal::layouts::Backend>::region(buf, offset, len)
            }

            fn region_mut(buf: &mut Self::OwnedBuf, offset: usize, len: usize) -> Self::BufMut<'_> {
                <$from as poulpy_hal::layouts::Backend>::region_mut(buf, offset, len)
            }

            fn region_ref<'a, 'b>(buf: &'a Self::BufRef<'b>, offset: usize, len: usize) -> Self::BufRef<'a>
            where
                Self: 'b,
            {
                <$from as poulpy_hal::layouts::Backend>::region_ref(buf, offset, len)
            }

            fn region_ref_mut<'a, 'b>(buf: &'a Self::BufMut<'b>, offset: usize, len: usize) -> Self::BufRef<'a>
            where
                Self: 'b,
            {
                <$from as poulpy_hal::layouts::Backend>::region_ref_mut(buf, offset, len)
            }

            fn region_mut_ref<'a, 'b>(buf: &'a mut Self::BufMut<'b>, offset: usize, len: usize) -> Self::BufMut<'a>
            where
                Self: 'b,
            {
                <$from as poulpy_hal::layouts::Backend>::region_mut_ref(buf, offset, len)
            }

            unsafe fn destroy(handle: std::ptr::NonNull<Self::Handle>) {
                <$from as poulpy_hal::layouts::Backend>::destroy(handle)
            }

            // Sizing must be forwarded explicitly: these are defaulted trait
            // methods, so without forwarding the delegate would silently get
            // the word-derived defaults instead of the source backend's
            // overrides (e.g. the packed IFMA `bytes_of_vmp_pmat`), breaking
            // the layout compatibility asserted by the markers below.
            const SCRATCH_ALIGN: usize = <$from as poulpy_hal::layouts::Backend>::SCRATCH_ALIGN;

            fn bytes_of_vec_znx_dft(n: usize, cols: usize, size: usize) -> usize {
                <$from as poulpy_hal::layouts::Backend>::bytes_of_vec_znx_dft(n, cols, size)
            }

            fn bytes_of_vec_znx_big(n: usize, cols: usize, size: usize) -> usize {
                <$from as poulpy_hal::layouts::Backend>::bytes_of_vec_znx_big(n, cols, size)
            }

            fn bytes_of_svp_ppol(n: usize, cols: usize) -> usize {
                <$from as poulpy_hal::layouts::Backend>::bytes_of_svp_ppol(n, cols)
            }

            fn bytes_of_vmp_pmat(n: usize, rows: usize, cols_in: usize, cols_out: usize, size: usize) -> usize {
                <$from as poulpy_hal::layouts::Backend>::bytes_of_vmp_pmat(n, rows, cols_in, cols_out, size)
            }

            fn bytes_of_cnv_pvec_left(n: usize, cols: usize, size: usize) -> usize {
                <$from as poulpy_hal::layouts::Backend>::bytes_of_cnv_pvec_left(n, cols, size)
            }

            fn bytes_of_cnv_pvec_right(n: usize, cols: usize, size: usize) -> usize {
                <$from as poulpy_hal::layouts::Backend>::bytes_of_cnv_pvec_right(n, cols, size)
            }
        }

        // A delegating backend forwards all storage behavior verbatim, so every
        // container layout is shared with the source backend by construction.
        unsafe impl poulpy_hal::layouts::VecZnxDftLayoutCompatible<$from> for $be {}
        unsafe impl poulpy_hal::layouts::VecZnxDftLayoutCompatible<$be> for $from {}
        unsafe impl poulpy_hal::layouts::VecZnxBigLayoutCompatible<$from> for $be {}
        unsafe impl poulpy_hal::layouts::VecZnxBigLayoutCompatible<$be> for $from {}
        unsafe impl poulpy_hal::layouts::SvpPPolLayoutCompatible<$from> for $be {}
        unsafe impl poulpy_hal::layouts::SvpPPolLayoutCompatible<$be> for $from {}
        unsafe impl poulpy_hal::layouts::VmpPMatLayoutCompatible<$from> for $be {}
        unsafe impl poulpy_hal::layouts::VmpPMatLayoutCompatible<$be> for $from {}
        unsafe impl poulpy_hal::layouts::CnvPVecLayoutCompatible<$from> for $be {}
        unsafe impl poulpy_hal::layouts::CnvPVecLayoutCompatible<$be> for $from {}
    };
}

#[derive(Clone, Copy, Debug)]
pub struct NoiseInfos {
    pub k: usize,
    pub sigma: f64,
    pub bound: f64,
}

impl NoiseInfos {
    pub fn new(k: usize, sigma: f64, bound: f64) -> Result<Self> {
        anyhow::ensure!(sigma.is_sign_positive(), "sigma must be positive");
        anyhow::ensure!(sigma >= 1.0, "sigma must be greater or equal to 1");
        anyhow::ensure!(bound >= sigma, "bound: {bound} must be greater or equal to sigma: {sigma}");
        Ok(Self { k, sigma, bound })
    }

    pub fn target_limb_and_scale(&self, base2k: usize) -> (usize, f64) {
        let limb: usize = self.k.div_ceil(base2k) - 1;
        let scale: f64 = (((limb + 1) * base2k - self.k) as f64).exp2();
        (limb, scale)
    }
}