molframe_core/io/input/
source_file.rs1use super::{ByteWindow, Compression, SourceBytes, SpillWindowedFile, WindowedFile};
4use crate::{Code, Diagnostic, ExecutionContext};
5use std::fs::File;
6#[cfg(any(feature = "gzip", feature = "zstd"))]
7use std::io::ErrorKind;
8use std::io::Read;
9use std::path::Path;
10#[cfg(any(feature = "gzip", feature = "zstd"))]
11use std::sync::atomic::{AtomicU64, Ordering};
12
13#[cfg(any(feature = "gzip", feature = "zstd"))]
14static SPOOL_SEQUENCE: AtomicU64 = AtomicU64::new(0);
15
16#[derive(Debug)]
19pub enum WindowedSourceFile {
20 Plain(WindowedFile),
22 Spilled(SpillWindowedFile),
24}
25
26impl WindowedSourceFile {
27 pub fn open(
36 path: impl AsRef<Path>,
37 max_window_bytes: usize,
38 context: &ExecutionContext,
39 ) -> Result<Self, Diagnostic> {
40 let path = path.as_ref();
41 let mut file = File::open(path).map_err(|error| io_error("open", &error))?;
42 let mut prefix = [0_u8; 4];
43 let prefix_len = file
44 .read(&mut prefix)
45 .map_err(|error| io_error("read compression prefix", &error))?;
46 match Compression::sniff(&prefix[..prefix_len]) {
47 Compression::None => {
48 WindowedFile::open(path, max_window_bytes, context).map(Self::Plain)
49 }
50 Compression::Gzip => open_gzip(path, max_window_bytes, context).map(Self::Spilled),
51 Compression::Zstd => open_zstd(path, max_window_bytes, context).map(Self::Spilled),
52 }
53 }
54
55 #[must_use]
57 pub const fn spill_bytes(&self) -> u64 {
58 match self {
59 Self::Plain(_) => 0,
60 Self::Spilled(source) => source.spill_bytes(),
61 }
62 }
63}
64
65impl SourceBytes for WindowedSourceFile {
66 fn window(&mut self, start: u64, len: usize) -> Result<ByteWindow<'_>, Diagnostic> {
67 match self {
68 Self::Plain(source) => source.window(start, len),
69 Self::Spilled(source) => source.window(start, len),
70 }
71 }
72
73 fn len_hint(&self) -> Option<u64> {
74 match self {
75 Self::Plain(source) => source.len_hint(),
76 Self::Spilled(source) => source.len_hint(),
77 }
78 }
79}
80
81#[cfg(feature = "gzip")]
82fn open_gzip(
83 path: &Path,
84 window_bytes: usize,
85 context: &ExecutionContext,
86) -> Result<SpillWindowedFile, Diagnostic> {
87 let file = File::open(path).map_err(|error| io_error("open gzip", &error))?;
88 spool(flate2::read::GzDecoder::new(file), window_bytes, context)
89}
90
91#[cfg(not(feature = "gzip"))]
92fn open_gzip(
93 _path: &Path,
94 _window_bytes: usize,
95 _context: &ExecutionContext,
96) -> Result<SpillWindowedFile, Diagnostic> {
97 Err(codec_error("gzip support is not enabled"))
98}
99
100#[cfg(feature = "zstd")]
101fn open_zstd(
102 path: &Path,
103 window_bytes: usize,
104 context: &ExecutionContext,
105) -> Result<SpillWindowedFile, Diagnostic> {
106 let file = File::open(path).map_err(|error| io_error("open zstd", &error))?;
107 let decoder = zstd::stream::read::Decoder::new(file)
108 .map_err(|error| codec_error(format!("zstd decoder failed: {error}")))?;
109 spool(decoder, window_bytes, context)
110}
111
112#[cfg(not(feature = "zstd"))]
113fn open_zstd(
114 _path: &Path,
115 _window_bytes: usize,
116 _context: &ExecutionContext,
117) -> Result<SpillWindowedFile, Diagnostic> {
118 Err(codec_error("zstd support is not enabled"))
119}
120
121#[cfg(any(feature = "gzip", feature = "zstd"))]
122fn spool(
123 mut reader: impl Read,
124 record_bytes: usize,
125 context: &ExecutionContext,
126) -> Result<SpillWindowedFile, Diagnostic> {
127 let key = format!(
128 "structure-input-{}",
129 SPOOL_SEQUENCE.fetch_add(1, Ordering::Relaxed)
130 );
131 let mut spill = context
132 .create_spill_file(&key)
133 .map_err(|error| spill_error(&error))?;
134 let reservation = context
135 .try_reserve(record_bytes)
136 .map_err(|error| resource_error(error.to_string()))?;
137 let mut buffer = Vec::new();
138 buffer
139 .try_reserve_exact(record_bytes)
140 .map_err(|error| resource_error(error.to_string()))?;
141 buffer.resize(record_bytes, 0);
142 let mut length = 0_u64;
143 loop {
144 let mut filled = 0usize;
145 while filled < record_bytes {
146 match reader.read(&mut buffer[filled..]) {
147 Ok(0) => break,
148 Ok(count) => filled += count,
149 Err(error) if error.kind() == ErrorKind::Interrupted => {}
150 Err(error) => return Err(codec_error(format!("decompression failed: {error}"))),
151 }
152 }
153 if filled == 0 {
154 break;
155 }
156 spill
157 .append(&buffer[..filled])
158 .map_err(|error| spill_error(&error))?;
159 length = length
160 .checked_add(
161 u64::try_from(filled)
162 .map_err(|_| resource_error("decompressed length exceeds u64"))?,
163 )
164 .ok_or_else(|| resource_error("decompressed length overflow"))?;
165 if filled < record_bytes {
166 break;
167 }
168 }
169 let artifact = spill.finish().map_err(|error| spill_error(&error))?;
170 drop(buffer);
171 drop(reservation);
172 SpillWindowedFile::open(artifact, record_bytes, length, record_bytes, context)
173}
174
175fn io_error(operation: &'static str, error: &std::io::Error) -> Diagnostic {
176 Diagnostic::new(Code::E7101)
177 .with_message("structural source I/O failed")
178 .with_context("operation", operation)
179 .with_context("reason", error.to_string())
180}
181
182fn codec_error(reason: impl Into<Box<str>>) -> Diagnostic {
183 Diagnostic::new(Code::E1401).with_message(reason)
184}
185
186#[cfg(any(feature = "gzip", feature = "zstd"))]
187fn spill_error(error: &crate::SpillError) -> Diagnostic {
188 resource_error(format!(
189 "compressed structure requires bounded spill: {error}"
190 ))
191}
192
193#[cfg(any(feature = "gzip", feature = "zstd"))]
194fn resource_error(reason: impl Into<Box<str>>) -> Diagnostic {
195 Diagnostic::new(Code::E1902).with_message(reason)
196}
197
198#[cfg(test)]
199#[path = "source_file_tests.rs"]
200mod tests;