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