From e7e82757c801c7e712ba9228daa2199b6299feb0 Mon Sep 17 00:00:00 2001
From: Toshiki Teramura <toshiki.teramura@gmail.com>
Date: Wed, 31 Aug 2022 22:50:44 +0900
Subject: [PATCH] Introduce alternative `vec_uninit2`, and replace `vec_uninit`
 in eig.rs

---
 lax/src/eig.rs | 56 +++++++++++++++++++++++++++++++-------------------
 lax/src/lib.rs | 12 +++++++++++
 2 files changed, 47 insertions(+), 21 deletions(-)

diff --git a/lax/src/eig.rs b/lax/src/eig.rs
index bf3b9f8c..2aeb2251 100644
--- a/lax/src/eig.rs
+++ b/lax/src/eig.rs
@@ -41,13 +41,14 @@ macro_rules! impl_eig_complex {
                 } else {
                     (EigenVectorFlag::Not, EigenVectorFlag::Not)
                 };
-                let mut eigs = unsafe { vec_uninit(n as usize) };
-                let mut rwork: Vec<Self::Real> = unsafe { vec_uninit(2 * n as usize) };
+                let mut eigs: Vec<MaybeUninit<Self>> = unsafe { vec_uninit2(n as usize) };
+                let mut rwork: Vec<MaybeUninit<Self::Real>> =
+                    unsafe { vec_uninit2(2 * n as usize) };
 
-                let mut vl: Option<Vec<Self>> =
-                    jobvl.then(|| unsafe { vec_uninit((n * n) as usize) });
-                let mut vr: Option<Vec<Self>> =
-                    jobvr.then(|| unsafe { vec_uninit((n * n) as usize) });
+                let mut vl: Option<Vec<MaybeUninit<Self>>> =
+                    jobvl.then(|| unsafe { vec_uninit2((n * n) as usize) });
+                let mut vr: Option<Vec<MaybeUninit<Self>>> =
+                    jobvr.then(|| unsafe { vec_uninit2((n * n) as usize) });
 
                 // calc work size
                 let mut info = 0;
@@ -74,7 +75,7 @@ macro_rules! impl_eig_complex {
 
                 // actal ev
                 let lwork = work_size[0].to_usize().unwrap();
-                let mut work: Vec<Self> = unsafe { vec_uninit(lwork) };
+                let mut work: Vec<MaybeUninit<Self>> = unsafe { vec_uninit2(lwork) };
                 let lwork = lwork as i32;
                 unsafe {
                     $ev(
@@ -99,11 +100,19 @@ macro_rules! impl_eig_complex {
                 // Hermite conjugate
                 if jobvl.is_calc() {
                     for c in vl.as_mut().unwrap().iter_mut() {
-                        c.im = -c.im
+                        let value = unsafe { c.assume_init_mut() };
+                        value.im = -value.im;
                     }
                 }
 
-                Ok((eigs, vr.or(vl).unwrap_or(Vec::new())))
+                unsafe {
+                    Ok((
+                        eigs.assume_init(),
+                        vr.map(|v| v.assume_init())
+                            .or(vl.map(|v| v.assume_init()))
+                            .unwrap_or(Vec::new()),
+                    ))
+                }
             }
         }
     };
@@ -145,13 +154,13 @@ macro_rules! impl_eig_real {
                 } else {
                     (EigenVectorFlag::Not, EigenVectorFlag::Not)
                 };
-                let mut eig_re: Vec<Self> = unsafe { vec_uninit(n as usize) };
-                let mut eig_im: Vec<Self> = unsafe { vec_uninit(n as usize) };
+                let mut eig_re: Vec<MaybeUninit<Self>> = unsafe { vec_uninit2(n as usize) };
+                let mut eig_im: Vec<MaybeUninit<Self>> = unsafe { vec_uninit2(n as usize) };
 
-                let mut vl: Option<Vec<Self>> =
-                    jobvl.then(|| unsafe { vec_uninit((n * n) as usize) });
-                let mut vr: Option<Vec<Self>> =
-                    jobvr.then(|| unsafe { vec_uninit((n * n) as usize) });
+                let mut vl: Option<Vec<MaybeUninit<Self>>> =
+                    jobvl.then(|| unsafe { vec_uninit2((n * n) as usize) });
+                let mut vr: Option<Vec<MaybeUninit<Self>>> =
+                    jobvr.then(|| unsafe { vec_uninit2((n * n) as usize) });
 
                 // calc work size
                 let mut info = 0;
@@ -178,7 +187,7 @@ macro_rules! impl_eig_real {
 
                 // actual ev
                 let lwork = work_size[0].to_usize().unwrap();
-                let mut work: Vec<Self> = unsafe { vec_uninit(lwork) };
+                let mut work: Vec<MaybeUninit<Self>> = unsafe { vec_uninit2(lwork) };
                 let lwork = lwork as i32;
                 unsafe {
                     $ev(
@@ -200,6 +209,11 @@ macro_rules! impl_eig_real {
                 };
                 info.as_lapack_result()?;
 
+                let eig_re = unsafe { eig_re.assume_init() };
+                let eig_im = unsafe { eig_im.assume_init() };
+                let vl = unsafe { vl.map(|v| v.assume_init()) };
+                let vr = unsafe { vr.map(|v| v.assume_init()) };
+
                 // reconstruct eigenvalues
                 let eigs: Vec<Self::Complex> = eig_re
                     .iter()
@@ -228,14 +242,14 @@ macro_rules! impl_eig_real {
 
                 let n = n as usize;
                 let v = vr.or(vl).unwrap();
-                let mut eigvecs = unsafe { vec_uninit(n * n) };
+                let mut eigvecs: Vec<MaybeUninit<Self::Complex>> = unsafe { vec_uninit2(n * n) };
                 let mut col = 0;
                 while col < n {
                     if eig_im[col] == 0. {
                         // The corresponding eigenvalue is real.
                         for row in 0..n {
                             let re = v[row + col * n];
-                            eigvecs[row + col * n] = Self::complex(re, 0.);
+                            eigvecs[row + col * n].write(Self::complex(re, 0.));
                         }
                         col += 1;
                     } else {
@@ -247,14 +261,14 @@ macro_rules! impl_eig_real {
                             if jobvl.is_calc() {
                                 im = -im;
                             }
-                            eigvecs[row + col * n] = Self::complex(re, im);
-                            eigvecs[row + (col + 1) * n] = Self::complex(re, -im);
+                            eigvecs[row + col * n].write(Self::complex(re, im));
+                            eigvecs[row + (col + 1) * n].write(Self::complex(re, -im));
                         }
                         col += 2;
                     }
                 }
 
-                Ok((eigs, eigvecs))
+                unsafe { Ok((eigs, eigvecs.assume_init())) }
             }
         }
     };
diff --git a/lax/src/lib.rs b/lax/src/lib.rs
index ac7b8256..370e02ed 100644
--- a/lax/src/lib.rs
+++ b/lax/src/lib.rs
@@ -278,3 +278,15 @@ unsafe fn vec_uninit<T: Sized>(n: usize) -> Vec<T> {
     v.set_len(n);
     v
 }
+
+/// Create a vector without initialization
+///
+/// Safety
+/// ------
+/// - Memory is not initialized. Do not read the memory before write.
+///
+unsafe fn vec_uninit2<T: Sized>(n: usize) -> Vec<MaybeUninit<T>> {
+    let mut v = Vec::with_capacity(n);
+    v.set_len(n);
+    v
+}