Support for continuations in execution (#111)

* Managing the stack from the JAAF family instructions.

* Code reuse.

* Frame allocator that can track sparse allocated ranges.

* Persisting WomController across segments.

* Carrying across segments only the values in active frames.

* Fixing memory reset.

* Sharing FP.

* Also saving newly allocated frames.

* Adapting the tests to use AllocateFrame.

* Commented out stack print.

* Small things.

Author: lvella

extensions/circuit/src/adapters/allocate_frame.rs

@@ -1,4 +1,7 @@
-use std::borrow::Borrow;
+use std::{
+    borrow::Borrow,
+    sync::{Arc, Mutex},
+};
 
 use openvm_circuit::{
     arch::{
@@ -23,15 +26,16 @@ use struct_reflection::{StructReflection, StructReflectionHelper};
 
 use crate::{
     FrameBridge, FrameBus, FrameState, VmAdapterChipWom, WomBridge, WomController, WomRecord,
-    adapters::{compose, decompose},
+    adapters::{compose, decompose, frame_allocator::FrameAllocator},
 };
 
 use super::RV32_REGISTER_NUM_LIMBS;
 
 #[derive(Debug)]
 pub struct AllocateFrameAdapterChipWom {
     pub air: AllocateFrameAdapterAirWom,
-    next_fp: u32,
+    frame_stack: Arc<Mutex<Vec<u32>>>,
+    frame_allocator: Arc<Mutex<FrameAllocator>>,
 }
 
 impl AllocateFrameAdapterChipWom {
@@ -41,6 +45,8 @@ impl AllocateFrameAdapterChipWom {
         frame_bus: FrameBus,
         memory_bridge: MemoryBridge,
         wom_bridge: WomBridge,
+        frame_allocator: Arc<Mutex<FrameAllocator>>,
+        frame_stack: Arc<Mutex<Vec<u32>>>,
     ) -> Self {
         Self {
             air: AllocateFrameAdapterAirWom {
@@ -49,8 +55,8 @@ impl AllocateFrameAdapterChipWom {
                 frame_bridge: FrameBridge::new(frame_bus),
                 _memory_bridge: memory_bridge,
             },
-            // Start from 8 because 0 and 4 are used by the startup code.
-            next_fp: 8,
+            frame_allocator,
+            frame_stack,
         }
     }
 }
@@ -214,9 +220,18 @@ impl<F: PrimeField32> VmAdapterChipWom<F> for AllocateFrameAdapterChipWom {
         };
         let amount_bytes = RV32_REGISTER_NUM_LIMBS as u32 * amount;
 
-        let allocated_ptr = self.next_fp;
-
-        self.next_fp += amount_bytes;
+        let allocated_ptr = self
+            .frame_allocator
+            .lock()
+            .unwrap()
+            .allocate(amount_bytes)
+            .expect("WOM frame allocation failed: not enough free contiguous space");
+
+        {
+            let mut frame_stack = self.frame_stack.lock().unwrap();
+            frame_stack.push(allocated_ptr);
+            //println!("A STACK: {frame_stack:?}");
+        }
 
         let amount_bytes = decompose(amount_bytes);
 
@@ -280,3 +295,90 @@ impl<F: PrimeField32> VmAdapterChipWom<F> for AllocateFrameAdapterChipWom {
         &self.air
     }
 }
+
+pub mod frame_allocator {
+    use std::collections::{BTreeMap, btree_map::Entry};
+
+    use serde::{Deserialize, Serialize};
+
+    #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
+    pub struct Range {
+        pub start: u32,
+        pub end: u32,
+    }
+
+    #[derive(Debug, Serialize, Deserialize)]
+    pub struct FrameAllocator {
+        /// The set of free frame ranges, indexed by its length.
+        free_ranges: BTreeMap<u32, Vec<Range>>,
+        /// The set of allocated frame ranges. Maps start address to end address.
+        allocated_ranges: BTreeMap<u32, u32>,
+    }
+
+    impl FrameAllocator {
+        pub fn new(largest_address: u32, existing_allocations: BTreeMap<u32, u32>) -> Self {
+            // Initialize free_ranges based on existing_allocations
+            let mut free_ranges: BTreeMap<u32, Vec<Range>> = BTreeMap::new();
+            let mut current_start = 0;
+
+            for (start, end) in &existing_allocations {
+                if current_start < *start {
+                    let range = Range {
+                        start: current_start,
+                        end: *start,
+                    };
+                    let length = range.end - range.start;
+                    free_ranges.entry(length).or_default().push(range);
+                }
+                current_start = *end;
+            }
+
+            if current_start < largest_address {
+                let range = Range {
+                    start: current_start,
+                    end: largest_address,
+                };
+                let length = range.end - range.start;
+                free_ranges.entry(length).or_default().push(range);
+            }
+
+            Self {
+                free_ranges,
+                allocated_ranges: existing_allocations,
+            }
+        }
+
+        pub fn allocate(&mut self, size: u32) -> Option<u32> {
+            // Find a free range that can accommodate the requested size
+            let fittest_len = *self.free_ranges.range_mut(size..).next()?.0;
+
+            // Apparently there is no way find this entry without searching again...
+            // TODO: change this when https://github.com/rust-lang/rust/issues/107540
+            // is stabilized.
+            let Entry::Occupied(mut entry) = self.free_ranges.entry(fittest_len) else {
+                unreachable!();
+            };
+            let mut range = entry.get_mut().pop().unwrap();
+            if entry.get().is_empty() {
+                entry.remove();
+            }
+
+            let allocated_start = range.start;
+            range.start += size;
+
+            if range.start < range.end {
+                let length = range.end - range.start;
+                self.free_ranges.entry(length).or_default().push(range);
+            }
+
+            self.allocated_ranges
+                .insert(allocated_start, allocated_start + size);
+
+            Some(allocated_start)
+        }
+
+        pub fn get_allocated_ranges(&self) -> &BTreeMap<u32, u32> {
+            &self.allocated_ranges
+        }
+    }
+}

extensions/circuit/src/adapters/jaaf.rs

@@ -1,4 +1,9 @@
-use std::{array, borrow::Borrow, marker::PhantomData};
+use std::{
+    array,
+    borrow::Borrow,
+    marker::PhantomData,
+    sync::{Arc, Mutex},
+};
 
 use openvm_circuit::{
     arch::{
@@ -51,6 +56,7 @@ pub struct JaafInstruction<T> {
 #[derive(Debug)]
 pub struct JaafAdapterChipWom<F: Field> {
     pub air: JaafAdapterAirWom,
+    frame_stack: Arc<Mutex<Vec<u32>>>,
     _marker: PhantomData<F>,
 }
 
@@ -59,6 +65,7 @@ impl<F: PrimeField32> JaafAdapterChipWom<F> {
         execution_bus: ExecutionBus,
         program_bus: ProgramBus,
         frame_bus: FrameBus,
+        frame_stack: Arc<Mutex<Vec<u32>>>,
         memory_bridge: MemoryBridge,
         wom_bridge: WomBridge,
     ) -> Self {
@@ -69,6 +76,7 @@ impl<F: PrimeField32> JaafAdapterChipWom<F> {
                 _memory_bridge: memory_bridge,
                 wom_bridge,
             },
+            frame_stack,
             _marker: PhantomData,
         }
     }
@@ -356,6 +364,41 @@ impl<F: PrimeField32> VmAdapterChipWom<F> for JaafAdapterChipWom<F> {
             (None, None)
         };
 
+        // Update the frame stack according to the opcode
+        {
+            let mut frame_stack = self.frame_stack.lock().unwrap();
+            match local_opcode {
+                RET | JAAF => {
+                    // The newly activated frame must be in the frame stack.
+                    let idx = frame_stack.iter().rposition(|x| x == &to_fp).unwrap();
+
+                    // Test if the activated frame has just been allocated with AllocateFrame
+                    // (i.e., is at the top of the stack)
+                    if idx == frame_stack.len() - 1 {
+                        // Yes, so this has to be a JAAF instruction jumping to a new frame.
+                        assert_eq!(local_opcode, JAAF);
+
+                        // The current frame can be safely dropped, since this opcode doesn't save it,
+                        // rendering it unreachable.
+                        let curr_frame_idx = idx - 1;
+                        let curr_frame = frame_stack.swap_remove(curr_frame_idx);
+
+                        // Since we are jumping to the last frame in the stack, which is the new
+                        // frame, the current frame must be the one before it, that we just removed.
+                        assert_eq!(curr_frame, from_frame.fp);
+                    } else {
+                        // Found an old frame, truncate the stack to that frame
+                        frame_stack.truncate(idx + 1);
+                    }
+                    //println!("J STACK: {frame_stack:?}");
+                }
+                CALL | CALL_INDIRECT | JAAF_SAVE => {
+                    // These calls always target a new frame, which is already on top of the stack.
+                    assert_eq!(frame_stack.last(), Some(&to_fp));
+                }
+            }
+        }
+
         Ok((
             ExecutionState {
                 pc: to_pc,