src/traits.rs

   1 use crate::{f16::F16, ieee754::FloatEncoding, scalar::Scalar};
   2 use core::ops::{
   3     Add, AddAssign, BitAnd, BitAndAssign, BitOr, BitOrAssign, BitXor, BitXorAssign, Div, DivAssign,
   4     Mul, MulAssign, Neg, Not, Rem, RemAssign, Shl, ShlAssign, Shr, ShrAssign, Sub, SubAssign,
   5 };
   6
   7 /// reference used to build IR for Kazan; an empty type for `core::simd`
   8 pub trait Context: Copy {
   9     vector_math_proc_macro::make_context_types!();
  10     fn make<T: Make<Context = Self>>(self, v: T::Prim) -> T {
  11         T::make(self, v)
  12     }
  13 }
  14
  15 pub trait Make: Copy {
  16     type Prim: Copy;
  17     type Context: Context;
  18     fn ctx(self) -> Self::Context;
  19     fn make(ctx: Self::Context, v: Self::Prim) -> Self;
  20 }
  21
  22 pub trait ConvertTo<T> {
  23     fn to(self) -> T;
  24 }
  25
  26 impl<T> ConvertTo<T> for T {
  27     fn to(self) -> T {
  28         self
  29     }
  30 }
  31
  32 macro_rules! impl_convert_to_using_as {
  33     ($first:ident $(, $ty:ident)*) => {
  34         $(
  35             impl ConvertTo<$first> for $ty {
  36                 fn to(self) -> $first {
  37                     self as $first
  38                 }
  39             }
  40             impl ConvertTo<$ty> for $first {
  41                 fn to(self) -> $ty {
  42                     self as $ty
  43                 }
  44             }
  45         )*
  46         impl_convert_to_using_as![$($ty),*];
  47     };
  48     () => {
  49     };
  50 }
  51
  52 impl_convert_to_using_as![u8, i8, u16, i16, u32, i32, u64, i64, f32, f64];
  53
  54 pub trait Number:
  55     Compare
  56     + Add<Output = Self>
  57     + Sub<Output = Self>
  58     + Mul<Output = Self>
  59     + Div<Output = Self>
  60     + Rem<Output = Self>
  61     + AddAssign
  62     + SubAssign
  63     + MulAssign
  64     + DivAssign
  65     + RemAssign
  66 {
  67 }
  68
  69 impl<T> Number for T where
  70     T: Compare
  71         + Add<Output = Self>
  72         + Sub<Output = Self>
  73         + Mul<Output = Self>
  74         + Div<Output = Self>
  75         + Rem<Output = Self>
  76         + AddAssign
  77         + SubAssign
  78         + MulAssign
  79         + DivAssign
  80         + RemAssign
  81 {
  82 }
  83
  84 pub trait BitOps:
  85     Copy
  86     + BitAnd<Output = Self>
  87     + BitOr<Output = Self>
  88     + BitXor<Output = Self>
  89     + Not<Output = Self>
  90     + BitAndAssign
  91     + BitOrAssign
  92     + BitXorAssign
  93 {
  94 }
  95
  96 impl<T> BitOps for T where
  97     T: Copy
  98         + BitAnd<Output = Self>
  99         + BitOr<Output = Self>
 100         + BitXor<Output = Self>
 101         + Not<Output = Self>
 102         + BitAndAssign
 103         + BitOrAssign
 104         + BitXorAssign
 105 {
 106 }
 107
 108 pub trait Int:
 109     Number + BitOps + Shl<Output = Self> + Shr<Output = Self> + ShlAssign + ShrAssign
 110 {
 111     fn leading_zeros(self) -> Self;
 112     fn leading_ones(self) -> Self {
 113         self.not().leading_zeros()
 114     }
 115     fn trailing_zeros(self) -> Self;
 116     fn trailing_ones(self) -> Self {
 117         self.not().trailing_zeros()
 118     }
 119     fn count_zeros(self) -> Self {
 120         self.not().count_ones()
 121     }
 122     fn count_ones(self) -> Self;
 123 }
 124
 125 pub trait UInt: Int {}
 126
 127 pub trait SInt: Int + Neg<Output = Self> {}
 128
 129 macro_rules! impl_int {
 130     ($ty:ident) => {
 131         impl Int for $ty {
 132             fn leading_zeros(self) -> Self {
 133                 self.leading_zeros() as Self
 134             }
 135             fn leading_ones(self) -> Self {
 136                 self.leading_ones() as Self
 137             }
 138             fn trailing_zeros(self) -> Self {
 139                 self.trailing_zeros() as Self
 140             }
 141             fn trailing_ones(self) -> Self {
 142                 self.trailing_ones() as Self
 143             }
 144             fn count_zeros(self) -> Self {
 145                 self.count_zeros() as Self
 146             }
 147             fn count_ones(self) -> Self {
 148                 self.count_ones() as Self
 149             }
 150         }
 151     };
 152 }
 153
 154 macro_rules! impl_uint {
 155     ($($ty:ident),*) => {
 156         $(
 157             impl_int!($ty);
 158             impl UInt for $ty {}
 159         )*
 160     };
 161 }
 162
 163 impl_uint![u8, u16, u32, u64];
 164
 165 macro_rules! impl_sint {
 166     ($($ty:ident),*) => {
 167         $(
 168             impl_int!($ty);
 169             impl SInt for $ty {}
 170         )*
 171     };
 172 }
 173
 174 impl_sint![i8, i16, i32, i64];
 175
 176 pub trait Float: Number + Neg<Output = Self> {
 177     type FloatEncoding: FloatEncoding + Make<Context = Scalar, Prim = <Self as Make>::Prim>;
 178     type BitsType: UInt
 179         + Make<Context = Self::Context, Prim = <Self::FloatEncoding as Float>::BitsType>
 180         + ConvertTo<Self::SignedBitsType>
 181         + Compare<Bool = Self::Bool>;
 182     type SignedBitsType: SInt
 183         + Make<Context = Self::Context, Prim = <Self::FloatEncoding as Float>::SignedBitsType>
 184         + ConvertTo<Self::BitsType>
 185         + Compare<Bool = Self::Bool>;
 186     fn abs(self) -> Self;
 187     fn trunc(self) -> Self;
 188     fn ceil(self) -> Self;
 189     fn floor(self) -> Self;
 190     fn round(self) -> Self;
 191     #[cfg(feature = "fma")]
 192     fn fma(self, a: Self, b: Self) -> Self;
 193     fn is_nan(self) -> Self::Bool {
 194         self.ne(self)
 195     }
 196     fn is_infinite(self) -> Self::Bool {
 197         self.abs().eq(Self::infinity(self.ctx()))
 198     }
 199     fn infinity(ctx: Self::Context) -> Self {
 200         Self::from_bits(ctx.make(Self::FloatEncoding::INFINITY_BITS))
 201     }
 202     fn nan(ctx: Self::Context) -> Self {
 203         Self::from_bits(ctx.make(Self::FloatEncoding::NAN_BITS))
 204     }
 205     fn is_finite(self) -> Self::Bool;
 206     fn is_zero_or_subnormal(self) -> Self::Bool {
 207         self.extract_exponent_field().eq(self
 208             .ctx()
 209             .make(Self::FloatEncoding::ZERO_SUBNORMAL_EXPONENT))
 210     }
 211     fn from_bits(v: Self::BitsType) -> Self;
 212     fn to_bits(self) -> Self::BitsType;
 213     fn extract_exponent_field(self) -> Self::BitsType {
 214         let mask = self.ctx().make(Self::FloatEncoding::EXPONENT_FIELD_MASK);
 215         let shift = self.ctx().make(Self::FloatEncoding::EXPONENT_FIELD_SHIFT);
 216         (self.to_bits() & mask) >> shift
 217     }
 218     fn extract_exponent_unbiased(self) -> Self::SignedBitsType {
 219         Self::sub_exponent_bias(self.extract_exponent_field())
 220     }
 221     fn extract_mantissa_field(self) -> Self::BitsType {
 222         let mask = self.ctx().make(Self::FloatEncoding::MANTISSA_FIELD_MASK);
 223         self.to_bits() & mask
 224     }
 225     fn sub_exponent_bias(exponent_field: Self::BitsType) -> Self::SignedBitsType {
 226         exponent_field.to()
 227             - exponent_field
 228                 .ctx()
 229                 .make(Self::FloatEncoding::EXPONENT_BIAS_SIGNED)
 230     }
 231     fn add_exponent_bias(exponent: Self::SignedBitsType) -> Self::BitsType {
 232         (exponent
 233             + exponent
 234                 .ctx()
 235                 .make(Self::FloatEncoding::EXPONENT_BIAS_SIGNED))
 236         .to()
 237     }
 238 }
 239
 240 macro_rules! impl_float {
 241     ($ty:ty, $bits:ty, $signed_bits:ty) => {
 242         impl Float for $ty {
 243             type FloatEncoding = $ty;
 244             type BitsType = $bits;
 245             type SignedBitsType = $signed_bits;
 246             fn abs(self) -> Self {
 247                 #[cfg(feature = "std")]
 248                 return self.abs();
 249                 #[cfg(not(feature = "std"))]
 250                 todo!();
 251             }
 252             fn trunc(self) -> Self {
 253                 #[cfg(feature = "std")]
 254                 return self.trunc();
 255                 #[cfg(not(feature = "std"))]
 256                 todo!();
 257             }
 258             fn ceil(self) -> Self {
 259                 #[cfg(feature = "std")]
 260                 return self.ceil();
 261                 #[cfg(not(feature = "std"))]
 262                 todo!();
 263             }
 264             fn floor(self) -> Self {
 265                 #[cfg(feature = "std")]
 266                 return self.floor();
 267                 #[cfg(not(feature = "std"))]
 268                 todo!();
 269             }
 270             fn round(self) -> Self {
 271                 #[cfg(feature = "std")]
 272                 return self.round();
 273                 #[cfg(not(feature = "std"))]
 274                 todo!();
 275             }
 276             #[cfg(feature = "fma")]
 277             fn fma(self, a: Self, b: Self) -> Self {
 278                 self.mul_add(a, b)
 279             }
 280             fn is_nan(self) -> Self::Bool {
 281                 self.is_nan()
 282             }
 283             fn is_infinite(self) -> Self::Bool {
 284                 self.is_infinite()
 285             }
 286             fn is_finite(self) -> Self::Bool {
 287                 self.is_finite()
 288             }
 289             fn from_bits(v: Self::BitsType) -> Self {
 290                 <$ty>::from_bits(v)
 291             }
 292             fn to_bits(self) -> Self::BitsType {
 293                 self.to_bits()
 294             }
 295         }
 296     };
 297 }
 298
 299 impl_float!(f32, u32, i32);
 300 impl_float!(f64, u64, i64);
 301
 302 pub trait Bool: Make + BitOps {}
 303
 304 impl Bool for bool {}
 305
 306 pub trait Select<T>: Bool {
 307     fn select(self, true_v: T, false_v: T) -> T;
 308 }
 309
 310 impl<T> Select<T> for bool {
 311     fn select(self, true_v: T, false_v: T) -> T {
 312         if self {
 313             true_v
 314         } else {
 315             false_v
 316         }
 317     }
 318 }
 319 pub trait Compare: Make {
 320     type Bool: Bool + Select<Self>;
 321     fn eq(self, rhs: Self) -> Self::Bool;
 322     fn ne(self, rhs: Self) -> Self::Bool;
 323     fn lt(self, rhs: Self) -> Self::Bool;
 324     fn gt(self, rhs: Self) -> Self::Bool;
 325     fn le(self, rhs: Self) -> Self::Bool;
 326     fn ge(self, rhs: Self) -> Self::Bool;
 327 }
 328
 329 macro_rules! impl_compare_using_partial_cmp {
 330     ($($ty:ty),*) => {
 331         $(
 332             impl Compare for $ty {
 333                 type Bool = bool;
 334                 fn eq(self, rhs: Self) -> Self::Bool {
 335                     self == rhs
 336                 }
 337                 fn ne(self, rhs: Self) -> Self::Bool {
 338                     self != rhs
 339                 }
 340                 fn lt(self, rhs: Self) -> Self::Bool {
 341                     self < rhs
 342                 }
 343                 fn gt(self, rhs: Self) -> Self::Bool {
 344                     self > rhs
 345                 }
 346                 fn le(self, rhs: Self) -> Self::Bool {
 347                     self <= rhs
 348                 }
 349                 fn ge(self, rhs: Self) -> Self::Bool {
 350                     self >= rhs
 351                 }
 352             }
 353         )*
 354     };
 355 }
 356
 357 impl_compare_using_partial_cmp![bool, u8, i8, u16, i16, F16, u32, i32, f32, u64, i64, f64];