dataset::values: check alignment rather than datatype size

gauteh · Oct 15, 2023 · 75195de · 75195de
1 parent 02cf097
commit 75195de
Showing 1 changed file with 7 additions and 7 deletions.
diff --git a/src/reader/dataset.rs b/src/reader/dataset.rs
@@ -46,18 +46,18 @@ pub trait ReaderExt: Reader {
         E::Error: Into<anyhow::Error>,
     {
         let dsz = self.dsize();
+        let extents = extents.try_into().map_err(|e| e.into())?;
+        let counts = extents.get_counts(self.shape())?;
+        let vsz = counts.product::<u64>() as usize * dsz / std::mem::size_of::<T>();
+
         ensure!(
             dsz % std::mem::size_of::<T>() == 0,
             "size of datatype ({}) not multiple of target {}",
             dsz,
             std::mem::size_of::<T>()
         );
 
-        ensure!(dsz == std::mem::size_of::<T>(), "size of datatype ({}) not same as target {}, alignment may not match and result in unsoundness", dsz, std::mem::size_of::<T>());
-
-        let extents = extents.try_into().map_err(|e| e.into())?;
-        let counts = extents.get_counts(self.shape())?;
-        let vsz = counts.product::<u64>() as usize * dsz / std::mem::size_of::<T>();
+        ensure!((dsz * vsz) % std::mem::align_of::<T>() == 0, "alignment of datatype ({}) not a multiple of datatype size and length {}*{}={}, alignment may not match and result in unsoundness", std::mem::align_of::<T>(), dsz, vsz, vsz * dsz);
 
         let values = Box::<[T]>::new_uninit_slice(vsz);
         let values = unsafe { values.assume_init() };
@@ -106,7 +106,7 @@ pub trait ParReaderExt: Reader + ParReader {
             std::mem::size_of::<T>()
         );
 
-        ensure!(dsz == std::mem::size_of::<T>(), "size of datatype ({}) not same as target {}, alignment may not match and result in unsoundness", dsz, std::mem::size_of::<T>());
+        ensure!((dsz * vsz) % std::mem::align_of::<T>() == 0, "alignment of datatype ({}) not a multiple of datatype size and length {}*{}={}, alignment may not match and result in unsoundness", std::mem::align_of::<T>(), dsz, vsz, vsz * dsz);
 
         let extents = extents.try_into().map_err(|e| e.into())?;
         let counts = extents.get_counts(self.shape())?;
@@ -135,7 +135,7 @@ pub trait ParReaderExt: Reader + ParReader {
             std::mem::size_of::<T>()
         );
 
-        ensure!(dsz == std::mem::size_of::<T>(), "size of datatype ({}) not same as target {}, alignment may not match and result in unsoundness", dsz, std::mem::size_of::<T>());
+        ensure!((dsz * vsz) % std::mem::align_of::<T>() == 0, "alignment of datatype ({}) not a multiple of datatype size and length {}*{}={}, alignment may not match and result in unsoundness", std::mem::align_of::<T>(), dsz, vsz, vsz * dsz);
 
         let extents = extents.try_into().map_err(|e| e.into())?;
         let counts = extents.get_counts(self.shape())?;