How to measure a functions stack usage in Rust?

Question

Is there a way I can measure how much stack memory a function uses?

This question isn't specific to recursive functions; however I was interested to know how much stack memory a function called recursively would take.

I was interested to optimize the function for stack memory usage; however, without knowing what optimizations the compiler is already making, it's just guess-work if this is making real improvements or not.

To be clear, this is not a question about how to optimize for better stack usage.

So is there some reliable way to find out how much stack memory a function uses in Rust?

---

Note that other compilers support this, GCC has -fstack-usage for example.

Answer 1

As a last resort, you can observe the stack pointer (using inline assembly) and infer the result from that. This method is definitely not something you'd use in production... but it works.

#![feature(asm)]

use std::cell::Cell;
use std::cmp;
use std::usize;

// This global variable tracks the highest point of the stack
thread_local!(static STACK_END: Cell<usize> = Cell::new(usize::MAX));

macro_rules! stack_ptr {
    () => ({
        // Grab a copy of the stack pointer
        let x: usize;
        unsafe {
            asm!("mov %rsp, $0" : "=r"(x) ::: "volatile");
        }
        x
    })
}

/// Saves the current position of the stack. Any function
/// being profiled must call this macro.
macro_rules! tick {
    () => ({
        // Save the current stack pointer in STACK_END
        let stack_end = stack_ptr!();
        STACK_END.with(|c| {
            // Since the stack grows down, the "tallest"
            // stack must have the least pointer value
            let best = cmp::min(c.get(), stack_end);
            c.set(best);
        });
    })
}

/// Runs the given callback, and returns its maximum stack usage
/// as reported by the `tick!()` macro.
fn measure<T, F: FnOnce() -> T>(callback: F) -> (T, usize) {
    STACK_END.with(|c| c.set(usize::MAX));
    let stack_start = stack_ptr!();
    let r = callback();
    let stack_end = STACK_END.with(|c| c.get());
    if stack_start < stack_end {
        panic!("tick!() was never called");
    }
    (r, stack_start - stack_end)
}

/// Example recursive function
fn fibonacci(n: i64) -> i64 {
    tick!();
    match n {
        0 => 0,
        1 => 1,
        _ => fibonacci(n-1) + fibonacci(n-2)
    }
}

fn main() {
    // Stack usage should grow linearly with `i`
    for i in 0 .. 10 {
        let (result, stack) = measure(|| fibonacci(i));
        println!("fibonacci({}) = {}: used {} bytes of stack", i, result, stack);
    }
}