1use std::fs::File;
4use std::io::{BufWriter, Write};
5use std::path::Path;
6#[cfg(test)]
7use std::path::PathBuf;
8
9use copc_core::{
10 deserialize_le, serialize_le, Bounds, Error, LasPointRecord, Result, StreamingLayout,
11};
12use memmap2::Mmap;
13use tempfile::{NamedTempFile, TempPath};
14
15pub struct SpillWriter {
17 #[cfg(test)]
18 path: PathBuf,
19 file: Option<BufWriter<NamedTempFile>>,
20 layout: StreamingLayout,
21 record_width: usize,
22 scratch: Vec<u8>,
23 count: u64,
24 bounds: Option<Bounds>,
25}
26
27impl SpillWriter {
28 pub fn create(spill_dir: &Path, layout: StreamingLayout) -> Result<Self> {
29 let file = tempfile::Builder::new()
30 .prefix(".copc-writer-spill.")
31 .suffix(".part")
32 .tempfile_in(spill_dir)
33 .map_err(|e| Error::io("create spill file", e))?;
34 #[cfg(test)]
35 let path = file.path().to_path_buf();
36 let record_width = layout.record_width();
37 Ok(Self {
38 #[cfg(test)]
39 path,
40 file: Some(BufWriter::new(file)),
41 layout,
42 record_width,
43 scratch: vec![0u8; record_width],
44 count: 0,
45 bounds: None,
46 })
47 }
48
49 pub fn push(&mut self, record: &LasPointRecord) -> Result<()> {
50 serialize_le(record, &self.layout, &mut self.scratch)
51 .map_err(|e| Error::InvalidInput(format!("encode spill record: {e}")))?;
52 let writer = self
53 .file
54 .as_mut()
55 .ok_or_else(|| Error::InvalidInput("spill writer already finalized".into()))?;
56 writer
57 .write_all(&self.scratch)
58 .map_err(|e| Error::io("write spill record", e))?;
59 match self.bounds.as_mut() {
60 Some(bounds) => bounds.extend(record.x, record.y, record.z),
61 None => self.bounds = Some(Bounds::point(record.x, record.y, record.z)),
62 }
63 self.count += 1;
64 Ok(())
65 }
66
67 pub fn count(&self) -> u64 {
68 self.count
69 }
70
71 pub fn finalize(mut self) -> Result<SpillReader> {
72 let mut writer = self
73 .file
74 .take()
75 .ok_or_else(|| Error::InvalidInput("spill writer already finalized".into()))?;
76 writer
77 .flush()
78 .map_err(|e| Error::io("flush spill writer", e))?;
79 let file = writer
80 .into_inner()
81 .map_err(|e| Error::io("unwrap spill writer", e.into_error()))?;
82 file.as_file()
83 .sync_all()
84 .map_err(|e| Error::io("sync spill file", e))?;
85 let mmap_file = file
86 .reopen()
87 .map_err(|e| Error::io("open spill for mmap", e))?;
88 let temp_path = file.into_temp_path();
89 let count = usize::try_from(self.count)
90 .map_err(|_| Error::InvalidInput("spill record count exceeds usize range".into()))?;
91 let bounds = self.bounds.unwrap_or_else(|| Bounds::point(0.0, 0.0, 0.0));
92 SpillReader::open(
93 temp_path,
94 mmap_file,
95 self.layout,
96 self.record_width,
97 count,
98 bounds,
99 )
100 }
101}
102
103pub struct SpillReader {
105 #[cfg(test)]
106 path: PathBuf,
107 mmap: Mmap,
108 _file: File,
109 _path: TempPath,
110 layout: StreamingLayout,
111 record_width: usize,
112 count: usize,
113 bounds: Bounds,
114}
115
116impl SpillReader {
117 fn open(
118 temp_path: TempPath,
119 file: File,
120 layout: StreamingLayout,
121 record_width: usize,
122 count: usize,
123 bounds: Bounds,
124 ) -> Result<Self> {
125 #[cfg(test)]
126 let path = temp_path.to_path_buf();
127 let mmap = unsafe { Mmap::map(&file) }.map_err(|e| Error::io("mmap spill file", e))?;
128 let expected = record_width
129 .checked_mul(count)
130 .ok_or_else(|| Error::InvalidInput("spill size overflow".into()))?;
131 if mmap.len() != expected {
132 return Err(Error::InvalidInput(format!(
133 "spill file is {} bytes, expected {}",
134 mmap.len(),
135 expected
136 )));
137 }
138 Ok(Self {
139 #[cfg(test)]
140 path,
141 mmap,
142 _file: file,
143 _path: temp_path,
144 layout,
145 record_width,
146 count,
147 bounds,
148 })
149 }
150
151 pub fn len(&self) -> usize {
152 self.count
153 }
154
155 pub fn is_empty(&self) -> bool {
156 self.count == 0
157 }
158
159 pub fn layout(&self) -> &StreamingLayout {
160 &self.layout
161 }
162
163 pub fn bounds(&self) -> Bounds {
164 self.bounds
165 }
166
167 #[inline]
168 fn record_bytes(&self, index: usize) -> &[u8] {
169 let start = index * self.record_width;
170 &self.mmap[start..start + self.record_width]
171 }
172
173 #[inline]
174 pub fn xyz_at(&self, index: usize) -> (f64, f64, f64) {
175 debug_assert!(index < self.count);
176 let bytes = self.record_bytes(index);
177 let x = f64::from_le_bytes(bytes[0..8].try_into().expect("spill x width"));
178 let y = f64::from_le_bytes(bytes[8..16].try_into().expect("spill y width"));
179 let z = f64::from_le_bytes(bytes[16..24].try_into().expect("spill z width"));
180 (x, y, z)
181 }
182
183 pub fn record_at(&self, index: usize) -> Result<LasPointRecord> {
184 if index >= self.count {
185 return Err(Error::InvalidInput(format!(
186 "spill index {index} out of range (len {})",
187 self.count
188 )));
189 }
190 deserialize_le(self.record_bytes(index), &self.layout)
191 .map_err(|e| Error::InvalidData(format!("decode spill record {index}: {e}")))
192 }
193}
194
195#[cfg(test)]
196mod tests {
197 use super::*;
198
199 fn layout_with_color() -> StreamingLayout {
200 StreamingLayout {
201 point_format: 3,
202 has_gps: true,
203 has_color: true,
204 has_nir: false,
205 has_waveform: false,
206 extra_bytes: 2,
207 extra_bytes_descriptors: Vec::new(),
208 }
209 }
210
211 fn record(seed: u32) -> LasPointRecord {
212 let f = f64::from(seed);
213 LasPointRecord {
214 x: f * 1.5,
215 y: -f * 2.25,
216 z: f * 0.125,
217 intensity: seed as u16,
218 return_number: (seed % 5) as u8,
219 number_of_returns: 5,
220 classification: (seed % 32) as u8,
221 scan_direction_flag: seed % 2 == 0,
222 edge_of_flight_line: seed % 3 == 0,
223 scan_angle: (seed as f32) - 100.25,
224 user_data: (seed % 256) as u8,
225 point_source_id: seed as u16,
226 synthetic: seed % 4 == 0,
227 key_point: seed % 4 == 1,
228 withheld: seed % 4 == 2,
229 overlap: false,
230 scan_channel: 0,
231 gps_time: 1.0e9 + f,
232 red: (seed * 7) as u16,
233 green: (seed * 11) as u16,
234 blue: (seed * 13) as u16,
235 nir: 0,
236 wave_packet_descriptor_index: 0,
237 byte_offset_to_waveform_data: 0,
238 waveform_packet_size: 0,
239 return_point_waveform_location: 0.0,
240 extra_bytes: vec![(seed & 0xff) as u8, ((seed >> 8) & 0xff) as u8],
241 }
242 }
243
244 #[test]
245 fn spill_round_trips_records_and_bounds() {
246 let dir = tempfile::tempdir().unwrap();
247 let layout = layout_with_color();
248 let mut writer = SpillWriter::create(dir.path(), layout).unwrap();
249 let originals: Vec<LasPointRecord> = (0..256).map(record).collect();
250 for rec in &originals {
251 writer.push(rec).unwrap();
252 }
253 assert_eq!(writer.count(), 256);
254 let reader = writer.finalize().unwrap();
255 assert_eq!(reader.len(), 256);
256 for (i, original) in originals.iter().enumerate() {
257 assert_eq!(reader.record_at(i).unwrap(), *original);
258 assert_eq!(reader.xyz_at(i), (original.x, original.y, original.z));
259 }
260 let bounds = reader.bounds();
261 assert_eq!(bounds.min, (0.0, -573.75, 0.0));
262 assert_eq!(bounds.max, (382.5, 0.0, 31.875));
263 }
264
265 #[test]
266 fn unfinalized_spill_writer_removes_file() {
267 let dir = tempfile::tempdir().unwrap();
268 let path = {
269 let mut writer = SpillWriter::create(dir.path(), layout_with_color()).unwrap();
270 writer.push(&record(1)).unwrap();
271 writer.path.clone()
272 };
273 assert!(!path.exists());
274 }
275
276 #[test]
277 fn finalized_spill_reader_removes_file() {
278 let dir = tempfile::tempdir().unwrap();
279 let mut writer = SpillWriter::create(dir.path(), layout_with_color()).unwrap();
280 writer.push(&record(1)).unwrap();
281 let reader = writer.finalize().unwrap();
282 let path = reader.path.clone();
283 assert!(path.exists());
284 drop(reader);
285 assert!(!path.exists());
286 }
287
288 #[cfg(unix)]
289 #[test]
290 fn spill_file_is_private_on_unix() {
291 use std::os::unix::fs::PermissionsExt;
292
293 let dir = tempfile::tempdir().unwrap();
294 let writer = SpillWriter::create(dir.path(), layout_with_color()).unwrap();
295 let mode = std::fs::metadata(&writer.path)
296 .unwrap()
297 .permissions()
298 .mode()
299 & 0o777;
300 assert_eq!(mode, 0o600);
301 }
302}