fixed io buffer overflow read and writes

src/io/reader.rs

@@ -18,45 +18,64 @@ pub struct Reader {
 impl Reader {
     pub fn new(filename: &str) -> Reader {
         let mut r = Reader {
-            file: BufReader::new(File::open(filename).expect("Can not create output file")),
+            file: BufReader::new(File::open(filename).expect("can not open input file")),
             buffer: 0,
             byte_buffer: [0; 16],
             read: 0,
         };
-        r.update_buffer();
+        r.fill_buffer();
         r
     }
 
     pub fn read_gamma(&mut self) -> u32 {
         let len = self.read_unary() - 1;
-        (self.read_len(len) | (1 << len)) - 1
+        (self.read_len(len) as u32 | (1 << len)) - 1
     }
 
     fn read_unary(&mut self) -> u32 {
+        let remaining = BUFFER_SIZE - self.read;
+
         let zeros = self.buffer.trailing_zeros();
 
+        if zeros >= remaining {
+            self.fill_buffer();
+            return remaining + self.read_unary();
+        }
+
         self.buffer >>= zeros + 1;
         self.read += zeros + 1;
 
         zeros + 1
     }
 
-    fn read_len(&mut self, len: u32) -> u32 {
+    fn read_len(&mut self, len: u32) -> u128 {
         let mask = (1 << len) - 1;
 
-        let res = self.buffer & mask;
+        let remaining = BUFFER_SIZE - self.read;
+
+        let mut res = self.buffer & mask;
         self.buffer >>= len;
-        self.read += len;
 
-        res as u32
+        if remaining <= len {
+            self.fill_buffer();
+
+            let delta = len - remaining;
+            res |= self.read_len(delta) << remaining;
+
+            return res;
+        }
+
+        self.read += len;
+        res
     }
 
-    fn update_buffer(&mut self) {
+    fn fill_buffer(&mut self) {
         self.file
             .read_exact(&mut self.byte_buffer)
-            .expect("erorr while filling byte buffer");
+            .expect("error while filling byte buffer");
 
         self.buffer = u128::from_be_bytes(self.byte_buffer);
+        self.read = 0;
     }
 }
 
@@ -73,15 +92,16 @@ mod test {
         create_dir_all("data/test/").expect("error while creating test dir");
 
         let mut w = Writer::new("data/test/writer.bin");
-        for i in 1..5 {
+        for i in 1..100 {
             w.write_gamma(i);
         }
         w.flush();
 
         let mut r = Reader::new("data/test/writer.bin");
 
-        for i in 1..5 {
-            assert_eq!(i, r.read_gamma());
+        for i in 1..100 {
+            let a = r.read_gamma();
+            assert_eq!(i, a);
         }
     }
 }

src/io/writer.rs

@@ -67,7 +67,7 @@ impl Writer {
 
         self.update_buffer();
         if len > free {
-            self.buffer |= payload >> (len - free);
+            self.buffer |= payload >> free;
             self.written += len - free;
         }
     }
@@ -118,4 +118,19 @@ mod test {
         assert_eq!(format!("{vb:b}"), "10000001");
         assert_eq!(l, 8);
     }
+
+    #[test]
+    fn test_buffer_overflow() {
+        let word = (1 << 10) - 1;
+        let len = 10;
+
+        let mut w = Writer::new("data/test/writer.bin");
+        w.written = 125;
+
+        w.write_internal(word, len);
+
+        let b = w.buffer;
+        println!("{:b}", b);
+        assert_eq!(b, (1 << 7) - 1)
+    }
 }