1use std::io::{Read, Result, Write};
 2
 3pub trait GetDistribution {
 4    fn dist(&self) -> &Distribution;
 5}
 6
 7pub trait GetDistributionMut {
 8    fn dist_mut(&mut self) -> &mut Distribution;
 9}
10
11#[derive(Clone, Copy, Debug)]
12pub enum Distribution {
13    TernaryFixed(usize), // Ternary with fixed Hamming weight
14    TernaryProb(f64),    // Ternary with probabilistic Hamming weight
15    BinaryFixed(usize),  // Binary with fixed Hamming weight
16    BinaryProb(f64),     // Binary with probabilistic Hamming weight
17    BinaryBlock(usize),  // Binary split in block of size 2^k
18    ZERO,                // Debug mod
19    NONE,                // Unitialized
20}
21
22const TAG_TERNARY_FIXED: u8 = 0;
23const TAG_TERNARY_PROB: u8 = 1;
24const TAG_BINARY_FIXED: u8 = 2;
25const TAG_BINARY_PROB: u8 = 3;
26const TAG_BINARY_BLOCK: u8 = 4;
27const TAG_ZERO: u8 = 5;
28const TAG_NONE: u8 = 6;
29
30use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
31
32impl Distribution {
33    pub fn write_to<W: Write>(&self, writer: &mut W) -> Result<()> {
34        let word: u64 = match self {
35            Distribution::TernaryFixed(v) => (TAG_TERNARY_FIXED as u64) << 56 | (*v as u64),
36            Distribution::TernaryProb(p) => {
37                let bits = p.to_bits(); // f64 -> u64 bit representation
38                (TAG_TERNARY_PROB as u64) << 56 | (bits & 0x00FF_FFFF_FFFF_FFFF)
39            }
40            Distribution::BinaryFixed(v) => (TAG_BINARY_FIXED as u64) << 56 | (*v as u64),
41            Distribution::BinaryProb(p) => {
42                let bits = p.to_bits();
43                (TAG_BINARY_PROB as u64) << 56 | (bits & 0x00FF_FFFF_FFFF_FFFF)
44            }
45            Distribution::BinaryBlock(v) => (TAG_BINARY_BLOCK as u64) << 56 | (*v as u64),
46            Distribution::ZERO => (TAG_ZERO as u64) << 56,
47            Distribution::NONE => (TAG_NONE as u64) << 56,
48        };
49        writer.write_u64::<LittleEndian>(word)
50    }
51
52    pub fn read_from<R: Read>(reader: &mut R) -> Result<Self> {
53        let word = reader.read_u64::<LittleEndian>()?;
54        let tag = (word >> 56) as u8;
55        let payload = word & 0x00FF_FFFF_FFFF_FFFF;
56
57        let dist = match tag {
58            TAG_TERNARY_FIXED => Distribution::TernaryFixed(payload as usize),
59            TAG_TERNARY_PROB => Distribution::TernaryProb(f64::from_bits(payload)),
60            TAG_BINARY_FIXED => Distribution::BinaryFixed(payload as usize),
61            TAG_BINARY_PROB => Distribution::BinaryProb(f64::from_bits(payload)),
62            TAG_BINARY_BLOCK => Distribution::BinaryBlock(payload as usize),
63            TAG_ZERO => Distribution::ZERO,
64            TAG_NONE => Distribution::NONE,
65            _ => {
66                return Err(std::io::Error::new(std::io::ErrorKind::InvalidData, "Invalid tag"));
67            }
68        };
69        Ok(dist)
70    }
71}
72
73impl PartialEq for Distribution {
74    fn eq(&self, other: &Self) -> bool {
75        use Distribution::*;
76        match (self, other) {
77            (TernaryFixed(a), TernaryFixed(b)) => a == b,
78            (TernaryProb(a), TernaryProb(b)) => a.to_bits() == b.to_bits(),
79            (BinaryFixed(a), BinaryFixed(b)) => a == b,
80            (BinaryProb(a), BinaryProb(b)) => a.to_bits() == b.to_bits(),
81            (BinaryBlock(a), BinaryBlock(b)) => a == b,
82            (ZERO, ZERO) => true,
83            (NONE, NONE) => true,
84            _ => false,
85        }
86    }
87}
88
89impl Eq for Distribution {}