Skip to main content

poulpy_core/
dist.rs

1use std::io::{Read, Result, Write};
2
3/// Read-only access to the [`Distribution`] associated with a secret key.
4pub trait GetDistribution {
5    /// Returns a reference to the distribution descriptor.
6    fn dist(&self) -> &Distribution;
7}
8
9/// Mutable access to the [`Distribution`] associated with a secret key.
10pub trait GetDistributionMut {
11    /// Returns a mutable reference to the distribution descriptor.
12    fn dist_mut(&mut self) -> &mut Distribution;
13}
14
15impl<T: GetDistribution + ?Sized> GetDistribution for &T {
16    fn dist(&self) -> &Distribution {
17        (*self).dist()
18    }
19}
20
21impl<T: GetDistribution + ?Sized> GetDistribution for &mut T {
22    fn dist(&self) -> &Distribution {
23        (**self).dist()
24    }
25}
26
27impl<T: GetDistributionMut + ?Sized> GetDistributionMut for &mut T {
28    fn dist_mut(&mut self) -> &mut Distribution {
29        (**self).dist_mut()
30    }
31}
32
33/// Describes the probability distribution used to sample secret-key
34/// coefficients.
35///
36/// Each variant encodes either a fixed Hamming weight or a per-coefficient
37/// probability. The enum is serialised as a single little-endian `u64`
38/// word via [`write_to`](Self::write_to) / [`read_from`](Self::read_from).
39///
40/// For probabilistic variants the `f64` payload is stored with a
41/// precision loss below 2^-44 (8 least-significant mantissa bits
42/// are discarded to fit the tag byte).
43#[derive(Clone, Copy, Debug)]
44pub enum Distribution {
45    /// Ternary in {-1, 0, 1} with exactly `h` non-zero coefficients.
46    TernaryFixed(usize),
47    /// Ternary in {-1, 0, 1} where each coefficient is non-zero with probability `p`.
48    TernaryProb(f64),
49    /// Binary in {0, 1} with exactly `h` ones.
50    BinaryFixed(usize),
51    /// Binary in {0, 1} where each coefficient is 1 with probability `p`.
52    BinaryProb(f64),
53    /// Binary in {0, 1} split into blocks of size 2^k, with one 1 per block.
54    BinaryBlock(usize),
55    /// All-zero secret (debug / testing only).
56    ZERO,
57    /// Uninitialized — no distribution has been set yet.
58    NONE,
59}
60
61const TAG_TERNARY_FIXED: u8 = 0;
62const TAG_TERNARY_PROB: u8 = 1;
63const TAG_BINARY_FIXED: u8 = 2;
64const TAG_BINARY_PROB: u8 = 3;
65const TAG_BINARY_BLOCK: u8 = 4;
66const TAG_ZERO: u8 = 5;
67const TAG_NONE: u8 = 6;
68
69use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
70
71impl Distribution {
72    /// Packs a tag (u8) and an f64 into a single u64.
73    /// The f64 is shifted right by 8, discarding the 8 least-significant
74    /// mantissa bits (precision loss < 2^-44), and the tag is placed
75    /// in the freed top byte.
76    #[inline]
77    fn pack_f64(tag: u8, p: f64) -> u64 {
78        (tag as u64) << 56 | (p.to_bits() >> 8)
79    }
80
81    /// Unpacks a tag-stripped 56-bit payload back into an f64
82    /// by shifting left by 8 (the 8 LSB mantissa bits become zero).
83    #[inline]
84    fn unpack_f64(payload: u64) -> f64 {
85        f64::from_bits(payload << 8)
86    }
87
88    /// Serialises this distribution as a single little-endian `u64` word.
89    ///
90    /// The top byte carries a variant tag; the lower 56 bits carry either
91    /// a `usize` payload (for fixed/block variants) or a truncated `f64`
92    /// (for probabilistic variants).
93    pub fn write_to<W: Write>(&self, writer: &mut W) -> Result<()> {
94        let word: u64 = match self {
95            Distribution::TernaryFixed(v) => (TAG_TERNARY_FIXED as u64) << 56 | (*v as u64),
96            Distribution::TernaryProb(p) => Self::pack_f64(TAG_TERNARY_PROB, *p),
97            Distribution::BinaryFixed(v) => (TAG_BINARY_FIXED as u64) << 56 | (*v as u64),
98            Distribution::BinaryProb(p) => Self::pack_f64(TAG_BINARY_PROB, *p),
99            Distribution::BinaryBlock(v) => (TAG_BINARY_BLOCK as u64) << 56 | (*v as u64),
100            Distribution::ZERO => (TAG_ZERO as u64) << 56,
101            Distribution::NONE => (TAG_NONE as u64) << 56,
102        };
103        writer.write_u64::<LittleEndian>(word)
104    }
105
106    /// Deserialises a [`Distribution`] from a single little-endian `u64` word.
107    ///
108    /// Returns [`std::io::ErrorKind::InvalidData`] if the tag byte is unrecognised.
109    pub fn read_from<R: Read>(reader: &mut R) -> Result<Self> {
110        let word = reader.read_u64::<LittleEndian>()?;
111        let tag = (word >> 56) as u8;
112        let payload = word & 0x00FF_FFFF_FFFF_FFFF;
113
114        let dist = match tag {
115            TAG_TERNARY_FIXED => Distribution::TernaryFixed(payload as usize),
116            TAG_TERNARY_PROB => Distribution::TernaryProb(Self::unpack_f64(payload)),
117            TAG_BINARY_FIXED => Distribution::BinaryFixed(payload as usize),
118            TAG_BINARY_PROB => Distribution::BinaryProb(Self::unpack_f64(payload)),
119            TAG_BINARY_BLOCK => Distribution::BinaryBlock(payload as usize),
120            TAG_ZERO => Distribution::ZERO,
121            TAG_NONE => Distribution::NONE,
122            _ => {
123                return Err(std::io::Error::new(std::io::ErrorKind::InvalidData, "Invalid tag"));
124            }
125        };
126        Ok(dist)
127    }
128}
129
130impl PartialEq for Distribution {
131    fn eq(&self, other: &Self) -> bool {
132        use Distribution::*;
133        match (self, other) {
134            (TernaryFixed(a), TernaryFixed(b)) => a == b,
135            (TernaryProb(a), TernaryProb(b)) => a.to_bits() == b.to_bits(),
136            (BinaryFixed(a), BinaryFixed(b)) => a == b,
137            (BinaryProb(a), BinaryProb(b)) => a.to_bits() == b.to_bits(),
138            (BinaryBlock(a), BinaryBlock(b)) => a == b,
139            (ZERO, ZERO) => true,
140            (NONE, NONE) => true,
141            _ => false,
142        }
143    }
144}
145
146impl Eq for Distribution {}