1use std::cmp::Ordering;
4use std::fmt::{Debug, Display, Formatter};
5use std::hash::Hash;
6use std::panic::RefUnwindSafe;
7
8use num_traits::bounds::UpperBounded;
9use num_traits::{FromPrimitive, Num, NumCast, ToPrimitive};
10use vortex_error::{VortexError, VortexResult, vortex_err};
11
12use crate::DType;
13use crate::DType::*;
14use crate::half::f16;
15use crate::nullability::Nullability::NonNullable;
16
17#[derive(Debug, Clone, Copy, PartialEq, PartialOrd, Eq, Hash)]
19#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
20#[cfg_attr(feature = "serde", serde(rename_all = "lowercase"))]
21#[cfg_attr(
22 feature = "rkyv",
23 derive(
24 rkyv::Archive,
25 rkyv::Portable,
26 rkyv::Serialize,
27 rkyv::Deserialize,
28 rkyv::bytecheck::CheckBytes,
29 ),
30 rkyv(as = PType),
31 bytecheck(crate = rkyv::bytecheck),
32)]
33#[repr(u8)]
34pub enum PType {
35 U8,
37 U16,
39 U32,
41 U64,
43 I8,
45 I16,
47 I32,
49 I64,
51 F16,
53 F32,
55 F64,
57}
58
59pub trait NativePType:
61 Send
62 + Sync
63 + Clone
64 + Copy
65 + Debug
66 + Display
67 + Default
68 + RefUnwindSafe
69 + Num
70 + NumCast
71 + FromPrimitive
72 + ToBytes
73 + TryFromBytes
74 + 'static
75{
76 const PTYPE: PType;
78
79 fn is_nan(self) -> bool;
82
83 fn total_compare(self, other: Self) -> Ordering;
85
86 #[inline]
88 fn is_le(self, other: Self) -> bool {
89 self.total_compare(other).is_le()
90 }
91
92 #[inline]
94 fn is_lt(self, other: Self) -> bool {
95 self.total_compare(other).is_lt()
96 }
97
98 #[inline]
100 fn is_ge(self, other: Self) -> bool {
101 self.total_compare(other).is_ge()
102 }
103
104 #[inline]
106 fn is_gt(self, other: Self) -> bool {
107 self.total_compare(other).is_gt()
108 }
109
110 fn is_eq(self, other: Self) -> bool;
112}
113
114macro_rules! native_ptype {
115 ($T:ty, $ptype:tt) => {
116 impl NativePType for $T {
117 const PTYPE: PType = PType::$ptype;
118
119 fn is_nan(self) -> bool {
120 false
121 }
122
123 fn total_compare(self, other: Self) -> Ordering {
124 self.cmp(&other)
125 }
126
127 fn is_eq(self, other: Self) -> bool {
128 self == other
129 }
130 }
131 };
132}
133
134macro_rules! native_float_ptype {
135 ($T:ty, $ptype:tt) => {
136 impl NativePType for $T {
137 const PTYPE: PType = PType::$ptype;
138
139 fn is_nan(self) -> bool {
140 <$T>::is_nan(self)
141 }
142
143 fn total_compare(self, other: Self) -> Ordering {
144 self.total_cmp(&other)
145 }
146
147 fn is_eq(self, other: Self) -> bool {
148 self.to_bits() == other.to_bits()
149 }
150 }
151 };
152}
153
154native_ptype!(u8, U8);
155native_ptype!(u16, U16);
156native_ptype!(u32, U32);
157native_ptype!(u64, U64);
158native_ptype!(i8, I8);
159native_ptype!(i16, I16);
160native_ptype!(i32, I32);
161native_ptype!(i64, I64);
162native_float_ptype!(f16, F16);
163native_float_ptype!(f32, F32);
164native_float_ptype!(f64, F64);
165
166#[macro_export]
168macro_rules! match_each_native_ptype {
169 ($self:expr, | $_:tt $enc:ident | $($body:tt)*) => ({
170 macro_rules! __with__ {( $_ $enc:ident ) => ( $($body)* )}
171 use $crate::PType;
172 use $crate::half::f16;
173 match $self {
174 PType::I8 => __with__! { i8 },
175 PType::I16 => __with__! { i16 },
176 PType::I32 => __with__! { i32 },
177 PType::I64 => __with__! { i64 },
178 PType::U8 => __with__! { u8 },
179 PType::U16 => __with__! { u16 },
180 PType::U32 => __with__! { u32 },
181 PType::U64 => __with__! { u64 },
182 PType::F16 => __with__! { f16 },
183 PType::F32 => __with__! { f32 },
184 PType::F64 => __with__! { f64 },
185 }
186 });
187 ($self:expr,
188 integral: | $_:tt $integral_enc:ident | { $($integral_body:tt)* }
189 floating_point: | $_2:tt $floating_point_enc:ident | { $($floating_point_body:tt)* }
190 ) => ({
191 macro_rules! __with_integer__ {( $_ $integral_enc:ident ) => ( { $($integral_body)* } )}
192 macro_rules! __with_floating_point__ {( $_ $floating_point_enc:ident ) => ( { $($floating_point_body)* } )}
193 use $crate::PType;
194 use $crate::half::f16;
195 match $self {
196 PType::I8 => __with_integer__! { i8 },
197 PType::I16 => __with_integer__! { i16 },
198 PType::I32 => __with_integer__! { i32 },
199 PType::I64 => __with_integer__! { i64 },
200 PType::U8 => __with_integer__! { u8 },
201 PType::U16 => __with_integer__! { u16 },
202 PType::U32 => __with_integer__! { u32 },
203 PType::U64 => __with_integer__! { u64 },
204 PType::F16 => __with_floating_point__! { f16 },
205 PType::F32 => __with_floating_point__! { f32 },
206 PType::F64 => __with_floating_point__! { f64 },
207 }
208 });
209 ($self:expr,
210 unsigned: | $_:tt $unsigned_enc:ident | { $($unsigned_body:tt)* }
211 signed: | $_1:tt $signed_enc:ident | { $($signed_body:tt)* }
212 floating: | $_2:tt $floating_point_enc:ident | { $($floating_point_body:tt)* }
213 ) => ({
214 macro_rules! __with_unsigned__ {( $_ $unsigned_enc:ident ) => ( { $($unsigned_body)* } )}
215 macro_rules! __with_signed__ {( $_ $signed_enc:ident ) => ( { $($signed_body)* } )}
216 macro_rules! __with_floating_point__ {( $_ $floating_point_enc:ident ) => ( { $($floating_point_body)* } )}
217 use $crate::PType;
218 use $crate::half::f16;
219 match $self {
220 PType::U8 => __with_unsigned__! { u8 },
221 PType::U16 => __with_unsigned__! { u16 },
222 PType::U32 => __with_unsigned__! { u32 },
223 PType::U64 => __with_unsigned__! { u64 },
224 PType::I8 => __with_signed__! { i8 },
225 PType::I16 => __with_signed__! { i16 },
226 PType::I32 => __with_signed__! { i32 },
227 PType::I64 => __with_signed__! { i64 },
228 PType::F16 => __with_floating_point__! { f16 },
229 PType::F32 => __with_floating_point__! { f32 },
230 PType::F64 => __with_floating_point__! { f64 },
231 }
232 })
233}
234
235#[macro_export]
237macro_rules! match_each_integer_ptype {
238 ($self:expr, | $_:tt $enc:ident | $($body:tt)*) => ({
239 macro_rules! __with__ {( $_ $enc:ident ) => ( $($body)* )}
240 use $crate::PType;
241 match $self {
242 PType::I8 => __with__! { i8 },
243 PType::I16 => __with__! { i16 },
244 PType::I32 => __with__! { i32 },
245 PType::I64 => __with__! { i64 },
246 PType::U8 => __with__! { u8 },
247 PType::U16 => __with__! { u16 },
248 PType::U32 => __with__! { u32 },
249 PType::U64 => __with__! { u64 },
250 other => panic!("Unsupported ptype {other}")
251 }
252 })
253}
254
255#[macro_export]
258macro_rules! match_each_integer_ptype_with_unsigned_type {
259 ($self:expr, | $_:tt $enc:ident, $_2:tt $unsigned:ident | $($body:tt)*) => ({
260 macro_rules! __with__ {( $_ $enc:ident, $_2 $unsigned:ident ) => ( $($body)* )}
261 use $crate::PType;
262 match $self {
263 PType::I8 => __with__! { i8, u8 },
264 PType::I16 => __with__! { i16, u16 },
265 PType::I32 => __with__! { i32, u32 },
266 PType::I64 => __with__! { i64, u64 },
267 PType::U8 => __with__! { u8, u8 },
268 PType::U16 => __with__! { u16, u16 },
269 PType::U32 => __with__! { u32, u32 },
270 PType::U64 => __with__! { u64, u64 },
271 other => panic!("Unsupported ptype {other}")
272 }
273 })
274}
275
276#[macro_export]
278macro_rules! match_each_unsigned_integer_ptype {
279 ($self:expr, | $_:tt $enc:ident | $($body:tt)*) => ({
280 macro_rules! __with__ {( $_ $enc:ident ) => ( $($body)* )}
281 use $crate::PType;
282 match $self {
283 PType::U8 => __with__! { u8 },
284 PType::U16 => __with__! { u16 },
285 PType::U32 => __with__! { u32 },
286 PType::U64 => __with__! { u64 },
287 other => panic!("Unsupported ptype {other}"),
288 }
289 })
290}
291
292#[macro_export]
294macro_rules! match_each_float_ptype {
295 ($self:expr, | $_:tt $enc:ident | $($body:tt)*) => ({
296 macro_rules! __with__ {( $_ $enc:ident ) => ( $($body)* )}
297 use $crate::PType;
298 use vortex_dtype::half::f16;
299 match $self {
300 PType::F16 => __with__! { f16 },
301 PType::F32 => __with__! { f32 },
302 PType::F64 => __with__! { f64 },
303 other => panic!("Unsupported ptype {other}"),
304 }
305 })
306}
307
308#[macro_export]
312macro_rules! match_each_native_simd_ptype {
313 ($self:expr, | $_:tt $enc:ident | $($body:tt)*) => ({
314 macro_rules! __with__ {( $_ $enc:ident ) => ( $($body)* )}
315 use $crate::PType;
316 match $self {
317 PType::I8 => __with__! { i8 },
318 PType::I16 => __with__! { i16 },
319 PType::I32 => __with__! { i32 },
320 PType::I64 => __with__! { i64 },
321 PType::U8 => __with__! { u8 },
322 PType::U16 => __with__! { u16 },
323 PType::U32 => __with__! { u32 },
324 PType::U64 => __with__! { u64 },
325 PType::F16 => panic!("f16 does not implement simd::SimdElement"),
326 PType::F32 => __with__! { f32 },
327 PType::F64 => __with__! { f64 },
328 }
329 })
330}
331
332impl PType {
333 pub const fn is_unsigned_int(self) -> bool {
335 matches!(self, Self::U8 | Self::U16 | Self::U32 | Self::U64)
336 }
337
338 pub const fn is_signed_int(self) -> bool {
340 matches!(self, Self::I8 | Self::I16 | Self::I32 | Self::I64)
341 }
342
343 pub const fn is_int(self) -> bool {
346 self.is_unsigned_int() || self.is_signed_int()
347 }
348
349 pub const fn is_float(self) -> bool {
351 matches!(self, Self::F16 | Self::F32 | Self::F64)
352 }
353
354 pub const fn byte_width(&self) -> usize {
356 match_each_native_ptype!(self, |$T| std::mem::size_of::<$T>())
357 }
358
359 pub const fn bit_width(&self) -> usize {
361 self.byte_width() * 8
362 }
363
364 pub fn max_value_as_u64(&self) -> u64 {
367 match_each_native_ptype!(self, |$T| <$T as UpperBounded>::max_value().to_u64().unwrap_or(u64::MAX))
368 }
369
370 pub const fn to_signed(self) -> Self {
372 match self {
373 Self::U8 => Self::I8,
374 Self::U16 => Self::I16,
375 Self::U32 => Self::I32,
376 Self::U64 => Self::I64,
377 _ => self,
378 }
379 }
380
381 pub const fn to_unsigned(self) -> Self {
384 match self {
385 Self::I8 => Self::U8,
386 Self::I16 => Self::U16,
387 Self::I32 => Self::U32,
388 Self::I64 => Self::U64,
389 _ => self,
390 }
391 }
392}
393
394impl Display for PType {
395 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
396 match self {
397 Self::U8 => write!(f, "u8"),
398 Self::U16 => write!(f, "u16"),
399 Self::U32 => write!(f, "u32"),
400 Self::U64 => write!(f, "u64"),
401 Self::I8 => write!(f, "i8"),
402 Self::I16 => write!(f, "i16"),
403 Self::I32 => write!(f, "i32"),
404 Self::I64 => write!(f, "i64"),
405 Self::F16 => write!(f, "f16"),
406 Self::F32 => write!(f, "f32"),
407 Self::F64 => write!(f, "f64"),
408 }
409 }
410}
411
412impl TryFrom<&DType> for PType {
413 type Error = VortexError;
414
415 fn try_from(value: &DType) -> VortexResult<Self> {
416 match value {
417 Primitive(p, _) => Ok(*p),
418 _ => Err(vortex_err!("Cannot convert DType {} into PType", value)),
419 }
420 }
421}
422
423impl From<PType> for &DType {
424 fn from(item: PType) -> Self {
425 match item {
427 PType::I8 => &Primitive(PType::I8, NonNullable),
428 PType::I16 => &Primitive(PType::I16, NonNullable),
429 PType::I32 => &Primitive(PType::I32, NonNullable),
430 PType::I64 => &Primitive(PType::I64, NonNullable),
431 PType::U8 => &Primitive(PType::U8, NonNullable),
432 PType::U16 => &Primitive(PType::U16, NonNullable),
433 PType::U32 => &Primitive(PType::U32, NonNullable),
434 PType::U64 => &Primitive(PType::U64, NonNullable),
435 PType::F16 => &Primitive(PType::F16, NonNullable),
436 PType::F32 => &Primitive(PType::F32, NonNullable),
437 PType::F64 => &Primitive(PType::F64, NonNullable),
438 }
439 }
440}
441
442impl From<PType> for DType {
443 fn from(item: PType) -> Self {
444 Primitive(item, NonNullable)
445 }
446}
447
448pub trait ToBytes: Sized {
450 fn to_le_bytes(&self) -> &[u8];
452}
453
454pub trait TryFromBytes: Sized {
456 fn try_from_le_bytes(bytes: &[u8]) -> VortexResult<Self>;
458}
459
460macro_rules! try_from_bytes {
461 ($T:ty) => {
462 impl ToBytes for $T {
463 #[inline]
464 #[allow(clippy::size_of_in_element_count)]
465 fn to_le_bytes(&self) -> &[u8] {
466 let raw_ptr = self as *const $T as *const u8;
469 unsafe { std::slice::from_raw_parts(raw_ptr, std::mem::size_of::<$T>()) }
470 }
471 }
472
473 impl TryFromBytes for $T {
474 fn try_from_le_bytes(bytes: &[u8]) -> VortexResult<Self> {
475 Ok(<$T>::from_le_bytes(bytes.try_into().map_err(|_| {
476 vortex_err!("Failed to convert bytes into {}", stringify!($T))
477 })?))
478 }
479 }
480 };
481}
482
483try_from_bytes!(u8);
484try_from_bytes!(u16);
485try_from_bytes!(u32);
486try_from_bytes!(u64);
487try_from_bytes!(i8);
488try_from_bytes!(i16);
489try_from_bytes!(i32);
490try_from_bytes!(i64);
491try_from_bytes!(f16);
492try_from_bytes!(f32);
493try_from_bytes!(f64);
494
495#[cfg(test)]
496mod tests {
497 use super::*;
498
499 #[test]
500 fn try_from_bytes() {
501 assert_eq!(u8::try_from_le_bytes(&[0x01]).unwrap(), 0x01);
502 assert_eq!(u16::try_from_le_bytes(&[0x01, 0x02]).unwrap(), 0x0201);
503 assert_eq!(
504 u32::try_from_le_bytes(&[0x01, 0x02, 0x03, 0x04]).unwrap(),
505 0x04030201
506 );
507 assert_eq!(
508 u64::try_from_le_bytes(&[0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08]).unwrap(),
509 0x0807060504030201
510 );
511 }
512
513 #[test]
514 fn to_bytes_rt() {
515 assert_eq!(&0x01u8.to_le_bytes(), &[0x01]);
516 assert_eq!(&0x0201u16.to_le_bytes(), &[0x01, 0x02]);
517
518 assert_eq!(u8::try_from_le_bytes(&42_u8.to_le_bytes()).unwrap(), 42);
519 assert_eq!(u16::try_from_le_bytes(&42_u16.to_le_bytes()).unwrap(), 42);
520 assert_eq!(u32::try_from_le_bytes(&42_u32.to_le_bytes()).unwrap(), 42);
521 assert_eq!(u64::try_from_le_bytes(&42_u64.to_le_bytes()).unwrap(), 42);
522 assert_eq!(i8::try_from_le_bytes(&42_i8.to_le_bytes()).unwrap(), 42);
523 assert_eq!(i16::try_from_le_bytes(&42_i16.to_le_bytes()).unwrap(), 42);
524 assert_eq!(i32::try_from_le_bytes(&42_i32.to_le_bytes()).unwrap(), 42);
525 assert_eq!(i64::try_from_le_bytes(&42_i64.to_le_bytes()).unwrap(), 42);
526 assert_eq!(
527 f16::try_from_le_bytes(&f16::from_f32(42.0).to_le_bytes()).unwrap(),
528 f16::from_f32(42.0)
529 );
530 assert_eq!(
531 f32::try_from_le_bytes(&42.0_f32.to_le_bytes()).unwrap(),
532 42.0
533 );
534 assert_eq!(
535 f64::try_from_le_bytes(&42.0_f64.to_le_bytes()).unwrap(),
536 42.0
537 );
538 }
539
540 #[test]
541 fn max_value_u64() {
542 assert_eq!(PType::U8.max_value_as_u64(), u8::MAX as u64);
543 assert_eq!(PType::U16.max_value_as_u64(), u16::MAX as u64);
544 assert_eq!(PType::U32.max_value_as_u64(), u32::MAX as u64);
545 assert_eq!(PType::U64.max_value_as_u64(), u64::MAX);
546 assert_eq!(PType::I8.max_value_as_u64(), i8::MAX as u64);
547 assert_eq!(PType::I16.max_value_as_u64(), i16::MAX as u64);
548 assert_eq!(PType::I32.max_value_as_u64(), i32::MAX as u64);
549 assert_eq!(PType::I64.max_value_as_u64(), i64::MAX as u64);
550 assert_eq!(PType::F16.max_value_as_u64(), 65504); assert_eq!(PType::F32.max_value_as_u64(), u64::MAX);
552 assert_eq!(PType::F64.max_value_as_u64(), u64::MAX);
553 }
554
555 #[test]
556 fn widths() {
557 assert_eq!(PType::U8.byte_width(), 1);
558 assert_eq!(PType::U16.byte_width(), 2);
559 assert_eq!(PType::U32.byte_width(), 4);
560 assert_eq!(PType::U64.byte_width(), 8);
561 assert_eq!(PType::I8.byte_width(), 1);
562 assert_eq!(PType::I16.byte_width(), 2);
563 assert_eq!(PType::I32.byte_width(), 4);
564 assert_eq!(PType::I64.byte_width(), 8);
565 assert_eq!(PType::F16.byte_width(), 2);
566 assert_eq!(PType::F32.byte_width(), 4);
567 assert_eq!(PType::F64.byte_width(), 8);
568
569 assert_eq!(PType::U8.bit_width(), 8);
570 assert_eq!(PType::U16.bit_width(), 16);
571 assert_eq!(PType::U32.bit_width(), 32);
572 assert_eq!(PType::U64.bit_width(), 64);
573 assert_eq!(PType::I8.bit_width(), 8);
574 assert_eq!(PType::I16.bit_width(), 16);
575 assert_eq!(PType::I32.bit_width(), 32);
576 assert_eq!(PType::I64.bit_width(), 64);
577 assert_eq!(PType::F16.bit_width(), 16);
578 assert_eq!(PType::F32.bit_width(), 32);
579 assert_eq!(PType::F64.bit_width(), 64);
580 }
581
582 #[test]
583 fn native_ptype_nan_handling() {
584 let a = f32::NAN;
585 let b = f32::NAN;
586 assert_ne!(a, b);
587 assert!(<f32 as NativePType>::is_nan(a));
588 assert!(<f32 as NativePType>::is_nan(b));
589 assert!(<f32 as NativePType>::is_eq(a, b));
590 assert!(<f32 as NativePType>::total_compare(a, b) == Ordering::Equal);
591 }
592
593 #[test]
594 fn to_signed() {
595 assert_eq!(PType::U8.to_signed(), PType::I8);
596 assert_eq!(PType::U16.to_signed(), PType::I16);
597 assert_eq!(PType::U32.to_signed(), PType::I32);
598 assert_eq!(PType::U64.to_signed(), PType::I64);
599 assert_eq!(PType::I8.to_signed(), PType::I8);
600 assert_eq!(PType::I16.to_signed(), PType::I16);
601 assert_eq!(PType::I32.to_signed(), PType::I32);
602 assert_eq!(PType::I64.to_signed(), PType::I64);
603 assert_eq!(PType::F16.to_signed(), PType::F16);
604 assert_eq!(PType::F32.to_signed(), PType::F32);
605 assert_eq!(PType::F64.to_signed(), PType::F64);
606 }
607
608 #[test]
609 fn to_unsigned() {
610 assert_eq!(PType::U8.to_unsigned(), PType::U8);
611 assert_eq!(PType::U16.to_unsigned(), PType::U16);
612 assert_eq!(PType::U32.to_unsigned(), PType::U32);
613 assert_eq!(PType::U64.to_unsigned(), PType::U64);
614 assert_eq!(PType::I8.to_unsigned(), PType::U8);
615 assert_eq!(PType::I16.to_unsigned(), PType::U16);
616 assert_eq!(PType::I32.to_unsigned(), PType::U32);
617 assert_eq!(PType::I64.to_unsigned(), PType::U64);
618 assert_eq!(PType::F16.to_unsigned(), PType::F16);
619 assert_eq!(PType::F32.to_unsigned(), PType::F32);
620 assert_eq!(PType::F64.to_unsigned(), PType::F64);
621 }
622
623 #[test]
624 fn to_dtype() {
625 assert_eq!(DType::from(PType::U8), Primitive(PType::U8, NonNullable));
626 assert_eq!(DType::from(PType::U16), Primitive(PType::U16, NonNullable));
627 assert_eq!(DType::from(PType::U32), Primitive(PType::U32, NonNullable));
628 assert_eq!(DType::from(PType::U64), Primitive(PType::U64, NonNullable));
629 assert_eq!(DType::from(PType::I8), Primitive(PType::I8, NonNullable));
630 assert_eq!(DType::from(PType::I16), Primitive(PType::I16, NonNullable));
631 assert_eq!(DType::from(PType::I32), Primitive(PType::I32, NonNullable));
632 assert_eq!(DType::from(PType::I64), Primitive(PType::I64, NonNullable));
633 assert_eq!(DType::from(PType::F16), Primitive(PType::F16, NonNullable));
634 assert_eq!(DType::from(PType::F32), Primitive(PType::F32, NonNullable));
635 assert_eq!(DType::from(PType::F64), Primitive(PType::F64, NonNullable));
636 }
637}