Skip to content

Commit fecf60a

Browse files
committed
feat(hpc): W3 SoA polish — iter_rows(), #[inline], Clone/Debug
Add SoaVec::iter_rows() (+ SoaRowIter, ExactSizeIterator), #[inline] on aos_to_soa/soa_to_aos, and Clone/Debug on SoaVec. No SoaBatch alias (not in the design doc). https://claude.ai/code/session_017GFLBnDy23AWBqvkbHHC41
1 parent cc93b19 commit fecf60a

1 file changed

Lines changed: 123 additions & 0 deletions

File tree

src/hpc/soa.rs

Lines changed: 123 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -84,6 +84,7 @@ use core::array;
8484
/// assert_eq!(soa.field(1), &[2.0, 5.0]);
8585
/// assert_eq!(soa.field(2), &[3.0, 6.0]);
8686
/// ```
87+
#[derive(Clone, Debug)]
8788
pub struct SoaVec<T, const N: usize> {
8889
fields: [Vec<T>; N],
8990
}
@@ -246,6 +247,33 @@ impl<T, const N: usize> SoaVec<T, N> {
246247
}
247248
}
248249

250+
impl<T: Copy, const N: usize> SoaVec<T, N> {
251+
/// Iterate over individual rows, yielding each as `[T; N]` (one value
252+
/// per field, reconstructed from the parallel field arrays).
253+
///
254+
/// Requires `T: Copy` so each field element can be copied out without
255+
/// cloning; the borrow on `self` lasts only as long as the iterator.
256+
/// For non-`Copy` types, use [`chunks`](Self::chunks) with `chunk_len = 1`
257+
/// and index into the single-element slices.
258+
///
259+
/// # Example
260+
///
261+
/// ```
262+
/// use ndarray::hpc::soa::SoaVec;
263+
/// let mut soa: SoaVec<f32, 3> = SoaVec::new();
264+
/// soa.push([1.0, 2.0, 3.0]);
265+
/// soa.push([4.0, 5.0, 6.0]);
266+
///
267+
/// let rows: Vec<[f32; 3]> = soa.iter_rows().collect();
268+
/// assert_eq!(rows[0], [1.0, 2.0, 3.0]);
269+
/// assert_eq!(rows[1], [4.0, 5.0, 6.0]);
270+
/// ```
271+
#[inline]
272+
pub fn iter_rows(&self) -> SoaRowIter<'_, T, N> {
273+
SoaRowIter { soa: self, cursor: 0 }
274+
}
275+
}
276+
249277
impl<T, const N: usize> Default for SoaVec<T, N> {
250278
fn default() -> Self {
251279
Self::new()
@@ -274,6 +302,37 @@ impl<'a, T, const N: usize> Iterator for SoaChunks<'a, T, N> {
274302
}
275303
}
276304

305+
/// Iterator yielded by [`SoaVec::iter_rows`].
306+
///
307+
/// Each call to [`next`](Iterator::next) copies one row (`[T; N]`) out of
308+
/// the parallel field arrays. Requires `T: Copy`.
309+
pub struct SoaRowIter<'a, T, const N: usize> {
310+
soa: &'a SoaVec<T, N>,
311+
cursor: usize,
312+
}
313+
314+
impl<'a, T: Copy, const N: usize> Iterator for SoaRowIter<'a, T, N> {
315+
type Item = [T; N];
316+
317+
#[inline]
318+
fn next(&mut self) -> Option<Self::Item> {
319+
if self.cursor >= self.soa.len() {
320+
return None;
321+
}
322+
let row: [T; N] = array::from_fn(|i| self.soa.fields[i][self.cursor]);
323+
self.cursor += 1;
324+
Some(row)
325+
}
326+
327+
#[inline]
328+
fn size_hint(&self) -> (usize, Option<usize>) {
329+
let remaining = self.soa.len().saturating_sub(self.cursor);
330+
(remaining, Some(remaining))
331+
}
332+
}
333+
334+
impl<T: Copy, const N: usize> ExactSizeIterator for SoaRowIter<'_, T, N> {}
335+
277336
/// Generate a named-field SoA struct from a struct-like declaration.
278337
///
279338
/// Each declared field `name: T` becomes `name: Vec<T>` on the generated
@@ -603,6 +662,7 @@ macro_rules! soa_struct {
603662
/// assert_eq!(soa.field(0), &[7u8, 3]);
604663
/// assert_eq!(soa.field(1), &[255u8, 128]);
605664
/// ```
665+
#[inline]
606666
pub fn aos_to_soa<T, U, const N: usize, F>(aos: &[T], extract: F) -> SoaVec<U, N>
607667
where
608668
F: Fn(&T) -> [U; N],
@@ -651,6 +711,7 @@ where
651711
/// assert_eq!(back[0], Pair { lo: 0x1234, hi: 0xABCD });
652712
/// assert_eq!(back[1], Pair { lo: 0x5678, hi: 0xEF01 });
653713
/// ```
714+
#[inline]
654715
pub fn soa_to_aos<T, U, const N: usize, F>(soa: &SoaVec<U, N>, build: F) -> Vec<T>
655716
where
656717
F: Fn([U; N]) -> T,
@@ -856,6 +917,68 @@ mod tests {
856917
assert!(chunks.is_empty());
857918
}
858919

920+
// -------------------------------------------------------------------
921+
// SoaVec::iter_rows
922+
// -------------------------------------------------------------------
923+
924+
#[test]
925+
fn soa_vec_iter_rows_basic() {
926+
let mut soa: SoaVec<f32, 3> = SoaVec::new();
927+
soa.push([1.0, 2.0, 3.0]);
928+
soa.push([4.0, 5.0, 6.0]);
929+
soa.push([7.0, 8.0, 9.0]);
930+
let rows: Vec<[f32; 3]> = soa.iter_rows().collect();
931+
assert_eq!(rows.len(), 3);
932+
assert_eq!(rows[0], [1.0, 2.0, 3.0]);
933+
assert_eq!(rows[1], [4.0, 5.0, 6.0]);
934+
assert_eq!(rows[2], [7.0, 8.0, 9.0]);
935+
}
936+
937+
#[test]
938+
fn soa_vec_iter_rows_empty_yields_nothing() {
939+
let soa: SoaVec<u32, 2> = SoaVec::new();
940+
let rows: Vec<[u32; 2]> = soa.iter_rows().collect();
941+
assert!(rows.is_empty());
942+
}
943+
944+
#[test]
945+
fn soa_vec_iter_rows_single_field() {
946+
let mut soa: SoaVec<i32, 1> = SoaVec::new();
947+
soa.push([10]);
948+
soa.push([20]);
949+
let rows: Vec<[i32; 1]> = soa.iter_rows().collect();
950+
assert_eq!(rows, [[10], [20]]);
951+
}
952+
953+
#[test]
954+
fn soa_vec_iter_rows_size_hint() {
955+
let mut soa: SoaVec<u8, 2> = SoaVec::new();
956+
soa.push([1, 2]);
957+
soa.push([3, 4]);
958+
soa.push([5, 6]);
959+
let mut it = soa.iter_rows();
960+
assert_eq!(it.size_hint(), (3, Some(3)));
961+
let _ = it.next();
962+
assert_eq!(it.size_hint(), (2, Some(2)));
963+
let _ = it.next();
964+
let _ = it.next();
965+
assert_eq!(it.size_hint(), (0, Some(0)));
966+
assert!(it.next().is_none());
967+
}
968+
969+
#[test]
970+
fn soa_vec_iter_rows_matches_push_order() {
971+
// Cross-check iter_rows against field() to ensure column order is preserved.
972+
let mut soa: SoaVec<u32, 4> = SoaVec::new();
973+
for i in 0..5u32 {
974+
soa.push([i, i * 10, i * 100, i * 1000]);
975+
}
976+
for (row_idx, row) in soa.iter_rows().enumerate() {
977+
let i = row_idx as u32;
978+
assert_eq!(row, [i, i * 10, i * 100, i * 1000]);
979+
}
980+
}
981+
859982
// -------------------------------------------------------------------
860983
// soa_struct! macro
861984
// -------------------------------------------------------------------

0 commit comments

Comments
 (0)