src/test/run-pass/allocator/custom.rs

   1 // Copyright 2017 The Rust Project Developers. See the COPYRIGHT
   2 // file at the top-level directory of this distribution and at
   3 // http://rust-lang.org/COPYRIGHT.
   4 //
   5 // Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
   6 // http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
   7 // <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
   8 // option. This file may not be copied, modified, or distributed
   9 // except according to those terms.
  10
  11 // aux-build:helper.rs
  12 // no-prefer-dynamic
  13
  14 #![feature(global_allocator, heap_api, allocator_api)]
  15
  16 extern crate helper;
  17
  18 use std::env;
  19 use std::heap::{Heap, Alloc, System, Layout, AllocErr};
  20 use std::sync::atomic::{AtomicUsize, Ordering, ATOMIC_USIZE_INIT};
  21
  22 static HITS: AtomicUsize = ATOMIC_USIZE_INIT;
  23
  24 struct A;
  25
  26 unsafe impl<'a> Alloc for &'a A {
  27     unsafe fn alloc(&mut self, layout: Layout) -> Result<*mut u8, AllocErr> {
  28         HITS.fetch_add(1, Ordering::SeqCst);
  29         System.alloc(layout)
  30     }
  31
  32     unsafe fn dealloc(&mut self, ptr: *mut u8, layout: Layout) {
  33         HITS.fetch_add(1, Ordering::SeqCst);
  34         System.dealloc(ptr, layout)
  35     }
  36 }
  37
  38 #[global_allocator]
  39 static GLOBAL: A = A;
  40
  41 fn main() {
  42     env::set_var("FOO", "bar");
  43     drop(env::var("FOO"));
  44
  45     let n = HITS.load(Ordering::SeqCst);
  46     assert!(n > 0);
  47     unsafe {
  48         let layout = Layout::from_size_align(4, 2).unwrap();
  49
  50         let ptr = Heap.alloc(layout.clone()).unwrap();
  51         helper::work_with(&ptr);
  52         assert_eq!(HITS.load(Ordering::SeqCst), n + 1);
  53         Heap.dealloc(ptr, layout.clone());
  54         assert_eq!(HITS.load(Ordering::SeqCst), n + 2);
  55
  56         let s = String::with_capacity(10);
  57         helper::work_with(&s);
  58         assert_eq!(HITS.load(Ordering::SeqCst), n + 3);
  59         drop(s);
  60         assert_eq!(HITS.load(Ordering::SeqCst), n + 4);
  61
  62         let ptr = System.alloc(layout.clone()).unwrap();
  63         assert_eq!(HITS.load(Ordering::SeqCst), n + 4);
  64         helper::work_with(&ptr);
  65         System.dealloc(ptr, layout);
  66         assert_eq!(HITS.load(Ordering::SeqCst), n + 4);
  67     }
  68 }