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}