compress/bwt/
mod.rs

1/*!
2
3BWT (Burrows-Wheeler Transform) forward and backward transformation. Requires `bwt` feature, enabled by default
4
5This module contains a bruteforce implementation of BWT encoding in Rust as well as standard decoding.
6These are exposed as a standard `Reader` and `Writer` interfaces wrapping an underlying stream.
7
8BWT output stream places together symbols with similar leading contexts. This reshaping of the entropy
9allows further stages to deal with repeated sequences of symbols for better compression.
10
11Typical compression schemes are:
12BWT + RLE (+ EC)
13RLE + BWT + MTF + RLE + EC  : bzip2
14BWT + DC + EC               : ybs
15
16Where the stage families are:
17BWT: BWT (Burrows-Wheeler Transform), ST (Shindler transform)
18RLE: RLE (Run-Length Encoding)
19MTF: MTF (Move-To-Front), WFC (Weighted Frequency Coding)
20DC: DC (Distance Coding), IF (Inverse Frequencies)
21EC (Entropy Coder): Huffman, Arithmetic, RC (Range Coder)
22
23
24# Example
25
26```rust
27use std::io::{BufWriter, BufReader, Read, Write};
28use compress::bwt;
29
30// Encode some text
31let text = "some text";
32let mut e = bwt::Encoder::new(BufWriter::new(Vec::new()), 4 << 20);
33e.write(text.as_bytes()).unwrap();
34let (encoded, _) = e.finish();
35let inner = encoded.into_inner().unwrap();
36
37// Decode the encoded text
38let mut d = bwt::Decoder::new(BufReader::new(&inner[..]), true);
39let mut decoded = Vec::new();
40d.read_to_end(&mut decoded).unwrap();
41
42assert_eq!(&decoded[..], text.as_bytes());
43```
44
45# Credit
46
47This is an original (mostly trivial) implementation.
48
49*/
50
51#![allow(missing_docs)]
52
53extern crate num;
54
55use std::{cmp, fmt, slice};
56use std::ptr;
57use std::iter::{self, Extend, repeat};
58use std::io::{self, Read, Write};
59use self::num::traits::{NumCast, ToPrimitive};
60
61use super::byteorder::{LittleEndian, WriteBytesExt, ReadBytesExt};
62use super::{byteorder_err_to_io, ReadExact};
63
64pub mod dc;
65pub mod mtf;
66
67/// A base element for the transformation
68pub type Symbol = u8;
69
70pub const ALPHABET_SIZE: usize = 0x100;
71
72/// Radix sorting primitive
73pub struct Radix    {
74    /// number of occurancies (frequency) per symbox
75    pub freq    : [usize; ALPHABET_SIZE+1],
76}
77
78impl Radix  {
79    /// create Radix sort instance
80    pub fn new() -> Radix   {
81        Radix   {
82            freq : [0; ALPHABET_SIZE+1],
83        }
84    }
85
86    /// reset counters
87    /// allows the struct to be re-used
88    pub fn reset(&mut self) {
89        for fr in self.freq.iter_mut()   {
90            *fr = 0;
91        }
92    }
93
94    /// count elements in the input
95    pub fn gather(&mut self, input: &[Symbol])  {
96        for &b in input.iter()  {
97            self.freq[b as usize] += 1;
98        }
99    }
100
101    /// build offset table
102    pub fn accumulate(&mut self)    {
103        let mut n = 0;
104        for freq in self.freq.iter_mut() {
105            let f = *freq;
106            *freq = n;
107            n += f;
108        }
109    }
110
111    /// return next byte position, advance it internally
112    pub fn place(&mut self, b: Symbol)-> usize   {
113        let pos = self.freq[b as usize];
114        assert!(self.freq[b as usize] < self.freq[(b as usize)+1],
115            "Unable to place symbol {} at offset {}",
116            b, pos);
117        self.freq[b as usize] += 1;
118        pos
119    }
120
121    /// shift frequences to the left
122    /// allows the offsets to be re-used after all positions are obtained
123    pub fn shift(&mut self) {
124        assert_eq!( self.freq[ALPHABET_SIZE-1], self.freq[ALPHABET_SIZE] );
125        for i in (0 .. ALPHABET_SIZE).rev()   {
126            self.freq[i+1] = self.freq[i];
127        }
128        self.freq[0] = 0;
129    }
130}
131
132
133/// Compute a suffix array from a given input string
134/// Resulting suffixes are guaranteed to be alphabetically sorted
135/// Run time: O(N^3), memory: N words (suf_array) + ALPHABET_SIZE words (Radix)
136pub fn compute_suffixes<SUF: NumCast + ToPrimitive + fmt::Debug>(input: &[Symbol], suf_array: &mut [SUF]) {
137    let mut radix = Radix::new();
138    radix.gather(input);
139    radix.accumulate();
140
141    debug!("SA compute input: {:?}", input);
142    debug!("radix offsets: {:?}", &radix.freq[..]);
143
144    for (i,&ch) in input.iter().enumerate() {
145        let p = radix.place(ch);
146        suf_array[p] = NumCast::from(i).unwrap();
147    }
148
149    // bring the original offsets back
150    radix.shift();
151
152    for i in 0..ALPHABET_SIZE {
153        let lo = radix.freq[i];
154        let hi = radix.freq[i+1];
155        if lo == hi {
156            continue;
157        }
158        let slice = &mut suf_array[lo..hi];
159        debug!("\tsorting group [{}-{}) for symbol {}", lo, hi, i);
160        slice.sort_by(|a,b| {
161            input[(a.to_usize().unwrap())..].cmp(&input[(b.to_usize().unwrap())..])
162        });
163    }
164
165    debug!("sorted SA: {:?}", suf_array);
166}
167
168/// An iterator over BWT output
169pub struct TransformIterator<'a, SUF: 'a> {
170    input      : &'a [Symbol],
171    suf_iter   : iter::Enumerate<slice::Iter<'a,SUF>>,
172    origin     : Option<usize>,
173}
174
175impl<'a, SUF> TransformIterator<'a, SUF> {
176    /// create a new BWT iterator from the suffix array
177    pub fn new(input: &'a [Symbol], suffixes: &'a [SUF]) -> TransformIterator<'a, SUF> {
178        TransformIterator {
179            input: input,
180            suf_iter: suffixes.iter().enumerate(),
181            origin: None,
182        }
183    }
184
185    /// return the index of the original string
186    pub fn get_origin(&self) -> usize {
187        self.origin.unwrap()
188    }
189}
190
191impl<'a, SUF: ToPrimitive + 'a> Iterator for TransformIterator<'a, SUF> {
192    type Item = Symbol;
193    fn next(&mut self) -> Option<Symbol> {
194        self.suf_iter.next().map(|(i,p)| {
195            if p.to_usize().unwrap() == 0 {
196                assert!( self.origin.is_none() );
197                self.origin = Some(i);
198                *self.input.last().unwrap()
199            }else {
200                self.input[p.to_usize().unwrap() - 1]
201            }
202        })
203    }
204}
205
206/// Encode BWT of a given input, using the 'suf_array'
207pub fn encode<'a, SUF: NumCast + ToPrimitive + fmt::Debug>(input: &'a [Symbol], suf_array: &'a mut [SUF]) -> TransformIterator<'a, SUF> {
208    compute_suffixes(input, suf_array);
209    TransformIterator::new(input, suf_array)
210}
211
212/// Transform an input block into the output slice, all-inclusive version.
213/// Returns the index of the original string in the output matrix.
214pub fn encode_simple(input: &[Symbol]) -> (Vec<Symbol>, usize) {
215    let mut suf_array: Vec<usize> = repeat(0).take(input.len()).collect();
216    let mut iter = encode(input, &mut suf_array[..]);
217    let output: Vec<Symbol> = iter.by_ref().collect();
218    (output, iter.get_origin())
219}
220
221
222/// Compute an inversion jump table, needed for BWT decoding
223pub fn compute_inversion_table<SUF: NumCast + fmt::Debug>(input: &[Symbol], origin: usize, table: &mut [SUF]) {
224    assert_eq!(input.len(), table.len());
225
226    let mut radix = Radix::new();
227    radix.gather(input);
228    radix.accumulate();
229
230    table[radix.place(input[origin])] = NumCast::from(0).unwrap();
231    for (i,&ch) in input[..origin].iter().enumerate() {
232        table[radix.place(ch)] = NumCast::from(i+1).unwrap();
233    }
234    for (i,&ch) in input[(origin+1)..].iter().enumerate() {
235        table[radix.place(ch)] = NumCast::from(origin+2+i).unwrap();
236    }
237    //table[-1] = origin;
238    debug!("inverse table: {:?}", table)
239}
240
241/// An iterator over inverse BWT
242/// Run time: O(N), memory: N words (table)
243pub struct InverseIterator<'a, SUF: 'a> {
244    input      : &'a [Symbol],
245    table      : &'a [SUF],
246    origin     : usize,
247    current    : usize,
248}
249
250impl<'a, SUF> InverseIterator<'a, SUF> {
251    /// create a new inverse BWT iterator with a given input, origin, and a jump table
252    pub fn new(input: &'a [Symbol], origin: usize, table: &'a [SUF]) -> InverseIterator<'a, SUF> {
253        debug!("inverse origin={:?}, input: {:?}", origin, input);
254        InverseIterator {
255            input: input,
256            table: table,
257            origin: origin,
258            current: origin,
259        }
260    }
261}
262
263impl<'a, SUF: ToPrimitive> Iterator for InverseIterator<'a, SUF> {
264    type Item = Symbol;
265
266    fn next(&mut self) -> Option<Symbol> {
267        if self.current == usize::max_value() {
268            None
269        } else {
270            self.current = self.table[self.current].to_usize().unwrap().wrapping_sub(1);
271            debug!("\tjumped to {}", self.current);
272
273            let p = if self.current != usize::max_value() {
274                self.current
275            } else {
276                self.origin
277            };
278
279            Some(self.input[p])
280        }
281    }
282}
283
284/// Decode a BWT block, given it's origin, and using 'table' temporarily
285pub fn decode<'a, SUF: NumCast + fmt::Debug>(input: &'a [Symbol], origin: usize, table: &'a mut [SUF]) -> InverseIterator<'a, SUF> {
286    compute_inversion_table(input, origin, table);
287    InverseIterator::new(input, origin, table)
288}
289
290/// A simplified BWT decode function, which allocates a temporary suffix array
291pub fn decode_simple(input: &[Symbol], origin: usize) -> Vec<Symbol> {
292    let mut suf: Vec<usize> = repeat(0).take(input.len()).collect();
293    decode(input, origin, &mut suf[..]).take(input.len()).collect()
294}
295
296/// Decode without additional memory, can be greatly optimized
297/// Run time: O(n^2), Memory: 0n
298fn decode_minimal(input: &[Symbol], origin: usize, output: &mut [Symbol]) {
299    assert_eq!(input.len(), output.len());
300    if input.len() == 0 {
301        assert_eq!(origin, 0);
302    }
303
304    let mut radix = Radix::new();
305    radix.gather(input);
306    radix.accumulate();
307
308    let n = input.len();
309    (0..n).fold(origin, |i,j| {
310        let ch = input[i];
311        output[n-j-1] = ch;
312        let offset = &input[..i].iter().filter(|&k| *k==ch).count();
313        radix.freq[ch as usize] + offset
314    });
315}
316
317
318/// This structure is used to decode a stream of BWT blocks. This wraps an
319/// internal reader which is read from when this decoder's read method is
320/// called.
321pub struct Decoder<R> {
322    /// The internally wrapped reader. This is exposed so it may be moved out
323    /// of. Note that if data is read from the reader while decoding is in
324    /// progress the output stream will get corrupted.
325    pub r: R,
326    start  : usize,
327
328    temp   : Vec<u8>,
329    output : Vec<u8>,
330    table  : Vec<usize>,
331
332    header         : bool,
333    max_block_size : usize,
334    extra_memory   : bool,
335}
336
337impl<R: Read> Decoder<R> {
338    /// Creates a new decoder which will read data from the given stream. The
339    /// inner stream can be re-acquired by moving out of the `r` field of this
340    /// structure.
341    /// 'extra_mem' switch allows allocating extra N words of memory for better performance
342    pub fn new(r: R, extra_mem: bool) -> Decoder<R> {
343        Decoder {
344            r: r,
345            start: 0,
346            temp: Vec::new(),
347            output: Vec::new(),
348            table: Vec::new(),
349            header: false,
350            max_block_size: 0,
351            extra_memory: extra_mem,
352        }
353    }
354
355    /// Resets this decoder back to its initial state. Note that the underlying
356    /// stream is not seeked on or has any alterations performed on it.
357    pub fn reset(&mut self) {
358        self.header = false;
359        self.start = 0;
360    }
361
362    fn read_header(&mut self) -> io::Result<()> {
363        match self.r.read_u32::<LittleEndian>() {
364            Ok(size) => {
365                self.max_block_size = size as usize;
366                debug!("max size: {}", self.max_block_size);
367                Ok(())
368            },
369            Err(e) => Err(byteorder_err_to_io(e)),
370        }
371    }
372
373    fn decode_block(&mut self) -> io::Result<bool> {
374        let n = match self.r.read_u32::<LittleEndian>() {
375            Ok(n) => n as usize,
376            Err(ref e) if e.kind() == io::ErrorKind::UnexpectedEof => return Ok(false), // EOF
377            Err(e) => return Err(e),
378        };
379
380        self.temp.truncate(0);
381        self.temp.reserve(n);
382        try!(self.r.push_exactly(n as u64, &mut self.temp));
383
384        let origin = try!(self.r.read_u32::<LittleEndian>()) as usize;
385        self.output.truncate(0);
386        self.output.reserve(n);
387
388        if self.extra_memory    {
389            self.table.truncate(0);
390            self.table.extend((0..n).map(|_| 0));
391            for ch in decode(&self.temp[..], origin, &mut self.table[..]) {
392                self.output.push(ch);
393            }
394        }else   {
395            self.output.extend((0..n).map(|_| 0));
396            decode_minimal(&self.temp[..], origin, &mut self.output[..]);
397        }
398
399        self.start = 0;
400        return Ok(true);
401    }
402}
403
404impl<R: Read> Read for Decoder<R> {
405    fn read(&mut self, dst: &mut [u8]) -> io::Result<usize> {
406        if !self.header {
407            try!(self.read_header());
408            self.header = true;
409        }
410        let mut amt = dst.len();
411        let dst_len = amt;
412
413        while amt > 0 {
414            if self.output.len() == self.start {
415                let keep_going = try!(self.decode_block());
416                if !keep_going {
417                   break
418                }
419            }
420            let n = cmp::min(amt, self.output.len() - self.start);
421            unsafe { ptr::copy_nonoverlapping(
422                &self.output[self.start],
423                &mut dst[dst_len - amt],
424                n,
425            )};
426            self.start += n;
427            amt -= n;
428        }
429
430        Ok(dst_len - amt)
431    }
432}
433
434
435/// This structure is used to compress a stream of bytes using the BWT.
436/// This is a wrapper around an internal writer which bytes will be written to.
437pub struct Encoder<W> {
438    w: W,
439    buf: Vec<u8>,
440    suf: Vec<usize>,
441    wrote_header: bool,
442    block_size: usize,
443}
444
445impl<W: Write> Encoder<W> {
446    /// Creates a new encoder which will have its output written to the given
447    /// output stream. The output stream can be re-acquired by calling
448    /// `finish()`
449    /// 'block_size' is idealy as big as your input, unless you know for sure that
450    /// the input consists of multiple parts of different nature. Often set as 4Mb.
451    pub fn new(w: W, block_size: usize) -> Encoder<W> {
452        Encoder {
453            w: w,
454            buf: Vec::new(),
455            suf: Vec::new(),
456            wrote_header: false,
457            block_size: block_size,
458        }
459    }
460
461    fn encode_block(&mut self) -> io::Result<()> {
462        let n = self.buf.len();
463        try!(self.w.write_u32::<LittleEndian>(n as u32));
464
465        self.suf.truncate(0);
466        self.suf.extend((0..n).map(|_| n));
467        let w = &mut self.w;
468
469        {
470            let mut iter = encode(&self.buf[..], &mut self.suf[..]);
471            for ch in iter.by_ref() {
472                try!(w.write_u8(ch));
473            }
474
475            try!(w.write_u32::<LittleEndian>(iter.get_origin() as u32));
476        }
477        self.buf.truncate(0);
478
479        Ok(())
480    }
481
482    /// This function is used to flag that this session of compression is done
483    /// with. The stream is finished up (final bytes are written), and then the
484    /// wrapped writer is returned.
485    pub fn finish(mut self) -> (W, io::Result<()>) {
486        let result = self.flush();
487        (self.w, result)
488    }
489}
490
491impl<W: Write> Write for Encoder<W> {
492    fn write(&mut self, mut buf: &[u8]) -> io::Result<usize> {
493        if !self.wrote_header {
494            try!(self.w.write_u32::<LittleEndian>(self.block_size as u32));
495            self.wrote_header = true;
496        }
497
498        while buf.len() > 0 {
499            let amt = cmp::min( self.block_size - self.buf.len(), buf.len() );
500            self.buf.extend(buf[..amt].iter().map(|b| *b));
501
502            if self.buf.len() == self.block_size {
503                try!(self.encode_block());
504            }
505            buf = &buf[amt..];
506        }
507        Ok(buf.len())
508    }
509
510    fn flush(&mut self) -> io::Result<()> {
511        let ret = if self.buf.len() > 0 {
512            self.encode_block()
513        } else {
514            Ok(())
515        };
516        ret.and(self.w.flush())
517    }
518}
519
520
521#[cfg(test)]
522mod test {
523    use std::io::{BufReader, BufWriter, Read, Write};
524    #[cfg(feature="unstable")]
525    use test::Bencher;
526    use super::{Decoder, Encoder};
527
528    fn roundtrip(bytes: &[u8], extra_mem: bool) {
529        let mut e = Encoder::new(BufWriter::new(Vec::new()), 1<<10);
530        e.write(bytes).unwrap();
531        let (e, err) = e.finish();
532        err.unwrap();
533        let encoded = e.into_inner().unwrap();
534
535        let mut d = Decoder::new(BufReader::new(&encoded[..]), extra_mem);
536        let mut decoded = Vec::new();
537        d.read_to_end(&mut decoded).unwrap();
538        assert_eq!(&decoded[..], bytes);
539    }
540
541    #[test]
542    fn some_roundtrips() {
543        roundtrip(b"test", true);
544        roundtrip(b"", true);
545        roundtrip(include_bytes!("../data/test.txt"), true);
546    }
547
548    #[test]
549    fn decode_minimal() {
550        roundtrip(b"abracadabra", false);
551    }
552
553    #[cfg(feature="unstable")]
554    #[bench]
555    fn decode_speed(bh: &mut Bencher) {
556        use std::iter::repeat;
557        use super::{encode, decode};
558
559        let input = include_bytes!("../data/test.txt");
560        let n = input.len();
561        let mut suf: Vec<u16> = repeat(0).take(n).collect();
562        let (output, origin) = {
563            let mut to_iter = encode(input, &mut suf[..]);
564            let out: Vec<u8> = to_iter.by_ref().collect();
565            (out, to_iter.get_origin())
566        };
567
568        bh.iter(|| {
569            let from_iter = decode(&output[..], origin, &mut suf[..]);
570            from_iter.last().unwrap();
571        });
572        bh.bytes = n as u64;
573    }
574}
compress/bwt/mod.rs

compress/bwt/
mod.rs
