From bfcbd93036fe87a7b6776b146811bde1d70fac90 Mon Sep 17 00:00:00 2001
From: Morgan Thomas <morgan.a.s.thomas@gmail.com>
Date: Tue, 28 Oct 2025 13:42:41 +0000
Subject: [PATCH] issue/83: function-level memory profiling

---
 .../src/b_compile_intermediate.rs             |  42 +++-
 crates/lean_compiler/src/c_compile_final.rs   |  43 +++-
 crates/lean_compiler/src/ir/bytecode.rs       |  18 +-
 crates/lean_compiler/src/ir/mod.rs            |   2 +-
 crates/lean_vm/src/core/constants.rs          |   2 +-
 crates/lean_vm/src/core/types.rs              |   3 +
 crates/lean_vm/src/diagnostics/error.rs       |   2 +
 crates/lean_vm/src/diagnostics/profiler.rs    | 235 +++++++++++++++++-
 crates/lean_vm/src/execution/runner.rs        |  37 ++-
 crates/lean_vm/src/isa/hint.rs                |  58 ++++-
 crates/lean_vm/tests/test_lean_vm.rs          |   4 +
 crates/poseidon_circuit/src/tests.rs          |   2 +-
 crates/rec_aggregation/src/recursion.rs       |   2 +-
 crates/rec_aggregation/src/xmss_aggregate.rs  |   2 +-
 14 files changed, 412 insertions(+), 40 deletions(-)

diff --git a/crates/lean_compiler/src/b_compile_intermediate.rs b/crates/lean_compiler/src/b_compile_intermediate.rs
index 92818757..cffa8f45 100644
--- a/crates/lean_compiler/src/b_compile_intermediate.rs
+++ b/crates/lean_compiler/src/b_compile_intermediate.rs
@@ -10,7 +10,7 @@ use utils::ToUsize;
 #[derive(Default)]
 struct Compiler {
     bytecode: BTreeMap<Label, Vec<IntermediateInstruction>>,
-    match_blocks: Vec<Vec<Vec<IntermediateInstruction>>>, // each match = many bytecode blocks, each bytecode block = many instructions
+    match_blocks: Vec<MatchBlock>,
     if_counter: usize,
     call_counter: usize,
     func_name: String,
@@ -152,10 +152,17 @@ fn compile_function(
     compiler.args_count = function.arguments.len();
 
     let mut declared_vars: BTreeSet<Var> = function.arguments.iter().cloned().collect();
-    compile_lines(&function.instructions, compiler, None, &mut declared_vars)
+    compile_lines(
+        &Label::function(function.name.clone()),
+        &function.instructions,
+        compiler,
+        None,
+        &mut declared_vars,
+    )
 }
 
 fn compile_lines(
+    function_name: &Label,
     lines: &[SimpleLine],
     compiler: &mut Compiler,
     final_jump: Option<Label>,
@@ -212,6 +219,7 @@ fn compile_lines(
                     let mut arm_declared_vars = declared_vars.clone();
                     compiler.stack_size = original_stack_size;
                     let arm_instructions = compile_lines(
+                        function_name,
                         arm,
                         compiler,
                         Some(end_label.clone()),
@@ -229,7 +237,10 @@ fn compile_lines(
                     };
                 }
                 compiler.stack_size = new_stack_size;
-                compiler.match_blocks.push(compiled_arms);
+                compiler.match_blocks.push(MatchBlock {
+                    function_name: function_name.clone(),
+                    match_cases: compiled_arms,
+                });
 
                 let value_scaled_offset = IntermediateValue::MemoryAfterFp {
                     offset: compiler.stack_size.into(),
@@ -261,8 +272,13 @@ fn compile_lines(
                     updated_fp: None,
                 });
 
-                let remaining =
-                    compile_lines(&lines[i + 1..], compiler, final_jump, declared_vars)?;
+                let remaining = compile_lines(
+                    function_name,
+                    &lines[i + 1..],
+                    compiler,
+                    final_jump,
+                    declared_vars,
+                )?;
                 compiler.bytecode.insert(end_label, remaining);
 
                 return Ok(instructions);
@@ -349,6 +365,7 @@ fn compile_lines(
 
                 let mut then_declared_vars = declared_vars.clone();
                 let then_instructions = compile_lines(
+                    function_name,
                     then_branch,
                     compiler,
                     Some(end_label.clone()),
@@ -359,6 +376,7 @@ fn compile_lines(
                 compiler.stack_size = original_stack;
                 let mut else_declared_vars = declared_vars.clone();
                 let else_instructions = compile_lines(
+                    function_name,
                     else_branch,
                     compiler,
                     Some(end_label.clone()),
@@ -375,8 +393,13 @@ fn compile_lines(
                 compiler.bytecode.insert(if_label, then_instructions);
                 compiler.bytecode.insert(else_label, else_instructions);
 
-                let remaining =
-                    compile_lines(&lines[i + 1..], compiler, final_jump, declared_vars)?;
+                let remaining = compile_lines(
+                    function_name,
+                    &lines[i + 1..],
+                    compiler,
+                    final_jump,
+                    declared_vars,
+                )?;
                 compiler.bytecode.insert(end_label, remaining);
 
                 return Ok(instructions);
@@ -399,7 +422,7 @@ fn compile_lines(
             }
 
             SimpleLine::FunctionCall {
-                function_name,
+                function_name: callee_function_name,
                 args,
                 return_data,
             } => {
@@ -411,7 +434,7 @@ fn compile_lines(
                 compiler.stack_size += 1;
 
                 instructions.extend(setup_function_call(
-                    function_name,
+                    callee_function_name,
                     args,
                     new_fp_pos,
                     &return_label,
@@ -436,6 +459,7 @@ fn compile_lines(
                     }
 
                     instructions.extend(compile_lines(
+                        function_name,
                         &lines[i + 1..],
                         compiler,
                         final_jump,
diff --git a/crates/lean_compiler/src/c_compile_final.rs b/crates/lean_compiler/src/c_compile_final.rs
index 57672790..350b9db7 100644
--- a/crates/lean_compiler/src/c_compile_final.rs
+++ b/crates/lean_compiler/src/c_compile_final.rs
@@ -52,18 +52,39 @@ pub fn compile_to_low_level_bytecode(
         .get("main")
         .unwrap();
 
+    let mut hints = BTreeMap::new();
     let mut label_to_pc = BTreeMap::new();
     label_to_pc.insert(Label::function("main"), 0);
     let entrypoint = intermediate_bytecode
         .bytecode
         .remove(&Label::function("main"))
         .ok_or("No main function found in the compiled program")?;
+    hints.insert(
+        0,
+        vec![Hint::StackFrame {
+            label: Label::function("main"),
+            size: starting_frame_memory,
+        }],
+    );
 
-    let mut code_blocks = vec![(0, entrypoint.clone())];
+    let mut code_blocks = vec![(Label::function("main"), 0, entrypoint.clone())];
     let mut pc = count_real_instructions(&entrypoint);
+
     for (label, instructions) in &intermediate_bytecode.bytecode {
         label_to_pc.insert(label.clone(), pc);
-        code_blocks.push((pc, instructions.clone()));
+        if let Label::Function(function_name) = label {
+            hints
+                .entry(pc)
+                .or_insert_with(Vec::new)
+                .push(Hint::StackFrame {
+                    label: label.clone(),
+                    size: *intermediate_bytecode
+                        .memory_size_per_function
+                        .get(function_name)
+                        .unwrap(),
+                });
+        }
+        code_blocks.push((label.clone(), pc, instructions.clone()));
         pc += count_real_instructions(instructions);
     }
 
@@ -71,8 +92,12 @@ pub fn compile_to_low_level_bytecode(
 
     let mut match_block_sizes = Vec::new();
     let mut match_first_block_starts = Vec::new();
-    for match_statement in intermediate_bytecode.match_blocks {
-        let max_block_size = match_statement
+    for MatchBlock {
+        function_name,
+        match_cases,
+    } in intermediate_bytecode.match_blocks
+    {
+        let max_block_size = match_cases
             .iter()
             .map(|block| count_real_instructions(block))
             .max()
@@ -80,19 +105,18 @@ pub fn compile_to_low_level_bytecode(
         match_first_block_starts.push(pc);
         match_block_sizes.push(max_block_size);
 
-        for mut block in match_statement {
+        for mut block in match_cases {
             // fill the end of block with unreachable instructions
             block.extend(vec![
                 IntermediateInstruction::Panic;
                 max_block_size - count_real_instructions(&block)
             ]);
-            code_blocks.push((pc, block));
+            code_blocks.push((function_name.clone(), pc, block));
             pc += max_block_size;
         }
     }
 
     let mut low_level_bytecode = Vec::new();
-    let mut hints = BTreeMap::new();
 
     for (label, pc) in label_to_pc.clone() {
         hints
@@ -108,9 +132,10 @@ pub fn compile_to_low_level_bytecode(
         match_first_block_starts,
     };
 
-    for (pc_start, block) in code_blocks {
+    for (function_name, pc_start, block) in code_blocks {
         compile_block(
             &compiler,
+            &function_name,
             &block,
             pc_start,
             &mut low_level_bytecode,
@@ -130,6 +155,7 @@ pub fn compile_to_low_level_bytecode(
 
 fn compile_block(
     compiler: &Compiler,
+    function_name: &Label,
     block: &[IntermediateInstruction],
     pc_start: CodeAddress,
     low_level_bytecode: &mut Vec<Instruction>,
@@ -359,6 +385,7 @@ fn compile_block(
                     .unwrap()
                     .to_usize();
                 let hint = Hint::RequestMemory {
+                    function_name: function_name.clone(),
                     offset: eval_const_expression_usize(&offset, compiler),
                     vectorized,
                     size,
diff --git a/crates/lean_compiler/src/ir/bytecode.rs b/crates/lean_compiler/src/ir/bytecode.rs
index dcfec1e9..1279d070 100644
--- a/crates/lean_compiler/src/ir/bytecode.rs
+++ b/crates/lean_compiler/src/ir/bytecode.rs
@@ -3,6 +3,16 @@ use lean_vm::Label;
 use std::collections::BTreeMap;
 use std::fmt::{Display, Formatter};
 
+/// A match statement bytecode block
+#[derive(Debug, Clone)]
+pub struct MatchBlock {
+    /// Name of the function containing the match block
+    pub function_name: Label,
+
+    /// Cases of the match block
+    pub match_cases: Vec<Vec<IntermediateInstruction>>,
+}
+
 /// Container for the complete intermediate representation of a program.
 ///
 /// This structure holds all the compiled intermediate bytecode along with
@@ -17,7 +27,7 @@ pub struct IntermediateBytecode {
     /// Match statement bytecode blocks.
     ///
     /// Each match statement produces multiple case blocks.
-    pub match_blocks: Vec<Vec<Vec<IntermediateInstruction>>>,
+    pub match_blocks: Vec<MatchBlock>,
 
     /// Memory requirements for each function.
     ///
@@ -29,13 +39,13 @@ impl Display for IntermediateBytecode {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
         for (label, instructions) in &self.bytecode {
             writeln!(f, "\n{label}:")?;
-            for instruction in instructions {
+            for instruction in instructions.iter() {
                 writeln!(f, "  {instruction}")?;
             }
         }
-        for (i, match_blocks) in self.match_blocks.iter().enumerate() {
+        for (i, MatchBlock { match_cases, .. }) in self.match_blocks.iter().enumerate() {
             writeln!(f, "\nMatch {i}:")?;
-            for (j, case) in match_blocks.iter().enumerate() {
+            for (j, case) in match_cases.iter().enumerate() {
                 writeln!(f, "  Case {j}:")?;
                 for instruction in case {
                     writeln!(f, "    {instruction}")?;
diff --git a/crates/lean_compiler/src/ir/mod.rs b/crates/lean_compiler/src/ir/mod.rs
index d81bc047..43d34bad 100644
--- a/crates/lean_compiler/src/ir/mod.rs
+++ b/crates/lean_compiler/src/ir/mod.rs
@@ -5,7 +5,7 @@ pub mod instruction;
 pub mod operation;
 pub mod value;
 
-pub use bytecode::IntermediateBytecode;
+pub use bytecode::{IntermediateBytecode, MatchBlock};
 pub use instruction::IntermediateInstruction;
 pub use operation::HighLevelOperation;
 pub use value::{IntermediaryMemOrFpOrConstant, IntermediateValue};
diff --git a/crates/lean_vm/src/core/constants.rs b/crates/lean_vm/src/core/constants.rs
index 1b8e4b15..7d6ca413 100644
--- a/crates/lean_vm/src/core/constants.rs
+++ b/crates/lean_vm/src/core/constants.rs
@@ -23,7 +23,7 @@ pub const MAX_RUNNER_MEMORY_SIZE: usize = 1 << 24;
 /// program_input: the input of the program we want to prove
 ///
 /// [reserved_area] = [00000000] [00000000] [10000000] [poseidon_16(0) (16 field elements)] [poseidon_24(0) (8 last field elements)]
-
+///
 /// Convention: vectorized pointer of size 2, pointing to 16 zeros
 pub const ZERO_VEC_PTR: usize = 0;
 
diff --git a/crates/lean_vm/src/core/types.rs b/crates/lean_vm/src/core/types.rs
index bfcaca23..8538abb6 100644
--- a/crates/lean_vm/src/core/types.rs
+++ b/crates/lean_vm/src/core/types.rs
@@ -11,3 +11,6 @@ pub type SourceLineNumber = usize;
 
 /// Bytecode address (i.e., a value of the program counter)
 pub type CodeAddress = usize;
+
+/// Memory address
+pub type MemoryAddress = usize;
diff --git a/crates/lean_vm/src/diagnostics/error.rs b/crates/lean_vm/src/diagnostics/error.rs
index 0a225733..cc8ed616 100644
--- a/crates/lean_vm/src/diagnostics/error.rs
+++ b/crates/lean_vm/src/diagnostics/error.rs
@@ -1,4 +1,5 @@
 use crate::core::F;
+use crate::diagnostics::profiler::MemoryProfile;
 use crate::execution::Memory;
 use crate::witness::{
     WitnessDotProduct, WitnessMultilinearEval, WitnessPoseidon16, WitnessPoseidon24,
@@ -46,4 +47,5 @@ pub struct ExecutionResult {
     pub dot_products: Vec<WitnessDotProduct>,
     pub multilinear_evals: Vec<WitnessMultilinearEval>,
     pub summary: String,
+    pub memory_profile: Option<MemoryProfile>,
 }
diff --git a/crates/lean_vm/src/diagnostics/profiler.rs b/crates/lean_vm/src/diagnostics/profiler.rs
index 182ee016..633ed3ad 100644
--- a/crates/lean_vm/src/diagnostics/profiler.rs
+++ b/crates/lean_vm/src/diagnostics/profiler.rs
@@ -1,10 +1,11 @@
-use std::collections::BTreeMap;
-use std::collections::HashMap;
+use std::collections::{BTreeMap, BTreeSet, HashMap};
+use std::ops::Range;
 
 use utils::pretty_integer;
 
-use crate::ExecutionHistory;
+use crate::core::{Label, MemoryAddress};
 use crate::stack_trace::find_function_for_line;
+use crate::{ExecutionHistory, NONRESERVED_PROGRAM_INPUT_START};
 
 pub(crate) fn profiling_report(
     instructions: &ExecutionHistory,
@@ -99,3 +100,231 @@ pub(crate) fn profiling_report(
 
     report
 }
+
+#[derive(Debug, Clone)]
+pub enum MemoryObjectType {
+    StackFrame,
+    NonVectorHeapObject,
+    VectorHeapObject,
+}
+
+/// An object (i.e., stack frame or heap object) in the memory profile
+#[derive(Debug, Clone)]
+pub struct MemoryObject {
+    pub object_type: MemoryObjectType,
+
+    /// The name of the function allocating this object
+    pub function_name: Label,
+
+    /// The number of words in this stack frame
+    pub size: usize,
+}
+
+/// Struct to accumulate memory profiling results
+#[derive(Debug, Clone)]
+pub struct MemoryProfile {
+    /// The set of used memory addresses
+    pub used: BTreeSet<MemoryAddress>,
+
+    /// The range of public input addresses
+    pub public_inputs: Range<MemoryAddress>,
+
+    /// The range of private input addresses
+    pub private_inputs: Range<MemoryAddress>,
+
+    /// The memory objects keyed by their start addresses
+    pub objects: BTreeMap<MemoryAddress, MemoryObject>,
+}
+
+pub(crate) fn memory_profiling_report(profile: &MemoryProfile) -> String {
+    let mut report = String::new();
+
+    let f_allocations = function_allocations(profile);
+    let all_allocated_mem = all_allocated_memory(profile);
+    let footprint = all_allocated_mem.last().unwrap() + 1;
+    let allocated = all_allocated_mem.len();
+    let allocated_percent = (allocated as f64 / footprint as f64) * 100.0;
+    let used = profile.used.len();
+    let used_but_not_allocated = count_used_but_not_allocated(&profile.used, &all_allocated_mem);
+    let used_but_not_allocated_percent = (used_but_not_allocated as f64 / footprint as f64) * 100.0;
+    let allocated_but_not_used = count_allocated_but_not_used(&profile.used, &all_allocated_mem);
+    let allocated_but_not_used_percent = (allocated_but_not_used as f64 / allocated as f64) * 100.0;
+
+    // TODO: add percentages used
+    report.push_str("============================================\n");
+    report.push_str("=        DETAILED MEMORY PROFILING         =\n");
+    report.push_str("============================================\n");
+    report.push('\n');
+    report.push_str(&format!(
+        "Total memory footprint: {}\n",
+        pretty_integer(footprint)
+    ));
+    report.push_str(&format!(
+        "Total allocated memory: {} ({:.2}% of footprint)\n",
+        pretty_integer(allocated),
+        allocated_percent
+    ));
+    report.push_str(&format!("Total used memory: {}\n", pretty_integer(used)));
+    report.push_str(&format!(
+        "Total used but not allocated: {} ({:.2}% of footprint)\n",
+        pretty_integer(used_but_not_allocated),
+        used_but_not_allocated_percent
+    ));
+    report.push_str(&format!(
+        "Total allocated but not used: {} ({:.2}% of allocated)\n\n",
+        pretty_integer(allocated_but_not_used),
+        allocated_but_not_used_percent
+    ));
+
+    for (function_name, f_allocations) in f_allocations {
+        let FunctionAllocationStats {
+            frame_words_allocated,
+            frame_words_used,
+            heap_nonvector_words_allocated,
+            heap_nonvector_words_used,
+            heap_vector_words_allocated,
+            heap_vector_words_used,
+        } = function_allocation_stats(&profile.used, &f_allocations);
+        report.push_str(&format!("Function {function_name}:\n"));
+        report.push_str(&format!(
+            " * Frame words allocated: {}\n",
+            pretty_integer(frame_words_allocated)
+        ));
+        report.push_str(&format!(
+            " * Frame words used: {} ({:.2}%)\n",
+            pretty_integer(frame_words_used),
+            (frame_words_used as f64 / frame_words_allocated as f64) * 100.0
+        ));
+        report.push_str(&format!(
+            " * Non-vector heap words allocated: {}\n",
+            pretty_integer(heap_nonvector_words_allocated)
+        ));
+        report.push_str(&format!(
+            " * Non-vector heap words used: {} ({:.2}%)\n",
+            pretty_integer(heap_nonvector_words_used),
+            (heap_nonvector_words_used as f64 / heap_nonvector_words_allocated as f64) * 100.0
+        ));
+        report.push_str(&format!(
+            " * Vector heap words allocated: {}\n",
+            pretty_integer(heap_vector_words_allocated)
+        ));
+        report.push_str(&format!(
+            " * Vector heap words used: {} ({:.2}%)\n\n",
+            pretty_integer(heap_vector_words_used),
+            (heap_vector_words_used as f64 / heap_vector_words_allocated as f64) * 100.0
+        ));
+    }
+
+    report
+}
+
+fn function_allocations(
+    profile: &MemoryProfile,
+) -> BTreeMap<Label, BTreeMap<MemoryAddress, MemoryObject>> {
+    let mut allocations = BTreeMap::new();
+
+    for (addr, object) in profile.objects.iter() {
+        let objects = if let Some(objects) = allocations.get_mut(&object.function_name) {
+            objects
+        } else {
+            allocations.insert(object.function_name.clone(), BTreeMap::new());
+            allocations.get_mut(&object.function_name).unwrap()
+        };
+        objects.insert(*addr, object.clone());
+    }
+
+    allocations
+}
+
+#[derive(Debug, Clone, Default)]
+struct FunctionAllocationStats {
+    frame_words_allocated: usize,
+    frame_words_used: usize,
+    heap_nonvector_words_allocated: usize,
+    heap_nonvector_words_used: usize,
+    heap_vector_words_allocated: usize,
+    heap_vector_words_used: usize,
+}
+
+/// Get the allocation stats for a particular function.
+fn function_allocation_stats(
+    used: &BTreeSet<MemoryAddress>,
+    function_allocations: &BTreeMap<MemoryAddress, MemoryObject>,
+) -> FunctionAllocationStats {
+    let mut stats = FunctionAllocationStats::default();
+
+    for (object_address, object) in function_allocations.iter() {
+        let address_range = *object_address..*object_address + object.size;
+        let object_words_used = used.range(address_range).count();
+
+        match object.object_type {
+            MemoryObjectType::StackFrame => {
+                stats.frame_words_allocated += object.size;
+                stats.frame_words_used += object_words_used;
+            }
+            MemoryObjectType::NonVectorHeapObject => {
+                stats.heap_nonvector_words_allocated += object.size;
+                stats.heap_nonvector_words_used += object_words_used;
+            }
+            MemoryObjectType::VectorHeapObject => {
+                stats.heap_vector_words_allocated += object.size;
+                stats.heap_vector_words_used += object_words_used;
+            }
+        }
+    }
+
+    stats
+}
+
+/// Get the set of all allocated memory addresses, including public and private input regions.
+fn all_allocated_memory(profile: &MemoryProfile) -> BTreeSet<MemoryAddress> {
+    let mut result = BTreeSet::new();
+
+    // Reserved public inputs
+    for addr in 0..NONRESERVED_PROGRAM_INPUT_START {
+        result.insert(addr);
+    }
+
+    // Non-reserved public inputs
+    for addr in profile.public_inputs.clone() {
+        result.insert(addr);
+    }
+
+    // Private inputs
+    for addr in profile.private_inputs.clone() {
+        result.insert(addr);
+    }
+
+    // Stack and heap objects
+    for (start_addr, object) in profile.objects.iter() {
+        for addr in *start_addr..start_addr + object.size {
+            result.insert(addr);
+        }
+    }
+
+    result
+}
+
+/// Get the number of used memory addresses which are not allocated.
+fn count_used_but_not_allocated(
+    used: &BTreeSet<MemoryAddress>,
+    allocated: &BTreeSet<MemoryAddress>,
+) -> usize {
+    let diff = BTreeSet::from_iter(used.difference(allocated));
+    let len = diff.len();
+    if len > 0 {
+        println!("!! Addresses defined but not used:");
+        for addr in diff {
+            println!("{addr}");
+        }
+    }
+    len
+}
+
+/// Get the number of allocated memory addresses which are not used.
+fn count_allocated_but_not_used(
+    used: &BTreeSet<MemoryAddress>,
+    allocated: &BTreeSet<MemoryAddress>,
+) -> usize {
+    allocated.difference(used).count()
+}
diff --git a/crates/lean_vm/src/execution/runner.rs b/crates/lean_vm/src/execution/runner.rs
index 37106d40..d19064f7 100644
--- a/crates/lean_vm/src/execution/runner.rs
+++ b/crates/lean_vm/src/execution/runner.rs
@@ -4,7 +4,7 @@ use crate::core::{
     DIMENSION, F, NONRESERVED_PROGRAM_INPUT_START, ONE_VEC_PTR, POSEIDON_16_NULL_HASH_PTR,
     POSEIDON_24_NULL_HASH_PTR, VECTOR_LEN, ZERO_VEC_PTR,
 };
-use crate::diagnostics::{ExecutionResult, RunnerError};
+use crate::diagnostics::{ExecutionResult, MemoryProfile, RunnerError, memory_profiling_report};
 use crate::execution::{ExecutionHistory, Memory};
 use crate::isa::Bytecode;
 use crate::isa::instruction::InstructionContext;
@@ -13,7 +13,7 @@ use crate::witness::{
 };
 use crate::{CodeAddress, HintExecutionContext, SourceLineNumber};
 use p3_field::PrimeCharacteristicRing;
-use std::collections::BTreeMap;
+use std::collections::{BTreeMap, BTreeSet};
 use utils::{poseidon16_permute, poseidon24_permute, pretty_integer};
 use xmss::{Poseidon16History, Poseidon24History};
 
@@ -67,7 +67,7 @@ pub fn execute_bytecode(
     bytecode: &Bytecode,
     (public_input, private_input): (&[F], &[F]),
     no_vec_runtime_memory: usize, // size of the "non-vectorized" runtime memory
-    profiler: bool,
+    profiling: bool,
     (poseidons_16_precomputed, poseidons_24_precomputed): (&Poseidon16History, &Poseidon24History),
 ) -> ExecutionResult {
     let mut std_out = String::new();
@@ -78,7 +78,7 @@ pub fn execute_bytecode(
         &mut std_out,
         &mut instruction_history,
         no_vec_runtime_memory,
-        profiler,
+        profiling,
         (poseidons_16_precomputed, poseidons_24_precomputed),
     )
     .unwrap_or_else(|err| {
@@ -101,9 +101,12 @@ pub fn execute_bytecode(
         }
         panic!("Error during bytecode execution: {err}");
     });
-    if profiler {
+    if profiling {
         print_line_cycle_counts(instruction_history);
         print_instruction_cycle_counts(bytecode, result.pcs.clone());
+        if let Some(ref mem_profile) = result.memory_profile {
+            print!("{}", memory_profiling_report(mem_profile));
+        }
     }
     result
 }
@@ -153,7 +156,7 @@ fn execute_bytecode_helper(
     std_out: &mut String,
     instruction_history: &mut ExecutionHistory,
     no_vec_runtime_memory: usize,
-    profiler: bool,
+    profiling: bool,
     (poseidons_16_precomputed, poseidons_24_precomputed): (&Poseidon16History, &Poseidon24History),
 ) -> Result<ExecutionResult, RunnerError> {
     // set public memory
@@ -166,6 +169,15 @@ fn execute_bytecode_helper(
     for (i, value) in private_input.iter().enumerate() {
         memory.set(fp + i, *value)?;
     }
+
+    let mut mem_profile = MemoryProfile {
+        used: BTreeSet::new(),
+        public_inputs: NONRESERVED_PROGRAM_INPUT_START
+            ..NONRESERVED_PROGRAM_INPUT_START + public_memory_size,
+        private_inputs: fp..fp + private_input.len(),
+        objects: BTreeMap::new(),
+    };
+
     fp += private_input.len();
     fp = fp.next_multiple_of(DIMENSION);
 
@@ -228,6 +240,8 @@ fn execute_bytecode_helper(
                 last_checkpoint_cpu_cycles: &mut last_checkpoint_cpu_cycles,
                 checkpoint_ap: &mut checkpoint_ap,
                 checkpoint_ap_vec: &mut checkpoint_ap_vec,
+                profiling,
+                memory_profile: &mut mem_profile,
             };
             hint.execute_hint(&mut hint_ctx)?;
         }
@@ -289,7 +303,7 @@ fn execute_bytecode_helper(
 
     let mut summary = String::new();
 
-    if profiler {
+    if profiling {
         let report =
             crate::diagnostics::profiling_report(instruction_history, &bytecode.function_locations);
         summary.push_str(&report);
@@ -392,6 +406,14 @@ fn execute_bytecode_helper(
     summary
         .push_str("──────────────────────────────────────────────────────────────────────────\n");
 
+    if profiling {
+        for (addr, val) in (0..).zip(memory.0.iter()) {
+            if val.is_some() {
+                mem_profile.used.insert(addr);
+            }
+        }
+    }
+
     Ok(ExecutionResult {
         no_vec_runtime_memory,
         public_memory_size,
@@ -403,5 +425,6 @@ fn execute_bytecode_helper(
         dot_products,
         multilinear_evals,
         summary,
+        memory_profile: if profiling { Some(mem_profile) } else { None },
     })
 }
diff --git a/crates/lean_vm/src/isa/hint.rs b/crates/lean_vm/src/isa/hint.rs
index 7b4f1c70..ddadc5d0 100644
--- a/crates/lean_vm/src/isa/hint.rs
+++ b/crates/lean_vm/src/isa/hint.rs
@@ -1,5 +1,5 @@
 use crate::core::{F, LOG_VECTOR_LEN, Label, SourceLineNumber, VECTOR_LEN};
-use crate::diagnostics::RunnerError;
+use crate::diagnostics::{MemoryObject, MemoryObjectType, MemoryProfile, RunnerError};
 use crate::execution::{ExecutionHistory, Memory};
 use crate::isa::operands::MemOrConstant;
 use p3_field::{Field, PrimeCharacteristicRing};
@@ -19,6 +19,8 @@ pub enum Hint {
     },
     /// Request memory allocation
     RequestMemory {
+        /// Function name this allocation is associated with (for profiling)
+        function_name: Label,
         /// Memory offset where hint will be stored: m[fp + offset]
         offset: usize,
         /// The requested memory size
@@ -65,6 +67,8 @@ pub enum Hint {
     },
     /// Jump destination label (for debugging purposes)
     Label { label: Label },
+    /// Stack frame size (for memory profiling)
+    StackFrame { label: Label, size: usize },
 }
 
 /// Execution state for hint processing
@@ -82,6 +86,8 @@ pub struct HintExecutionContext<'a> {
     pub last_checkpoint_cpu_cycles: &'a mut usize,
     pub checkpoint_ap: &'a mut usize,
     pub checkpoint_ap_vec: &'a mut usize,
+    pub profiling: bool,
+    pub memory_profile: &'a mut MemoryProfile,
 }
 
 impl Hint {
@@ -90,6 +96,7 @@ impl Hint {
     pub fn execute_hint(&self, ctx: &mut HintExecutionContext<'_>) -> Result<(), RunnerError> {
         match self {
             Self::RequestMemory {
+                function_name,
                 offset,
                 size,
                 vectorized,
@@ -104,14 +111,41 @@ impl Hint {
                     while !(*ctx.ap_vec * VECTOR_LEN).is_multiple_of(1 << *vectorized_len) {
                         *ctx.ap_vec += 1;
                     }
+                    let allocation_start_addr = *ctx.ap_vec;
                     ctx.memory.set(
                         ctx.fp + *offset,
-                        F::from_usize(*ctx.ap_vec >> (*vectorized_len - LOG_VECTOR_LEN)),
+                        F::from_usize(allocation_start_addr >> (*vectorized_len - LOG_VECTOR_LEN)),
                     )?;
-                    *ctx.ap_vec += size << (*vectorized_len - LOG_VECTOR_LEN);
+                    let size_vectors = size << (*vectorized_len - LOG_VECTOR_LEN);
+                    let size_words = size_vectors * VECTOR_LEN;
+                    *ctx.ap_vec += size_vectors;
+
+                    if ctx.profiling {
+                        ctx.memory_profile.objects.insert(
+                            allocation_start_addr * VECTOR_LEN,
+                            MemoryObject {
+                                object_type: MemoryObjectType::VectorHeapObject,
+                                function_name: function_name.clone(),
+                                size: size_words,
+                            },
+                        );
+                    }
                 } else {
-                    ctx.memory.set(ctx.fp + *offset, F::from_usize(*ctx.ap))?;
+                    let allocation_start_addr = *ctx.ap;
+                    ctx.memory
+                        .set(ctx.fp + *offset, F::from_usize(allocation_start_addr))?;
                     *ctx.ap += size;
+
+                    if ctx.profiling {
+                        ctx.memory_profile.objects.insert(
+                            allocation_start_addr,
+                            MemoryObject {
+                                object_type: MemoryObjectType::NonVectorHeapObject,
+                                function_name: function_name.clone(),
+                                size,
+                            },
+                        );
+                    }
                 }
             }
             Self::DecomposeBits {
@@ -193,6 +227,18 @@ impl Hint {
                 *ctx.cpu_cycles_before_new_line = 0;
             }
             Self::Label { .. } => {}
+            Self::StackFrame { label, size } => {
+                if ctx.profiling {
+                    ctx.memory_profile.objects.insert(
+                        ctx.fp,
+                        MemoryObject {
+                            object_type: MemoryObjectType::StackFrame,
+                            function_name: label.clone(),
+                            size: *size,
+                        },
+                    );
+                }
+            }
         }
         Ok(())
     }
@@ -202,6 +248,7 @@ impl Display for Hint {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
         match self {
             Self::RequestMemory {
+                function_name: _,
                 offset,
                 size,
                 vectorized,
@@ -266,6 +313,9 @@ impl Display for Hint {
             Self::Label { label } => {
                 write!(f, "label: {label}")
             }
+            Self::StackFrame { label, size } => {
+                write!(f, "stack frame for {label} size {size}")
+            }
         }
     }
 }
diff --git a/crates/lean_vm/tests/test_lean_vm.rs b/crates/lean_vm/tests/test_lean_vm.rs
index e58cfd6c..3b6841cd 100644
--- a/crates/lean_vm/tests/test_lean_vm.rs
+++ b/crates/lean_vm/tests/test_lean_vm.rs
@@ -122,6 +122,7 @@ fn build_test_case() -> (Bytecode, Vec<F>) {
     hints.insert(
         0,
         vec![Hint::RequestMemory {
+            function_name: Label::function("main"),
             offset: POSEIDON16_RES_OFFSET,
             size: MemOrConstant::Constant(f(2)),
             vectorized: true,
@@ -131,6 +132,7 @@ fn build_test_case() -> (Bytecode, Vec<F>) {
     hints.insert(
         1,
         vec![Hint::RequestMemory {
+            function_name: Label::function("main"),
             offset: POSEIDON24_RES_OFFSET,
             size: MemOrConstant::Constant(f(1)),
             vectorized: true,
@@ -140,6 +142,7 @@ fn build_test_case() -> (Bytecode, Vec<F>) {
     hints.insert(
         2,
         vec![Hint::RequestMemory {
+            function_name: Label::function("main"),
             offset: DOT_RES_OFFSET,
             size: MemOrConstant::Constant(f(1)),
             vectorized: false,
@@ -149,6 +152,7 @@ fn build_test_case() -> (Bytecode, Vec<F>) {
     hints.insert(
         3,
         vec![Hint::RequestMemory {
+            function_name: Label::function("main"),
             offset: MLE_RES_OFFSET,
             size: MemOrConstant::Constant(f(1)),
             vectorized: true,
diff --git a/crates/poseidon_circuit/src/tests.rs b/crates/poseidon_circuit/src/tests.rs
index 17e92b68..ee053ac6 100644
--- a/crates/poseidon_circuit/src/tests.rs
+++ b/crates/poseidon_circuit/src/tests.rs
@@ -314,7 +314,7 @@ pub fn run_poseidon_benchmark<
             .collect::<Vec<_>>()
     );
 
-    println!("2^{} Poseidon2", log_n_poseidons);
+    println!("2^{log_n_poseidons} Poseidon2");
     println!(
         "Plaintext (no proof) time: {:.3}s ({:.2}M Poseidons / s)",
         plaintext_duration.as_secs_f64(),
diff --git a/crates/rec_aggregation/src/recursion.rs b/crates/rec_aggregation/src/recursion.rs
index 607aa5ce..a01dfed2 100644
--- a/crates/rec_aggregation/src/recursion.rs
+++ b/crates/rec_aggregation/src/recursion.rs
@@ -202,7 +202,7 @@ pub fn run_whir_recursion_benchmark() {
     let proving_time = time.elapsed();
     verify_execution(&bytecode, &public_input, proof_data, whir_config_builder()).unwrap();
 
-    println!("{}", summary);
+    println!("{summary}");
     println!(
         "WHIR recursion, proving time: {} ms, proof size: {} KiB (not optimized)",
         proving_time.as_millis(),
diff --git a/crates/rec_aggregation/src/xmss_aggregate.rs b/crates/rec_aggregation/src/xmss_aggregate.rs
index 535182a2..833bc0e2 100644
--- a/crates/rec_aggregation/src/xmss_aggregate.rs
+++ b/crates/rec_aggregation/src/xmss_aggregate.rs
@@ -285,7 +285,7 @@ pub fn run_xmss_benchmark(n_xmss: usize) {
     let proving_time = time.elapsed();
     verify_execution(&bytecode, &public_input, proof_data, whir_config_builder()).unwrap();
 
-    println!("{}", summary);
+    println!("{summary}");
     println!(
         "XMSS aggregation, proving time: {:.3} s ({:.1} XMSS/s), proof size: {} KiB (not optimized)",
         proving_time.as_secs_f64(),