1use crate::attributes::NDAttributeList;
2use crate::codec::Codec;
3use crate::error::{ADError, ADResult};
4use crate::timestamp::EpicsTimestamp;
5
6#[derive(Debug, Clone, Copy, PartialEq, Eq)]
8#[repr(u8)]
9pub enum NDDataType {
10 Int8 = 0,
11 UInt8 = 1,
12 Int16 = 2,
13 UInt16 = 3,
14 Int32 = 4,
15 UInt32 = 5,
16 Int64 = 6,
17 UInt64 = 7,
18 Float32 = 8,
19 Float64 = 9,
20}
21
22impl NDDataType {
23 pub fn element_size(&self) -> usize {
24 match self {
25 Self::Int8 | Self::UInt8 => 1,
26 Self::Int16 | Self::UInt16 => 2,
27 Self::Int32 | Self::UInt32 | Self::Float32 => 4,
28 Self::Int64 | Self::UInt64 | Self::Float64 => 8,
29 }
30 }
31
32 pub fn from_ordinal(v: u8) -> Option<Self> {
33 match v {
34 0 => Some(Self::Int8),
35 1 => Some(Self::UInt8),
36 2 => Some(Self::Int16),
37 3 => Some(Self::UInt16),
38 4 => Some(Self::Int32),
39 5 => Some(Self::UInt32),
40 6 => Some(Self::Int64),
41 7 => Some(Self::UInt64),
42 8 => Some(Self::Float32),
43 9 => Some(Self::Float64),
44 _ => None,
45 }
46 }
47}
48
49#[derive(Debug, Clone)]
51pub enum NDDataBuffer {
52 I8(Vec<i8>),
53 U8(Vec<u8>),
54 I16(Vec<i16>),
55 U16(Vec<u16>),
56 I32(Vec<i32>),
57 U32(Vec<u32>),
58 I64(Vec<i64>),
59 U64(Vec<u64>),
60 F32(Vec<f32>),
61 F64(Vec<f64>),
62}
63
64impl NDDataBuffer {
65 pub fn zeros(data_type: NDDataType, count: usize) -> Self {
66 match data_type {
67 NDDataType::Int8 => Self::I8(vec![0; count]),
68 NDDataType::UInt8 => Self::U8(vec![0; count]),
69 NDDataType::Int16 => Self::I16(vec![0; count]),
70 NDDataType::UInt16 => Self::U16(vec![0; count]),
71 NDDataType::Int32 => Self::I32(vec![0; count]),
72 NDDataType::UInt32 => Self::U32(vec![0; count]),
73 NDDataType::Int64 => Self::I64(vec![0; count]),
74 NDDataType::UInt64 => Self::U64(vec![0; count]),
75 NDDataType::Float32 => Self::F32(vec![0.0; count]),
76 NDDataType::Float64 => Self::F64(vec![0.0; count]),
77 }
78 }
79
80 pub fn data_type(&self) -> NDDataType {
81 match self {
82 Self::I8(_) => NDDataType::Int8,
83 Self::U8(_) => NDDataType::UInt8,
84 Self::I16(_) => NDDataType::Int16,
85 Self::U16(_) => NDDataType::UInt16,
86 Self::I32(_) => NDDataType::Int32,
87 Self::U32(_) => NDDataType::UInt32,
88 Self::I64(_) => NDDataType::Int64,
89 Self::U64(_) => NDDataType::UInt64,
90 Self::F32(_) => NDDataType::Float32,
91 Self::F64(_) => NDDataType::Float64,
92 }
93 }
94
95 pub fn len(&self) -> usize {
96 match self {
97 Self::I8(v) => v.len(),
98 Self::U8(v) => v.len(),
99 Self::I16(v) => v.len(),
100 Self::U16(v) => v.len(),
101 Self::I32(v) => v.len(),
102 Self::U32(v) => v.len(),
103 Self::I64(v) => v.len(),
104 Self::U64(v) => v.len(),
105 Self::F32(v) => v.len(),
106 Self::F64(v) => v.len(),
107 }
108 }
109
110 pub fn is_empty(&self) -> bool {
111 self.len() == 0
112 }
113
114 pub fn total_bytes(&self) -> usize {
115 self.len() * self.data_type().element_size()
116 }
117
118 pub fn capacity_bytes(&self) -> usize {
120 let cap = match self {
121 Self::I8(v) => v.capacity(),
122 Self::U8(v) => v.capacity(),
123 Self::I16(v) => v.capacity(),
124 Self::U16(v) => v.capacity(),
125 Self::I32(v) => v.capacity(),
126 Self::U32(v) => v.capacity(),
127 Self::I64(v) => v.capacity(),
128 Self::U64(v) => v.capacity(),
129 Self::F32(v) => v.capacity(),
130 Self::F64(v) => v.capacity(),
131 };
132 cap * self.data_type().element_size()
133 }
134
135 pub fn resize(&mut self, new_len: usize) {
137 match self {
138 Self::I8(v) => v.resize(new_len, 0),
139 Self::U8(v) => v.resize(new_len, 0),
140 Self::I16(v) => v.resize(new_len, 0),
141 Self::U16(v) => v.resize(new_len, 0),
142 Self::I32(v) => v.resize(new_len, 0),
143 Self::U32(v) => v.resize(new_len, 0),
144 Self::I64(v) => v.resize(new_len, 0),
145 Self::U64(v) => v.resize(new_len, 0),
146 Self::F32(v) => v.resize(new_len, 0.0),
147 Self::F64(v) => v.resize(new_len, 0.0),
148 }
149 }
150
151 pub fn as_u8_slice(&self) -> &[u8] {
153 match self {
154 Self::I8(v) => unsafe {
155 std::slice::from_raw_parts(v.as_ptr() as *const u8, v.len())
156 },
157 Self::U8(v) => v.as_slice(),
158 Self::I16(v) => unsafe {
159 std::slice::from_raw_parts(v.as_ptr() as *const u8, v.len() * 2)
160 },
161 Self::U16(v) => unsafe {
162 std::slice::from_raw_parts(v.as_ptr() as *const u8, v.len() * 2)
163 },
164 Self::I32(v) => unsafe {
165 std::slice::from_raw_parts(v.as_ptr() as *const u8, v.len() * 4)
166 },
167 Self::U32(v) => unsafe {
168 std::slice::from_raw_parts(v.as_ptr() as *const u8, v.len() * 4)
169 },
170 Self::I64(v) => unsafe {
171 std::slice::from_raw_parts(v.as_ptr() as *const u8, v.len() * 8)
172 },
173 Self::U64(v) => unsafe {
174 std::slice::from_raw_parts(v.as_ptr() as *const u8, v.len() * 8)
175 },
176 Self::F32(v) => unsafe {
177 std::slice::from_raw_parts(v.as_ptr() as *const u8, v.len() * 4)
178 },
179 Self::F64(v) => unsafe {
180 std::slice::from_raw_parts(v.as_ptr() as *const u8, v.len() * 8)
181 },
182 }
183 }
184
185 pub fn get_as_f64(&self, index: usize) -> Option<f64> {
187 match self {
188 Self::I8(v) => v.get(index).map(|&x| x as f64),
189 Self::U8(v) => v.get(index).map(|&x| x as f64),
190 Self::I16(v) => v.get(index).map(|&x| x as f64),
191 Self::U16(v) => v.get(index).map(|&x| x as f64),
192 Self::I32(v) => v.get(index).map(|&x| x as f64),
193 Self::U32(v) => v.get(index).map(|&x| x as f64),
194 Self::I64(v) => v.get(index).map(|&x| x as f64),
195 Self::U64(v) => v.get(index).map(|&x| x as f64),
196 Self::F32(v) => v.get(index).map(|&x| x as f64),
197 Self::F64(v) => v.get(index).copied(),
198 }
199 }
200
201 pub fn set_from_f64(&mut self, index: usize, value: f64) {
203 match self {
204 Self::I8(v) => { if let Some(e) = v.get_mut(index) { *e = value as i8; } }
205 Self::U8(v) => { if let Some(e) = v.get_mut(index) { *e = value as u8; } }
206 Self::I16(v) => { if let Some(e) = v.get_mut(index) { *e = value as i16; } }
207 Self::U16(v) => { if let Some(e) = v.get_mut(index) { *e = value as u16; } }
208 Self::I32(v) => { if let Some(e) = v.get_mut(index) { *e = value as i32; } }
209 Self::U32(v) => { if let Some(e) = v.get_mut(index) { *e = value as u32; } }
210 Self::I64(v) => { if let Some(e) = v.get_mut(index) { *e = value as i64; } }
211 Self::U64(v) => { if let Some(e) = v.get_mut(index) { *e = value as u64; } }
212 Self::F32(v) => { if let Some(e) = v.get_mut(index) { *e = value as f32; } }
213 Self::F64(v) => { if let Some(e) = v.get_mut(index) { *e = value; } }
214 }
215 }
216}
217
218#[derive(Debug, Clone)]
220pub struct NDDimension {
221 pub size: usize,
222 pub offset: usize,
223 pub binning: usize,
224 pub reverse: bool,
225}
226
227impl NDDimension {
228 pub fn new(size: usize) -> Self {
229 Self {
230 size,
231 offset: 0,
232 binning: 1,
233 reverse: false,
234 }
235 }
236}
237
238#[derive(Debug, Clone)]
240pub struct NDArrayInfo {
241 pub total_bytes: usize,
242 pub bytes_per_element: usize,
243 pub num_elements: usize,
244 pub x_size: usize,
245 pub y_size: usize,
246 pub color_size: usize,
247}
248
249#[derive(Debug, Clone)]
251pub struct NDArray {
252 pub unique_id: i32,
253 pub timestamp: EpicsTimestamp,
254 pub dims: Vec<NDDimension>,
255 pub data: NDDataBuffer,
256 pub attributes: NDAttributeList,
257 pub codec: Option<Codec>,
258}
259
260impl NDArray {
261 pub fn new(dims: Vec<NDDimension>, data_type: NDDataType) -> Self {
263 let num_elements: usize = dims.iter().map(|d| d.size).product();
264 Self {
265 unique_id: 0,
266 timestamp: EpicsTimestamp::default(),
267 dims,
268 data: NDDataBuffer::zeros(data_type, num_elements),
269 attributes: NDAttributeList::new(),
270 codec: None,
271 }
272 }
273
274 pub fn info(&self) -> NDArrayInfo {
276 let bytes_per_element = self.data.data_type().element_size();
277 let num_elements = self.data.len();
278 let total_bytes = num_elements * bytes_per_element;
279
280 let ndims = self.dims.len();
281 let (x_size, y_size, color_size) = match ndims {
282 0 => (0, 0, 0),
283 1 => (self.dims[0].size, 1, 1),
284 2 => (self.dims[0].size, self.dims[1].size, 1),
285 _ => {
286 (self.dims[1].size, self.dims[2].size, self.dims[0].size)
289 }
290 };
291
292 NDArrayInfo {
293 total_bytes,
294 bytes_per_element,
295 num_elements,
296 x_size,
297 y_size,
298 color_size,
299 }
300 }
301
302 pub fn validate(&self) -> ADResult<()> {
304 let expected: usize = if self.dims.is_empty() {
305 0
306 } else {
307 self.dims.iter().map(|d| d.size).product()
308 };
309 if self.data.len() != expected {
310 return Err(ADError::BufferSizeMismatch {
311 expected,
312 actual: self.data.len(),
313 });
314 }
315 Ok(())
316 }
317}
318
319#[cfg(test)]
320mod tests {
321 use super::*;
322
323 #[test]
324 fn test_element_size_all_types() {
325 assert_eq!(NDDataType::Int8.element_size(), 1);
326 assert_eq!(NDDataType::UInt8.element_size(), 1);
327 assert_eq!(NDDataType::Int16.element_size(), 2);
328 assert_eq!(NDDataType::UInt16.element_size(), 2);
329 assert_eq!(NDDataType::Int32.element_size(), 4);
330 assert_eq!(NDDataType::UInt32.element_size(), 4);
331 assert_eq!(NDDataType::Int64.element_size(), 8);
332 assert_eq!(NDDataType::UInt64.element_size(), 8);
333 assert_eq!(NDDataType::Float32.element_size(), 4);
334 assert_eq!(NDDataType::Float64.element_size(), 8);
335 }
336
337 #[test]
338 fn test_from_ordinal_roundtrip() {
339 for i in 0..10u8 {
340 let dt = NDDataType::from_ordinal(i).unwrap();
341 assert_eq!(dt as u8, i);
342 }
343 assert!(NDDataType::from_ordinal(10).is_none());
344 }
345
346 #[test]
347 fn test_buffer_zeros_type_and_len() {
348 let buf = NDDataBuffer::zeros(NDDataType::UInt16, 100);
349 assert_eq!(buf.data_type(), NDDataType::UInt16);
350 assert_eq!(buf.len(), 100);
351 assert_eq!(buf.total_bytes(), 200);
352 }
353
354 #[test]
355 fn test_buffer_zeros_all_types() {
356 for i in 0..10u8 {
357 let dt = NDDataType::from_ordinal(i).unwrap();
358 let buf = NDDataBuffer::zeros(dt, 10);
359 assert_eq!(buf.data_type(), dt);
360 assert_eq!(buf.len(), 10);
361 assert_eq!(buf.total_bytes(), 10 * dt.element_size());
362 }
363 }
364
365 #[test]
366 fn test_buffer_as_u8_slice() {
367 let buf = NDDataBuffer::U8(vec![1, 2, 3]);
368 assert_eq!(buf.as_u8_slice(), &[1, 2, 3]);
369 }
370
371 #[test]
372 fn test_ndarray_new_allocates() {
373 let dims = vec![NDDimension::new(256), NDDimension::new(256)];
374 let arr = NDArray::new(dims, NDDataType::UInt8);
375 assert_eq!(arr.data.len(), 256 * 256);
376 assert_eq!(arr.data.data_type(), NDDataType::UInt8);
377 }
378
379 #[test]
380 fn test_ndarray_validate_ok() {
381 let dims = vec![NDDimension::new(10), NDDimension::new(20)];
382 let arr = NDArray::new(dims, NDDataType::Float64);
383 arr.validate().unwrap();
384 }
385
386 #[test]
387 fn test_ndarray_validate_mismatch() {
388 let mut arr = NDArray::new(
389 vec![NDDimension::new(10), NDDimension::new(20)],
390 NDDataType::UInt8,
391 );
392 arr.data = NDDataBuffer::U8(vec![0; 100]);
393 assert!(arr.validate().is_err());
394 }
395
396 #[test]
397 fn test_ndarray_info_2d_mono() {
398 let dims = vec![NDDimension::new(640), NDDimension::new(480)];
399 let arr = NDArray::new(dims, NDDataType::UInt16);
400 let info = arr.info();
401 assert_eq!(info.x_size, 640);
402 assert_eq!(info.y_size, 480);
403 assert_eq!(info.color_size, 1);
404 assert_eq!(info.num_elements, 640 * 480);
405 assert_eq!(info.bytes_per_element, 2);
406 assert_eq!(info.total_bytes, 640 * 480 * 2);
407 }
408
409 #[test]
410 fn test_ndarray_info_3d_rgb() {
411 let dims = vec![
412 NDDimension::new(3),
413 NDDimension::new(640),
414 NDDimension::new(480),
415 ];
416 let arr = NDArray::new(dims, NDDataType::UInt8);
417 let info = arr.info();
418 assert_eq!(info.color_size, 3);
419 assert_eq!(info.x_size, 640);
420 assert_eq!(info.y_size, 480);
421 assert_eq!(info.num_elements, 3 * 640 * 480);
422 }
423
424 #[test]
425 fn test_ndarray_info_1d() {
426 let dims = vec![NDDimension::new(1024)];
427 let arr = NDArray::new(dims, NDDataType::Float64);
428 let info = arr.info();
429 assert_eq!(info.x_size, 1024);
430 assert_eq!(info.y_size, 1);
431 assert_eq!(info.color_size, 1);
432 }
433
434 #[test]
435 fn test_buffer_is_empty() {
436 let buf = NDDataBuffer::zeros(NDDataType::UInt8, 0);
437 assert!(buf.is_empty());
438 let buf2 = NDDataBuffer::zeros(NDDataType::UInt8, 1);
439 assert!(!buf2.is_empty());
440 }
441
442 #[test]
443 fn test_codec_field_preserved() {
444 let mut arr = NDArray::new(vec![NDDimension::new(10)], NDDataType::UInt8);
445 arr.codec = Some(Codec { name: crate::codec::CodecName::JPEG, compressed_size: 42 });
446 let cloned = arr.clone();
447 assert_eq!(cloned.codec.as_ref().unwrap().compressed_size, 42);
448 }
449}