Skip to content

Commit

Permalink
feat: VecDeque (#11)
Browse files Browse the repository at this point in the history
  • Loading branch information
amunra authored Nov 11, 2024
1 parent 8535088 commit c8c08bd
Show file tree
Hide file tree
Showing 5 changed files with 984 additions and 79 deletions.
32 changes: 32 additions & 0 deletions .github/workflows/ci.yml
Original file line number Diff line number Diff line change
Expand Up @@ -59,3 +59,35 @@ jobs:
# Test
- name: Test
run: cargo test ${{ matrix.profile }} ${{ matrix.features }}

miri_test:
name: Miri Test
runs-on: ubuntu-latest
steps:
- name: Checkout code
uses: actions/checkout@v3
- name: Set up Rust
uses: actions-rs/toolchain@v1
with:
profile: minimal
toolchain: nightly
- name: Install miri
run: rustup +nightly component add miri
- name: Run miri tests
run: cargo miri test

asan_test:
name: Address Sanitizer Test
runs-on: ubuntu-latest
steps:
- name: Checkout code
uses: actions/checkout@v3
- name: Set up Rust
uses: actions-rs/toolchain@v1
with:
profile: minimal
toolchain: nightly
- name: Install miri
run: rustup +nightly component add miri
- name: Run ASan tests
run: RUSTFLAGS="-Z sanitizer=address" cargo test
10 changes: 9 additions & 1 deletion src/lib.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
#![cfg_attr(feature = "no_std", no_std)]
#![cfg_attr(not(test), cfg_attr(feature = "no_std", no_std))]
#![feature(allocator_api)]
#![cfg_attr(not(doctest), doc = include_str!("../README.md"))]

Expand All @@ -8,3 +8,11 @@ extern crate core;
pub mod claim;
pub mod try_clone;
pub mod vec;
pub mod vec_deque;

#[cfg(test)]
pub(crate) mod testing;

#[cfg(test)]
#[global_allocator]
static GLOBAL: testing::GlobalAllocTestGuardAllocator = testing::GlobalAllocTestGuardAllocator;
199 changes: 199 additions & 0 deletions src/testing.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,199 @@
use crate::claim::Claim;
use alloc::sync::Arc;
use core::ptr::NonNull;
use core::sync::atomic::{AtomicUsize, Ordering};
use std::alloc::{AllocError, Allocator, Global, GlobalAlloc, Layout, System};

thread_local! {
static GLOBAL_ALLOC_ALLOWED: std::cell::RefCell<bool> = const { std::cell::RefCell::new(true) };
}

struct NoPubCtor;

/// A guard that temporarily error if a test performs global allocation in the current thread.
pub struct NoGlobalAllocGuard(NoPubCtor);

impl NoGlobalAllocGuard {
pub fn new() -> Self {
GLOBAL_ALLOC_ALLOWED.with(|alloc_allowed| {
let mut alloc_allowed = alloc_allowed.borrow_mut();
if !*alloc_allowed {
panic!("NoGlobalAllocGuard is not re-entrant.");
}
*alloc_allowed = false; // Disable global allocation
});

Self(NoPubCtor)
}
}

impl Drop for NoGlobalAllocGuard {
fn drop(&mut self) {
GLOBAL_ALLOC_ALLOWED.with(|alloc_allowed| {
let mut alloc_allowed = alloc_allowed.borrow_mut();
*alloc_allowed = true;
});
}
}

pub struct AllowGlobalAllocGuard {
was_allowed: bool,
}

impl AllowGlobalAllocGuard {
pub fn new() -> Self {
let was_allowed = GLOBAL_ALLOC_ALLOWED.with(|alloc_allowed| {
let was_allowed = *alloc_allowed.borrow();
if !was_allowed {
let mut alloc_allowed = alloc_allowed.borrow_mut();
*alloc_allowed = true;
}
was_allowed
});

Self { was_allowed }
}
}

impl Drop for AllowGlobalAllocGuard {
fn drop(&mut self) {
GLOBAL_ALLOC_ALLOWED.with(|alloc_allowed| {
let mut alloc_allowed = alloc_allowed.borrow_mut();
*alloc_allowed = self.was_allowed;
});
}
}

/// Enables the `NoGlobalAllocGuard` by acting as a global allocator.
pub struct GlobalAllocTestGuardAllocator;

impl GlobalAllocTestGuardAllocator {
fn is_allowed(&self) -> bool {
GLOBAL_ALLOC_ALLOWED.with(|alloc_allowed| {
*alloc_allowed.borrow() // Check if allocation is allowed for the current thread
})
}

fn guard(&self) {
if !self.is_allowed() {
panic!("Caught unexpected global allocation with the NoGlobalAllocGuard. Run tests under debugger.");
}
}
}

unsafe impl GlobalAlloc for GlobalAllocTestGuardAllocator {
unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
self.guard();
System.alloc(layout)
}

unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
self.guard();
System.dealloc(ptr, layout)
}

unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
self.guard();
System.alloc_zeroed(layout)
}

unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
self.guard();
System.realloc(ptr, layout, new_size)
}
}

#[derive(Clone)]
pub struct WatermarkAllocator {
watermark: usize,
in_use: Option<Arc<AtomicUsize>>,
}

impl Drop for WatermarkAllocator {
fn drop(&mut self) {
let in_use = self.in_use.take().unwrap();
let _g = AllowGlobalAllocGuard::new();
drop(in_use);
}
}

impl WatermarkAllocator {
pub fn new(watermark: usize) -> Self {
let in_use = Some({
let _g = AllowGlobalAllocGuard::new();
AtomicUsize::new(0).into()
});
Self { watermark, in_use }
}

pub fn in_use(&self) -> usize {
self.in_use.as_ref().unwrap().load(Ordering::SeqCst)
}
}

impl Claim for WatermarkAllocator {}

unsafe impl Allocator for WatermarkAllocator {
fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
let current_in_use = self.in_use.as_ref().unwrap().load(Ordering::SeqCst);
let new_in_use = current_in_use + layout.size();
if new_in_use > self.watermark {
return Err(AllocError);
}
let allocated = {
let _g = AllowGlobalAllocGuard::new();
Global.allocate(layout)?
};
let true_new_in_use = self
.in_use
.as_ref()
.unwrap()
.fetch_add(allocated.len(), Ordering::SeqCst);
unsafe {
if true_new_in_use > self.watermark {
let ptr = allocated.as_ptr() as *mut u8;
let to_dealloc = NonNull::new_unchecked(ptr);
{
let _g = AllowGlobalAllocGuard::new();
Global.deallocate(to_dealloc, layout);
}
Err(AllocError)
} else {
Ok(allocated)
}
}
}

unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
let _g = AllowGlobalAllocGuard::new();
Global.deallocate(ptr, layout);
self.in_use
.as_ref()
.unwrap()
.fetch_sub(layout.size(), Ordering::SeqCst);
}
}

/// A second watermark allocator. This is just to test cases where we need generic types
/// to interoperate, even when their allocator differs. E.g. `lhs: Vec<T, A1> == rhs: Vec<T, A2>`.
#[derive(Clone)]
pub struct WatermarkAllocator2(WatermarkAllocator);

impl WatermarkAllocator2 {
pub fn new(watermark: usize) -> Self {
let inner = WatermarkAllocator::new(watermark);
Self(inner)
}
}

impl Claim for WatermarkAllocator2 {}

unsafe impl Allocator for WatermarkAllocator2 {
fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
self.0.allocate(layout)
}

unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
self.0.deallocate(ptr, layout)
}
}
Loading

0 comments on commit c8c08bd

Please sign in to comment.