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 { std::slice::from_raw_parts(v.as_ptr() as *const u8, v.len()) },
155 Self::U8(v) => v.as_slice(),
156 Self::I16(v) => unsafe {
157 std::slice::from_raw_parts(v.as_ptr() as *const u8, v.len() * 2)
158 },
159 Self::U16(v) => unsafe {
160 std::slice::from_raw_parts(v.as_ptr() as *const u8, v.len() * 2)
161 },
162 Self::I32(v) => unsafe {
163 std::slice::from_raw_parts(v.as_ptr() as *const u8, v.len() * 4)
164 },
165 Self::U32(v) => unsafe {
166 std::slice::from_raw_parts(v.as_ptr() as *const u8, v.len() * 4)
167 },
168 Self::I64(v) => unsafe {
169 std::slice::from_raw_parts(v.as_ptr() as *const u8, v.len() * 8)
170 },
171 Self::U64(v) => unsafe {
172 std::slice::from_raw_parts(v.as_ptr() as *const u8, v.len() * 8)
173 },
174 Self::F32(v) => unsafe {
175 std::slice::from_raw_parts(v.as_ptr() as *const u8, v.len() * 4)
176 },
177 Self::F64(v) => unsafe {
178 std::slice::from_raw_parts(v.as_ptr() as *const u8, v.len() * 8)
179 },
180 }
181 }
182
183 pub fn get_as_f64(&self, index: usize) -> Option<f64> {
185 match self {
186 Self::I8(v) => v.get(index).map(|&x| x as f64),
187 Self::U8(v) => v.get(index).map(|&x| x as f64),
188 Self::I16(v) => v.get(index).map(|&x| x as f64),
189 Self::U16(v) => v.get(index).map(|&x| x as f64),
190 Self::I32(v) => v.get(index).map(|&x| x as f64),
191 Self::U32(v) => v.get(index).map(|&x| x as f64),
192 Self::I64(v) => v.get(index).map(|&x| x as f64),
193 Self::U64(v) => v.get(index).map(|&x| x as f64),
194 Self::F32(v) => v.get(index).map(|&x| x as f64),
195 Self::F64(v) => v.get(index).copied(),
196 }
197 }
198
199 pub fn set_from_f64(&mut self, index: usize, value: f64) {
201 match self {
202 Self::I8(v) => {
203 if let Some(e) = v.get_mut(index) {
204 *e = value as i8;
205 }
206 }
207 Self::U8(v) => {
208 if let Some(e) = v.get_mut(index) {
209 *e = value as u8;
210 }
211 }
212 Self::I16(v) => {
213 if let Some(e) = v.get_mut(index) {
214 *e = value as i16;
215 }
216 }
217 Self::U16(v) => {
218 if let Some(e) = v.get_mut(index) {
219 *e = value as u16;
220 }
221 }
222 Self::I32(v) => {
223 if let Some(e) = v.get_mut(index) {
224 *e = value as i32;
225 }
226 }
227 Self::U32(v) => {
228 if let Some(e) = v.get_mut(index) {
229 *e = value as u32;
230 }
231 }
232 Self::I64(v) => {
233 if let Some(e) = v.get_mut(index) {
234 *e = value as i64;
235 }
236 }
237 Self::U64(v) => {
238 if let Some(e) = v.get_mut(index) {
239 *e = value as u64;
240 }
241 }
242 Self::F32(v) => {
243 if let Some(e) = v.get_mut(index) {
244 *e = value as f32;
245 }
246 }
247 Self::F64(v) => {
248 if let Some(e) = v.get_mut(index) {
249 *e = value;
250 }
251 }
252 }
253 }
254}
255
256#[derive(Debug, Clone)]
258pub struct NDDimension {
259 pub size: usize,
260 pub offset: usize,
261 pub binning: usize,
262 pub reverse: bool,
263}
264
265impl NDDimension {
266 pub fn new(size: usize) -> Self {
267 Self {
268 size,
269 offset: 0,
270 binning: 1,
271 reverse: false,
272 }
273 }
274}
275
276#[derive(Debug, Clone)]
278pub struct NDArrayInfo {
279 pub total_bytes: usize,
280 pub bytes_per_element: usize,
281 pub num_elements: usize,
282 pub x_size: usize,
283 pub y_size: usize,
284 pub color_size: usize,
285 pub x_dim: usize,
287 pub y_dim: usize,
289 pub color_dim: usize,
291 pub x_stride: usize,
293 pub y_stride: usize,
295 pub color_stride: usize,
297 pub color_mode: crate::color::NDColorMode,
299}
300
301#[derive(Debug, Clone)]
303pub struct NDArray {
304 pub unique_id: i32,
305 pub timestamp: EpicsTimestamp,
306 pub time_stamp: f64,
308 pub dims: Vec<NDDimension>,
309 pub data: NDDataBuffer,
310 pub attributes: NDAttributeList,
311 pub codec: Option<Codec>,
312}
313
314impl NDArray {
315 pub fn new(dims: Vec<NDDimension>, data_type: NDDataType) -> Self {
317 let num_elements: usize = if dims.is_empty() {
318 0
319 } else {
320 dims.iter().map(|d| d.size).product()
321 };
322 Self {
323 unique_id: 0,
324 timestamp: EpicsTimestamp::default(),
325 time_stamp: 0.0,
326 dims,
327 data: NDDataBuffer::zeros(data_type, num_elements),
328 attributes: NDAttributeList::new(),
329 codec: None,
330 }
331 }
332
333 pub fn info(&self) -> NDArrayInfo {
338 use crate::color::NDColorMode;
339
340 let bytes_per_element = self.data.data_type().element_size();
341 let num_elements = self.data.len();
342 let total_bytes = num_elements * bytes_per_element;
343
344 let ndims = self.dims.len();
345
346 let color_mode = self
348 .attributes
349 .get("ColorMode")
350 .and_then(|a| a.value.as_i64())
351 .map(|v| NDColorMode::from_i32(v as i32))
352 .unwrap_or(NDColorMode::Mono);
353
354 let (x_size, y_size, color_size, x_dim, y_dim, color_dim, x_stride, y_stride, color_stride) =
355 match ndims {
356 0 => (0, 0, 0, 0, 0, 0, 0, 0, 0),
357 1 => (self.dims[0].size, 1, 1, 0, 0, 0, 1, self.dims[0].size, 0),
358 2 => {
359 let xs = self.dims[0].size;
360 let ys = self.dims[1].size;
361 (xs, ys, 1, 0, 1, 0, 1, xs, 0)
362 }
363 _ => {
364 match color_mode {
366 NDColorMode::RGB1 => {
367 let cs = self.dims[0].size;
369 let xs = self.dims[1].size;
370 let ys = self.dims[2].size;
371 (xs, ys, cs, 1, 2, 0, cs, xs * cs, 1)
372 }
373 NDColorMode::RGB2 => {
374 let xs = self.dims[0].size;
376 let cs = self.dims[1].size;
377 let ys = self.dims[2].size;
378 (xs, ys, cs, 0, 2, 1, 1, xs * cs, xs)
379 }
380 NDColorMode::RGB3 => {
381 let xs = self.dims[0].size;
383 let ys = self.dims[1].size;
384 let cs = self.dims[2].size;
385 (xs, ys, cs, 0, 1, 2, 1, xs, xs * ys)
386 }
387 _ => {
388 let xs = self.dims[0].size;
390 let ys = self.dims[1].size;
391 let cs = self.dims[2].size;
392 (xs, ys, cs, 0, 1, 2, 1, xs, xs * ys)
393 }
394 }
395 }
396 };
397
398 NDArrayInfo {
399 total_bytes,
400 bytes_per_element,
401 num_elements,
402 x_size,
403 y_size,
404 color_size,
405 x_dim,
406 y_dim,
407 color_dim,
408 x_stride,
409 y_stride,
410 color_stride,
411 color_mode,
412 }
413 }
414
415 pub fn validate(&self) -> ADResult<()> {
417 let expected: usize = if self.dims.is_empty() {
418 0
419 } else {
420 self.dims.iter().map(|d| d.size).product()
421 };
422 if self.data.len() != expected {
423 return Err(ADError::BufferSizeMismatch {
424 expected,
425 actual: self.data.len(),
426 });
427 }
428 Ok(())
429 }
430}
431
432#[cfg(test)]
433mod tests {
434 use super::*;
435
436 #[test]
437 fn test_element_size_all_types() {
438 assert_eq!(NDDataType::Int8.element_size(), 1);
439 assert_eq!(NDDataType::UInt8.element_size(), 1);
440 assert_eq!(NDDataType::Int16.element_size(), 2);
441 assert_eq!(NDDataType::UInt16.element_size(), 2);
442 assert_eq!(NDDataType::Int32.element_size(), 4);
443 assert_eq!(NDDataType::UInt32.element_size(), 4);
444 assert_eq!(NDDataType::Int64.element_size(), 8);
445 assert_eq!(NDDataType::UInt64.element_size(), 8);
446 assert_eq!(NDDataType::Float32.element_size(), 4);
447 assert_eq!(NDDataType::Float64.element_size(), 8);
448 }
449
450 #[test]
451 fn test_from_ordinal_roundtrip() {
452 for i in 0..10u8 {
453 let dt = NDDataType::from_ordinal(i).unwrap();
454 assert_eq!(dt as u8, i);
455 }
456 assert!(NDDataType::from_ordinal(10).is_none());
457 }
458
459 #[test]
460 fn test_buffer_zeros_type_and_len() {
461 let buf = NDDataBuffer::zeros(NDDataType::UInt16, 100);
462 assert_eq!(buf.data_type(), NDDataType::UInt16);
463 assert_eq!(buf.len(), 100);
464 assert_eq!(buf.total_bytes(), 200);
465 }
466
467 #[test]
468 fn test_buffer_zeros_all_types() {
469 for i in 0..10u8 {
470 let dt = NDDataType::from_ordinal(i).unwrap();
471 let buf = NDDataBuffer::zeros(dt, 10);
472 assert_eq!(buf.data_type(), dt);
473 assert_eq!(buf.len(), 10);
474 assert_eq!(buf.total_bytes(), 10 * dt.element_size());
475 }
476 }
477
478 #[test]
479 fn test_buffer_as_u8_slice() {
480 let buf = NDDataBuffer::U8(vec![1, 2, 3]);
481 assert_eq!(buf.as_u8_slice(), &[1, 2, 3]);
482 }
483
484 #[test]
485 fn test_ndarray_new_allocates() {
486 let dims = vec![NDDimension::new(256), NDDimension::new(256)];
487 let arr = NDArray::new(dims, NDDataType::UInt8);
488 assert_eq!(arr.data.len(), 256 * 256);
489 assert_eq!(arr.data.data_type(), NDDataType::UInt8);
490 }
491
492 #[test]
493 fn test_ndarray_validate_ok() {
494 let dims = vec![NDDimension::new(10), NDDimension::new(20)];
495 let arr = NDArray::new(dims, NDDataType::Float64);
496 arr.validate().unwrap();
497 }
498
499 #[test]
500 fn test_ndarray_validate_mismatch() {
501 let mut arr = NDArray::new(
502 vec![NDDimension::new(10), NDDimension::new(20)],
503 NDDataType::UInt8,
504 );
505 arr.data = NDDataBuffer::U8(vec![0; 100]);
506 assert!(arr.validate().is_err());
507 }
508
509 #[test]
510 fn test_ndarray_info_2d_mono() {
511 let dims = vec![NDDimension::new(640), NDDimension::new(480)];
512 let arr = NDArray::new(dims, NDDataType::UInt16);
513 let info = arr.info();
514 assert_eq!(info.x_size, 640);
515 assert_eq!(info.y_size, 480);
516 assert_eq!(info.color_size, 1);
517 assert_eq!(info.num_elements, 640 * 480);
518 assert_eq!(info.bytes_per_element, 2);
519 assert_eq!(info.total_bytes, 640 * 480 * 2);
520 }
521
522 #[test]
523 fn test_ndarray_info_3d_rgb() {
524 use crate::attributes::{NDAttrSource, NDAttrValue, NDAttribute};
525 use crate::color::NDColorMode;
526
527 let dims = vec![
529 NDDimension::new(3),
530 NDDimension::new(640),
531 NDDimension::new(480),
532 ];
533 let arr = NDArray::new(dims, NDDataType::UInt8);
534 let info = arr.info();
535 assert_eq!(info.x_size, 3);
536 assert_eq!(info.y_size, 640);
537 assert_eq!(info.color_size, 480);
538
539 let dims = vec![
541 NDDimension::new(3),
542 NDDimension::new(640),
543 NDDimension::new(480),
544 ];
545 let mut arr = NDArray::new(dims, NDDataType::UInt8);
546 arr.attributes.add(NDAttribute {
547 name: "ColorMode".into(),
548 description: "Color Mode".into(),
549 source: NDAttrSource::Driver,
550 value: NDAttrValue::Int32(NDColorMode::RGB1 as i32),
551 });
552 let info = arr.info();
553 assert_eq!(info.color_size, 3);
554 assert_eq!(info.x_size, 640);
555 assert_eq!(info.y_size, 480);
556 assert_eq!(info.x_dim, 1);
557 assert_eq!(info.y_dim, 2);
558 assert_eq!(info.color_dim, 0);
559 assert_eq!(info.num_elements, 3 * 640 * 480);
560 }
561
562 #[test]
563 fn test_ndarray_info_1d() {
564 let dims = vec![NDDimension::new(1024)];
565 let arr = NDArray::new(dims, NDDataType::Float64);
566 let info = arr.info();
567 assert_eq!(info.x_size, 1024);
568 assert_eq!(info.y_size, 1);
569 assert_eq!(info.color_size, 1);
570 }
571
572 #[test]
573 fn test_buffer_is_empty() {
574 let buf = NDDataBuffer::zeros(NDDataType::UInt8, 0);
575 assert!(buf.is_empty());
576 let buf2 = NDDataBuffer::zeros(NDDataType::UInt8, 1);
577 assert!(!buf2.is_empty());
578 }
579
580 #[test]
581 fn test_codec_field_preserved() {
582 let mut arr = NDArray::new(vec![NDDimension::new(10)], NDDataType::UInt8);
583 arr.codec = Some(Codec {
584 name: crate::codec::CodecName::JPEG,
585 compressed_size: 42,
586 level: 0,
587 shuffle: 0,
588 compressor: 0,
589 });
590 let cloned = arr.clone();
591 assert_eq!(cloned.codec.as_ref().unwrap().compressed_size, 42);
592 }
593}