Handle near-converged state

.gitignore

@@ -1 +1,2 @@
 /target
+Cargo.lock

src/bit_buffer.rs

@@ -0,0 +1,87 @@
+#[derive(Debug)]
+pub struct BitWriter {
+    data: Vec<u8>,
+    bits: u8,
+    nextbit: usize,
+}
+impl BitWriter {
+    pub fn new() -> Self {
+        return BitWriter {
+            data: vec![],
+            bits: 0,
+            nextbit: 7,
+        };
+    }
+    pub fn write(&mut self, bit: bool) {
+        if bit {
+            self.bits = 1 << self.nextbit | self.bits;
+        }
+        if self.nextbit == 0 {
+            self.data.push(self.bits);
+            self.bits = 0;
+            self.nextbit = 7;
+        } else {
+            self.nextbit -= 1;
+        }
+    }
+    pub fn flush(mut self) -> Vec<u8> {
+        if self.bits != 0 {
+            self.data.push(self.bits);
+        }
+        return self.data;
+    }
+}
+
+trait Poppable {
+    fn pop(&mut self) -> Option<u8>;
+}
+impl Poppable for &[u8] {
+    fn pop(&mut self) -> Option<u8> {
+        if self.len() == 0 {
+            return None;
+        }
+        let out = self[0];
+        *self = &self[1..];
+        return Some(out);
+    }
+}
+
+pub struct BitReader<'a> {
+    data: &'a [u8],
+    // bits: u8,
+    nextbit: usize,
+}
+
+impl<'a> BitReader<'a> {
+    pub fn new(data: &'a [u8]) -> Self {
+        BitReader { data, nextbit: 7 }
+    }
+}
+
+impl<'a> From<&'a [u8]> for BitReader<'a> {
+    fn from(value: &'a [u8]) -> Self {
+        BitReader::new(value)
+    }
+}
+
+impl BitReader<'_> {
+    pub fn pop(&mut self) -> Option<bool> {
+        if self.data.len() == 0 {
+            return None;
+        }
+        let bit = (self.data[0] >> self.nextbit) & 1;
+        if self.nextbit == 0 {
+            self.data.pop();
+            self.nextbit = 8;
+        }
+        self.nextbit -= 1;
+        return Some(bit > 0);
+    }
+}
+
+impl Iterator for BitReader<'_> {
+    type Item = bool;
+    fn next(&mut self) -> Option<Self::Item> {
+        self.pop()
+    }
+}

src/main.rs

@@ -1,53 +1,97 @@
 // https://marknelson.us/posts/2014/10/19/data-compression-with-arithmetic-coding.html
 use std::collections::HashMap;
 
+mod bit_buffer;
+use bit_buffer::{BitReader, BitWriter};
+
 type Model = HashMap<u8, (f64, f64)>;
 
-fn get_symbol(model: &Model, d: f64) -> Option<u8> {
-    // Brute force
+fn get_symbol(model: &Model, low: f64, high: f64) -> Option<u8> {
     for (symbol, (start, end)) in model {
-        if d >= *start && d < *end {
+        if low >= *start && high < *end {
             return Some(*symbol);
         }
     }
     return None;
 }
 
-fn encode(data: &[u8], model: &Model) -> f64 {
-    let mut high: f64 = 1.0;
-    let mut low: f64 = 0.0;
-    for symbol in data {
-        let p = model.get(symbol).expect("Invalid/Unsupported data");
+fn encode(input: &[u8], model: &Model) -> Vec<u8> {
+    const HALF: u64 = 1 << (u64::BITS - 1);
+    const LOW_CONVERGE: u64 = 0b10 << (u64::BITS - 2);
+    const HIGH_CONVERGE: u64 = 0b01 << (u64::BITS - 2);
+
+    let mut output = BitWriter::new();
+
+    let mut high = u64::MAX;
+    let mut low = u64::MIN;
+    let mut pending_bits = 0;
+
+    for symbol in input {
         let range = high - low;
-        high = low + range * p.1;
-        low = low + range * p.0;
+        let p = model.get(symbol).expect("Invalid/Unsupported data");
+        high = low + (range as f64 * p.1) as u64;
+        low = low + (range as f64 * p.0) as u64;
+        loop {
+            if high < HALF {
+                output.write(false);
+                print!("0");
+                while pending_bits > 0 {
+                    output.write(true);
+                    print!("1");
+                    pending_bits -= 1;
                 }
-    return low + (high - low) / 2.0;
+            } else if low >= HALF {
+                output.write(true);
+                print!("1");
+                while pending_bits > 0 {
+                    output.write(true);
+                    print!("0");
+                    pending_bits -= 1;
+                }
+            } else if low >= LOW_CONVERGE && high < HIGH_CONVERGE {
+                println!("BET");
+                pending_bits += 1;
+                low <<= 1;
+                low &= HALF - 1;
+                high <<= 1;
+                high &= HALF + 1;
+                continue;
+            } else {
+                break;
+            }
+            low <<= 1;
+            high <<= 1;
+            high |= 1;
+        }
+    }
+    println!("");
+    return output.flush();
 }
 
-fn decode(message: f64, model: &Model) {
-    let mut high: f64 = 1.0;
-    let mut low: f64 = 0.0;
-    loop {
-        let range = high - low;
-        let d = (message - low) / range;
-        let c = match get_symbol(&model, d) {
-            Some(c) => c,
-            None => {
-                println!("");
-                eprintln!("Decode error: d={d}");
-                return;
+fn decode(input: &[u8], model: &Model) -> Vec<u8> {
+    let mut high = 1.0;
+    let mut low = 0.0;
+    let mut output = vec![];
+    for bit in BitReader::new(input) {
+        let diff = high - low;
+        if bit {
+            //print!("1");
+            low = low + (diff / 2.0);
+        } else {
+            high = high - (diff / 2.0);
+            //print!("0");
         }
-        };
-        if c == b'-' {
-            println!("");
-            return;
+        if let Some(symbol) = get_symbol(model, low, high) {
+            //println!("\nGot sym: {} from [{}, {})", symbol as char, low, high);
+            output.push(symbol);
+            let (slow, shigh) = model.get(&symbol).unwrap();
+            let symdiff = *shigh - *slow;
+            high = (high - *slow) / symdiff;
+            low = (low - *slow) / symdiff;
         }
-        print!("{}", c as char);
-        let p = model.get(&c).expect("Decode error");
-        high = low + range * p.1;
-        low = low + range * p.0;
     }
+
+    return output;
 }
 
 fn make_model(probabilities: &[(u8, f64)]) -> Model {
@@ -92,8 +136,17 @@ const ENGLISH: &[(u8, f64)] = &[
     (b'-', 0.02),
 ];
 fn main() {
+    let data = b"hello world-";
+    println!("MODEL:");
     let model: Model = make_model(ENGLISH);
-    let message = encode(b"hello world-", &model);
-    println!("{message}");
-    decode(message, &model);
+    println!("");
+
+    let _enc = encode(data, &model);
+    let _dec = decode(&_enc, &model);
+
+    println!("{}", String::from_utf8(_dec).unwrap());
+    println!(
+        "Compression Ratio: {}",
+        data.len() as f64 / _enc.len() as f64
+    );
 }