ferritin_core/columns/
buffer.rs

1//! Typed numeric column views and their little-endian byte representation.
2
3use std::borrow::Cow;
4
5/// Scalar element type of an exported column.
6///
7/// Discriminants are part of the boundary contract and will not be renumbered.
8#[repr(u8)]
9#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
10pub enum ScalarType {
11    U8 = 1,
12    I8 = 2,
13    U32 = 3,
14    I32 = 4,
15    F32 = 5,
16}
17
18impl ScalarType {
19    /// Size of one scalar in bytes.
20    pub const fn size_bytes(self) -> usize {
21        match self {
22            ScalarType::U8 | ScalarType::I8 => 1,
23            ScalarType::U32 | ScalarType::I32 | ScalarType::F32 => 4,
24        }
25    }
26
27    /// Stable numeric code for this type.
28    pub const fn code(self) -> u8 {
29        self as u8
30    }
31}
32
33/// Plain-old-data scalars that may be reinterpreted as bytes.
34///
35/// Sealed: implemented only for padding-free primitive types.
36pub(crate) trait Scalar: Copy + 'static {
37    fn extend_le(self, out: &mut Vec<u8>);
38}
39
40macro_rules! impl_scalar {
41    ($($t:ty),*) => {$(
42        impl Scalar for $t {
43            fn extend_le(self, out: &mut Vec<u8>) {
44                out.extend_from_slice(&self.to_le_bytes());
45            }
46        }
47    )*};
48}
49impl_scalar!(u8, i8, u32, i32, f32);
50
51/// Little-endian bytes of `values`: borrowed on little-endian targets (wasm32,
52/// x86_64, aarch64), copied on big-endian ones.
53pub(crate) fn le_bytes<T: Scalar>(values: &[T]) -> Cow<'_, [u8]> {
54    if cfg!(target_endian = "little") {
55        // SAFETY: `Scalar` is only implemented for primitive integer/float types,
56        // which have no padding and no invalid byte patterns when read as `u8`.
57        // `u8` has alignment 1 and the length is the exact byte size of the slice.
58        Cow::Borrowed(unsafe {
59            std::slice::from_raw_parts(values.as_ptr().cast::<u8>(), std::mem::size_of_val(values))
60        })
61    } else {
62        let mut out = Vec::with_capacity(std::mem::size_of_val(values));
63        for &value in values {
64            value.extend_le(&mut out);
65        }
66        Cow::Owned(out)
67    }
68}
69
70/// A borrowed, contiguous, typed numeric column.
71///
72/// Every variant borrows storage owned elsewhere — a [`Model`](crate::Model)'s
73/// conformation, its [`AtomicHierarchy`](crate::AtomicHierarchy), or a
74/// [`TopologyColumns`](super::TopologyColumns) snapshot — so creating one never
75/// copies.
76#[derive(Clone, Copy, Debug, PartialEq)]
77pub enum ColumnData<'a> {
78    U8(&'a [u8]),
79    I8(&'a [i8]),
80    U32(&'a [u32]),
81    I32(&'a [i32]),
82    F32(&'a [f32]),
83}
84
85impl<'a> ColumnData<'a> {
86    /// Element type of the column.
87    pub fn scalar_type(&self) -> ScalarType {
88        match self {
89            ColumnData::U8(_) => ScalarType::U8,
90            ColumnData::I8(_) => ScalarType::I8,
91            ColumnData::U32(_) => ScalarType::U32,
92            ColumnData::I32(_) => ScalarType::I32,
93            ColumnData::F32(_) => ScalarType::F32,
94        }
95    }
96
97    /// Number of scalars (rows).
98    pub fn len(&self) -> usize {
99        match self {
100            ColumnData::U8(v) => v.len(),
101            ColumnData::I8(v) => v.len(),
102            ColumnData::U32(v) => v.len(),
103            ColumnData::I32(v) => v.len(),
104            ColumnData::F32(v) => v.len(),
105        }
106    }
107
108    /// Returns `true` if the column has no rows.
109    pub fn is_empty(&self) -> bool {
110        self.len() == 0
111    }
112
113    /// Size of the column in bytes.
114    pub fn byte_len(&self) -> usize {
115        self.len() * self.scalar_type().size_bytes()
116    }
117
118    /// Pointer to the first scalar, for FFI hand-off together with [`Self::len`].
119    ///
120    /// The pointer is valid exactly as long as the borrow `'a`.
121    pub fn as_ptr(&self) -> *const u8 {
122        match self {
123            ColumnData::U8(v) => v.as_ptr(),
124            ColumnData::I8(v) => v.as_ptr().cast(),
125            ColumnData::U32(v) => v.as_ptr().cast(),
126            ColumnData::I32(v) => v.as_ptr().cast(),
127            ColumnData::F32(v) => v.as_ptr().cast(),
128        }
129    }
130
131    /// Little-endian bytes of the column; zero-copy on little-endian targets.
132    pub fn le_bytes(&self) -> Cow<'a, [u8]> {
133        match *self {
134            ColumnData::U8(v) => Cow::Borrowed(v),
135            ColumnData::I8(v) => le_bytes(v),
136            ColumnData::U32(v) => le_bytes(v),
137            ColumnData::I32(v) => le_bytes(v),
138            ColumnData::F32(v) => le_bytes(v),
139        }
140    }
141}
142
143/// A named column, as listed by [`TopologyColumns::columns`](super::TopologyColumns::columns)
144/// and [`FrameColumns::columns`](super::FrameColumns::columns).
145#[derive(Clone, Copy, Debug, PartialEq)]
146pub struct Column<'a> {
147    /// Stable column name, e.g. `"x"`, `"atom.element"`, `"bond.atom_a"`.
148    pub name: &'static str,
149    /// The column's values.
150    pub data: ColumnData<'a>,
151}
152
153#[cfg(test)]
154mod tests {
155    use super::*;
156
157    #[test]
158    fn test_le_bytes_matches_to_le_bytes() {
159        let values = [1.5f32, -2.0, 3.25];
160        let expected: Vec<u8> = values.iter().flat_map(|v| v.to_le_bytes()).collect();
161        assert_eq!(&*ColumnData::F32(&values).le_bytes(), expected.as_slice());
162
163        let ints = [-1i32, 7, i32::MIN];
164        let expected: Vec<u8> = ints.iter().flat_map(|v| v.to_le_bytes()).collect();
165        assert_eq!(&*ColumnData::I32(&ints).le_bytes(), expected.as_slice());
166    }
167
168    #[test]
169    fn test_le_bytes_is_zero_copy_on_little_endian() {
170        let values = [1u32, 2, 3];
171        let bytes = ColumnData::U32(&values).le_bytes();
172        if cfg!(target_endian = "little") {
173            assert!(matches!(bytes, Cow::Borrowed(_)));
174            assert_eq!(bytes.as_ptr(), values.as_ptr().cast());
175        }
176        assert_eq!(bytes.len(), 12);
177    }
178
179    #[test]
180    fn test_column_data_metadata() {
181        let values = [0i8, -1, 2, 3];
182        let column = ColumnData::I8(&values);
183        assert_eq!(column.scalar_type(), ScalarType::I8);
184        assert_eq!(column.len(), 4);
185        assert_eq!(column.byte_len(), 4);
186        assert_eq!(column.as_ptr(), values.as_ptr().cast());
187        assert!(ColumnData::F32(&[]).is_empty());
188    }
189}
ferritin_core/columns/buffer.rs

ferritin_core/columns/
buffer.rs
