timetrap/
lib.rs

1#[cfg(test)]
2mod tests;
3
4/// Log macro is intended to expand beyond println
5/// to use trace lib
6macro_rules! log {
7    ($s:literal $(,$i:expr)+) => {
8        println!($s $(, $i)+);
9    };
10}
11
12#[derive(Debug)]
13pub enum MemUnits {
14    Bytes,
15    Kb,
16    Mb,
17    Gb,
18}
19
20/// `trap!` method to wrap your code with
21/// to measure the time taken for the wrapped section execution.
22///
23/// `trap!` will return value of anything passed from within closure
24///
25/// # Example
26/// ```
27/// use timetrap::*;
28/// trap!("add", {
29///     let a = 0;
30///     let b = a + 1;
31///     assert_eq!(b, a + 1);
32/// });
33/// ```
34#[macro_export]
35macro_rules! trap {
36    ($s:literal, $e:expr) => {{ measure_time($s, || $e) }};
37    ($e:expr) => {{ measure_time("", || $e) }};
38}
39
40/// `trap_mem!` method to wrap your code with
41/// to measure the time taken and the memory consumed for the wrapped section execution.
42///
43/// `trap_mem!` will return value of anything passed from within closure
44///
45/// # Example 1
46/// ```
47/// use timetrap::*;
48/// trap_mem!("add", {
49///     let mut map = std::collections::HashMap::with_capacity(1_000_000);
50///         for i in 0..1_000_000u64 {
51///             map.insert(i, i);
52///         }
53///         map
54/// });
55/// ```
56///
57/// # Example 2
58/// ```
59/// use timetrap::*;
60/// trap_mem!("add", MemUnits::Mb, {
61///     let mut map = std::collections::HashMap::with_capacity(1_000_000);
62///         for i in 0..1_000_000u64 {
63///             map.insert(i, i);
64///         }
65///         map
66/// });
67/// ```
68#[macro_export]
69macro_rules! trap_mem {
70    ($s:literal, $u:expr, $e:expr) => {{ measure_time_and_memory($s, $u, || $e) }};
71    ($s:literal, $e:expr) => {{ measure_time_and_memory($s, MemUnits::Bytes, || $e) }};
72    ($u:expr, $e:expr) => {{ measure_time_and_memory("", $u, || $e) }};
73    ($e:expr) => {{ measure_time_and_memory("", MemUnits::Bytes, || $e) }};
74}
75
76pub fn measure_time<F, R>(name: &str, f: F) -> R
77where
78    F: FnOnce() -> R,
79{
80    let start = std::time::Instant::now();
81
82    let result = f();
83
84    let duration = start.elapsed();
85    if name.is_empty() {
86        log!("Took {:?}", duration);
87    } else {
88        log!("{} took {:?}", name, duration);
89    }
90
91    result
92}
93
94pub fn measure_time_and_memory<F, R>(name: &str, units: MemUnits, f: F) -> R
95where
96    F: FnOnce() -> R,
97{
98    // TODO: check devider is correct
99    let divider = (match units {
100        MemUnits::Bytes => 1,
101        MemUnits::Kb => 1024,
102        MemUnits::Mb => 1048576,
103        MemUnits::Gb => 1073741824,
104    }) as f64;
105
106    use sysinfo::System;
107    let mut sys: System = System::new_all();
108    sys.refresh_all();
109    let init_used_memory = sys.used_memory();
110    let init_used_swap = sys.used_swap();
111
112    let result = measure_time(name, f);
113
114    sys.refresh_all();
115    let final_used_memory = sys.used_memory();
116    let final_used_swap = sys.used_swap();
117    if name.is_empty() {
118        log!(
119            "Consumed memory: {:.2}{:?}",
120            ((final_used_memory as f64) - (init_used_memory as f64)) / divider,
121            units
122        );
123        log!(
124            "Consumed swap: {:.2}{:?}",
125            ((final_used_swap as f64) - (init_used_swap as f64)) / divider,
126            units
127        );
128    } else {
129        log!(
130            "{} consumed memory: {:.2}{:?}",
131            name,
132            ((final_used_memory as f64) - (init_used_memory as f64)) / divider,
133            units
134        );
135        log!(
136            "{} consumed swap: {:.2}{:?}",
137            name,
138            ((final_used_swap as f64) - (init_used_swap as f64)) / divider,
139            units
140        );
141    }
142
143    result
144}