Rust Stack and Heap Size Demonstration

calc_alloc.rs

// This program demonstrates how to get a rough estimate of stack usage
// and how to track heap allocations in Rust.
//
// NOTE: This approach for stack measurement is not reliable or portable.
// It is for demonstration purposes only. The heap tracking method is a
// more robust way to monitor memory usage.

use std::alloc::{GlobalAlloc, Layout, System};
use std::sync::atomic::{AtomicUsize, Ordering};

// A custom allocator that wraps the system allocator to track all
// heap allocations.
struct TrackingAllocator;

static ALLOCATED: AtomicUsize = AtomicUsize::new(0);

// We implement the GlobalAlloc trait for our custom allocator.
// This allows it to be used as the default allocator for the entire program.
unsafe impl GlobalAlloc for TrackingAllocator {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        // Use the system allocator to perform the actual allocation.
        let ptr = System.alloc(layout);
        if !ptr.is_null() {
            // If allocation was successful, add the size of the allocation
            // to our total count. We use relaxed ordering for performance.
            ALLOCATED.fetch_add(layout.size(), Ordering::Relaxed);
        }
        ptr
    }

    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        // Use the system allocator to perform the deallocation.
        System.dealloc(ptr, layout);
        // Subtract the size of the deallocated block from our total count.
        ALLOCATED.fetch_sub(layout.size(), Ordering::Relaxed);
    }
}

// Tell Rust to use our custom allocator for all heap allocations.
#[global_allocator]
static A: TrackingAllocator = TrackingAllocator;

// This function recursively calls itself to simulate stack usage.
// It stops when it hits a depth of 1000.
fn stack_eater(mut depth: i32, canary: &u8) {
    let _on_stack = [0_u8; 1024]; // 1KB array on the stack
    depth += 1;

    if depth < 1000 {
        stack_eater(depth, canary);
    } else {
        // Find the difference in memory addresses between the start of the
        // stack (where the canary is) and the current location.
        let current_stack_addr = &_on_stack as *const _ as usize;
        let start_stack_addr = canary as *const _ as usize;

        // The size is the difference between the two addresses.
        let stack_size_estimate = start_stack_addr - current_stack_addr;

        println!("--- Stack Size Estimation ---");
        println!("Estimated stack usage: {} KB", stack_size_estimate / 1024);
    }
}

fn main() {
    // --- Heap demonstration ---
    println!("--- Heap Allocation Tracking ---");

    // Initially, the heap is empty, so our counter is 0.
    println!("Total heap allocated before: {} bytes", ALLOCATED.load(Ordering::Relaxed));

    // Allocate some memory on the heap.
    let mut data = Vec::new();
    for i in 0..100 {
        data.push(vec![i; 1000]); // Allocate 100 * 1000 bytes.
    }

    // Now our allocator has tracked the new memory.
    println!("Total heap allocated after: {} bytes", ALLOCATED.load(Ordering::Relaxed));

    // When `data` goes out of scope, it will be deallocated and
    // our counter will decrease. But for this example, we keep it.
    std::mem::drop(data);
    println!("Total heap allocated after dealloc: {} bytes", ALLOCATED.load(Ordering::Relaxed));

    // --- Stack demonstration ---
    // Place a canary value at the start of the main function's stack frame.
    let canary = 0_u8;
    stack_eater(0, &canary);
}