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 = 0,
37 U16 = 1,
39 U32 = 2,
41 U64 = 3,
43 I8 = 4,
45 I16 = 5,
47 I32 = 6,
49 I64 = 7,
51 F16 = 8,
53 F32 = 9,
55 F64 = 10,
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 #[inline]
120 fn is_nan(self) -> bool {
121 false
122 }
123
124 #[inline]
125 fn total_compare(self, other: Self) -> Ordering {
126 self.cmp(&other)
127 }
128
129 #[inline]
130 fn is_eq(self, other: Self) -> bool {
131 self == other
132 }
133 }
134 };
135}
136
137macro_rules! native_float_ptype {
138 ($T:ty, $ptype:tt) => {
139 impl NativePType for $T {
140 const PTYPE: PType = PType::$ptype;
141
142 #[inline]
143 fn is_nan(self) -> bool {
144 <$T>::is_nan(self)
145 }
146
147 #[inline]
148 fn total_compare(self, other: Self) -> Ordering {
149 self.total_cmp(&other)
150 }
151
152 #[inline]
153 fn is_eq(self, other: Self) -> bool {
154 self.to_bits() == other.to_bits()
155 }
156 }
157 };
158}
159
160native_ptype!(u8, U8);
161native_ptype!(u16, U16);
162native_ptype!(u32, U32);
163native_ptype!(u64, U64);
164native_ptype!(i8, I8);
165native_ptype!(i16, I16);
166native_ptype!(i32, I32);
167native_ptype!(i64, I64);
168native_float_ptype!(f16, F16);
169native_float_ptype!(f32, F32);
170native_float_ptype!(f64, F64);
171
172#[macro_export]
174macro_rules! match_each_native_ptype {
175 ($self:expr, | $_:tt $enc:ident | $($body:tt)*) => ({
176 macro_rules! __with__ {( $_ $enc:ident ) => ( $($body)* )}
177 use $crate::PType;
178 use $crate::half::f16;
179 match $self {
180 PType::I8 => __with__! { i8 },
181 PType::I16 => __with__! { i16 },
182 PType::I32 => __with__! { i32 },
183 PType::I64 => __with__! { i64 },
184 PType::U8 => __with__! { u8 },
185 PType::U16 => __with__! { u16 },
186 PType::U32 => __with__! { u32 },
187 PType::U64 => __with__! { u64 },
188 PType::F16 => __with__! { f16 },
189 PType::F32 => __with__! { f32 },
190 PType::F64 => __with__! { f64 },
191 }
192 });
193 ($self:expr,
194 integral: | $_:tt $integral_enc:ident | { $($integral_body:tt)* }
195 floating_point: | $_2:tt $floating_point_enc:ident | { $($floating_point_body:tt)* }
196 ) => ({
197 macro_rules! __with_integer__ {( $_ $integral_enc:ident ) => ( { $($integral_body)* } )}
198 macro_rules! __with_floating_point__ {( $_ $floating_point_enc:ident ) => ( { $($floating_point_body)* } )}
199 use $crate::PType;
200 use $crate::half::f16;
201 match $self {
202 PType::I8 => __with_integer__! { i8 },
203 PType::I16 => __with_integer__! { i16 },
204 PType::I32 => __with_integer__! { i32 },
205 PType::I64 => __with_integer__! { i64 },
206 PType::U8 => __with_integer__! { u8 },
207 PType::U16 => __with_integer__! { u16 },
208 PType::U32 => __with_integer__! { u32 },
209 PType::U64 => __with_integer__! { u64 },
210 PType::F16 => __with_floating_point__! { f16 },
211 PType::F32 => __with_floating_point__! { f32 },
212 PType::F64 => __with_floating_point__! { f64 },
213 }
214 });
215 ($self:expr,
216 unsigned: | $_:tt $unsigned_enc:ident | { $($unsigned_body:tt)* }
217 signed: | $_1:tt $signed_enc:ident | { $($signed_body:tt)* }
218 floating: | $_2:tt $floating_point_enc:ident | { $($floating_point_body:tt)* }
219 ) => ({
220 macro_rules! __with_unsigned__ {( $_ $unsigned_enc:ident ) => ( { $($unsigned_body)* } )}
221 macro_rules! __with_signed__ {( $_ $signed_enc:ident ) => ( { $($signed_body)* } )}
222 macro_rules! __with_floating_point__ {( $_ $floating_point_enc:ident ) => ( { $($floating_point_body)* } )}
223 use $crate::PType;
224 use $crate::half::f16;
225 match $self {
226 PType::U8 => __with_unsigned__! { u8 },
227 PType::U16 => __with_unsigned__! { u16 },
228 PType::U32 => __with_unsigned__! { u32 },
229 PType::U64 => __with_unsigned__! { u64 },
230 PType::I8 => __with_signed__! { i8 },
231 PType::I16 => __with_signed__! { i16 },
232 PType::I32 => __with_signed__! { i32 },
233 PType::I64 => __with_signed__! { i64 },
234 PType::F16 => __with_floating_point__! { f16 },
235 PType::F32 => __with_floating_point__! { f32 },
236 PType::F64 => __with_floating_point__! { f64 },
237 }
238 })
239}
240
241#[macro_export]
243macro_rules! match_each_integer_ptype {
244 ($self:expr, | $_:tt $enc:ident | $($body:tt)*) => ({
245 macro_rules! __with__ {( $_ $enc:ident ) => ( $($body)* )}
246 use $crate::PType;
247 match $self {
248 PType::I8 => __with__! { i8 },
249 PType::I16 => __with__! { i16 },
250 PType::I32 => __with__! { i32 },
251 PType::I64 => __with__! { i64 },
252 PType::U8 => __with__! { u8 },
253 PType::U16 => __with__! { u16 },
254 PType::U32 => __with__! { u32 },
255 PType::U64 => __with__! { u64 },
256 other => panic!("Unsupported ptype {other}")
257 }
258 })
259}
260
261#[macro_export]
263macro_rules! match_each_unsigned_integer_ptype {
264 ($self:expr, | $_:tt $enc:ident | $($body:tt)*) => ({
265 macro_rules! __with__ {( $_ $enc:ident ) => ( $($body)* )}
266 use $crate::PType;
267 match $self {
268 PType::U8 => __with__! { u8 },
269 PType::U16 => __with__! { u16 },
270 PType::U32 => __with__! { u32 },
271 PType::U64 => __with__! { u64 },
272 other => panic!("Unsupported ptype {other}"),
273 }
274 })
275}
276
277#[macro_export]
279macro_rules! match_each_float_ptype {
280 ($self:expr, | $_:tt $enc:ident | $($body:tt)*) => ({
281 macro_rules! __with__ {( $_ $enc:ident ) => ( $($body)* )}
282 use $crate::PType;
283 use vortex_dtype::half::f16;
284 match $self {
285 PType::F16 => __with__! { f16 },
286 PType::F32 => __with__! { f32 },
287 PType::F64 => __with__! { f64 },
288 other => panic!("Unsupported ptype {other}"),
289 }
290 })
291}
292
293#[macro_export]
297macro_rules! match_each_native_simd_ptype {
298 ($self:expr, | $_:tt $enc:ident | $($body:tt)*) => ({
299 macro_rules! __with__ {( $_ $enc:ident ) => ( $($body)* )}
300 use $crate::PType;
301 match $self {
302 PType::I8 => __with__! { i8 },
303 PType::I16 => __with__! { i16 },
304 PType::I32 => __with__! { i32 },
305 PType::I64 => __with__! { i64 },
306 PType::U8 => __with__! { u8 },
307 PType::U16 => __with__! { u16 },
308 PType::U32 => __with__! { u32 },
309 PType::U64 => __with__! { u64 },
310 PType::F16 => panic!("f16 does not implement simd::SimdElement"),
311 PType::F32 => __with__! { f32 },
312 PType::F64 => __with__! { f64 },
313 }
314 })
315}
316
317impl PType {
318 pub const fn is_unsigned_int(self) -> bool {
320 matches!(self, Self::U8 | Self::U16 | Self::U32 | Self::U64)
321 }
322
323 pub const fn is_signed_int(self) -> bool {
325 matches!(self, Self::I8 | Self::I16 | Self::I32 | Self::I64)
326 }
327
328 pub const fn is_int(self) -> bool {
331 self.is_unsigned_int() || self.is_signed_int()
332 }
333
334 pub const fn is_float(self) -> bool {
336 matches!(self, Self::F16 | Self::F32 | Self::F64)
337 }
338
339 pub const fn byte_width(&self) -> usize {
341 match_each_native_ptype!(self, |$T| std::mem::size_of::<$T>())
342 }
343
344 pub const fn bit_width(&self) -> usize {
346 self.byte_width() * 8
347 }
348
349 pub fn max_value_as_u64(&self) -> u64 {
352 match_each_native_ptype!(self, |$T| <$T as UpperBounded>::max_value().to_u64().unwrap_or(u64::MAX))
353 }
354
355 pub const fn to_signed(self) -> Self {
357 match self {
358 Self::U8 => Self::I8,
359 Self::U16 => Self::I16,
360 Self::U32 => Self::I32,
361 Self::U64 => Self::I64,
362 _ => self,
363 }
364 }
365
366 pub const fn to_unsigned(self) -> Self {
369 match self {
370 Self::I8 => Self::U8,
371 Self::I16 => Self::U16,
372 Self::I32 => Self::U32,
373 Self::I64 => Self::U64,
374 _ => self,
375 }
376 }
377}
378
379impl Display for PType {
380 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
381 match self {
382 Self::U8 => write!(f, "u8"),
383 Self::U16 => write!(f, "u16"),
384 Self::U32 => write!(f, "u32"),
385 Self::U64 => write!(f, "u64"),
386 Self::I8 => write!(f, "i8"),
387 Self::I16 => write!(f, "i16"),
388 Self::I32 => write!(f, "i32"),
389 Self::I64 => write!(f, "i64"),
390 Self::F16 => write!(f, "f16"),
391 Self::F32 => write!(f, "f32"),
392 Self::F64 => write!(f, "f64"),
393 }
394 }
395}
396
397impl TryFrom<&DType> for PType {
398 type Error = VortexError;
399
400 fn try_from(value: &DType) -> VortexResult<Self> {
401 match value {
402 Primitive(p, _) => Ok(*p),
403 _ => Err(vortex_err!("Cannot convert DType {} into PType", value)),
404 }
405 }
406}
407
408impl From<PType> for &DType {
409 fn from(item: PType) -> Self {
410 match item {
412 PType::I8 => &Primitive(PType::I8, NonNullable),
413 PType::I16 => &Primitive(PType::I16, NonNullable),
414 PType::I32 => &Primitive(PType::I32, NonNullable),
415 PType::I64 => &Primitive(PType::I64, NonNullable),
416 PType::U8 => &Primitive(PType::U8, NonNullable),
417 PType::U16 => &Primitive(PType::U16, NonNullable),
418 PType::U32 => &Primitive(PType::U32, NonNullable),
419 PType::U64 => &Primitive(PType::U64, NonNullable),
420 PType::F16 => &Primitive(PType::F16, NonNullable),
421 PType::F32 => &Primitive(PType::F32, NonNullable),
422 PType::F64 => &Primitive(PType::F64, NonNullable),
423 }
424 }
425}
426
427impl From<PType> for DType {
428 fn from(item: PType) -> Self {
429 Primitive(item, NonNullable)
430 }
431}
432
433pub trait ToBytes: Sized {
435 fn to_le_bytes(&self) -> &[u8];
437}
438
439pub trait TryFromBytes: Sized {
441 fn try_from_le_bytes(bytes: &[u8]) -> VortexResult<Self>;
443}
444
445macro_rules! try_from_bytes {
446 ($T:ty) => {
447 impl ToBytes for $T {
448 #[inline]
449 #[allow(clippy::size_of_in_element_count)]
450 fn to_le_bytes(&self) -> &[u8] {
451 let raw_ptr = self as *const $T as *const u8;
454 unsafe { std::slice::from_raw_parts(raw_ptr, std::mem::size_of::<$T>()) }
455 }
456 }
457
458 impl TryFromBytes for $T {
459 fn try_from_le_bytes(bytes: &[u8]) -> VortexResult<Self> {
460 Ok(<$T>::from_le_bytes(bytes.try_into().map_err(|_| {
461 vortex_err!("Failed to convert bytes into {}", stringify!($T))
462 })?))
463 }
464 }
465 };
466}
467
468try_from_bytes!(u8);
469try_from_bytes!(u16);
470try_from_bytes!(u32);
471try_from_bytes!(u64);
472try_from_bytes!(i8);
473try_from_bytes!(i16);
474try_from_bytes!(i32);
475try_from_bytes!(i64);
476try_from_bytes!(f16);
477try_from_bytes!(f32);
478try_from_bytes!(f64);
479
480#[cfg(test)]
481mod tests {
482 use super::*;
483
484 #[test]
485 fn try_from_bytes() {
486 assert_eq!(u8::try_from_le_bytes(&[0x01]).unwrap(), 0x01);
487 assert_eq!(u16::try_from_le_bytes(&[0x01, 0x02]).unwrap(), 0x0201);
488 assert_eq!(
489 u32::try_from_le_bytes(&[0x01, 0x02, 0x03, 0x04]).unwrap(),
490 0x04030201
491 );
492 assert_eq!(
493 u64::try_from_le_bytes(&[0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08]).unwrap(),
494 0x0807060504030201
495 );
496 }
497
498 #[test]
499 fn to_bytes_rt() {
500 assert_eq!(&0x01u8.to_le_bytes(), &[0x01]);
501 assert_eq!(&0x0201u16.to_le_bytes(), &[0x01, 0x02]);
502
503 assert_eq!(u8::try_from_le_bytes(&42_u8.to_le_bytes()).unwrap(), 42);
504 assert_eq!(u16::try_from_le_bytes(&42_u16.to_le_bytes()).unwrap(), 42);
505 assert_eq!(u32::try_from_le_bytes(&42_u32.to_le_bytes()).unwrap(), 42);
506 assert_eq!(u64::try_from_le_bytes(&42_u64.to_le_bytes()).unwrap(), 42);
507 assert_eq!(i8::try_from_le_bytes(&42_i8.to_le_bytes()).unwrap(), 42);
508 assert_eq!(i16::try_from_le_bytes(&42_i16.to_le_bytes()).unwrap(), 42);
509 assert_eq!(i32::try_from_le_bytes(&42_i32.to_le_bytes()).unwrap(), 42);
510 assert_eq!(i64::try_from_le_bytes(&42_i64.to_le_bytes()).unwrap(), 42);
511 assert_eq!(
512 f16::try_from_le_bytes(&f16::from_f32(42.0).to_le_bytes()).unwrap(),
513 f16::from_f32(42.0)
514 );
515 assert_eq!(
516 f32::try_from_le_bytes(&42.0_f32.to_le_bytes()).unwrap(),
517 42.0
518 );
519 assert_eq!(
520 f64::try_from_le_bytes(&42.0_f64.to_le_bytes()).unwrap(),
521 42.0
522 );
523 }
524
525 #[test]
526 fn max_value_u64() {
527 assert_eq!(PType::U8.max_value_as_u64(), u8::MAX as u64);
528 assert_eq!(PType::U16.max_value_as_u64(), u16::MAX as u64);
529 assert_eq!(PType::U32.max_value_as_u64(), u32::MAX as u64);
530 assert_eq!(PType::U64.max_value_as_u64(), u64::MAX);
531 assert_eq!(PType::I8.max_value_as_u64(), i8::MAX as u64);
532 assert_eq!(PType::I16.max_value_as_u64(), i16::MAX as u64);
533 assert_eq!(PType::I32.max_value_as_u64(), i32::MAX as u64);
534 assert_eq!(PType::I64.max_value_as_u64(), i64::MAX as u64);
535 assert_eq!(PType::F16.max_value_as_u64(), 65504); assert_eq!(PType::F32.max_value_as_u64(), u64::MAX);
537 assert_eq!(PType::F64.max_value_as_u64(), u64::MAX);
538 }
539
540 #[test]
541 fn widths() {
542 assert_eq!(PType::U8.byte_width(), 1);
543 assert_eq!(PType::U16.byte_width(), 2);
544 assert_eq!(PType::U32.byte_width(), 4);
545 assert_eq!(PType::U64.byte_width(), 8);
546 assert_eq!(PType::I8.byte_width(), 1);
547 assert_eq!(PType::I16.byte_width(), 2);
548 assert_eq!(PType::I32.byte_width(), 4);
549 assert_eq!(PType::I64.byte_width(), 8);
550 assert_eq!(PType::F16.byte_width(), 2);
551 assert_eq!(PType::F32.byte_width(), 4);
552 assert_eq!(PType::F64.byte_width(), 8);
553
554 assert_eq!(PType::U8.bit_width(), 8);
555 assert_eq!(PType::U16.bit_width(), 16);
556 assert_eq!(PType::U32.bit_width(), 32);
557 assert_eq!(PType::U64.bit_width(), 64);
558 assert_eq!(PType::I8.bit_width(), 8);
559 assert_eq!(PType::I16.bit_width(), 16);
560 assert_eq!(PType::I32.bit_width(), 32);
561 assert_eq!(PType::I64.bit_width(), 64);
562 assert_eq!(PType::F16.bit_width(), 16);
563 assert_eq!(PType::F32.bit_width(), 32);
564 assert_eq!(PType::F64.bit_width(), 64);
565 }
566
567 #[test]
568 fn native_ptype_nan_handling() {
569 let a = f32::NAN;
570 let b = f32::NAN;
571 assert_ne!(a, b);
572 assert!(<f32 as NativePType>::is_nan(a));
573 assert!(<f32 as NativePType>::is_nan(b));
574 assert!(<f32 as NativePType>::is_eq(a, b));
575 assert!(<f32 as NativePType>::total_compare(a, b) == Ordering::Equal);
576 }
577
578 #[test]
579 fn to_signed() {
580 assert_eq!(PType::U8.to_signed(), PType::I8);
581 assert_eq!(PType::U16.to_signed(), PType::I16);
582 assert_eq!(PType::U32.to_signed(), PType::I32);
583 assert_eq!(PType::U64.to_signed(), PType::I64);
584 assert_eq!(PType::I8.to_signed(), PType::I8);
585 assert_eq!(PType::I16.to_signed(), PType::I16);
586 assert_eq!(PType::I32.to_signed(), PType::I32);
587 assert_eq!(PType::I64.to_signed(), PType::I64);
588 assert_eq!(PType::F16.to_signed(), PType::F16);
589 assert_eq!(PType::F32.to_signed(), PType::F32);
590 assert_eq!(PType::F64.to_signed(), PType::F64);
591 }
592
593 #[test]
594 fn to_unsigned() {
595 assert_eq!(PType::U8.to_unsigned(), PType::U8);
596 assert_eq!(PType::U16.to_unsigned(), PType::U16);
597 assert_eq!(PType::U32.to_unsigned(), PType::U32);
598 assert_eq!(PType::U64.to_unsigned(), PType::U64);
599 assert_eq!(PType::I8.to_unsigned(), PType::U8);
600 assert_eq!(PType::I16.to_unsigned(), PType::U16);
601 assert_eq!(PType::I32.to_unsigned(), PType::U32);
602 assert_eq!(PType::I64.to_unsigned(), PType::U64);
603 assert_eq!(PType::F16.to_unsigned(), PType::F16);
604 assert_eq!(PType::F32.to_unsigned(), PType::F32);
605 assert_eq!(PType::F64.to_unsigned(), PType::F64);
606 }
607
608 #[test]
609 fn to_dtype() {
610 assert_eq!(DType::from(PType::U8), Primitive(PType::U8, NonNullable));
611 assert_eq!(DType::from(PType::U16), Primitive(PType::U16, NonNullable));
612 assert_eq!(DType::from(PType::U32), Primitive(PType::U32, NonNullable));
613 assert_eq!(DType::from(PType::U64), Primitive(PType::U64, NonNullable));
614 assert_eq!(DType::from(PType::I8), Primitive(PType::I8, NonNullable));
615 assert_eq!(DType::from(PType::I16), Primitive(PType::I16, NonNullable));
616 assert_eq!(DType::from(PType::I32), Primitive(PType::I32, NonNullable));
617 assert_eq!(DType::from(PType::I64), Primitive(PType::I64, NonNullable));
618 assert_eq!(DType::from(PType::F16), Primitive(PType::F16, NonNullable));
619 assert_eq!(DType::from(PType::F32), Primitive(PType::F32, NonNullable));
620 assert_eq!(DType::from(PType::F64), Primitive(PType::F64, NonNullable));
621 }
622}