src/tools/miri/tests/pass/box-custom-alloc.rs

   1 #![allow(incomplete_features)] // for trait upcasting
   2 #![feature(allocator_api, trait_upcasting)]
   3
   4 use std::alloc::Layout;
   5 use std::alloc::{AllocError, Allocator};
   6 use std::cell::Cell;
   7 use std::mem::MaybeUninit;
   8 use std::ptr::{self, NonNull};
   9
  10 struct OnceAlloc<'a> {
  11     space: Cell<&'a mut [MaybeUninit<u8>]>,
  12 }
  13
  14 unsafe impl<'shared, 'a: 'shared> Allocator for &'shared OnceAlloc<'a> {
  15     fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
  16         let space = self.space.replace(&mut []);
  17
  18         let (ptr, len) = (space.as_mut_ptr(), space.len());
  19
  20         if ptr.align_offset(layout.align()) != 0 || len < layout.size() {
  21             return Err(AllocError);
  22         }
  23
  24         let slice_ptr = ptr::slice_from_raw_parts_mut(ptr as *mut u8, len);
  25         unsafe { Ok(NonNull::new_unchecked(slice_ptr)) }
  26     }
  27
  28     unsafe fn deallocate(&self, _ptr: NonNull<u8>, _layout: Layout) {}
  29 }
  30
  31 trait MyTrait {
  32     fn hello(&self) -> u8;
  33 }
  34
  35 impl MyTrait for [u8; 1] {
  36     fn hello(&self) -> u8 {
  37         self[0]
  38     }
  39 }
  40
  41 trait TheTrait: MyTrait {}
  42
  43 impl TheTrait for [u8; 1] {}
  44
  45 /// `Box<T, G>` is a `ScalarPair` where the 2nd component is the allocator.
  46 fn test1() {
  47     let mut space = vec![MaybeUninit::new(0); 1];
  48     let once_alloc = OnceAlloc { space: Cell::new(&mut space[..]) };
  49
  50     let boxed = Box::new_in([42u8; 1], &once_alloc);
  51     let _val = *boxed;
  52     let with_dyn: Box<dyn TheTrait, &OnceAlloc> = boxed;
  53     assert_eq!(42, with_dyn.hello());
  54     let with_dyn: Box<dyn MyTrait, &OnceAlloc> = with_dyn; // upcast
  55     assert_eq!(42, with_dyn.hello());
  56 }
  57
  58 // Make the allocator itself so big that the Box is not even a ScalarPair any more.
  59 struct OnceAllocRef<'s, 'a>(&'s OnceAlloc<'a>, u64);
  60
  61 unsafe impl<'shared, 'a: 'shared> Allocator for OnceAllocRef<'shared, 'a> {
  62     fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
  63         self.0.allocate(layout)
  64     }
  65
  66     unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
  67         self.0.deallocate(ptr, layout)
  68     }
  69 }
  70
  71 /// `Box<T, G>` is an `Aggregate`.
  72 fn test2() {
  73     let mut space = vec![MaybeUninit::new(0); 1];
  74     let once_alloc = OnceAlloc { space: Cell::new(&mut space[..]) };
  75
  76     let boxed = Box::new_in([42u8; 1], OnceAllocRef(&once_alloc, 0));
  77     let _val = *boxed;
  78     let with_dyn: Box<dyn TheTrait, OnceAllocRef> = boxed;
  79     assert_eq!(42, with_dyn.hello());
  80     let with_dyn: Box<dyn MyTrait, OnceAllocRef> = with_dyn; // upcast
  81     assert_eq!(42, with_dyn.hello());
  82 }
  83
  84 fn main() {
  85     test1();
  86     test2();
  87 }