1use std::io::{Read, Result, Write};
2
3pub trait GetDistribution {
5 fn dist(&self) -> &Distribution;
7}
8
9pub trait GetDistributionMut {
11 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#[derive(Clone, Copy, Debug)]
44pub enum Distribution {
45 TernaryFixed(usize),
47 TernaryProb(f64),
49 BinaryFixed(usize),
51 BinaryProb(f64),
53 BinaryBlock(usize),
55 ZERO,
57 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 #[inline]
77 fn pack_f64(tag: u8, p: f64) -> u64 {
78 (tag as u64) << 56 | (p.to_bits() >> 8)
79 }
80
81 #[inline]
84 fn unpack_f64(payload: u64) -> f64 {
85 f64::from_bits(payload << 8)
86 }
87
88 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 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 {}