simple_arithmetic_coding/
lib.rs

1#![forbid(unsafe_code)]
2
3use std::{collections::VecDeque, io::{BufRead, Write}};
4
5use bitvec::BitVec;
6use io::{Next, Push};
7
8mod codec;
9mod fenwick;
10mod io;
11mod bitvec;
12
13const SIZE: usize = 257;
14const EOF: usize = 256;
15
16const BUFFER_SIZE: usize = 16000;
17
18/**
19 * The encoding routine for arithmetic coding. Takes an input and an output.
20 */
21pub fn encode_routine<I,O> (input_handle: &mut I, output_handle: &mut O) -> Result<(), std::io::Error>
22where I: BufRead, O: Write
23{
24    let mut encoder = codec::ArithmeticEncoder::new(SIZE);
25    let mut bits_out = io::Output::new(output_handle, BUFFER_SIZE);
26    let mut bits_in = io::Input::new(input_handle)?;
27
28    while let Some(byte) = bits_in.next_byte()? {
29        encoder.encode(byte as usize, &mut bits_out)?;
30        encoder.discover(byte as usize);
31    }
32    encoder.encode(EOF, &mut bits_out)?;
33    encoder.finish(&mut bits_out)?;
34    bits_out.flush()?;
35
36    Ok(())
37}
38
39/**
40 * The decoding routine for arithmetic coding. Takes an input and an output.
41 */
42pub fn decode_routine<I,O> (input_handle: &mut I, output_handle: &mut O) -> Result<(), std::io::Error>
43where I: BufRead, O: Write
44{
45    let mut decoder = codec::ArithmeticDecoder::new(SIZE);
46    let mut bits_out = io::Output::new(output_handle, BUFFER_SIZE);
47    let mut bits_in = io::Input::new(input_handle)?;
48
49    decoder.begin(&mut bits_in)?;
50    loop {
51        let s = decoder.decode(&mut bits_in)?;
52        if s == EOF {
53            break;
54        }
55        bits_out.push_byte(s as u8)?;
56        decoder.discover(s);
57    }
58    bits_out.flush()?;
59    Ok(())
60}
61
62/**
63 * Convenience iterator for arithmetic decoding. Panics if IO errors occur during _next_. Take care that you know what you're doing when using iterator methods on this, such as _filter_ or _step_by_, etc.
64 */
65pub struct ArithmeticStreamDecoder<'a,I>
66where I: BufRead
67{
68    decoder: codec::ArithmeticDecoder,
69    bits_in: io::Input<&'a mut I>,
70}
71
72impl<'a,I> ArithmeticStreamDecoder<'a,I>
73where I: BufRead
74{
75    pub fn new(input: &'a mut I) -> Result<Self, std::io::Error>
76    {
77        let mut decoder = codec::ArithmeticDecoder::new(SIZE);
78        let mut bits_in = io::Input::new(input)?;
79
80        decoder.begin(&mut bits_in)?;
81
82        Ok(Self { decoder, bits_in })
83    }
84}
85
86impl<I> Iterator for ArithmeticStreamDecoder<'_,I>
87where I: BufRead
88{
89    type Item = u8;
90
91    fn next(&mut self) -> Option<Self::Item>
92    {
93        let s = self.decoder.decode(&mut self.bits_in).unwrap();
94        if s == EOF {
95            return None;
96        }
97        self.decoder.discover(s);
98        Some(s as u8)
99    }
100}
101
102struct SmallBuffer {
103    data: BitVec,
104    q: VecDeque<u8>,
105    eof: bool,
106}
107
108impl SmallBuffer {
109    fn new(capacity: usize) -> Self {
110        Self {
111            data: BitVec::with_capacity(capacity),
112            q: VecDeque::new(),
113            eof: false,
114        }
115    }
116}
117
118impl Push for SmallBuffer {
119    fn push_bit(&mut self, bit: bool) -> Result<(), std::io::Error> {
120        if self.data.len() == self.data.capacity() {
121            self.flush()?;
122        }
123        self.data.push(bit);
124        Ok(())
125    }
126    fn push_byte(&mut self, byte: u8) -> Result<(), std::io::Error> {
127        for i in (0..8).rev() {
128            self.data.push(((byte >> i) & 1) != 0);
129        }
130        Ok(())
131    }
132    fn flush (&mut self) -> Result<(), std::io::Error> {
133        if !self.data.is_empty() {
134            self.q.write_all(&self.data.get_bytes())?;
135            self.data.clear();
136        }
137        Ok(())
138    }
139}
140
141/**
142 * Convenience iterator for arithmetic encoding. Panics if IO errors occur during _next_. Take care that you know what you're doing when using iterator methods on this, such as _filter_ or _step_by_, etc.
143 */
144pub struct ArithmeticStreamEncoder<'a, I>
145where I: BufRead
146{
147    encoder: codec::ArithmeticEncoder,
148    bits_in: io::Input<&'a mut I>,
149    buf: SmallBuffer,
150}
151
152impl<'a, I> ArithmeticStreamEncoder<'a, I>
153where I: BufRead
154{
155    pub fn new(input: &'a mut I) -> Result<Self, std::io::Error>
156    {
157        let encoder = codec::ArithmeticEncoder::new(SIZE);
158        let bits_in = io::Input::new(input)?;
159        let buf = SmallBuffer::new(BUFFER_SIZE/4);
160
161        Ok(Self { encoder, bits_in, buf })
162    }
163}
164
165impl<I> Iterator for ArithmeticStreamEncoder<'_, I>
166where I: BufRead
167{
168    type Item = u8;
169
170    fn next(&mut self) -> Option<Self::Item>
171    {
172        while self.buf.q.is_empty() && !self.buf.eof {
173            if let Some(byte) = self.bits_in.next_byte().unwrap() {
174                self.encoder.encode(byte as usize, &mut self.buf).unwrap();
175                self.encoder.discover(byte as usize);
176            } else {
177                self.encoder.encode(EOF, &mut self.buf).unwrap();
178                self.encoder.finish(&mut self.buf).unwrap();
179                self.buf.flush().unwrap();
180                self.buf.eof = true;
181            }
182        }
183
184        self.buf.q.pop_front()
185    }
186}
187
188
189
190
191
simple_arithmetic_coding/lib.rs

simple_arithmetic_coding/
lib.rs
