Skip to main content

locus_core/
batch.rs

1/// The maximum number of candidates in a single batch.
2pub(crate) const MAX_CANDIDATES: usize = 1024;
3
4/// A 2D point with subpixel precision (f32).
5#[derive(Debug, Clone, Copy, Default)]
6#[repr(C)]
7pub struct Point2f {
8    /// X coordinate.
9    pub x: f32,
10    /// Y coordinate.
11    pub y: f32,
12}
13
14/// A 3x3 homography matrix (f32).
15#[derive(Debug, Clone, Copy, Default)]
16#[repr(C, align(32))]
17pub struct Matrix3x3 {
18    /// The matrix elements in row-major or column-major format (internal use).
19    pub data: [f32; 9],
20    /// Padding to ensure 64-byte size (cache line) and alignment for SIMD.
21    pub padding: [f32; 7],
22}
23
24/// A 6D pose representing translation and rotation (unit quaternion).
25#[derive(Debug, Clone, Copy, Default)]
26#[repr(C, align(32))]
27pub struct Pose6D {
28    /// Translation (x, y, z) and Rotation as a unit quaternion (x, y, z, w).
29    pub data: [f32; 7],
30    /// Padding to 32-byte alignment.
31    pub padding: f32,
32}
33
34/// The lifecycle state of a candidate in the detection pipeline.
35#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
36#[repr(u8)]
37pub enum CandidateState {
38    /// No candidate at this index.
39    #[default]
40    Empty = 0,
41    /// Quad vertices have been extracted.
42    Active = 1,
43    /// Failed to decode the marker payload.
44    FailedDecode = 2,
45    /// Successfully decoded and mathematically verified.
46    Valid = 3,
47}
48
49/// The status of a candidate in the fast-path decoding funnel.
50#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
51#[repr(u8)]
52pub enum FunnelStatus {
53    /// Candidate has not yet been processed by the funnel.
54    #[default]
55    None = 0,
56    /// Candidate passed the O(1) contrast gate.
57    PassedContrast = 1,
58    /// Candidate rejected by the O(1) contrast gate.
59    RejectedContrast = 2,
60    /// Candidate rejected during homography DDA or SIMD sampling.
61    RejectedSampling = 3,
62}
63
64/// A batched state container for fiducial marker detections using a Structure of Arrays (SoA) layout.
65/// This structure is designed for high-performance SIMD processing and zero heap allocations.
66#[repr(C, align(32))]
67pub struct DetectionBatch {
68    /// Flattened array of sub-pixel quad vertices (4 corners per candidate).
69    pub corners: [[Point2f; 4]; MAX_CANDIDATES],
70    /// The 3x3 projection matrices.
71    pub homographies: [Matrix3x3; MAX_CANDIDATES],
72    /// The decoded IDs of the tags.
73    pub ids: [u32; MAX_CANDIDATES],
74    /// The extracted bitstrings.
75    pub payloads: [u64; MAX_CANDIDATES],
76    /// The MSE or Log-Likelihood Ratio confidence scores.
77    pub error_rates: [f32; MAX_CANDIDATES],
78    /// Translation vectors and unit quaternions.
79    pub poses: [Pose6D; MAX_CANDIDATES],
80    /// A dense byte-array tracking the lifecycle of each candidate.
81    pub status_mask: [CandidateState; MAX_CANDIDATES],
82    /// Detailed status from the fast-path funnel.
83    pub funnel_status: [FunnelStatus; MAX_CANDIDATES],
84    /// Four 2x2 corner covariance matrices per quad (16 floats).
85    /// Layout: [c0_xx, c0_xy, c0_yx, c0_yy, c1_xx, ...]
86    pub corner_covariances: [[f32; 16]; MAX_CANDIDATES],
87}
88
89impl DetectionBatch {
90    /// Creates a new DetectionBatch with all fields initialized to zero (Empty state).
91    #[must_use]
92    #[allow(clippy::large_stack_arrays)]
93    pub fn new() -> Self {
94        *Box::new(Self {
95            corners: [[Point2f { x: 0.0, y: 0.0 }; 4]; MAX_CANDIDATES],
96            homographies: [Matrix3x3 {
97                data: [0.0; 9],
98                padding: [0.0; 7],
99            }; MAX_CANDIDATES],
100            ids: [0; MAX_CANDIDATES],
101            payloads: [0; MAX_CANDIDATES],
102            error_rates: [0.0; MAX_CANDIDATES],
103            poses: [Pose6D {
104                data: [0.0; 7],
105                padding: 0.0,
106            }; MAX_CANDIDATES],
107            status_mask: [CandidateState::Empty; MAX_CANDIDATES],
108            funnel_status: [FunnelStatus::None; MAX_CANDIDATES],
109            corner_covariances: [[0.0; 16]; MAX_CANDIDATES],
110        })
111    }
112    /// Returns the maximum capacity of the batch.
113    #[must_use]
114    pub fn capacity(&self) -> usize {
115        MAX_CANDIDATES
116    }
117
118    /// Partitions the batch so that all `Valid` candidates are at the front `[0..V]`.
119    /// Returns the number of valid candidates `V`.
120    pub fn partition(&mut self, n: usize) -> usize {
121        let mut v = 0;
122        let n_clamped = n.min(MAX_CANDIDATES);
123
124        // Two-pointer partition:
125        // Move Valid candidates to [0..v]
126        // Move everything else to [v..n_clamped]
127        for i in 0..n_clamped {
128            if self.status_mask[i] == CandidateState::Valid {
129                if i != v {
130                    // Swap index i with index v across all parallel arrays.
131                    self.corners.swap(i, v);
132                    self.homographies.swap(i, v);
133                    self.ids.swap(i, v);
134                    self.payloads.swap(i, v);
135                    self.error_rates.swap(i, v);
136                    self.poses.swap(i, v);
137                    self.status_mask.swap(i, v);
138                    self.funnel_status.swap(i, v);
139                    self.corner_covariances.swap(i, v);
140                }
141                v += 1;
142            }
143        }
144        v
145    }
146
147    /// Reassemble the batched SoA data into a list of discrete `Detection` objects.
148    #[must_use]
149    #[allow(clippy::cast_sign_loss)]
150    pub fn reassemble(&self, v: usize) -> Vec<crate::Detection> {
151        let mut detections = Vec::with_capacity(v);
152        for i in 0..v {
153            let corners = [
154                [
155                    f64::from(self.corners[i][0].x),
156                    f64::from(self.corners[i][0].y),
157                ],
158                [
159                    f64::from(self.corners[i][1].x),
160                    f64::from(self.corners[i][1].y),
161                ],
162                [
163                    f64::from(self.corners[i][2].x),
164                    f64::from(self.corners[i][2].y),
165                ],
166                [
167                    f64::from(self.corners[i][3].x),
168                    f64::from(self.corners[i][3].y),
169                ],
170            ];
171
172            let center = [
173                (corners[0][0] + corners[1][0] + corners[2][0] + corners[3][0]) / 4.0,
174                (corners[0][1] + corners[1][1] + corners[2][1] + corners[3][1]) / 4.0,
175            ];
176
177            // Reconstruct Pose if available (Z translation > 0)
178            let pose = if self.poses[i].data[2] > 0.0 {
179                let d = self.poses[i].data;
180                // layout: [tx, ty, tz, qx, qy, qz, qw]
181                let t = nalgebra::Vector3::new(f64::from(d[0]), f64::from(d[1]), f64::from(d[2]));
182                let q = nalgebra::UnitQuaternion::from_quaternion(nalgebra::Quaternion::new(
183                    f64::from(d[6]), // w
184                    f64::from(d[3]), // x
185                    f64::from(d[4]), // y
186                    f64::from(d[5]), // z
187                ));
188                Some(crate::pose::Pose {
189                    rotation: q.to_rotation_matrix().into_inner(),
190                    translation: t,
191                })
192            } else {
193                None
194            };
195
196            detections.push(crate::Detection {
197                id: self.ids[i],
198                center,
199                corners,
200                hamming: self.error_rates[i] as u32,
201                rotation: 0,
202                decision_margin: 0.0,
203                bits: self.payloads[i],
204                pose,
205                pose_covariance: None,
206            });
207        }
208        detections
209    }
210}
211
212/// Helper function to partition a batch, moving all valid candidates to the front.
213pub fn partition_batch_soa(batch: &mut DetectionBatch, n: usize) -> usize {
214    batch.partition(n)
215}
216
217/// Metadata for zero-copy telemetry extraction.
218#[derive(Debug, Clone, Copy)]
219pub struct TelemetryPayload {
220    /// Pointer to the binarized image buffer.
221    pub binarized_ptr: *const u8,
222    /// Pointer to the threshold map buffer.
223    pub threshold_map_ptr: *const u8,
224    /// Pointer to subpixel jitter data [4 corners * 2 (dx, dy)] per valid candidate.
225    /// Allocated in the detector's arena.
226    pub subpixel_jitter_ptr: *const f32,
227    /// Number of valid candidates jitter data is available for.
228    pub num_jitter: usize,
229    /// Pointer to reprojection RMSE values per valid candidate.
230    /// Allocated in the detector's arena.
231    pub reprojection_errors_ptr: *const f32,
232    /// Number of valid candidates reprojection data is available for.
233    pub num_reprojection: usize,
234    /// Number of quads that fell back to coarse corners during GWLF.
235    pub gwlf_fallback_count: usize,
236    /// Average Euclidean distance (delta) of GWLF refinement (pixels).
237    pub gwlf_avg_delta: f32,
238    /// Width of the buffers.
239    pub width: usize,
240    /// Height of the buffers.
241    pub height: usize,
242    /// Stride of the buffers.
243    pub stride: usize,
244}
245
246// SAFETY: TelemetryPayload contains raw pointers to the detector's arena.
247// These pointers are stable for the duration of the frame processing and
248// are only accessed from the Python thread while the arena is still alive.
249#[allow(unsafe_code)]
250unsafe impl Send for TelemetryPayload {}
251#[allow(unsafe_code)]
252unsafe impl Sync for TelemetryPayload {}
253
254/// A lightweight, borrowed view of the detection results.
255///
256/// This struct holds slices to the active elements in a [`DetectionBatch`].
257/// It avoids heap allocations and provides efficient access to detection data.
258#[derive(Debug, Clone, Copy)]
259pub struct DetectionBatchView<'a> {
260    /// Decoded IDs of the markers.
261    pub ids: &'a [u32],
262    /// Refined corners in image coordinates.
263    pub corners: &'a [[Point2f; 4]],
264    /// Computed homography matrices.
265    pub homographies: &'a [Matrix3x3],
266    /// Decoded bitstrings (paylods).
267    pub payloads: &'a [u64],
268    /// Confidence scores or error rates.
269    pub error_rates: &'a [f32],
270    /// 3D poses (rotation + translation).
271    pub poses: &'a [Pose6D],
272    /// Optional telemetry data for intermediate images.
273    pub telemetry: Option<TelemetryPayload>,
274    /// Corners of quads that were extracted but rejected during decoding or verification.
275    pub rejected_corners: &'a [[Point2f; 4]],
276    /// Error rates (e.g. Hamming distance) for rejected quads.
277    pub rejected_error_rates: &'a [f32],
278    /// Detailed status from the fast-path funnel for rejected quads.
279    pub rejected_funnel_status: &'a [FunnelStatus],
280}
281
282impl DetectionBatchView<'_> {
283    /// Returns the number of detections in the view.
284    #[must_use]
285    pub fn len(&self) -> usize {
286        self.ids.len()
287    }
288
289    /// Returns true if the view contains no detections.
290    #[must_use]
291    pub fn is_empty(&self) -> bool {
292        self.ids.is_empty()
293    }
294}
295
296impl DetectionBatch {
297    /// Returns a borrowed view of the first `v` candidates in the batch.
298    #[must_use]
299    pub fn view(&self, v: usize) -> DetectionBatchView<'_> {
300        let n = v.min(MAX_CANDIDATES);
301        DetectionBatchView {
302            ids: &self.ids[..n],
303            corners: &self.corners[..n],
304            homographies: &self.homographies[..n],
305            payloads: &self.payloads[..n],
306            error_rates: &self.error_rates[..n],
307            poses: &self.poses[..n],
308            telemetry: None,
309            rejected_corners: &[],
310            rejected_error_rates: &[],
311            rejected_funnel_status: &[],
312        }
313    }
314
315    /// Returns a borrowed view with telemetry data.
316    #[must_use]
317    pub fn view_with_telemetry(
318        &self,
319        v: usize,
320        n: usize,
321        telemetry: Option<TelemetryPayload>,
322    ) -> DetectionBatchView<'_> {
323        let v_clamped = v.min(MAX_CANDIDATES);
324        let n_clamped = n.min(MAX_CANDIDATES);
325        DetectionBatchView {
326            ids: &self.ids[..v_clamped],
327            corners: &self.corners[..v_clamped],
328            homographies: &self.homographies[..v_clamped],
329            payloads: &self.payloads[..v_clamped],
330            error_rates: &self.error_rates[..v_clamped],
331            poses: &self.poses[..v_clamped],
332            telemetry,
333            rejected_corners: &self.corners[v_clamped..n_clamped],
334            rejected_error_rates: &self.error_rates[v_clamped..n_clamped],
335            rejected_funnel_status: &self.funnel_status[v_clamped..n_clamped],
336        }
337    }
338}
339
340impl Default for DetectionBatch {
341    fn default() -> Self {
342        Self::new()
343    }
344}