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), TernaryProb(f64), BinaryFixed(usize), BinaryProb(f64), BinaryBlock(usize), ZERO, NONE, }
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(); (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(
67 std::io::ErrorKind::InvalidData,
68 "Invalid tag",
69 ));
70 }
71 };
72 Ok(dist)
73 }
74}
75
76impl PartialEq for Distribution {
77 fn eq(&self, other: &Self) -> bool {
78 use Distribution::*;
79 match (self, other) {
80 (TernaryFixed(a), TernaryFixed(b)) => a == b,
81 (TernaryProb(a), TernaryProb(b)) => a.to_bits() == b.to_bits(),
82 (BinaryFixed(a), BinaryFixed(b)) => a == b,
83 (BinaryProb(a), BinaryProb(b)) => a.to_bits() == b.to_bits(),
84 (BinaryBlock(a), BinaryBlock(b)) => a == b,
85 (ZERO, ZERO) => true,
86 (NONE, NONE) => true,
87 _ => false,
88 }
89 }
90}
91
92impl Eq for Distribution {}