wolfram_expr/
array_buf.rs

1//! Shared implementation backing both [`NumericArray`][crate::NumericArray] and
2//! [`PackedArray`][crate::PackedArray].
3//!
4//! Both types are dense N-dimensional buffers with an element-type tag — the only
5//! difference is the *set* of valid element types (PackedArray supports a strict
6//! subset).
7
8use crate::wxf::NumericArrayEnum;
9use crate::ByteArray;
10
11/// Sealed marker for Rust primitives valid as an array element. The set is
12/// fixed (the C ABI only knows these widths), so external types can't
13/// implement [`ArrayElement`].
14mod sealed {
15    use crate::complex::{Complex32, Complex64};
16    pub trait Sealed {}
17    impl Sealed for i8 {}
18    impl Sealed for i16 {}
19    impl Sealed for i32 {}
20    impl Sealed for i64 {}
21    impl Sealed for u8 {}
22    impl Sealed for u16 {}
23    impl Sealed for u32 {}
24    impl Sealed for u64 {}
25    impl Sealed for f32 {}
26    impl Sealed for f64 {}
27    impl Sealed for Complex32 {}
28    impl Sealed for Complex64 {}
29}
30
31/// Connects a Rust primitive to its element-type discriminant. Implemented
32/// once per `(type, tag)` pair: e.g. `i32: ArrayElement<NumericArrayEnum>`
33/// (with `TAG = Integer32`) and `i32: ArrayElement<PackedArrayEnum>` (with
34/// `TAG = Integer32`). Sealed — only the primitives listed above can
35/// satisfy the `Sealed` super-bound.
36pub trait ArrayElement<Tag: Copy + PartialEq>: Copy + 'static + sealed::Sealed {
37    /// The element-type tag for `Self` under this array kind.
38    const TAG: Tag;
39}
40
41/// Generic dense N-dimensional buffer parameterized by an element-type tag.
42///
43/// `NumericArray = ArrayBuf<NumericArrayEnum>` and
44/// `PackedArray   = ArrayBuf<PackedArrayEnum>`. Each provides specialized
45/// constructors (`from_slice<T: …Element>`) and a typed slice view; shape,
46/// byte access, and element count are shared via this struct.
47#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
48pub struct ArrayBuf<Tag> {
49    pub(crate) data_type: Tag,
50    pub(crate) dimensions: Vec<usize>,
51    pub(crate) bytes: ByteArray,
52}
53
54impl<Tag: Copy + PartialEq> ArrayBuf<Tag> {
55    /// Construct from raw parts. Caller is responsible for ensuring
56    /// `bytes.len() == prod(dimensions) * element_size`.
57    pub fn new(data_type: Tag, dimensions: Vec<usize>, bytes: ByteArray) -> Self {
58        ArrayBuf {
59            data_type,
60            dimensions,
61            bytes,
62        }
63    }
64
65    /// The concrete element-type tag.
66    pub fn data_type(&self) -> Tag {
67        self.data_type
68    }
69
70    /// Multi-dimensional shape.
71    pub fn dimensions(&self) -> &[usize] {
72        &self.dimensions
73    }
74
75    /// Raw byte buffer.
76    pub(crate) fn as_bytes(&self) -> &[u8] {
77        &self.bytes
78    }
79
80    /// Number of dimensions.
81    pub fn rank(&self) -> usize {
82        self.dimensions.len()
83    }
84
85    /// Total element count = product of dimensions.
86    pub fn element_count(&self) -> usize {
87        self.dimensions.iter().product()
88    }
89
90    /// Construct from a typed slice. Dimensions must satisfy
91    /// `prod(dimensions) == slice.len()`.
92    pub fn from_slice<T: ArrayElement<Tag>>(dimensions: Vec<usize>, slice: &[T]) -> Self {
93        assert_eq!(
94            dimensions.iter().product::<usize>(),
95            slice.len(),
96            "ArrayBuf::from_slice: dims product must equal slice length"
97        );
98        let bytes: &[u8] = unsafe {
99            std::slice::from_raw_parts(
100                slice.as_ptr() as *const u8,
101                std::mem::size_of_val(slice),
102            )
103        };
104        ArrayBuf::new(T::TAG, dimensions, ByteArray::from(bytes))
105    }
106
107    /// Try to view the buffer as a slice of `T`. Returns `None` if `T`'s tag
108    /// doesn't match this array's [`data_type`][Self::data_type].
109    pub fn try_as_slice<T: ArrayElement<Tag>>(&self) -> Option<&[T]> {
110        if self.data_type != T::TAG {
111            return None;
112        }
113        let bytes = self.as_bytes();
114        let elem_size = std::mem::size_of::<T>();
115        debug_assert_eq!(bytes.len() % elem_size, 0);
116        if bytes.is_empty() {
117            return Some(&[]);
118        }
119        // SAFETY: tag matches T, so the bytes were produced from a `[T]`.
120        Some(unsafe {
121            std::slice::from_raw_parts(
122                bytes.as_ptr() as *const T,
123                bytes.len() / elem_size,
124            )
125        })
126    }
127}
128
129/// Common read API implemented by both the owned [`crate::NumericArray`] /
130/// [`crate::PackedArray`] and the runtime-handle `NumericArray<T>` in
131/// `wolfram-library-link`.
132pub trait NumericArrayRead {
133    /// The element-type tag.
134    fn data_type(&self) -> NumericArrayEnum;
135    /// The multi-dimensional shape (row-major).
136    fn dimensions(&self) -> &[usize];
137    #[doc(hidden)]
138    fn as_bytes(&self) -> &[u8];
139
140    /// Number of dimensions.
141    fn rank(&self) -> usize {
142        self.dimensions().len()
143    }
144    /// Total element count = product of dimensions.
145    fn element_count(&self) -> usize {
146        self.dimensions().iter().product()
147    }
148    #[doc(hidden)]
149    fn byte_count(&self) -> usize {
150        self.as_bytes().len()
151    }
152    #[doc(hidden)]
153    fn element_size(&self) -> usize {
154        self.data_type().size_in_bytes()
155    }
156
157    /// View the buffer as `&[T]` if `T`'s element type matches; else `None`.
158    fn try_as_slice<T: ArrayElement<NumericArrayEnum>>(&self) -> Option<&[T]> {
159        if self.data_type() != T::TAG {
160            return None;
161        }
162        let bytes = self.as_bytes();
163        let elem_size = std::mem::size_of::<T>();
164        debug_assert_eq!(bytes.len() % elem_size, 0);
165        // SAFETY: tag matches, alignment guaranteed by construction.
166        Some(unsafe {
167            std::slice::from_raw_parts(
168                bytes.as_ptr() as *const T,
169                bytes.len() / elem_size,
170            )
171        })
172    }
173}
174
175impl<Tag: Into<NumericArrayEnum> + Copy + PartialEq> NumericArrayRead for ArrayBuf<Tag> {
176    fn data_type(&self) -> NumericArrayEnum {
177        self.data_type.into()
178    }
179    fn dimensions(&self) -> &[usize] {
180        &self.dimensions
181    }
182    fn as_bytes(&self) -> &[u8] {
183        &self.bytes
184    }
185}
wolfram_expr/array_buf.rs

wolfram_expr/
array_buf.rs
