Add neg for host vector, as well as x/y/z/w axis functions.

LuisaGroup · Feb 25, 2024 · f023ae5 · f023ae5
1 parent 716c603
commit f023ae5
Show file tree

Hide file tree

Showing 2 changed files with 96 additions and 2 deletions.
diff --git a/luisa_compute/src/lang/types/core.rs b/luisa_compute/src/lang/types/core.rs
@@ -1,3 +1,5 @@
+use std::ops::Neg;
+
 use crate::lang::soa::SoaBuffer;
 
 use super::*;
@@ -496,5 +498,5 @@ impls!(Numeric for f16, f32, f64, i8, i16, i32, i64, u8, u16, u32, u64);
 pub trait Floating: Numeric {}
 impls!(Floating for f16, f32, f64);
 
-pub trait Signed: Numeric {}
+pub trait Signed: Numeric + Neg<Output = Self> {}
 impls!(Signed for f16, f32, f64, i8, i16, i32, i64);
diff --git a/luisa_compute/src/lang/types/vector/impls.rs b/luisa_compute/src/lang/types/vector/impls.rs
@@ -1,6 +1,6 @@
 use super::*;
 use crate::lang::index::IntoIndex;
-use std::ops::Index;
+use std::ops::{Index, Neg};
 
 impl<T: VectorAlign<N>, const N: usize> From<[T; N]> for Vector<T, N> {
     fn from(elements: [T; N]) -> Self {
@@ -115,6 +115,98 @@ impl_sized!(Vec2(2), VectorExprProxy2, VectorVarProxy2: x, y);
 impl_sized!(Vec3(3), VectorExprProxy3, VectorVarProxy3: x, y, z);
 impl_sized!(Vec4(4), VectorExprProxy4, VectorVarProxy4: x, y, z, w);
 
+pub trait ZeroOne {
+    fn zero() -> Self;
+    fn one() -> Self;
+}
+impl ZeroOne for bool {
+    fn zero() -> Self {
+        false
+    }
+    fn one() -> Self {
+        true
+    }
+}
+macro_rules! zero_one {
+    ($($t:ty),*) => {
+        $(
+            impl ZeroOne for $t {
+                fn zero() -> Self {
+                    0
+                }
+                fn one() -> Self {
+                    1
+                }
+            }
+        )*
+    }
+}
+zero_one!(u8, u16, u32, u64, i8, i16, i32, i64);
+impl ZeroOne for f16 {
+    fn zero() -> Self {
+        f16::ZERO
+    }
+    fn one() -> Self {
+        f16::ONE
+    }
+}
+impl ZeroOne for f32 {
+    fn zero() -> Self {
+        0.0
+    }
+    fn one() -> Self {
+        1.0
+    }
+}
+impl ZeroOne for f64 {
+    fn zero() -> Self {
+        0.0
+    }
+    fn one() -> Self {
+        1.0
+    }
+}
+
+impl<T: ZeroOne + VectorAlign<2>> Vector<T, 2> {
+    pub fn x() -> Self {
+        Self::new(T::one(), T::zero())
+    }
+    pub fn y() -> Self {
+        Self::new(T::zero(), T::one())
+    }
+}
+impl<T: ZeroOne + VectorAlign<3>> Vector<T, 3> {
+    pub fn x() -> Self {
+        Self::new(T::one(), T::zero(), T::zero())
+    }
+    pub fn y() -> Self {
+        Self::new(T::zero(), T::one(), T::zero())
+    }
+    pub fn z() -> Self {
+        Self::new(T::zero(), T::zero(), T::one())
+    }
+}
+impl<T: ZeroOne + VectorAlign<4>> Vector<T, 4> {
+    pub fn x() -> Self {
+        Self::new(T::one(), T::zero(), T::zero(), T::zero())
+    }
+    pub fn y() -> Self {
+        Self::new(T::zero(), T::one(), T::zero(), T::zero())
+    }
+    pub fn z() -> Self {
+        Self::new(T::zero(), T::zero(), T::one(), T::zero())
+    }
+    pub fn w() -> Self {
+        Self::new(T::zero(), T::zero(), T::zero(), T::one())
+    }
+}
+impl<T: Neg<Output = T> + VectorAlign<N>, const N: usize> Neg for Vector<T, N> {
+    type Output = Self;
+    fn neg(self) -> Self {
+        Self::from(self.elements.map(|x| -x))
+    }
+}
+
 pub trait VectorExprProxy {
     const N: usize;
     type T: Primitive;