diff --git a/src/compression/mod.rs b/src/compression/mod.rs index 9cecfbd0..eb4afc6e 100644 --- a/src/compression/mod.rs +++ b/src/compression/mod.rs @@ -74,9 +74,8 @@ impl<'a, R: Read + Seek + Send> PointReader for CompressedPointReader<'a, R> { .decompress_one(&mut self.decompressor_output.get_mut()); if let Err(e) = res { Some(Err(e.into())) - } else if let Err(e) = self.decompressor_output.seek(SeekFrom::Start(0)) { - Some(Err(e.into())) } else { + self.decompressor_output.set_position(0); Some(read_point_from(&mut self.decompressor_output, &self.header)) } } else { @@ -84,6 +83,27 @@ impl<'a, R: Read + Seek + Send> PointReader for CompressedPointReader<'a, R> { } } + fn read_into_vec(&mut self, points: &mut Vec, n: u64) -> Result { + let points_left = self.header().number_of_points() - self.last_point_idx; + let num_points_to_read = points_left.min(n); + + self.decompressor_output.get_mut().resize( + num_points_to_read as usize * self.header.point_format().len() as usize, + 0u8, + ); + self.decompressor + .decompress_many(self.decompressor_output.get_mut())?; + self.decompressor_output.set_position(0); + points.reserve(num_points_to_read as usize); + + for _ in 0..num_points_to_read { + let point = read_point_from(&mut self.decompressor_output, &self.header)?; + self.last_point_idx += 1; + points.push(point); + } + Ok(num_points_to_read) + } + fn seek(&mut self, position: u64) -> Result<()> { self.last_point_idx = position; self.decompressor.seek(position)?; diff --git a/src/reader.rs b/src/reader.rs index 691d1cfe..fd6b5efd 100644 --- a/src/reader.rs +++ b/src/reader.rs @@ -74,9 +74,15 @@ pub(crate) fn read_point_from( point } -/// Trait to specify behaviour a a PointReader +/// Trait to specify behaviour of a PointReader pub(crate) trait PointReader: Debug + Send { fn read_next(&mut self) -> Option>; + fn read_into_vec(&mut self, points: &mut Vec, n: u64) -> Result; + fn read_next_points(&mut self, n: u64) -> Result> { + let mut points = Vec::with_capacity(n as usize); + self.read_into_vec(&mut points, n)?; + Ok(points) + } fn seek(&mut self, position: u64) -> Result<()>; fn header(&self) -> &Header; } @@ -116,6 +122,18 @@ impl PointReader for UncompressedPointRe } } + fn read_into_vec(&mut self, points: &mut Vec, n: u64) -> Result { + let points_left = self.header().number_of_points() - self.last_point_idx; + let num_points_to_read = points_left.min(n); + points.reserve(num_points_to_read as usize); + for _ in 0..num_points_to_read { + let point = read_point_from(&mut self.source, &self.header)?; + self.last_point_idx += 1; + points.push(point); + } + Ok(num_points_to_read) + } + fn seek(&mut self, position: u64) -> Result<()> { self.last_point_idx = position; self.source.seek(SeekFrom::Start( @@ -153,6 +171,15 @@ pub trait Read { /// ``` fn read(&mut self) -> Option>; + /// Reads n points. + fn read_n(&mut self, n: u64) -> Result>; + + /// Reads n points into the vec + fn read_n_into(&mut self, n: u64, points: &mut Vec) -> Result; + + /// Reads all points left into the vec + fn read_all_points(&mut self, points: &mut Vec) -> Result; + /// Seeks to the given point number, zero-indexed. /// /// Note that seeking on compressed (LAZ) data can be expensive as the reader @@ -301,6 +328,19 @@ impl<'a> Read for Reader<'a> { self.point_reader.read_next() } + fn read_n(&mut self, n: u64) -> Result> { + self.point_reader.read_next_points(n) + } + + fn read_n_into(&mut self, n: u64, points: &mut Vec) -> Result { + self.point_reader.read_into_vec(points, n) + } + + fn read_all_points(&mut self, points: &mut Vec) -> Result { + let point_count = self.point_reader.header().number_of_points(); + self.point_reader.read_into_vec(points, point_count) + } + /// Seeks to the given point number, zero-indexed. fn seek(&mut self, position: u64) -> Result<()> { self.point_reader.seek(position) diff --git a/tests/autzen.rs b/tests/autzen.rs index 96a0ec63..a4a0c1f4 100644 --- a/tests/autzen.rs +++ b/tests/autzen.rs @@ -89,3 +89,100 @@ fn test_seek_0_works_on_las() { fn test_seek_0_works_on_laz() { test_seek_0_works_on("tests/data/autzen.laz"); } + +fn test_read_n_on(path: &str) { + use las::{Point, Read, Reader}; + + let ground_truth_points = { + let mut ground_truth_reader = Reader::from_path(path).unwrap(); + ground_truth_reader + .points() + .collect::>>() + .unwrap() + }; + + let mut reader = Reader::from_path(path).unwrap(); + let n = 7; + let mut all_points = Vec::with_capacity(reader.header().number_of_points() as usize); + loop { + let mut points = reader.read_n(n).unwrap(); + if points.len() == 0 { + break; + } + all_points.append(&mut points); + } + + assert_eq!(all_points, ground_truth_points); +} + +fn test_read_n_into_on(path: &str) { + use las::{Point, Read, Reader}; + + let ground_truth_points = { + let mut ground_truth_reader = Reader::from_path(path).unwrap(); + ground_truth_reader + .points() + .collect::>>() + .unwrap() + }; + + let mut reader = Reader::from_path(path).unwrap(); + let n = 7; + let mut all_points = Vec::with_capacity(reader.header().number_of_points() as usize); + let mut points_buffer = Vec::with_capacity(n as usize); + while reader.read_n_into(n, &mut points_buffer).unwrap() != 0 { + all_points.append(&mut points_buffer); + } + + assert_eq!(all_points, ground_truth_points); +} + +#[test] +fn test_las_read_n() { + test_read_n_on("tests/data/autzen.las"); +} + +#[cfg(feature = "laz")] +#[test] +fn test_laz_read_n() { + test_read_n_on("tests/data/autzen.laz"); +} + +#[test] +fn test_las_read_n_into() { + test_read_n_into_on("tests/data/autzen.las"); +} + +#[cfg(feature = "laz")] +#[test] +fn test_laz_read_n_into() { + test_read_n_into_on("tests/data/autzen.laz"); +} + +fn test_read_all_points_on(path: &str) { + use las::{Point, Read, Reader}; + + let ground_truth_points = { + let mut ground_truth_reader = Reader::from_path(path).unwrap(); + ground_truth_reader + .points() + .collect::>>() + .unwrap() + }; + + let mut reader = Reader::from_path(path).unwrap(); + let mut all_points = vec![]; + reader.read_all_points(&mut all_points).unwrap(); + assert_eq!(all_points, ground_truth_points); +} + +#[test] +fn test_las_read_all_points() { + test_read_all_points_on("tests/data/autzen.las"); +} + +#[cfg(feature = "laz")] +#[test] +fn test_laz_read_all_points() { + test_read_all_points_on("tests/data/autzen.laz"); +}