From 306e57a5508f6cd1ba9c21194daab4b594c0bad5 Mon Sep 17 00:00:00 2001
From: "Jorge C. Leitao" <jorgecarleitao@gmail.com>
Date: Fri, 24 Sep 2021 15:42:40 +0000
Subject: [PATCH] Added MIRI checks and coverage for custom allocator

---
 .github/workflows/test.yml | 36 +++++++++++++++++++++++++++++++++---
 Cargo.toml                 |  2 +-
 src/buffer/mutable.rs      |  4 ++--
 src/vec.rs                 | 10 ++++++++--
 4 files changed, 44 insertions(+), 8 deletions(-)

diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 87765f68541..e57b03c941f 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -67,7 +67,7 @@ jobs:
           rustup default stable
           rustup component add rustfmt
       - name: Setup parquet files
-        run: | 
+        run: |
           apt update && apt install python3-pip python3-venv -y -q
           python3 -m venv venv
           venv/bin/pip install pip --upgrade
@@ -158,13 +158,42 @@ jobs:
       - name: Run
         env:
           RUST_BACKTRACE: full
-          RUST_LOG: 'trace'
+          RUST_LOG: "trace"
         # --skip io: miri can't handle opening of files, so we skip those
         run: |
           cargo miri setup
           cargo clean
           cargo miri test --features full -- --skip io::parquet --skip io::ipc
 
+  miri-checks-custom-allocator:
+    name: MIRI with custom allocator
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v2
+        with:
+          submodules: true
+      - uses: actions/cache@v2
+        with:
+          path: |
+            ~/.cargo/registry
+            ~/.cargo/git
+            target
+          key: ${{ runner.os }}-cargo-miri-${{ hashFiles('**/Cargo.lock') }}
+      - name: Setup toolchain
+        run: |
+          rustup toolchain install nightly-2021-07-03
+          rustup default nightly-2021-07-03
+          rustup component add rustfmt miri
+      - name: Run
+        env:
+          RUST_BACKTRACE: full
+          RUST_LOG: "trace"
+        # --skip io: miri can't handle opening of files, so we skip those
+        run: |
+          cargo miri setup
+          cargo clean
+          cargo miri test --features full,cache_aligned -- --skip io::parquet --skip io::ipc
+
   coverage:
     name: Coverage
     runs-on: ubuntu-latest
@@ -185,7 +214,7 @@ jobs:
           # this key is not equal because coverage uses different compilation flags.
           key: ${{ runner.os }}-amd64-target-coverage-cache-stable-
       - name: Setup parquet files
-        run: | 
+        run: |
           apt update && apt install python3-pip python3-venv -y -q
           python3 -m venv venv
           venv/bin/pip install pip --upgrade
@@ -199,6 +228,7 @@ jobs:
           # 2020-11-15: There is a cargo-tarpaulin regression in 0.17.0
           # see https://github.com/xd009642/tarpaulin/issues/618
           cargo install --version 0.16.0 cargo-tarpaulin
+          cargo tarpaulin --features cache_aligned --out Xml
           cargo tarpaulin --features full --out Xml
       - name: Report coverage
         continue-on-error: true
diff --git a/Cargo.toml b/Cargo.toml
index 6d5684dda3d..fa4963318c2 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -98,7 +98,7 @@ full = [
     "merge_sort",
     "compute",
     # parses timezones used in timestamp conversions
-    "chrono-tz",
+    "chrono-tz"
 ]
 merge_sort = ["itertools"]
 io_csv = ["io_csv_read", "io_csv_write"]
diff --git a/src/buffer/mutable.rs b/src/buffer/mutable.rs
index f4ac85e1996..70c4e4415d0 100644
--- a/src/buffer/mutable.rs
+++ b/src/buffer/mutable.rs
@@ -241,7 +241,7 @@ impl<T: NativeType> MutableBuffer<T> {
     }
 
     /// Shrinks the capacity of the [`MutableBuffer`] to fit its current length.
-    /// The new capacity will be a multiple of 64 bytes.
+    /// When the feature `cache_aligned`, the new capacity will be a multiple of 64 bytes.
     ///
     /// # Example
     /// ```
@@ -253,7 +253,7 @@ impl<T: NativeType> MutableBuffer<T> {
     /// buffer.push(2);
     ///
     /// buffer.shrink_to_fit();
-    /// assert!(buffer.capacity() == 2);
+    /// assert!(buffer.capacity() < 16);  // 2 or 8 depending on feature `cache_aligned`
     /// ```
     pub fn shrink_to_fit(&mut self) {
         self.data.shrink_to_fit();
diff --git a/src/vec.rs b/src/vec.rs
index 8ce7e7c785a..48966c0c616 100644
--- a/src/vec.rs
+++ b/src/vec.rs
@@ -36,8 +36,8 @@ unsafe fn reallocate<T: NativeType>(
     (ptr, new_capacity)
 }
 
-/// An interface equivalent to `std::vec::Vec` with an allocator aligned along cache-lines.
-pub struct AlignedVec<T: NativeType> {
+/// An interface equivalent to `std::vec::Vec` with an allocator aligned along cache lines.
+pub(crate) struct AlignedVec<T: NativeType> {
     // dangling iff capacity = 0
     ptr: NonNull<T>,
     // invariant: len <= capacity
@@ -45,6 +45,12 @@ pub struct AlignedVec<T: NativeType> {
     capacity: usize,
 }
 
+impl<T: NativeType> Drop for AlignedVec<T> {
+    fn drop(&mut self) {
+        unsafe { alloc::free_aligned(self.ptr, self.capacity) }
+    }
+}
+
 impl<T: NativeType> AlignedVec<T> {
     #[inline]
     pub fn new() -> Self {