1use crate::diagnostic::{Code, Diagnostic};
4use std::fmt;
5use std::fs::File;
6use std::io::{self, BufWriter, Write};
7use std::path::Path;
8
9pub const DEFAULT_OUTPUT_MEMORY_LIMIT_BYTES: usize = 100_000_000;
14const OUTPUT_BUFFER_BYTES: usize = 64 * 1_024;
15
16#[derive(Clone, Copy, Debug, PartialEq, Eq)]
18pub struct OutputOptions {
19 pub memory_limit_bytes: usize,
21}
22
23impl OutputOptions {
24 #[must_use]
26 pub const fn with_memory_limit(mut self, bytes: usize) -> Self {
27 self.memory_limit_bytes = bytes;
28 self
29 }
30}
31
32impl Default for OutputOptions {
33 fn default() -> Self {
34 Self {
35 memory_limit_bytes: DEFAULT_OUTPUT_MEMORY_LIMIT_BYTES,
36 }
37 }
38}
39
40pub struct TextOutput<'a, W> {
45 output: &'a mut W,
46 error: Option<io::Error>,
47}
48
49impl<W> fmt::Debug for TextOutput<'_, W> {
50 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
51 formatter
52 .debug_struct("TextOutput")
53 .field("failed", &self.error.is_some())
54 .finish_non_exhaustive()
55 }
56}
57
58impl<'a, W> TextOutput<'a, W> {
59 pub const fn new(output: &'a mut W) -> Self {
61 Self {
62 output,
63 error: None,
64 }
65 }
66
67 pub fn finish(self) -> io::Result<()> {
73 match self.error {
74 Some(error) => Err(error),
75 None => Ok(()),
76 }
77 }
78}
79
80impl<W: Write> std::fmt::Write for TextOutput<'_, W> {
81 fn write_str(&mut self, text: &str) -> std::fmt::Result {
82 if self.error.is_some() {
83 return Err(std::fmt::Error);
84 }
85 if let Err(error) = self.output.write_all(text.as_bytes()) {
86 self.error = Some(error);
87 return Err(std::fmt::Error);
88 }
89 Ok(())
90 }
91}
92
93#[must_use = "compressed streams must be finished explicitly"]
98pub struct OutputSink<W: Write> {
99 encoder: Option<Encoder<W>>,
100 destination: Box<str>,
101}
102
103impl<W: Write> fmt::Debug for OutputSink<W> {
104 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
105 formatter
106 .debug_struct("OutputSink")
107 .field("destination", &self.destination)
108 .field("finished", &self.encoder.is_none())
109 .finish_non_exhaustive()
110 }
111}
112
113enum Encoder<W: Write> {
114 Plain(BufWriter<W>),
115 #[cfg(feature = "gzip")]
116 Gzip(flate2::write::GzEncoder<BufWriter<W>>),
117 #[cfg(feature = "zstd")]
118 Zstd(zstd::stream::write::Encoder<'static, BufWriter<W>>),
119}
120
121impl OutputSink<File> {
122 pub fn create(path: impl AsRef<Path>, options: OutputOptions) -> Result<Self, Diagnostic> {
129 let path = path.as_ref();
130 validate_options(path, options)?;
131 let compression = compression_for(path);
132 ensure_available(path, compression)?;
133 let file = File::create(path).map_err(|error| output_error(path, error))?;
134 Self::from_writer(
135 file,
136 compression,
137 options,
138 path.display().to_string().into(),
139 )
140 .map_err(|error| output_error(path, error))
141 }
142}
143
144impl<W: Write> OutputSink<W> {
145 fn from_writer(
146 writer: W,
147 compression: OutputCompression,
148 options: OutputOptions,
149 destination: Box<str>,
150 ) -> io::Result<Self> {
151 let capacity = OUTPUT_BUFFER_BYTES.min(options.memory_limit_bytes);
152 let writer = BufWriter::with_capacity(capacity, writer);
153 let encoder = match compression {
154 OutputCompression::None => Encoder::Plain(writer),
155 #[cfg(feature = "gzip")]
156 OutputCompression::Gzip => Encoder::Gzip(
157 flate2::GzBuilder::new()
158 .mtime(0)
159 .write(writer, flate2::Compression::default()),
160 ),
161 #[cfg(feature = "zstd")]
162 OutputCompression::Zstd => Encoder::Zstd(zstd::stream::write::Encoder::new(writer, 0)?),
163 #[cfg(any(not(feature = "gzip"), not(feature = "zstd")))]
164 _ => return Err(io::Error::other("compression support is not enabled")),
165 };
166 Ok(Self {
167 encoder: Some(encoder),
168 destination,
169 })
170 }
171
172 pub fn finish(mut self) -> Result<(), Diagnostic> {
178 let Some(encoder) = self.encoder.take() else {
179 return Err(output_error_name(
180 &self.destination,
181 "output stream was already finished",
182 ));
183 };
184 finish_encoder(encoder).map_err(|error| output_error_name(&self.destination, error))
185 }
186
187 #[cfg(test)]
188 fn buffer_capacity(&self) -> Option<usize> {
189 match self.encoder.as_ref() {
190 Some(Encoder::Plain(writer)) => Some(writer.capacity()),
191 #[cfg(feature = "gzip")]
192 Some(Encoder::Gzip(writer)) => Some(writer.get_ref().capacity()),
193 #[cfg(feature = "zstd")]
194 Some(Encoder::Zstd(writer)) => Some(writer.get_ref().capacity()),
195 None => None,
196 }
197 }
198}
199
200impl<W: Write> Write for OutputSink<W> {
201 fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
202 match self.encoder.as_mut() {
203 Some(Encoder::Plain(writer)) => writer.write(bytes),
204 #[cfg(feature = "gzip")]
205 Some(Encoder::Gzip(writer)) => writer.write(bytes),
206 #[cfg(feature = "zstd")]
207 Some(Encoder::Zstd(writer)) => writer.write(bytes),
208 None => Err(io::Error::new(
209 io::ErrorKind::BrokenPipe,
210 "output stream is finished",
211 )),
212 }
213 }
214
215 fn flush(&mut self) -> io::Result<()> {
216 match self.encoder.as_mut() {
217 Some(Encoder::Plain(writer)) => writer.flush(),
218 #[cfg(feature = "gzip")]
219 Some(Encoder::Gzip(writer)) => writer.flush(),
220 #[cfg(feature = "zstd")]
221 Some(Encoder::Zstd(writer)) => writer.flush(),
222 None => Err(io::Error::new(
223 io::ErrorKind::BrokenPipe,
224 "output stream is finished",
225 )),
226 }
227 }
228}
229
230pub fn write_output(path: impl AsRef<Path>, bytes: &[u8]) -> Result<(), Diagnostic> {
239 let path = path.as_ref();
240 let mut output = OutputSink::create(path, OutputOptions::default())?;
241 output
242 .write_all(bytes)
243 .map_err(|error| output_error(path, error))?;
244 output.finish()
245}
246
247#[cfg(any(feature = "gzip", feature = "zstd"))]
248fn finish_encoder<W: Write>(encoder: Encoder<W>) -> io::Result<()> {
249 let mut writer = match encoder {
250 Encoder::Plain(writer) => writer,
251 #[cfg(feature = "gzip")]
252 Encoder::Gzip(writer) => writer.finish()?,
253 #[cfg(feature = "zstd")]
254 Encoder::Zstd(writer) => writer.finish()?,
255 };
256 writer.flush()
257}
258
259#[cfg(not(any(feature = "gzip", feature = "zstd")))]
260fn finish_encoder<W: Write>(encoder: Encoder<W>) -> io::Result<()> {
261 let Encoder::Plain(mut writer) = encoder;
262 writer.flush()
263}
264
265#[derive(Clone, Copy)]
266enum OutputCompression {
267 None,
268 Gzip,
269 Zstd,
270}
271
272fn compression_for(path: &Path) -> OutputCompression {
273 match path.extension().and_then(|suffix| suffix.to_str()) {
274 Some(suffix) if suffix.eq_ignore_ascii_case("gz") => OutputCompression::Gzip,
275 Some(suffix) if suffix.eq_ignore_ascii_case("zst") => OutputCompression::Zstd,
276 _ => OutputCompression::None,
277 }
278}
279
280#[cfg_attr(
284 all(feature = "gzip", feature = "zstd"),
285 allow(clippy::unnecessary_wraps)
286)]
287fn ensure_available(path: &Path, compression: OutputCompression) -> Result<(), Diagnostic> {
288 let _ = path;
289 match compression {
290 #[cfg(not(feature = "gzip"))]
291 OutputCompression::Gzip => Err(unavailable(path, "gzip")),
292 #[cfg(not(feature = "zstd"))]
293 OutputCompression::Zstd => Err(unavailable(path, "zstandard")),
294 _ => Ok(()),
295 }
296}
297
298fn validate_options(path: &Path, options: OutputOptions) -> Result<(), Diagnostic> {
299 if options.memory_limit_bytes == 0 {
300 return Err(Diagnostic::new(Code::E7901)
301 .with_context("path", path.display().to_string())
302 .with_context("memory limit", options.memory_limit_bytes.to_string()));
303 }
304 Ok(())
305}
306
307#[cfg(any(not(feature = "gzip"), not(feature = "zstd")))]
308fn unavailable(path: &Path, compression: &str) -> Diagnostic {
309 Diagnostic::new(Code::E7901)
310 .with_context("path", path.display().to_string())
311 .with_context("compression", compression)
312 .with_context("reason", "support is not enabled in this build")
313}
314
315fn output_error(path: &Path, error: impl std::fmt::Display) -> Diagnostic {
316 output_error_name(&path.display().to_string(), error)
317}
318
319fn output_error_name(destination: &str, error: impl std::fmt::Display) -> Diagnostic {
320 Diagnostic::new(Code::E7901)
321 .with_context("path", destination)
322 .with_context("reason", error.to_string())
323}
324
325#[cfg(test)]
326#[path = "output_tests.rs"]
327mod tests;