1use std::io::{self, Write};
2use bzip2::write::BzEncoder;
3use bzip2::Compression as BzCompression;
4
5const BSDIFF_MAGIC: &[u8; 8] = b"BSDIFF40";
6const BSDF2_MAGIC: &[u8; 5] = b"BSDF2";
7
8#[derive(Debug, Clone, Copy, PartialEq, Eq)]
9pub enum CompressionAlgorithm {
10 None = 0,
11 Bz2 = 1,
12 Brotli = 2,
13}
14
15fn compress(alg: CompressionAlgorithm, data: &[u8]) -> io::Result<Vec<u8>> {
16 match alg {
17 CompressionAlgorithm::None => Ok(data.to_vec()),
18 CompressionAlgorithm::Bz2 => {
19 let mut encoder = BzEncoder::new(Vec::new(), BzCompression::best());
20 encoder.write_all(data)?;
21 encoder.finish()
22 }
23 CompressionAlgorithm::Brotli => {
24 let mut compressed = Vec::new();
25 {
26 let mut encoder = brotli::CompressorWriter::new(
27 &mut compressed,
28 4096, 11, 20, );
32 encoder.write_all(data)?;
33 encoder.flush()?;
34 } Ok(compressed)
36 }
37 }
38}
39
40#[inline]
42fn encode_int64(x: i64, buf: &mut [u8]) {
43 if x >= 0 {
44 buf.copy_from_slice(&x.to_le_bytes());
45 } else {
46 let tmp = ((-x) as u64) | (1u64 << 63);
47 buf.copy_from_slice(&tmp.to_le_bytes());
48 }
49}
50
51#[derive(Debug, Clone, Copy)]
53pub struct ControlEntry {
54 pub diff_size: i64,
55 pub extra_size: i64,
56 pub offset_increment: i64,
57}
58
59pub struct Bsdf2Writer {
61 ctrl_data: Vec<u8>,
62 diff_data: Vec<u8>,
63 extra_data: Vec<u8>,
64 ctrl_alg: CompressionAlgorithm,
65 diff_alg: CompressionAlgorithm,
66 extra_alg: CompressionAlgorithm,
67 written_output: u64,
68}
69
70impl Bsdf2Writer {
71 pub fn new(
73 ctrl_alg: CompressionAlgorithm,
74 diff_alg: CompressionAlgorithm,
75 extra_alg: CompressionAlgorithm,
76 ) -> Self {
77 Self {
78 ctrl_data: Vec::new(),
79 diff_data: Vec::new(),
80 extra_data: Vec::new(),
81 ctrl_alg,
82 diff_alg,
83 extra_alg,
84 written_output: 0,
85 }
86 }
87 pub fn new_legacy() -> Self {
88 Self::new(
89 CompressionAlgorithm::Bz2,
90 CompressionAlgorithm::Bz2,
91 CompressionAlgorithm::Bz2,
92 )
93 }
94 pub fn add_control_entry(&mut self, entry: ControlEntry) -> io::Result<()> {
95 let mut buf = [0u8; 24];
96 encode_int64(entry.diff_size, &mut buf[0..8]);
97 encode_int64(entry.extra_size, &mut buf[8..16]);
98 encode_int64(entry.offset_increment, &mut buf[16..24]);
99
100 self.ctrl_data.extend_from_slice(&buf);
101 self.written_output += (entry.diff_size + entry.extra_size) as u64;
102 Ok(())
103 }
104
105 pub fn write_diff_stream(&mut self, data: &[u8]) -> io::Result<()> {
107 self.diff_data.extend_from_slice(data);
108 Ok(())
109 }
110 pub fn write_extra_stream(&mut self, data: &[u8]) -> io::Result<()> {
111 self.extra_data.extend_from_slice(data);
112 Ok(())
113 }
114 pub fn close<W: Write>(&mut self, writer: &mut W) -> io::Result<()> {
115 let ctrl_compressed = compress(self.ctrl_alg, &self.ctrl_data)?;
117 let diff_compressed = compress(self.diff_alg, &self.diff_data)?;
118 let extra_compressed = compress(self.extra_alg, &self.extra_data)?;
119
120 let is_legacy = self.ctrl_alg == CompressionAlgorithm::Bz2
122 && self.diff_alg == CompressionAlgorithm::Bz2
123 && self.extra_alg == CompressionAlgorithm::Bz2;
124
125 self.write_header(
126 writer,
127 is_legacy,
128 ctrl_compressed.len() as u64,
129 diff_compressed.len() as u64,
130 )?;
131
132 writer.write_all(&ctrl_compressed)?;
134 writer.write_all(&diff_compressed)?;
135 writer.write_all(&extra_compressed)?;
136
137 Ok(())
138 }
139
140 fn write_header<W: Write>(
141 &self,
142 writer: &mut W,
143 is_legacy: bool,
144 ctrl_size: u64,
145 diff_size: u64,
146 ) -> io::Result<()> {
147 let mut header = [0u8; 32];
148
149 if is_legacy {
150 header[0..8].copy_from_slice(BSDIFF_MAGIC);
151 } else {
152 header[0..5].copy_from_slice(BSDF2_MAGIC);
153 header[5] = self.ctrl_alg as u8;
154 header[6] = self.diff_alg as u8;
155 header[7] = self.extra_alg as u8;
156 }
157
158 encode_int64(ctrl_size as i64, &mut header[8..16]);
159 encode_int64(diff_size as i64, &mut header[16..24]);
160 encode_int64(self.written_output as i64, &mut header[24..32]);
161
162 writer.write_all(&header)?;
163 Ok(())
164 }
165}
166
167#[cfg(test)]
168mod tests {
169 use super::*;
170
171 #[test]
172 fn test_encode_int64_positive() {
173 let mut buf = [0u8; 8];
174 encode_int64(42, &mut buf);
175 assert_eq!(buf, [42, 0, 0, 0, 0, 0, 0, 0]);
176 }
177
178 #[test]
179 fn test_encode_int64_negative() {
180 let mut buf = [0u8; 8];
181 encode_int64(-42, &mut buf);
182 assert_eq!(buf, [42, 0, 0, 0, 0, 0, 0, 0x80]);
183 }
184
185 #[test]
186 fn test_encode_int64_zero() {
187 let mut buf = [0u8; 8];
188 encode_int64(0, &mut buf);
189 assert_eq!(buf, [0; 8]);
190 }
191
192 #[test]
193 fn test_writer_creation() {
194 let writer = Bsdf2Writer::new(
195 CompressionAlgorithm::Brotli,
196 CompressionAlgorithm::Brotli,
197 CompressionAlgorithm::Bz2,
198 );
199 assert_eq!(writer.ctrl_alg, CompressionAlgorithm::Brotli);
200 assert_eq!(writer.written_output, 0);
201 }
202
203 #[test]
204 fn test_legacy_writer() {
205 let writer = Bsdf2Writer::new_legacy();
206 assert_eq!(writer.ctrl_alg, CompressionAlgorithm::Bz2);
207 assert_eq!(writer.diff_alg, CompressionAlgorithm::Bz2);
208 assert_eq!(writer.extra_alg, CompressionAlgorithm::Bz2);
209 }
210}