trx_rs/
dtype.rs

1use bytemuck::{Pod, Zeroable};
2use half::f16;
3use std::fmt;
4
5use crate::error::{Result, TrxError};
6
7/// Supported element data types in TRX files.
8#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
9pub enum DType {
10    Float16,
11    Float32,
12    Float64,
13    Int8,
14    Int16,
15    Int32,
16    Int64,
17    UInt8,
18    UInt16,
19    UInt32,
20    UInt64,
21}
22
23impl DType {
24    /// Size of one element in bytes.
25    pub fn size_of(self) -> usize {
26        match self {
27            DType::Int8 | DType::UInt8 => 1,
28            DType::Float16 | DType::Int16 | DType::UInt16 => 2,
29            DType::Float32 | DType::Int32 | DType::UInt32 => 4,
30            DType::Float64 | DType::Int64 | DType::UInt64 => 8,
31        }
32    }
33
34    /// Canonical string name as used in TRX filenames.
35    pub fn name(self) -> &'static str {
36        match self {
37            DType::Float16 => "float16",
38            DType::Float32 => "float32",
39            DType::Float64 => "float64",
40            DType::Int8 => "int8",
41            DType::Int16 => "int16",
42            DType::Int32 => "int32",
43            DType::Int64 => "int64",
44            DType::UInt8 => "uint8",
45            DType::UInt16 => "uint16",
46            DType::UInt32 => "uint32",
47            DType::UInt64 => "uint64",
48        }
49    }
50
51    /// Parse a dtype string (e.g. `"float32"`).
52    pub fn parse(s: &str) -> Result<Self> {
53        match s {
54            "float16" => Ok(DType::Float16),
55            "float32" => Ok(DType::Float32),
56            "float64" => Ok(DType::Float64),
57            "int8" => Ok(DType::Int8),
58            "int16" => Ok(DType::Int16),
59            "int32" => Ok(DType::Int32),
60            "int64" => Ok(DType::Int64),
61            "uint8" => Ok(DType::UInt8),
62            "uint16" => Ok(DType::UInt16),
63            "uint32" => Ok(DType::UInt32),
64            "uint64" => Ok(DType::UInt64),
65            _ => Err(TrxError::DType(s.to_string())),
66        }
67    }
68
69    /// Whether this is a floating-point type.
70    pub fn is_float(self) -> bool {
71        matches!(self, DType::Float16 | DType::Float32 | DType::Float64)
72    }
73}
74
75impl fmt::Display for DType {
76    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
77        f.write_str(self.name())
78    }
79}
80
81/// Trait for scalar types that can be stored as TRX position coordinates.
82///
83/// Implementors must be [`Pod`] + [`Zeroable`] (for zero-copy casts) and
84/// carry their [`DType`] as an associated constant.
85pub trait TrxScalar: Pod + Zeroable + Copy + 'static + fmt::Debug {
86    const DTYPE: DType;
87
88    /// Convert to f32 for operations that need floating-point comparisons.
89    fn to_f32(self) -> f32;
90
91    /// Convert to f64 for operations that need higher precision.
92    fn to_f64(self) -> f64;
93}
94
95impl TrxScalar for f32 {
96    const DTYPE: DType = DType::Float32;
97    fn to_f32(self) -> f32 {
98        self
99    }
100    fn to_f64(self) -> f64 {
101        self as f64
102    }
103}
104
105impl TrxScalar for f64 {
106    const DTYPE: DType = DType::Float64;
107    fn to_f32(self) -> f32 {
108        self as f32
109    }
110    fn to_f64(self) -> f64 {
111        self
112    }
113}
114
115impl TrxScalar for f16 {
116    const DTYPE: DType = DType::Float16;
117    fn to_f32(self) -> f32 {
118        f16::to_f32(self)
119    }
120    fn to_f64(self) -> f64 {
121        f16::to_f64(self)
122    }
123}
124
125impl TrxScalar for i8 {
126    const DTYPE: DType = DType::Int8;
127    fn to_f32(self) -> f32 {
128        self as f32
129    }
130    fn to_f64(self) -> f64 {
131        self as f64
132    }
133}
134impl TrxScalar for i16 {
135    const DTYPE: DType = DType::Int16;
136    fn to_f32(self) -> f32 {
137        self as f32
138    }
139    fn to_f64(self) -> f64 {
140        self as f64
141    }
142}
143impl TrxScalar for i32 {
144    const DTYPE: DType = DType::Int32;
145    fn to_f32(self) -> f32 {
146        self as f32
147    }
148    fn to_f64(self) -> f64 {
149        self as f64
150    }
151}
152impl TrxScalar for i64 {
153    const DTYPE: DType = DType::Int64;
154    fn to_f32(self) -> f32 {
155        self as f32
156    }
157    fn to_f64(self) -> f64 {
158        self as f64
159    }
160}
161impl TrxScalar for u8 {
162    const DTYPE: DType = DType::UInt8;
163    fn to_f32(self) -> f32 {
164        self as f32
165    }
166    fn to_f64(self) -> f64 {
167        self as f64
168    }
169}
170impl TrxScalar for u16 {
171    const DTYPE: DType = DType::UInt16;
172    fn to_f32(self) -> f32 {
173        self as f32
174    }
175    fn to_f64(self) -> f64 {
176        self as f64
177    }
178}
179impl TrxScalar for u32 {
180    const DTYPE: DType = DType::UInt32;
181    fn to_f32(self) -> f32 {
182        self as f32
183    }
184    fn to_f64(self) -> f64 {
185        self as f64
186    }
187}
188impl TrxScalar for u64 {
189    const DTYPE: DType = DType::UInt64;
190    fn to_f32(self) -> f32 {
191        self as f32
192    }
193    fn to_f64(self) -> f64 {
194        self as f64
195    }
196}
197
198#[cfg(test)]
199mod tests {
200    use super::*;
201
202    #[test]
203    fn dtype_round_trip() {
204        for dt in [
205            DType::Float16,
206            DType::Float32,
207            DType::Float64,
208            DType::Int8,
209            DType::Int16,
210            DType::Int32,
211            DType::Int64,
212            DType::UInt8,
213            DType::UInt16,
214            DType::UInt32,
215            DType::UInt64,
216        ] {
217            assert_eq!(DType::parse(dt.name()).unwrap(), dt);
218        }
219    }
220
221    #[test]
222    fn dtype_sizes() {
223        assert_eq!(DType::Float32.size_of(), 4);
224        assert_eq!(DType::Float64.size_of(), 8);
225        assert_eq!(DType::Float16.size_of(), 2);
226        assert_eq!(DType::UInt8.size_of(), 1);
227    }
228
229    #[test]
230    fn dtype_parse_invalid() {
231        assert!(DType::parse("complex128").is_err());
232    }
233}
trx_rs/dtype.rs

trx_rs/
dtype.rs
