Refactor EVM memory planning

crates/codegen/src/isa/evm/backend.rs

@@ -21,7 +21,7 @@ use super::{
     exact_func_merge::run_exact_private_func_merge,
     high_alias::compute_high_evm_aliases,
     late_section_merge::run_late_section_terminal_outline,
-    memory_plan::{ArenaCostModel, ObjLoc, PreserveMode, StableMode, WORD_BYTES},
+    memory_plan::{ArenaCostModel, ObjLoc, StableMode, WORD_BYTES},
     opcode::OpCode,
     prepare::{self, EvmPreparedSection},
 };
@@ -161,11 +161,11 @@ impl EvmBackend {
         let mut out = String::new();
         writeln!(
             &mut out,
-            "evm mem plan: global_dyn_base=0x{:x} arena_base=0x{:x} scratch_peak_words={} static_chain_peak_words={}",
+            "evm mem plan: global_dyn_base=0x{:x} arena_base=0x{:x} scratch_peak_words={} stable_chain_peak_words={}",
             prepared.section_plan().dyn_base,
             prepared.section_plan().arena_base,
             prepared.section_plan().scratch_peak_words,
-            prepared.section_plan().static_chain_peak_words
+            prepared.section_plan().stable_chain_peak_words
         )
         .expect("mem plan write failed");
 
@@ -178,9 +178,9 @@ impl EvmBackend {
             let stable = match func_plan.stable_mode {
                 StableMode::None => "None".to_string(),
                 StableMode::DynamicFrame => "DynamicFrame".to_string(),
-                StableMode::StaticAbs { base_word } => {
+                StableMode::StableAbs { base_word } => {
                     format!(
-                        "StaticAbs(base=0x{:x})",
+                        "StableAbs(base=0x{:x})",
                         func_plan.abs_addr_for_word(base_word)
                     )
                 }
@@ -218,7 +218,6 @@ impl EvmBackend {
                         .sp_relative_bytes(frame_layout.expect_local_word(off, "debug object"));
                     format!("sp-0x{addr_bytes:x}")
                 }
-                ObjLoc::StackPinned(depth) => format!("stack[{depth}]"),
             };
 
             if detail {
@@ -249,31 +248,17 @@ impl EvmBackend {
                         )
                     });
                     let callee_name = module.ctx.func_sig(callee, |sig| sig.name().to_string());
-                    match &plan.mode {
-                        PreserveMode::None => {}
-                        PreserveMode::ShadowRuns { shadow_obj, runs } => {
-                            let save_words: u32 = runs.iter().map(|run| run.len_words).sum();
-                            writeln!(
-                                &mut out,
-                                "  call inst{} callee=%{callee_name} preserve=ShadowRuns shadow_obj={} result_count={} save_words={} runs={runs:?}",
-                                inst.as_u32(),
-                                shadow_obj.as_u32(),
-                                plan.result_count,
-                                save_words,
-                            )
-                            .expect("mem plan write failed");
-                        }
-                        PreserveMode::TinyStackLift { word_offsets } => {
-                            writeln!(
-                                &mut out,
-                                "  call inst{} callee=%{callee_name} preserve=TinyStackLift result_count={} save_words={} save_offsets={word_offsets:?}",
-                                inst.as_u32(),
-                                plan.result_count,
-                                word_offsets.len(),
-                            )
-                            .expect("mem plan write failed");
-                        }
-                    }
+                    let runs = &plan.runs;
+                    let save_words: u32 = runs.iter().map(|run| run.len_words).sum();
+                    writeln!(
+                        &mut out,
+                        "  call inst{} callee=%{callee_name} preserve=ShadowRuns shadow_obj={} result_count={} save_words={} runs={runs:?}",
+                        inst.as_u32(),
+                        plan.shadow_obj.as_u32(),
+                        plan.result_count,
+                        save_words,
+                    )
+                    .expect("mem plan write failed");
                 }
 
                 let mut scratch_spills: Vec<(ValueId, u32)> = Vec::new();

crates/codegen/src/isa/evm/dyn_sp.rs

@@ -1,12 +1,9 @@
 use rustc_hash::{FxHashMap, FxHashSet};
 use sonatina_ir::{InstId, Module, module::FuncRef};
 
-use crate::module_analysis::{CallGraph, SccBuilder};
+use crate::module_analysis::CallGraphSchedule;
 
-use super::{
-    machine::lazy_frame::FrameSummary,
-    memory_plan::{self, FuncMemPlan},
-};
+use super::{machine::lazy_frame::FrameSummary, memory_plan::MachineFuncPlan};
 
 #[derive(Clone, Default)]
 pub(crate) struct FuncDynSpPlan {
@@ -60,43 +57,28 @@ impl EntryFrameState {
 
 pub(crate) fn compute_machine_dyn_sp_plan(
     module: &Module,
-    funcs: &[FuncRef],
+    schedule: &CallGraphSchedule,
     section_entry: FuncRef,
-    mem_plans: &FxHashMap<FuncRef, FuncMemPlan>,
+    mem_plans: &FxHashMap<FuncRef, MachineFuncPlan>,
     frame_summaries: &FxHashMap<FuncRef, FrameSummary>,
 ) -> FxHashMap<FuncRef, FuncDynSpPlan> {
-    let funcs_set: FxHashSet<FuncRef> = funcs.iter().copied().collect();
-    let call_graph = CallGraph::build_graph_subset(module, &funcs_set);
-    let mut reachable_funcs = FxHashSet::default();
-    let mut worklist = vec![section_entry];
-    while let Some(func) = worklist.pop() {
-        if !reachable_funcs.insert(func) {
-            continue;
-        }
-        for &callee in call_graph.callee_of(func) {
-            worklist.push(callee);
-        }
-    }
+    let reachable_funcs = schedule.reachable_from(section_entry);
 
-    let mut ordered_funcs = funcs.to_vec();
-    ordered_funcs.sort_unstable_by_key(|func| func.as_u32());
+    let ordered_funcs = schedule.funcs().to_vec();
     let mut ordered_reachable: Vec<FuncRef> = reachable_funcs.iter().copied().collect();
     ordered_reachable.sort_unstable_by_key(|func| func.as_u32());
 
-    let scc = SccBuilder::new().compute_scc(&call_graph);
-    let topo = memory_plan::topo_sort_sccs(&reachable_funcs, &call_graph, &scc);
     let mut ready_sccs = FxHashSet::default();
-    for &scc_ref in &topo {
-        if scc
-            .scc_info(scc_ref)
-            .components
+    for &scc_ref in &schedule.topo {
+        if schedule
+            .members(scc_ref)
             .iter()
             .copied()
             .filter(|func| reachable_funcs.contains(func))
             .any(|func| {
                 mem_plans
                     .get(&func)
-                    .is_some_and(FuncMemPlan::uses_dynamic_frame)
+                    .is_some_and(MachineFuncPlan::uses_dynamic_frame)
             })
         {
             ready_sccs.insert(scc_ref);
@@ -156,7 +138,7 @@ pub(crate) fn compute_machine_dyn_sp_plan(
         })
         .collect();
 
-    if ready_sccs.contains(&scc.scc_ref(section_entry)) {
+    if ready_sccs.contains(&schedule.sccs.scc_ref(section_entry)) {
         let entry_live_frame = plans
             .get(&section_entry)
             .is_some_and(|plan| plan.entry_live_frame);
@@ -171,7 +153,7 @@ pub(crate) fn compute_machine_dyn_sp_plan(
     }
 
     for &caller in &ordered_reachable {
-        let caller_scc = scc.scc_ref(caller);
+        let caller_scc = schedule.sccs.scc_ref(caller);
         if ready_sccs.contains(&caller_scc) {
             continue;
         }
@@ -190,7 +172,7 @@ pub(crate) fn compute_machine_dyn_sp_plan(
                     if !reachable_funcs.contains(&callee) {
                         continue;
                     }
-                    let callee_scc = scc.scc_ref(callee);
+                    let callee_scc = schedule.sccs.scc_ref(callee);
                     if caller_scc == callee_scc
                         || !ready_sccs.contains(&callee_scc)
                         || caller_summary.local_frame_active_before_inst(inst)

crates/codegen/src/isa/evm/emit/alloc.rs

@@ -2,17 +2,15 @@ use sonatina_ir::{BlockId, Function, I256, Immediate, InstId};
 
 use crate::stackalloc::{Action, Actions, Allocator, StackifyAlloc};
 
-use super::super::{
-    DynamicFrameLayout, FuncMemPlan, ObjLoc, PreserveMode, static_arena_alloc::StackObjId,
-};
+use super::super::{DynamicFrameLayout, MachineFuncPlan, ObjLoc, static_arena_alloc::StackObjId};
 
 pub(crate) struct FinalAlloc {
     inner: StackifyAlloc,
-    pub(crate) mem_plan: FuncMemPlan,
+    pub(crate) mem_plan: MachineFuncPlan,
 }
 
 impl FinalAlloc {
-    pub(crate) fn new(inner: StackifyAlloc, mem_plan: FuncMemPlan) -> Self {
+    pub(crate) fn new(inner: StackifyAlloc, mem_plan: MachineFuncPlan) -> Self {
         let alloc = Self { inner, mem_plan };
         alloc.validate_object_actions();
         alloc
@@ -84,9 +82,6 @@ impl FinalAlloc {
                 ))
             }
             ObjLoc::StableFrame(off) => frame(self.frame_action_offset(off, extra_words, kind)),
-            ObjLoc::StackPinned(depth) => {
-                panic!("stack-pinned objects are not supported in EVM lowering (depth={depth})")
-            }
         }
     }
 
@@ -151,9 +146,6 @@ impl FinalAlloc {
                 offset_words: self.frame_action_offset(off, 0, "exact local addr"),
                 extra_bytes: offset_bytes,
             },
-            ObjLoc::StackPinned(depth) => {
-                panic!("stack-pinned exact local addresses are not supported (depth={depth})")
-            }
         }
     }
 
@@ -166,9 +158,7 @@ impl FinalAlloc {
         let Some(plan) = self.mem_plan.call_preserve.get(&inst) else {
             return actions;
         };
-        let PreserveMode::ShadowRuns { shadow_obj, runs } = &plan.mode else {
-            return actions;
-        };
+        let (shadow_obj, runs) = (&plan.shadow_obj, &plan.runs);
         if runs.is_empty() {
             return actions;
         }
@@ -212,9 +202,7 @@ impl FinalAlloc {
         let Some(plan) = self.mem_plan.call_preserve.get(&inst) else {
             return actions;
         };
-        let PreserveMode::ShadowRuns { shadow_obj, runs } = &plan.mode else {
-            return actions;
-        };
+        let (shadow_obj, runs) = (&plan.shadow_obj, &plan.runs);
         if runs.is_empty() {
             return actions;
         }
@@ -295,19 +283,19 @@ impl Allocator for FinalAlloc {
 #[cfg(test)]
 mod tests {
     use cranelift_entity::SecondaryMap;
-    use rustc_hash::{FxHashMap, FxHashSet};
+    use rustc_hash::FxHashMap;
     use sonatina_ir::InstId;
 
     use super::{FinalAlloc, StackifyAlloc};
     use crate::{
-        isa::evm::{FuncMemPlan, ObjLoc, STATIC_BASE, memory_plan::StableMode},
+        isa::evm::{MachineFuncPlan, ObjLoc, STATIC_BASE, memory_plan::StableMode},
         stackalloc::Action,
     };
 
-    fn mem_plan_with_alloca(alloca: InstId, loc: ObjLoc, stable_words: u32) -> FuncMemPlan {
+    fn mem_plan_with_alloca(alloca: InstId, loc: ObjLoc, stable_words: u32) -> MachineFuncPlan {
         let mut alloca_loc = FxHashMap::default();
         alloca_loc.insert(alloca, loc);
-        FuncMemPlan {
+        MachineFuncPlan {
             arena_base: STATIC_BASE,
             scratch_words: 0,
             stable_words,
@@ -317,10 +305,6 @@ mod tests {
             alloca_loc,
             spill_obj: SecondaryMap::new(),
             call_preserve: FxHashMap::default(),
-            malloc_future_abs_words: FxHashMap::default(),
-            transient_mallocs: FxHashSet::default(),
-            malloc_escape_kinds: FxHashMap::default(),
-            return_escape_caller_abs_words: 0,
         }
     }
 

crates/codegen/src/isa/evm/escape_scan.rs

@@ -8,7 +8,7 @@ use sonatina_ir::{
     module::{FuncRef, ModuleCtx},
 };
 
-use super::{provenance::Provenance, ptr_escape::PtrEscapeSummary};
+use super::{ptr_escape::PtrEscapeSummary, ptr_provenance::Provenance};
 
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
 pub(crate) enum PtrWriteKind {