Rust Vec Fill Test

iterators.txt

S 100 0
A 1500 40000
A 4100 40000
A 14300 40000
E 23500 0

loops.txt

S 400 0
A 25100 16
A 25700 32
F 25900 16
A 26900 64
F 27100 32
A 27600 128
F 27900 64
A 28800 256
F 29000 128
A 29400 512
F 29700 256
A 30400 1024
F 30700 512
A 33400 2048
F 33700 1024
A 34600 4096
F 34800 2048
A 42300 8192
F 44500 4096
A 51500 16384
F 54000 8192
A 64100 32768
F 70500 16384
A 86900 65536
F 101200 32768
A 111800 16
A 112700 32
F 112900 16
A 113600 64
F 113800 32
A 114100 128
F 114300 64
A 115100 256
F 115300 128
A 115700 512
F 115900 256
A 116400 1024
F 116600 512
A 118900 2048
F 119200 1024
A 121500 4096
F 121800 2048
A 123100 8192
F 125100 4096
A 127300 16384
F 127700 8192
A 131700 32768
F 132100 16384
A 144600 65536
F 152400 32768
A 170000 16
A 170400 32
F 170500 16
A 170900 64
F 171100 32
A 171400 128
F 171500 64
A 172100 256
F 172200 128
A 172600 512
F 172800 256
A 173400 1024
F 173600 512
A 177400 2048
F 177600 1024
A 180000 4096
F 180200 2048
A 182900 8192
F 184800 4096
A 191000 16384
F 193100 8192
A 199500 32768
F 205300 16384
A 221800 65536
F 228300 32768
E 234100 0

main.rs

use once_cell::sync::Lazy;
use std::{
    alloc::{GlobalAlloc, System},
    fs::File,
    io::{BufWriter, Write},
    sync::atomic::{AtomicBool, Ordering},
    sync::Mutex,
    time::Instant,
};

#[global_allocator]
static ALLOC: TracingAlloc = TracingAlloc::new();

fn loops(n: i32) {
    let mut memory_trace = BufWriter::new(File::create("loops.txt").unwrap());
    ALLOC.set_file(memory_trace);
    ALLOC.enable_tracing();

    let mut a = Vec::new();
    for i in 0..n {
        a.push(i);
    }

    let mut b = Vec::new();
    for i in 0..n {
        b.push(i);
    }

    let mut c = Vec::new();
    for i in 0..n {
        c.push(a[i as usize] + b[i as usize]);
    }

    ALLOC.disable_tracing();
    ALLOC.clear_file();

    dbg!(c[n as usize - 1]);
}

fn iterators(n: i32) {
    let mut memory_trace = BufWriter::new(File::create("iterators.txt").unwrap());
    ALLOC.set_file(memory_trace);
    ALLOC.enable_tracing();

    let a: Vec<_> = (0..n).collect();
    let b: Vec<_> = (0..n).collect();
    let c: Vec<_> = a.iter().zip(&b).map(|(a, b)| a + b).collect();

    ALLOC.disable_tracing();
    ALLOC.clear_file();

    dbg!(c[n as usize - 1]);
}
fn main() {
    let mut stdin = std::io::stdin();
    let mut n = String::new();
    stdin.read_line(&mut n).unwrap();
    let n: i32 = n.trim().parse().unwrap();

    loops(n);
    iterators(n);
}

pub enum Event {
    Alloc { addr: usize, size: usize },
    Free { addr: usize, size: usize },
    Start,
    End,
}

struct TraceData {
    output_file: Option<BufWriter<File>>,
    timestamp: Option<Instant>,
}

static OUTPUT_FILE: Lazy<Mutex<TraceData>> = Lazy::new(|| {
    Mutex::new(TraceData {
        output_file: None,
        timestamp: None,
    })
});

pub struct TracingAlloc {
    inner: System,
    active: AtomicBool,
}

unsafe impl Sync for TracingAlloc {}

impl TracingAlloc {
    pub const fn new() -> Self {
        Self {
            inner: System,
            active: AtomicBool::new(false),
        }
    }

    pub fn enable_tracing(&self) {
        let mut lock = OUTPUT_FILE.lock().unwrap();
        lock.timestamp = Some(Instant::now());
        std::mem::drop(lock); // Must drop the lock before writing the start event.

        self.write_ev(Event::Start);
        self.active.store(true, Ordering::SeqCst);
    }

    pub fn disable_tracing(&self) {
        self.write_ev(Event::End);
        self.active.store(false, Ordering::SeqCst);
    }

    pub fn set_file(&self, file: BufWriter<File>) -> Option<BufWriter<File>> {
        let mut lock = OUTPUT_FILE.lock().unwrap();
        lock.output_file.replace(file)
    }

    pub fn clear_file(&self) -> Option<BufWriter<File>> {
        let mut lock = OUTPUT_FILE.lock().unwrap();
        lock.output_file.take()
    }

    fn write_ev(&self, ev: Event) {
        let mut lock = OUTPUT_FILE.lock().unwrap();

        if let (Some(ts), Some(file)) = (lock.timestamp, lock.output_file.as_mut()) {
            let elapsed = ts.elapsed();
            let (symbol, size) = match &ev {
                Event::Alloc { size, .. } => ('A', *size),
                Event::Free { size, .. } => ('F', *size),
                Event::Start => ('S', 0),
                Event::End => ('E', 0),
            };

            // Just eat the error so we don't get a panic during allocation.
            let _ = writeln!(file, "{} {} {}", symbol, elapsed.as_nanos(), size);
        }
    }
}

unsafe impl GlobalAlloc for TracingAlloc {
    unsafe fn alloc(&self, layout: std::alloc::Layout) -> *mut u8 {
        let res = self.inner.alloc(layout);

        if self.active.load(Ordering::SeqCst) {
            self.write_ev(Event::Alloc {
                addr: res as _,
                size: layout.size(),
            });
        }

        res
    }

    unsafe fn dealloc(&self, ptr: *mut u8, layout: std::alloc::Layout) {
        if self.active.load(Ordering::SeqCst) {
            self.write_ev(Event::Free {
                addr: ptr as _,
                size: layout.size(),
            });
        }

        self.inner.dealloc(ptr, layout)
    }
}