extern crate rand;
use core::alloc::{GlobalAlloc, Layout};
use std::alloc::System;
use std::collections::BTreeMap;
use std::fmt;
use std::sync::{Arc, Mutex};
pub type GAlloc = TestAlloc<System>;
pub struct TestAlloc<A>
where
A: GlobalAlloc,
{
alloc: A,
allocatings: Arc<Mutex<BTreeMap<*mut u8, Layout>>>,
}
impl<A> Default for TestAlloc<A>
where
A: GlobalAlloc + Default,
{
fn default() -> Self {
Self::from(A::default())
}
}
impl<A> From<A> for TestAlloc<A>
where
A: GlobalAlloc,
{
fn from(inner: A) -> Self {
Self {
alloc: inner,
allocatings: Arc::default(),
}
}
}
impl<A> Clone for TestAlloc<A>
where
A: GlobalAlloc + Clone,
{
fn clone(&self) -> Self {
Self {
alloc: self.alloc.clone(),
allocatings: self.allocatings.clone(),
}
}
}
impl<A> Drop for TestAlloc<A>
where
A: GlobalAlloc,
{
fn drop(&mut self) {
if Arc::strong_count(&self.allocatings) == 1 {
let allocatings = self.allocatings.lock().unwrap();
if allocatings.is_empty() == false {
let message0 = "Memory leak is detected";
let message1 =
"The allocator is dropped before the allocated pointer is deallocated";
panic!("{}: {}", message0, message1);
}
}
}
}
unsafe impl<A> GlobalAlloc for TestAlloc<A>
where
A: GlobalAlloc,
{
unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
let ptr = self.alloc.alloc(layout);
if !ptr.is_null() {
let mut allocatings = self.allocatings.lock().unwrap();
let prev = allocatings.insert(ptr, layout);
assert_eq!(true, prev.is_none());
}
ptr
}
unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
if ptr.is_null() {
panic!("Null pointer is passed to method GlobalAlloc.dealloc().");
}
{
let mut allocatings = self.allocatings.lock().unwrap();
let prev = allocatings.remove(&ptr).unwrap();
if layout != prev {
panic!(
"GlobalAlloc.dealloc() is passed a different layout from GlobalAlloc.alloc()"
);
}
}
self.alloc.dealloc(ptr, layout);
}
}
impl<A> TestAlloc<A>
where
A: GlobalAlloc,
{
pub fn providing_pointers(&self) -> Vec<(*mut u8, Layout)> {
self.allocatings
.lock()
.unwrap()
.iter()
.map(|(&k, &v)| (k, v))
.collect()
}
}
impl<A> fmt::Debug for TestAlloc<A>
where
A: GlobalAlloc + fmt::Debug,
{
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("TestAlloc")
.field("alloc", &self.alloc)
.field("info", &format!("{:p}", self.allocatings))
.finish()
}
}
unsafe impl<A> Send for TestAlloc<A> where A: GlobalAlloc + Send {}
unsafe impl<A> Sync for TestAlloc<A> where A: GlobalAlloc + Send + Sync {}
#[derive(Clone, Copy, Debug)]
pub struct NeverAlloc;
impl Default for NeverAlloc {
fn default() -> Self {
Self
}
}
unsafe impl GlobalAlloc for NeverAlloc {
unsafe fn alloc(&self, _layout: Layout) -> *mut u8 {
core::ptr::null_mut()
}
unsafe fn dealloc(&self, _ptr: *mut u8, _layout: Layout) {
panic!("Method NeverAlloc.dealloc() is called.");
}
}
#[derive(Debug)]
pub struct MaybeAlloc<A = TestAlloc<System>>
where
A: GlobalAlloc,
{
alloc: A,
}
impl<A> Default for MaybeAlloc<A>
where
A: GlobalAlloc + Default,
{
fn default() -> Self {
Self::from(A::default())
}
}
impl<A> From<A> for MaybeAlloc<A>
where
A: GlobalAlloc,
{
fn from(alloc: A) -> Self {
Self { alloc }
}
}
impl<A> Clone for MaybeAlloc<A>
where
A: GlobalAlloc + Clone,
{
fn clone(&self) -> Self {
Self::from(self.alloc.clone())
}
}
unsafe impl<A> GlobalAlloc for MaybeAlloc<A>
where
A: GlobalAlloc,
{
unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
if rand::random::<u8>() % 16 == 0 {
core::ptr::null_mut()
} else {
self.alloc.alloc(layout)
}
}
unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
if ptr.is_null() {
panic!("Null pointer is passed to method GlobalAlloc.dealloc().");
}
self.alloc.dealloc(ptr, layout);
}
}