microsoft/qdk
Publicmirrored fromhttps://github.com/microsoft/qdkAvailable
compiler/qsc/src/bin/memtest.rs
61lines · modecode
| 1 | // Copyright (c) Microsoft Corporation. |
| 2 | // Licensed under the MIT License. |
| 3 | |
| 4 | //! Records the memory usage of the compiler. |
| 5 | |
| 6 | use qsc::{compile, CompileUnit}; |
| 7 | use qsc_data_structures::target::TargetCapabilityFlags; |
| 8 | use qsc_frontend::compile::PackageStore; |
| 9 | use std::{ |
| 10 | alloc::{GlobalAlloc, Layout, System}, |
| 11 | sync::atomic::{AtomicU64, Ordering}, |
| 12 | }; |
| 13 | |
| 14 | /// A wrapper around a memory allocator that tracks allocation amounts. |
| 15 | pub struct AllocationCounter<A: GlobalAlloc> { |
| 16 | pub allocator: A, |
| 17 | pub counter: AtomicU64, |
| 18 | } |
| 19 | |
| 20 | unsafe impl<A: GlobalAlloc> GlobalAlloc for AllocationCounter<A> { |
| 21 | unsafe fn alloc(&self, l: Layout) -> *mut u8 { |
| 22 | self.counter.fetch_add(l.size() as u64, Ordering::SeqCst); |
| 23 | self.allocator.alloc(l) |
| 24 | } |
| 25 | unsafe fn dealloc(&self, ptr: *mut u8, l: Layout) { |
| 26 | self.allocator.dealloc(ptr, l); |
| 27 | self.counter.fetch_sub(l.size() as u64, Ordering::SeqCst); |
| 28 | } |
| 29 | } |
| 30 | |
| 31 | impl<A: GlobalAlloc> AllocationCounter<A> { |
| 32 | pub const fn new(allocator: A) -> Self { |
| 33 | AllocationCounter { |
| 34 | allocator, |
| 35 | counter: AtomicU64::new(0), |
| 36 | } |
| 37 | } |
| 38 | pub fn reset(&self) { |
| 39 | self.counter.store(0, Ordering::SeqCst); |
| 40 | } |
| 41 | pub fn read(&self) -> u64 { |
| 42 | self.counter.load(Ordering::SeqCst) |
| 43 | } |
| 44 | } |
| 45 | |
| 46 | #[global_allocator] |
| 47 | static ALLOCATOR: AllocationCounter<System> = AllocationCounter::new(System); |
| 48 | |
| 49 | #[must_use] |
| 50 | pub fn compile_stdlib() -> CompileUnit { |
| 51 | let store = PackageStore::new(compile::core()); |
| 52 | compile::std(&store, TargetCapabilityFlags::all()) |
| 53 | } |
| 54 | |
| 55 | fn main() { |
| 56 | let _stdlib = compile_stdlib(); |
| 57 | let std = ALLOCATOR.read(); |
| 58 | |
| 59 | ALLOCATOR.reset(); |
| 60 | println!("{std}"); |
| 61 | } |
| 62 | |