Reduce code duplication for allocating a storage from a raw iterator.

dimforge · Jul 30, 2022 · 89767ee · 89767ee
1 parent 59b01e9
commit 89767ee
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 86 deletions.
diff --git a/src/base/allocator.rs b/src/base/allocator.rs
@@ -42,12 +42,40 @@ pub trait Allocator<T, R: Dim, C: Dim = U1>: Any + Sized {
         iter: I,
     ) -> Self::Buffer;
 
+    #[inline]
     /// Allocates a buffer initialized with the content of the given row-major order iterator.
     fn allocate_from_row_iterator<I: IntoIterator<Item = T>>(
         nrows: R,
         ncols: C,
         iter: I,
-    ) -> Self::Buffer;
+    ) -> Self::Buffer {
+        let mut res = Self::allocate_uninit(nrows, ncols);
+        let mut count = 0;
+
+        unsafe {
+            // OK because the allocated buffer is guaranteed to be contiguous.
+            let res_ptr = res.as_mut_slice_unchecked();
+
+            for (k, e) in iter
+                .into_iter()
+                .take(ncols.value() * nrows.value())
+                .enumerate()
+            {
+                let i = k / ncols.value();
+                let j = k % ncols.value();
+                // result[(i, j)] = e;
+                *res_ptr.get_unchecked_mut(i + j * nrows.value()) = MaybeUninit::new(e);
+                count += 1;
+            }
+
+            assert!(
+                count == nrows.value() * ncols.value(),
+                "Matrix init. from row iterator: iterator not long enough."
+            );
+
+            <Self as Allocator<T, R, C>>::assume_init(res)
+        }
+    }
 }
 
 /// A matrix reallocator. Changes the size of the memory buffer that initially contains (`RFrom` ×

diff --git a/src/base/default_allocator.rs b/src/base/default_allocator.rs
@@ -80,38 +80,6 @@ impl<T: Scalar, const R: usize, const C: usize> Allocator<T, Const<R>, Const<C>>
         //         yielded enough elements to initialize our matrix.
         unsafe { <Self as Allocator<T, Const<R>, Const<C>>>::assume_init(res) }
     }
-
-    #[inline]
-    fn allocate_from_row_iterator<I: IntoIterator<Item = T>>(
-        nrows: Const<R>,
-        ncols: Const<C>,
-        iter: I,
-    ) -> Self::Buffer {
-        let mut res = Self::allocate_uninit(nrows, ncols);
-        let mut count = 0;
-        let res_ptr = res.as_mut_slice();
-
-        for (i, e) in iter
-            .into_iter()
-            .take(ncols.value() * nrows.value())
-            .enumerate()
-        {
-            unsafe {
-                *res_ptr
-                    .get_unchecked_mut((i % ncols.value()) * nrows.value() + i / ncols.value()) =
-                    MaybeUninit::new(e);
-            }
-            // res_ptr[(i % ncols.value()) * nrows.value() + i / ncols.value()] = e;
-            count += 1;
-        }
-
-        assert!(
-            count == nrows.value() * ncols.value(),
-            "Matrix init. from row iterator: iterator not long enough."
-        );
-
-        unsafe { <Self as Allocator<T, Const<R>, Const<C>>>::assume_init(res) }
-    }
 }
 
 // Dynamic - Static
@@ -160,32 +128,6 @@ impl<T: Scalar, C: Dim> Allocator<T, Dynamic, C> for DefaultAllocator {
 
         VecStorage::new(nrows, ncols, res)
     }
-
-    #[inline]
-    fn allocate_from_row_iterator<I: IntoIterator<Item = T>>(
-        nrows: Dynamic,
-        ncols: C,
-        iter: I,
-    ) -> Self::Buffer {
-        let it = iter.into_iter().take(nrows.value() * ncols.value());
-        let mut res: Vec<T> = Vec::with_capacity(nrows.value() * ncols.value());
-        let res_ptr = res.as_mut_ptr();
-        let mut count = 0;
-
-        unsafe {
-            for (i, e) in it.enumerate() {
-                *res_ptr.add((i % ncols.value()) * nrows.value() + i / ncols.value()) = e;
-                count += 1;
-            }
-            res.set_len(nrows.value() * ncols.value());
-        }
-        assert!(
-            count == nrows.value() * ncols.value(),
-            "Matrix init. from row iterator: iterator not long enough."
-        );
-
-        VecStorage::new(nrows, ncols, res)
-    }
 }
 
 // Static - Dynamic
@@ -234,33 +176,6 @@ impl<T: Scalar, R: DimName> Allocator<T, R, Dynamic> for DefaultAllocator {
 
         VecStorage::new(nrows, ncols, res)
     }
-
-    #[inline]
-    fn allocate_from_row_iterator<I: IntoIterator<Item = T>>(
-        nrows: R,
-        ncols: Dynamic,
-        iter: I,
-    ) -> Self::Buffer {
-        let it = iter.into_iter().take(nrows.value() * ncols.value());
-        let mut res: Vec<T> = Vec::with_capacity(nrows.value() * ncols.value());
-        let res_ptr = res.as_mut_ptr();
-        let mut count = 0;
-
-        unsafe {
-            for (i, e) in it.enumerate() {
-                *res_ptr.add((i % ncols.value()) * nrows.value() + i / ncols.value()) =
-                    MaybeUninit::new(e).assume_init();
-                count += 1;
-            }
-            res.set_len(nrows.value() * ncols.value());
-        }
-        assert!(
-            count == nrows.value() * ncols.value(),
-            "Matrix init. from row iterator: iterator not long enough."
-        );
-
-        VecStorage::new(nrows, ncols, res)
-    }
 }
 
 /*