leetcode

Log | Files | Refs | README

benchmark.rs (3233B)

      1 use std::alloc::{GlobalAlloc, Layout, System};
      2 use std::fmt::Debug;
      3 use std::sync::Once;
      4 use std::sync::atomic::{AtomicUsize, Ordering};
      5 use std::time::{Duration, Instant};
      6 use tracing::info;
      7 
      8 static ALLOCATED: AtomicUsize = AtomicUsize::new(0);
      9 static TRACING_INIT: Once = Once::new();
     10 
     11 struct TrackingAllocator;
     12 
     13 unsafe impl GlobalAlloc for TrackingAllocator {
     14     unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
     15         ALLOCATED.fetch_add(layout.size(), Ordering::SeqCst);
     16         unsafe { System.alloc(layout) }
     17     }
     18 
     19     unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
     20         ALLOCATED.fetch_sub(layout.size(), Ordering::SeqCst);
     21         unsafe { System.dealloc(ptr, layout) }
     22     }
     23 }
     24 
     25 #[global_allocator]
     26 static TRACKING_ALLOCATOR: TrackingAllocator = TrackingAllocator;
     27 
     28 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
     29 pub struct BenchmarkResult {
     30     pub elapsed: Duration,
     31     pub heap_delta_bytes: isize,
     32 }
     33 
     34 pub fn current_heap_allocated_bytes() -> usize {
     35     ALLOCATED.load(Ordering::SeqCst)
     36 }
     37 
     38 pub fn init_benchmark_tracing() {
     39     TRACING_INIT.call_once(|| {
     40         let _ = tracing_subscriber::fmt()
     41             .with_test_writer()
     42             .with_target(false)
     43             .with_level(false)
     44             .without_time()
     45             .compact()
     46             .try_init();
     47     });
     48 }
     49 
     50 fn format_duration(elapsed: Duration) -> String {
     51     let nanos = elapsed.as_nanos();
     52     if nanos >= 1_000_000_000 {
     53         format!("{:.3} s", nanos as f64 / 1_000_000_000.0)
     54     } else if nanos >= 1_000_000 {
     55         format!("{:.3} ms", nanos as f64 / 1_000_000.0)
     56     } else if nanos >= 1_000 {
     57         format!("{:.3} us", nanos as f64 / 1_000.0)
     58     } else {
     59         format!("{nanos} ns")
     60     }
     61 }
     62 
     63 fn format_signed_bytes(bytes: isize) -> String {
     64     let sign = if bytes >= 0 { "+" } else { "-" };
     65     let abs = (bytes as i128).abs() as f64;
     66     if abs >= 1024.0 * 1024.0 {
     67         format!("{sign}{:.2} MiB", abs / (1024.0 * 1024.0))
     68     } else if abs >= 1024.0 {
     69         format!("{sign}{:.2} KiB", abs / 1024.0)
     70     } else {
     71         format!("{sign}{} B", abs as u64)
     72     }
     73 }
     74 
     75 pub fn run_with_metrics<T, F>(label: &str, test: F) -> (T, BenchmarkResult)
     76 where
     77     F: FnOnce() -> T,
     78 {
     79     init_benchmark_tracing();
     80 
     81     let mem_before = current_heap_allocated_bytes() as isize;
     82     let start = Instant::now();
     83 
     84     let output = test();
     85 
     86     let elapsed = start.elapsed();
     87     let mem_after = current_heap_allocated_bytes() as isize;
     88     let result = BenchmarkResult {
     89         elapsed,
     90         heap_delta_bytes: mem_after - mem_before,
     91     };
     92 
     93     info!(
     94         "\n[benchmark] {label}\n  time       : {}\n  heap delta : {}",
     95         format_duration(result.elapsed),
     96         format_signed_bytes(result.heap_delta_bytes)
     97     );
     98 
     99     (output, result)
    100 }
    101 
    102 pub fn assert_vec_any_order<T>(mut actual: Vec<T>, mut expected: Vec<T>)
    103 where
    104     T: Ord + Debug,
    105 {
    106     actual.sort_unstable();
    107     expected.sort_unstable();
    108     assert_eq!(actual, expected);
    109 }
    110 
    111 pub fn run_and_assert_vec_any_order<T, F>(label: &str, expected: Vec<T>, test: F) -> BenchmarkResult
    112 where
    113     T: Ord + Debug,
    114     F: FnOnce() -> Vec<T>,
    115 {
    116     let (actual, metrics) = run_with_metrics(label, test);
    117     assert_vec_any_order(actual, expected);
    118     metrics
    119 }