Eig, EigRef

Author: termoshtt

lax/src/eig.rs

@@ -57,26 +57,28 @@ pub struct EigWork<T: Scalar> {
     pub rwork: Option<Vec<MaybeUninit<T::Real>>>,
 }
 
+#[derive(Debug, Clone, PartialEq)]
+pub struct Eig<T: Scalar> {
+    pub eigs: Vec<T::Complex>,
+    pub vr: Option<Vec<T::Complex>>,
+    pub vl: Option<Vec<T::Complex>>,
+}
+
+#[derive(Debug, Clone, PartialEq)]
+pub struct EigRef<'work, T: Scalar> {
+    pub eigs: &'work [T::Complex],
+    pub vr: Option<&'work [T::Complex]>,
+    pub vl: Option<&'work [T::Complex]>,
+}
+
 pub trait EigWorkImpl: Sized {
     type Elem: Scalar;
     /// Create new working memory for eigenvalues compution.
     fn new(calc_v: bool, l: MatrixLayout) -> Result<Self>;
     /// Compute eigenvalues and vectors on this working memory.
-    fn calc<'work>(
-        &'work mut self,
-        a: &mut [Self::Elem],
-    ) -> Result<(
-        &'work [<Self::Elem as Scalar>::Complex],
-        Option<&'work [<Self::Elem as Scalar>::Complex]>,
-    )>;
+    fn calc<'work>(&'work mut self, a: &mut [Self::Elem]) -> Result<EigRef<'work, Self::Elem>>;
     /// Compute eigenvalues and vectors by consuming this working memory.
-    fn eval(
-        self,
-        a: &mut [Self::Elem],
-    ) -> Result<(
-        Vec<<Self::Elem as Scalar>::Complex>,
-        Option<Vec<<Self::Elem as Scalar>::Complex>>,
-    )>;
+    fn eval(self, a: &mut [Self::Elem]) -> Result<Eig<Self::Elem>>;
 }
 
 impl EigWorkImpl for EigWork<c64> {
@@ -137,7 +139,7 @@ impl EigWorkImpl for EigWork<c64> {
         })
     }
 
-    fn calc<'work>(&'work mut self, a: &mut [c64]) -> Result<(&'work [c64], Option<&'work [c64]>)> {
+    fn calc<'work>(&'work mut self, a: &mut [c64]) -> Result<EigRef<'work, c64>> {
         let lwork = self.work.len().to_i32().unwrap();
         let mut info = 0;
         unsafe {
@@ -173,15 +175,20 @@ impl EigWorkImpl for EigWork<c64> {
                 value.im = -value.im;
             }
         }
-        let v = match (self.vl.as_ref(), self.vr.as_ref()) {
-            (Some(v), None) | (None, Some(v)) => Some(unsafe { v.slice_assume_init_ref() }),
-            (None, None) => None,
-            _ => unreachable!(),
-        };
-        Ok((eigs, v))
+        Ok(EigRef {
+            eigs,
+            vl: self
+                .vl
+                .as_ref()
+                .map(|v| unsafe { v.slice_assume_init_ref() }),
+            vr: self
+                .vr
+                .as_ref()
+                .map(|v| unsafe { v.slice_assume_init_ref() }),
+        })
     }
 
-    fn eval(mut self, a: &mut [c64]) -> Result<(Vec<c64>, Option<Vec<c64>>)> {
+    fn eval(mut self, a: &mut [c64]) -> Result<Eig<c64>> {
         let lwork = self.work.len().to_i32().unwrap();
         let mut info = 0;
         unsafe {
@@ -212,12 +219,11 @@ impl EigWorkImpl for EigWork<c64> {
                 value.im = -value.im;
             }
         }
-        let v = match (self.vl, self.vr) {
-            (Some(v), None) | (None, Some(v)) => Some(unsafe { v.assume_init() }),
-            (None, None) => None,
-            _ => unreachable!(),
-        };
-        Ok((eigs, v))
+        Ok(Eig {
+            eigs,
+            vl: self.vl.map(|v| unsafe { v.assume_init() }),
+            vr: self.vr.map(|v| unsafe { v.assume_init() }),
+        })
     }
 }