Skip to main content

locus_core/simd/
roi.rs

1//! Region of Interest (ROI) caching for fast sampling.
2
3use crate::image::ImageView;
4use bumpalo::Bump;
5
6/// A cache for a small region of the image to improve L1 cache hits during sampling.
7#[allow(clippy::large_enum_variant)]
8pub enum RoiCache<'a> {
9    /// Small ROI stored on the stack.
10    Stack {
11        /// The cached pixel data.
12        data: [u8; 1024],
13        /// The bounding box of the ROI in the original image.
14        min_x: usize,
15        /// The bounding box of the ROI in the original image.
16        min_y: usize,
17        /// Width of the ROI.
18        width: usize,
19        /// Height of the ROI.
20        height: usize,
21    },
22    /// Larger ROI stored in the arena.
23    Arena {
24        /// The cached pixel data.
25        data: &'a [u8],
26        /// The bounding box of the ROI in the original image.
27        min_x: usize,
28        /// The bounding box of the ROI in the original image.
29        min_y: usize,
30        /// Width of the ROI.
31        width: usize,
32        /// Height of the ROI.
33        height: usize,
34    },
35}
36
37impl<'a> RoiCache<'a> {
38    /// Create a new ROI cache by copying a region from the image.
39    ///
40    /// If the region fits in 1024 bytes, it is stored on the stack.
41    /// Otherwise, it is allocated from the provided arena.
42    #[must_use]
43    pub fn new(
44        img: &ImageView,
45        arena: &'a Bump,
46        min_x: usize,
47        min_y: usize,
48        max_x: usize,
49        max_y: usize,
50    ) -> Self {
51        let width = (max_x - min_x + 1).min(img.width - min_x);
52        let height = (max_y - min_y + 1).min(img.height - min_y);
53        let size = width * height;
54
55        if size <= 1024 {
56            let mut data = [0u8; 1024];
57            for y in 0..height {
58                let src_offset = (min_y + y) * img.stride + min_x;
59                let dst_offset = y * width;
60                data[dst_offset..dst_offset + width]
61                    .copy_from_slice(&img.data[src_offset..src_offset + width]);
62            }
63            RoiCache::Stack {
64                data,
65                min_x,
66                min_y,
67                width,
68                height,
69            }
70        } else {
71            let dst = arena.alloc_slice_fill_default(size);
72            for y in 0..height {
73                let src_offset = (min_y + y) * img.stride + min_x;
74                let dst_offset = y * width;
75                dst[dst_offset..dst_offset + width]
76                    .copy_from_slice(&img.data[src_offset..src_offset + width]);
77            }
78            RoiCache::Arena {
79                data: dst,
80                min_x,
81                min_y,
82                width,
83                height,
84            }
85        }
86    }
87
88    /// Get a pixel from the cache using global coordinates.
89    #[must_use]
90    pub fn get(&self, x: usize, y: usize) -> u8 {
91        match self {
92            RoiCache::Stack {
93                data,
94                min_x,
95                min_y,
96                width,
97                height,
98                ..
99            } => {
100                let lx = x.saturating_sub(*min_x).min(width.saturating_sub(1));
101                let ly = y.saturating_sub(*min_y).min(height.saturating_sub(1));
102                data[ly * width + lx]
103            },
104            RoiCache::Arena {
105                data,
106                min_x,
107                min_y,
108                width,
109                height,
110                ..
111            } => {
112                let lx = x.saturating_sub(*min_x).min(width.saturating_sub(1));
113                let ly = y.saturating_sub(*min_y).min(height.saturating_sub(1));
114                data[ly * width + lx]
115            },
116        }
117    }
118}
119
120#[cfg(test)]
121mod tests {
122    use super::*;
123    use crate::image::ImageView;
124    use bumpalo::Bump;
125
126    #[test]
127    #[allow(clippy::cast_sign_loss)]
128    fn test_roi_cache_stack() {
129        let data: Vec<u8> = (0..100).map(|i| i as u8).collect();
130        let img = ImageView::new(&data, 10, 10, 10).expect("valid view");
131        let arena = Bump::new();
132
133        // 3x3 = 9 bytes, should be Stack
134        let cache = RoiCache::new(&img, &arena, 2, 2, 4, 4);
135        assert!(matches!(cache, RoiCache::Stack { .. }));
136        assert_eq!(cache.get(2, 2), 22);
137        assert_eq!(cache.get(4, 4), 44);
138    }
139
140    #[test]
141    fn test_roi_cache_arena() {
142        let mut data = vec![0u8; 40 * 40];
143        data[20 * 40 + 20] = 255;
144        let img = ImageView::new(&data, 40, 40, 40).expect("valid view");
145        let arena = Bump::new();
146
147        // 33x33 = 1089 bytes, should be Arena
148        let cache = RoiCache::new(&img, &arena, 0, 0, 32, 32);
149        assert!(matches!(cache, RoiCache::Arena { .. }));
150        assert_eq!(cache.get(20, 20), 255);
151    }
152
153    #[test]
154    #[allow(clippy::cast_sign_loss)]
155    fn test_roi_cache_clamping() {
156        let data: Vec<u8> = (0..100).map(|i| i as u8).collect();
157        let img = ImageView::new(&data, 10, 10, 10).expect("valid view");
158        let arena = Bump::new();
159        let cache = RoiCache::new(&img, &arena, 2, 2, 4, 4);
160
161        // Should clamp to edges instead of panic
162        assert_eq!(cache.get(1, 1), 22); // Clamps to (2,2)
163        assert_eq!(cache.get(10, 10), 44); // Clamps to (4,4)
164    }
165}