1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130

//! This module implements `LzmaWriter`.
//!
//! `LzmaWriter` implements the LZMA (XZ) compression/decompression algorithm as a generic
//! Writer.  In other words, it behaves similar to BufWriter.  Instead of buffering the `Write`
//! object passed to it, `LzmaWriter` applies compression or decompression to it as it's write.
//!
//!
//! # Examples
//!
//! ```no_run
//! use lzma::LzmaWriter;
//! use std::io::prelude::*;
//! use std::fs::File;
//!
//! let f = File::create("foo.xz").unwrap();
//! let mut f = LzmaWriter::new_compressor(f, 6).unwrap();
//!
//! write!(f, "It's a small world!").unwrap();
//! f.finish().unwrap();
//! ```

use std::io::{self, Write};
use lzma_sys::*;
use std;
use error::LzmaError;
use ::Direction;
use lzma_stream_wrapper::{LzmaStreamWrapper, LzmaCodeResult};


const DEFAULT_BUF_SIZE: usize = 4 * 1024;


pub struct LzmaWriter<T> {
	inner: T,
	stream: LzmaStreamWrapper,
	buffer: Vec<u8>,
	direction: Direction,
}


impl<T: Write> LzmaWriter<T> {
	pub fn new_compressor(inner: T, preset: u32) -> Result<LzmaWriter<T>, LzmaError> {
		LzmaWriter::with_capacity(DEFAULT_BUF_SIZE, inner, Direction::Compress, preset)
	}

	pub fn new_decompressor(inner: T) -> Result<LzmaWriter<T>, LzmaError> {
		LzmaWriter::with_capacity(DEFAULT_BUF_SIZE, inner, Direction::Decompress, 0)
	}

	pub fn with_capacity(capacity: usize, inner: T, direction: Direction, preset: u32) -> Result<LzmaWriter<T>, LzmaError> {
		let mut writer = LzmaWriter {
			inner,
			stream: LzmaStreamWrapper::new(),
			buffer: vec![0; capacity],
			direction,
		};

		match writer.direction {
			Direction::Compress => {
				writer.stream.easy_encoder(preset, lzma_check::LzmaCheckCrc64)?
			},
			Direction::Decompress => {
				writer.stream.stream_decoder(std::u64::MAX, 0)?
			},
		}

		Ok(writer)
	}
}

impl<W: Write> LzmaWriter<W> {
	/// Finalizes the LZMA stream so that it finishes compressing or decompressing.
	///
	/// This *must* be called after all writing is done to ensure the last pieces of the compressed
	/// or decompressed stream get written out.
	pub fn finish(mut self) -> Result<W, LzmaError> {
		loop {
			match self.lzma_code_and_write(&[], lzma_action::LzmaFinish) {
				Ok(LzmaCodeResult {
					ret: Ok(lzma_ret::LzmaStreamEnd),
					bytes_read: _,
					bytes_written: _,
				}) => break,
				Ok(_) => continue,
				Err(err) => return Err(err),
			}
		}

		Ok(self.inner)
	}

	#[allow(clippy::question_mark)]
	fn lzma_code_and_write(&mut self, input: &[u8], action: lzma_action) -> Result<LzmaCodeResult, LzmaError> {
		let result = self.stream.code(input, &mut self.buffer, action);
		if let Err(err) = result.ret {
			return Err(err);
		}

		if result.bytes_written > 0 {
			Write::write_all(&mut self.inner, &self.buffer[..result.bytes_written])?;
		}

		Ok(result)
	}
}


impl<W: Write> Write for LzmaWriter<W> {
	fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
		// Loop until at least one byte from buf was consumed in order to be
		// compliant with std::io::Write trait API.
		loop {
			match self.lzma_code_and_write(buf, lzma_action::LzmaRun) {
				Ok(result) => if result.bytes_read == 0 && result.bytes_written > 0 {
					continue
				} else {
					// If result.bytes_read is zero, then neither was something
					// written nor read. This indicates, something went wrong.
					return Ok(result.bytes_read)
				},
				Err(LzmaError::Io(err)) => return Err(err),
				Err(err) => return Err(io::Error::new(io::ErrorKind::Other, err)),
			}
		}
	}

	fn flush(&mut self) -> io::Result<()> {
		self.inner.flush()
	}
}