Skip to main content

oxigdal_algorithms/raster/streaming/
iter.rs

1use super::ChunkedRaster;
2use super::chunk::{Chunk, RasterError, RasterSource};
3
4/// An iterator over the chunks of a [`ChunkedRaster`].
5///
6/// Yields `Result<Chunk, RasterError>` values in raster-scan order (row by row,
7/// left to right).  Each chunk's `data` includes the surrounding halo cells,
8/// zero-padded where the window would extend beyond the source raster boundary.
9///
10/// Edge chunks (those at the right or bottom boundary of the raster) will be
11/// smaller than `chunk_size` in the truncated dimension.
12///
13/// # Halo data layout
14///
15/// `chunk.data` has dimensions `(chunk.width + 2 * halo) × (chunk.height + 2 * halo)`.
16/// Row 0 is the top halo row, row `height + halo` is the last halo row.  Within
17/// each row column 0 is the left halo column.  Pixels that fall outside the source
18/// raster are `0.0`.  In particular, when the chunk is at the top-left corner of
19/// the raster (`x=0, y=0`) the entire top halo row and left halo column contain
20/// zeros.
21pub struct ChunkIterator<'r, R: RasterSource> {
22    raster: &'r ChunkedRaster<R>,
23    current_x: usize,
24    current_y: usize,
25}
26
27impl<'r, R: RasterSource> ChunkIterator<'r, R> {
28    /// Creates a new `ChunkIterator` positioned at the top-left corner.
29    pub fn new(raster: &'r ChunkedRaster<R>) -> Self {
30        Self {
31            raster,
32            current_x: 0,
33            current_y: 0,
34        }
35    }
36}
37
38impl<'r, R: RasterSource> Iterator for ChunkIterator<'r, R> {
39    type Item = Result<Chunk, RasterError>;
40
41    fn next(&mut self) -> Option<Self::Item> {
42        // Exhausted once we've advanced past the bottom of the raster.
43        if self.current_y >= self.raster.total_height {
44            return None;
45        }
46
47        let x = self.current_x;
48        let y = self.current_y;
49        let halo = self.raster.halo;
50
51        // Core region dimensions — smaller for right/bottom edge chunks.
52        let w = self.raster.chunk_size.min(self.raster.total_width - x);
53        let h = self.raster.chunk_size.min(self.raster.total_height - y);
54
55        // Full data buffer dimensions (core + halo on both sides).
56        let full_w = w + 2 * halo;
57        let full_h = h + 2 * halo;
58
59        // Build the halo-padded data buffer manually.
60        //
61        // For each position (row, col) in the full_h × full_w buffer, compute
62        // the corresponding source coordinate.  Positions that map outside the
63        // source raster [0, total_width) × [0, total_height) are filled with 0.0.
64        //
65        // We do this by asking the source for the largest contiguous read window
66        // that stays in-bounds, then scatter-write it into the correct position of
67        // the output buffer (filling the rest with zeros).
68        let mut data = vec![0.0_f32; full_w * full_h];
69
70        // Source window: the slice of the source raster that overlaps with the
71        // halo-extended chunk window.
72        //
73        // For each side: the halo extends `halo` pixels beyond the core origin.
74        // If the core origin is < halo (e.g. x=0, halo=1) then part of the halo
75        // falls outside the raster and must stay zero.
76        //
77        // `src_left` = actual source x of the first in-bounds halo pixel.
78        // `left_pad` = number of zero columns to the left of `src_left` in `data`.
79        let src_left = x.saturating_sub(halo);
80        let src_top = y.saturating_sub(halo);
81        let left_pad = halo.saturating_sub(x); // zero-pad columns on the left
82        let top_pad = halo.saturating_sub(y); // zero-pad rows on the top
83
84        // How many source columns / rows are actually available?
85        let src_right = (x + w + halo).min(self.raster.total_width);
86        let src_bottom = (y + h + halo).min(self.raster.total_height);
87        let src_w = src_right.saturating_sub(src_left);
88        let src_h = src_bottom.saturating_sub(src_top);
89
90        if src_w > 0 && src_h > 0 {
91            // Read the in-bounds source region.
92            let src_data = match self
93                .raster
94                .source
95                .read_window(src_left, src_top, src_w, src_h)
96            {
97                Ok(d) => d,
98                Err(e) => return Some(Err(e)),
99            };
100
101            // Scatter-write into the correct position in `data`.
102            for row in 0..src_h {
103                let dst_row = row + top_pad;
104                for col in 0..src_w {
105                    let dst_col = col + left_pad;
106                    data[dst_row * full_w + dst_col] = src_data[row * src_w + col];
107                }
108            }
109        }
110
111        // Advance the position cursor for the next call.
112        self.current_x += w;
113        if self.current_x >= self.raster.total_width {
114            self.current_x = 0;
115            self.current_y += h;
116        }
117
118        Some(Ok(Chunk {
119            x,
120            y,
121            width: w,
122            height: h,
123            halo,
124            data,
125        }))
126    }
127}