Split slice

Author: termoshtt

src/traits.rs

@@ -45,27 +45,68 @@ impl<A, S, Sq, Sr> QR<ArrayBase<Sq, Ix2>, ArrayBase<Sr, Ix2>> for ArrayBase<S, I
           S: DataMut<Elem = A>,
           Sq: DataOwned<Elem = A> + DataMut,
           Sr: DataOwned<Elem = A> + DataMut
+{
+    fn qr(mut self) -> Result<(ArrayBase<Sq, Ix2>, ArrayBase<Sr, Ix2>)> {
+        (&mut self).qr()
+    }
+}
+
+fn take_slice<A, S1, S2>(a: &ArrayBase<S1, Ix2>, n: usize, m: usize) -> ArrayBase<S2, Ix2>
+    where A: Copy,
+          S1: Data<Elem = A>,
+          S2: DataMut<Elem = A> + DataOwned
+{
+    let av = a.slice(s![..n as isize, ..m as isize]);
+    let mut a = unsafe { ArrayBase::uninitialized((n, m)) };
+    a.assign(&av);
+    a
+}
+
+fn take_slice_upper<A, S1, S2>(a: &ArrayBase<S1, Ix2>, n: usize, m: usize) -> ArrayBase<S2, Ix2>
+    where A: Copy + Zero,
+          S1: Data<Elem = A>,
+          S2: DataMut<Elem = A> + DataOwned
+{
+    let av = a.slice(s![..n as isize, ..m as isize]);
+    let mut a = unsafe { ArrayBase::uninitialized((n, m)) };
+    for ((i, j), val) in a.indexed_iter_mut() {
+        *val = if i <= j { av[(i, j)] } else { A::zero() };
+    }
+    a
+}
+
+impl<'a, A, S, Sq, Sr> QR<ArrayBase<Sq, Ix2>, ArrayBase<Sr, Ix2>> for &'a mut ArrayBase<S, Ix2>
+    where A: LapackScalar + Copy + Zero,
+          S: DataMut<Elem = A>,
+          Sq: DataOwned<Elem = A> + DataMut,
+          Sr: DataOwned<Elem = A> + DataMut
 {
     fn qr(mut self) -> Result<(ArrayBase<Sq, Ix2>, ArrayBase<Sr, Ix2>)> {
         let n = self.rows();
         let m = self.cols();
         let k = ::std::cmp::min(n, m);
         let l = self.layout()?;
-        // calc QR decomposition
         let r = A::qr(l, self.as_allocated_mut()?)?;
         let r: Array2<_> = reconstruct(l, r)?;
         let q = self;
-        // get slice
-        let qv = q.slice(s![..n as isize, ..k as isize]);
-        let mut q = unsafe { ArrayBase::uninitialized((n, k)) };
-        q.assign(&qv);
-        let rv = r.slice(s![..k as isize, ..m as isize]);
-        let mut r = ArrayBase::zeros((k, m));
-        for ((i, j), val) in r.indexed_iter_mut() {
-            if i <= j {
-                *val = rv[(i, j)];
-            }
-        }
-        Ok((q, r))
+        Ok((take_slice(q, n, k), take_slice_upper(&r, k, m)))
+    }
+}
+
+impl<'a, A, S, Sq, Sr> QR<ArrayBase<Sq, Ix2>, ArrayBase<Sr, Ix2>> for &'a ArrayBase<S, Ix2>
+    where A: LapackScalar + Copy + Zero,
+          S: Data<Elem = A>,
+          Sq: DataOwned<Elem = A> + DataMut,
+          Sr: DataOwned<Elem = A> + DataMut
+{
+    fn qr(self) -> Result<(ArrayBase<Sq, Ix2>, ArrayBase<Sr, Ix2>)> {
+        let n = self.rows();
+        let m = self.cols();
+        let k = ::std::cmp::min(n, m);
+        let l = self.layout()?;
+        let mut q = self.to_owned();
+        let r = A::qr(l, q.as_allocated_mut()?)?;
+        let r: Array2<_> = reconstruct(l, r)?;
+        Ok((take_slice(&q, n, k), take_slice_upper(&r, k, m)))
     }
 }