safer reinterpret; panic in case of field reordering (#2841)

pola-rs · Mar 7, 2022 · 5bb514a · 5bb514a
1 parent 1202408
commit 5bb514a
Show file tree

Hide file tree

Showing 5 changed files with 64 additions and 12 deletions.
diff --git a/polars/polars-core/src/chunked_array/drop.rs b/polars/polars-core/src/chunked_array/drop.rs
@@ -4,8 +4,14 @@ use crate::prelude::*;
 impl<T> Drop for ChunkedArray<T> {
     fn drop(&mut self) {
         if matches!(self.dtype(), DataType::List(_)) {
+            // Safety
             // guarded by the type system
-            unsafe { drop_list(std::mem::transmute(self)) }
+            // the transmute only convinces the type system that we are a list
+            // (which we are)
+            #[allow(clippy::transmute_undefined_repr)]
+            unsafe {
+                drop_list(std::mem::transmute(self))
+            }
         }
     }
 }
diff --git a/polars/polars-core/src/chunked_array/ops/bit_repr.rs b/polars/polars-core/src/chunked_array/ops/bit_repr.rs
@@ -18,10 +18,19 @@ where
                     let buf = array.values().clone();
                     // Safety:
                     // we just check the size of T::Native to be 64 bits
-                    let buf = unsafe { std::mem::transmute::<_, Buffer<u64>>(buf) };
+                    // The fields can still be reordered between generic types
+                    // so we do some extra assertions
+                    let len = buf.len();
+                    let offset = buf.offset();
+                    let ptr = buf.as_slice().as_ptr() as usize;
+                    #[allow(clippy::transmute_undefined_repr)]
+                    let reinterpretted_buf = unsafe { std::mem::transmute::<_, Buffer<u64>>(buf) };
+                    assert_eq!(reinterpretted_buf.len(), len);
+                    assert_eq!(reinterpretted_buf.offset(), offset);
+                    assert_eq!(reinterpretted_buf.as_slice().as_ptr() as usize, ptr);
                     Arc::new(PrimitiveArray::from_data(
                         ArrowDataType::UInt64,
-                        buf,
+                        reinterpretted_buf,
                         array.validity().cloned(),
                     )) as Arc<dyn Array>
                 })
@@ -40,10 +49,19 @@ where
                     let buf = array.values().clone();
                     // Safety:
                     // we just check the size of T::Native to be 32 bits
-                    let buf = unsafe { std::mem::transmute::<_, Buffer<u32>>(buf) };
+                    // The fields can still be reordered between generic types
+                    // so we do some extra assertions
+                    let len = buf.len();
+                    let offset = buf.offset();
+                    let ptr = buf.as_slice().as_ptr() as usize;
+                    #[allow(clippy::transmute_undefined_repr)]
+                    let reinterpretted_buf = unsafe { std::mem::transmute::<_, Buffer<u32>>(buf) };
+                    assert_eq!(reinterpretted_buf.len(), len);
+                    assert_eq!(reinterpretted_buf.offset(), offset);
+                    assert_eq!(reinterpretted_buf.as_slice().as_ptr() as usize, ptr);
                     Arc::new(PrimitiveArray::from_data(
                         ArrowDataType::UInt32,
-                        buf,
+                        reinterpretted_buf,
                         array.validity().cloned(),
                     )) as Arc<dyn Array>
                 })
@@ -64,10 +82,19 @@ impl Reinterpret for UInt64Chunked {
                 let buf = array.values().clone();
                 // Safety
                 // same bit length u64 <-> i64
-                let buf = unsafe { std::mem::transmute::<_, Buffer<i64>>(buf) };
+                // The fields can still be reordered between generic types
+                // so we do some extra assertions
+                let len = buf.len();
+                let offset = buf.offset();
+                let ptr = buf.as_slice().as_ptr() as usize;
+                #[allow(clippy::transmute_undefined_repr)]
+                let reinterpretted_buf = unsafe { std::mem::transmute::<_, Buffer<i64>>(buf) };
+                assert_eq!(reinterpretted_buf.len(), len);
+                assert_eq!(reinterpretted_buf.offset(), offset);
+                assert_eq!(reinterpretted_buf.as_slice().as_ptr() as usize, ptr);
                 Arc::new(PrimitiveArray::from_data(
                     ArrowDataType::Int64,
-                    buf,
+                    reinterpretted_buf,
                     array.validity().cloned(),
                 )) as Arc<dyn Array>
             })
@@ -98,10 +125,19 @@ impl UInt64Chunked {
                 let buf = array.values().clone();
                 // Safety
                 // same bit length u64 <-> f64
-                let buf = unsafe { std::mem::transmute::<_, Buffer<f64>>(buf) };
+                // The fields can still be reordered between generic types
+                // so we do some extra assertions
+                let len = buf.len();
+                let offset = buf.offset();
+                let ptr = buf.as_slice().as_ptr() as usize;
+                #[allow(clippy::transmute_undefined_repr)]
+                let reinterpretted_buf = unsafe { std::mem::transmute::<_, Buffer<f64>>(buf) };
+                assert_eq!(reinterpretted_buf.len(), len);
+                assert_eq!(reinterpretted_buf.offset(), offset);
+                assert_eq!(reinterpretted_buf.as_slice().as_ptr() as usize, ptr);
                 Arc::new(PrimitiveArray::from_data(
                     ArrowDataType::Float64,
-                    buf,
+                    reinterpretted_buf,
                     array.validity().cloned(),
                 )) as Arc<dyn Array>
             })
@@ -117,10 +153,19 @@ impl UInt32Chunked {
                 let buf = array.values().clone();
                 // Safety
                 // same bit length u32 <-> f32
-                let buf = unsafe { std::mem::transmute::<_, Buffer<f32>>(buf) };
+                // The fields can still be reordered between generic types
+                // so we do some extra assertions
+                let len = buf.len();
+                let offset = buf.offset();
+                let ptr = buf.as_slice().as_ptr() as usize;
+                #[allow(clippy::transmute_undefined_repr)]
+                let reinterpretted_buf = unsafe { std::mem::transmute::<_, Buffer<f32>>(buf) };
+                assert_eq!(reinterpretted_buf.len(), len);
+                assert_eq!(reinterpretted_buf.offset(), offset);
+                assert_eq!(reinterpretted_buf.as_slice().as_ptr() as usize, ptr);
                 Arc::new(PrimitiveArray::from_data(
                     ArrowDataType::Float32,
-                    buf,
+                    reinterpretted_buf,
                     array.validity().cloned(),
                 )) as Arc<dyn Array>
             })

diff --git a/polars/polars-core/src/frame/mod.rs b/polars/polars-core/src/frame/mod.rs
@@ -179,6 +179,7 @@ impl DataFrame {
         let series_cols = if S::is_series() {
             // Safety:
             // we are guarded by the type system here.
+            #[allow(clippy::transmute_undefined_repr)]
             let series_cols = unsafe { std::mem::transmute::<Vec<S>, Vec<Series>>(columns) };
             let mut names = PlHashSet::with_capacity(series_cols.len());
 

diff --git a/polars/polars-core/src/lib.rs b/polars/polars-core/src/lib.rs
@@ -1,4 +1,3 @@
-#![allow(clippy::transmute_undefined_repr)]
 #![cfg_attr(docsrs, feature(doc_cfg))]
 extern crate core;
 #[macro_use]

diff --git a/polars/polars-core/src/testing.rs b/polars/polars-core/src/testing.rs
@@ -35,6 +35,7 @@ impl Series {
         // A fat pointer consists of a data ptr and a ptr to the vtable.
         // we specifically check that we only transmute &dyn SeriesTrait e.g.
         // a trait object, therefore this is sound.
+        #[allow(clippy::transmute_undefined_repr)]
         let (data_ptr, _vtable_ptr) =
             unsafe { std::mem::transmute::<&dyn SeriesTrait, (usize, usize)>(object) };
         data_ptr