From 5c025fb4951c5c1466994c101b4c5dc289190f7b Mon Sep 17 00:00:00 2001
From: guillaume didier <guillaume.didier@inria.fr>
Date: Tue, 7 Apr 2020 14:25:13 +0200
Subject: [PATCH] Make the MMappedMemory generic

---
 cache_utils/src/bin/l3.rs |  2 +-
 cache_utils/src/lib.rs    |  1 +
 cache_utils/src/main.rs   |  2 +-
 cache_utils/src/mmap.rs   | 79 +++++++++++++++++++++++----------------
 4 files changed, 49 insertions(+), 35 deletions(-)

diff --git a/cache_utils/src/bin/l3.rs b/cache_utils/src/bin/l3.rs
index 829a8dc..b1437fe 100644
--- a/cache_utils/src/bin/l3.rs
+++ b/cache_utils/src/bin/l3.rs
@@ -2,7 +2,7 @@ use cache_utils::flush;
 use cache_utils::mmap::MMappedMemory;
 
 pub fn main() {
-    let m = unsafe { MMappedMemory::new(2 << 20) };
+    let m = MMappedMemory::new(2 << 20);
     let array = m.slice();
     loop {
         unsafe {
diff --git a/cache_utils/src/lib.rs b/cache_utils/src/lib.rs
index f3546a2..4532fea 100644
--- a/cache_utils/src/lib.rs
+++ b/cache_utils/src/lib.rs
@@ -1,4 +1,5 @@
 #![cfg_attr(feature = "no_std", no_std)]
+#![feature(ptr_internals)]
 
 use static_assertions::assert_cfg;
 
diff --git a/cache_utils/src/main.rs b/cache_utils/src/main.rs
index 756e2af..6781bea 100644
--- a/cache_utils/src/main.rs
+++ b/cache_utils/src/main.rs
@@ -34,7 +34,7 @@ struct Page {
 }
 */
 pub fn main() {
-    let m = unsafe { MMappedMemory::new(SIZE) };
+    let m = MMappedMemory::new(SIZE);
     let array = m.slice();
 
     let old = sched_getaffinity(Pid::from_raw(0)).unwrap();
diff --git a/cache_utils/src/mmap.rs b/cache_utils/src/mmap.rs
index f959a69..b22b915 100644
--- a/cache_utils/src/mmap.rs
+++ b/cache_utils/src/mmap.rs
@@ -2,9 +2,12 @@
 
 use core::borrow::{Borrow, BorrowMut};
 use core::ffi::c_void;
+use core::mem::size_of;
 use core::ops::{Deref, DerefMut};
 use core::ptr::null_mut;
+use core::ptr::Unique;
 use core::slice::{from_raw_parts, from_raw_parts_mut};
+use nix::errno::Errno::EINVAL;
 use nix::sys::mman;
 
 /* from linux kernel headers.
@@ -17,78 +20,88 @@ use nix::sys::mman;
 #define HUGETLB_FLAG_ENCODE_2MB         (21 << HUGETLB_FLAG_ENCODE_SHIFT)
 */
 
-pub struct MMappedMemory {
-    pointer: *mut u8,
+pub struct MMappedMemory<T> {
+    pointer: Unique<T>,
     size: usize,
 }
 
-impl MMappedMemory {
-    pub unsafe fn new(size: usize) -> MMappedMemory {
-        let p: *mut u8 = mman::mmap(
-            null_mut(),
-            size,
-            mman::ProtFlags::PROT_READ | mman::ProtFlags::PROT_WRITE,
-            mman::MapFlags::MAP_PRIVATE
-                | mman::MapFlags::MAP_ANONYMOUS
-                | mman::MapFlags::MAP_HUGETLB,
-            -1,
-            0,
-        )
-        .unwrap() as *mut u8;
-        MMappedMemory { pointer: p, size }
+impl<T> MMappedMemory<T> {
+    pub fn try_new(size: usize) -> Result<MMappedMemory<T>, nix::Error> {
+        assert_ne!(size_of::<T>(), 0);
+        if let Some(p) = unsafe {
+            let p = mman::mmap(
+                null_mut(),
+                size * size_of::<T>(),
+                mman::ProtFlags::PROT_READ | mman::ProtFlags::PROT_WRITE,
+                mman::MapFlags::MAP_PRIVATE
+                    | mman::MapFlags::MAP_ANONYMOUS
+                    | mman::MapFlags::MAP_HUGETLB,
+                -1,
+                0,
+            )?;
+            Unique::new(p as *mut T)
+        } {
+            Ok(MMappedMemory { pointer: p, size })
+        } else {
+            Err(nix::Error::Sys(EINVAL))
+        }
     }
 
-    pub fn slice(&self) -> &[u8] {
-        unsafe { from_raw_parts(self.pointer, self.size) }
+    pub fn new(size: usize) -> MMappedMemory<T> {
+        Self::try_new(size).unwrap()
     }
 
-    pub fn slice_mut(&mut self) -> &mut [u8] {
-        unsafe { from_raw_parts_mut(self.pointer, self.size) }
+    pub fn slice(&self) -> &[T] {
+        unsafe { from_raw_parts(self.pointer.as_ptr(), self.size) }
+    }
+
+    pub fn slice_mut(&mut self) -> &mut [T] {
+        unsafe { from_raw_parts_mut(self.pointer.as_ptr(), self.size) }
     }
 }
 
-impl Drop for MMappedMemory {
+impl<T> Drop for MMappedMemory<T> {
     fn drop(&mut self) {
         unsafe {
-            mman::munmap(self.pointer as *mut c_void, self.size).unwrap();
+            mman::munmap(self.pointer.as_ptr() as *mut c_void, self.size).unwrap();
         }
     }
 }
 
-impl Deref for MMappedMemory {
-    type Target = [u8];
+impl<T> Deref for MMappedMemory<T> {
+    type Target = [T];
 
     fn deref(&self) -> &Self::Target {
         self.slice()
     }
 }
 
-impl DerefMut for MMappedMemory {
+impl<T> DerefMut for MMappedMemory<T> {
     fn deref_mut(&mut self) -> &mut Self::Target {
         self.slice_mut()
     }
 }
 
-impl AsRef<[u8]> for MMappedMemory {
-    fn as_ref(&self) -> &[u8] {
+impl<T> AsRef<[T]> for MMappedMemory<T> {
+    fn as_ref(&self) -> &[T] {
         unimplemented!()
     }
 }
 
-impl AsMut<[u8]> for MMappedMemory {
-    fn as_mut(&mut self) -> &mut [u8] {
+impl<T> AsMut<[T]> for MMappedMemory<T> {
+    fn as_mut(&mut self) -> &mut [T] {
         self.slice_mut()
     }
 }
 
-impl Borrow<[u8]> for MMappedMemory {
-    fn borrow(&self) -> &[u8] {
+impl<T> Borrow<[T]> for MMappedMemory<T> {
+    fn borrow(&self) -> &[T] {
         self.slice()
     }
 }
 
-impl BorrowMut<[u8]> for MMappedMemory {
-    fn borrow_mut(&mut self) -> &mut [u8] {
+impl<T> BorrowMut<[T]> for MMappedMemory<T> {
+    fn borrow_mut(&mut self) -> &mut [T] {
         self.slice_mut()
     }
 }