more efficient save/restore of registers

Author: wtholliday

src/vm.rs

@@ -1,6 +1,7 @@
 // Register-based virtual machine for when we can't JIT,
 // like on iOS.
 
+use std::convert::TryInto;
 use std::fmt;
 
 /// Register index (0-255)
@@ -285,6 +286,15 @@ pub enum Opcode {
     /// Zero memory region
     MemZero { dst: Reg, size: u32 },
 
+    // ============ Register Save/Restore ============
+
+    /// Save registers [start_reg..start_reg+count] to locals at slot offset.
+    /// Each register is 8 bytes.
+    SaveRegs { start_reg: Reg, count: u8, slot: u32 },
+
+    /// Restore registers [start_reg..start_reg+count] from locals at slot offset.
+    RestoreRegs { start_reg: Reg, count: u8, slot: u32 },
+
     // ============ Debugging ============
 
     /// Print integer value (for debugging)
@@ -394,9 +404,6 @@ struct CallFrame {
 
     /// Register snapshot (for return value handling)
     return_reg: Reg,
-
-    /// Saved registers - stored as boxed array to avoid stack overflow
-    saved_registers: Box<[u64; 256]>,
 }
 
 /// The virtual machine state
@@ -888,13 +895,12 @@ impl VM {
                 }
 
                 Opcode::Call { func, args_start, arg_count } => {
-                    // Save current frame with registers
+                    // Save current frame
                     let frame = CallFrame {
                         func_idx: self.current_func,
                         ip: self.ip,
                         locals_base: self.locals_base,
                         return_reg: 0,
-                        saved_registers: Box::new(self.registers),
                     };
                     self.call_stack.push(frame);
 
@@ -925,13 +931,12 @@ impl VM {
                 Opcode::CallIndirect { func_reg, args_start, arg_count } => {
                     let func_idx = self.get_u64(func_reg) as FuncIdx;
 
-                    // Save current frame with registers
+                    // Save current frame
                     let frame = CallFrame {
                         func_idx: self.current_func,
                         ip: self.ip,
                         locals_base: self.locals_base,
                         return_reg: 0,
-                        saved_registers: Box::new(self.registers),
                     };
                     self.call_stack.push(frame);
 
@@ -953,36 +958,26 @@ impl VM {
                 }
 
                 Opcode::Return => {
-                    // Save return value (implicitly in r0) before restoring registers
-                    let return_value = self.registers[0];
-
                     if self.call_stack.is_empty() {
                         return self.get_i64(0);
                     }
                     let frame = self.call_stack.pop().unwrap();
                     self.current_func = frame.func_idx;
                     self.ip = frame.ip;
                     self.locals_base = frame.locals_base;
-                    // Restore caller's registers and put return value in r0
-                    self.registers = *frame.saved_registers;
-                    self.registers[0] = return_value;
                 }
 
                 Opcode::ReturnReg { src } => {
-                    // Save return value before restoring registers
-                    let return_value = self.registers[src as usize];
+                    // Move return value to r0
+                    self.registers[0] = self.registers[src as usize];
 
                     if self.call_stack.is_empty() {
-                        self.registers[0] = return_value;
                         return self.get_i64(0);
                     }
                     let frame = self.call_stack.pop().unwrap();
                     self.current_func = frame.func_idx;
                     self.ip = frame.ip;
                     self.locals_base = frame.locals_base;
-                    // Restore caller's registers and put return value in r0
-                    self.registers = *frame.saved_registers;
-                    self.registers[0] = return_value;
                 }
 
                 Opcode::AllocLocals { size } => {
@@ -1011,6 +1006,24 @@ impl VM {
                     }
                 }
 
+                Opcode::SaveRegs { start_reg, count, slot } => {
+                    let base = self.locals_base + slot as usize;
+                    for i in 0..count as usize {
+                        let reg_val = self.registers[start_reg as usize + i];
+                        let offset = base + i * 8;
+                        self.locals[offset..offset + 8].copy_from_slice(&reg_val.to_le_bytes());
+                    }
+                }
+
+                Opcode::RestoreRegs { start_reg, count, slot } => {
+                    let base = self.locals_base + slot as usize;
+                    for i in 0..count as usize {
+                        let offset = base + i * 8;
+                        let bytes: [u8; 8] = self.locals[offset..offset + 8].try_into().unwrap();
+                        self.registers[start_reg as usize + i] = u64::from_le_bytes(bytes);
+                    }
+                }
+
                 // Debugging
                 Opcode::PrintI32 { src } => {
                     println!("{}", self.get_i64(src) as i32);

src/vm_codegen.rs

@@ -217,6 +217,9 @@ struct FunctionTranslator<'a> {
 
     /// Tracks if a return has been emitted (so we don't emit epilogue).
     has_returned: bool,
+
+    /// Byte offset in locals where registers are saved.
+    save_regs_offset: u32,
 }
 
 /// Check if a type should be returned via output pointer.
@@ -246,6 +249,7 @@ impl<'a> FunctionTranslator<'a> {
             globals,
             output_ptr: None,
             has_returned: false,
+            save_regs_offset: 0,
         }
     }
 
@@ -263,6 +267,17 @@ impl<'a> FunctionTranslator<'a> {
             self.alloc_reg();
         }
 
+        // Reserve space for saving registers. We'll patch this after translation
+        // when we know how many registers are used.
+        // Reserve 256 * 8 = 2048 bytes maximum, but we'll only use what we need.
+        let save_regs_slot = self.alloc_local(256 * 8);
+        self.save_regs_offset = (save_regs_slot as u32) * 8; // Convert slot to byte offset
+        func.emit(Opcode::SaveRegs {
+            start_reg: 0,
+            count: 0, // Placeholder - will be patched
+            slot: self.save_regs_offset,
+        });
+
         // Save parameters to local slots so they persist across recursive calls.
         // Registers are shared across all call frames, but locals are per-frame.
         for (i, param) in self.decl.params.iter().enumerate() {
@@ -295,17 +310,58 @@ impl<'a> FunctionTranslator<'a> {
                         src: result_reg,
                         size,
                     });
+                    // Restore all registers before return.
+                    func.emit(Opcode::RestoreRegs {
+                        start_reg: 0,
+                        count: 0, // Placeholder - will be patched
+                        slot: self.save_regs_offset,
+                    });
                     func.emit(Opcode::Return);
-                } else if result_reg != 0 {
-                    func.emit(Opcode::ReturnReg { src: result_reg });
                 } else {
+                    // Move result to r0, restore r1..N, return
+                    if result_reg != 0 {
+                        func.emit(Opcode::Move { dst: 0, src: result_reg });
+                    }
+                    func.emit(Opcode::RestoreRegs {
+                        start_reg: 1, // Skip r0 which has the return value
+                        count: 0,     // Placeholder - will be patched
+                        slot: self.save_regs_offset + 8,
+                    });
                     func.emit(Opcode::Return);
                 }
             }
         } else {
+            // No body - restore all registers and return
+            func.emit(Opcode::RestoreRegs {
+                start_reg: 0,
+                count: 0, // Placeholder - will be patched
+                slot: self.save_regs_offset,
+            });
             func.emit(Opcode::Return);
         }
 
+        // Patch SaveRegs/RestoreRegs instructions with actual register count.
+        let reg_count = self.next_reg;
+        for op in &mut func.code {
+            match op {
+                Opcode::SaveRegs { count, .. } => {
+                    *count = reg_count;
+                }
+                Opcode::RestoreRegs { start_reg, count, .. } => {
+                    // If start_reg is 1, we're skipping r0 (for returns with value)
+                    if *start_reg == 1 {
+                        *count = if reg_count > 1 { reg_count - 1 } else { 0 };
+                    } else {
+                        *count = reg_count;
+                    }
+                }
+                _ => {}
+            }
+        }
+
+        // Adjust locals_size to account for saved registers.
+        self.locals_size = self.save_regs_offset + (reg_count as u32) * 8;
+
         // Set function metadata.
         func.locals_size = self.locals_size;
         func.local_slots = self.next_slot;
@@ -539,9 +595,25 @@ impl<'a> FunctionTranslator<'a> {
                         src: result,
                         size,
                     });
+                    // Restore caller's registers before returning.
+                    func.emit(Opcode::RestoreRegs {
+                        start_reg: 0,
+                        count: 0, // Will be patched
+                        slot: self.save_regs_offset,
+                    });
                     func.emit(Opcode::Return);
                 } else {
-                    func.emit(Opcode::ReturnReg { src: result });
+                    // Move result to r0 before restoring (it will be preserved).
+                    if result != 0 {
+                        func.emit(Opcode::Move { dst: 0, src: result });
+                    }
+                    // Restore caller's registers, but skip r0 which has return value.
+                    func.emit(Opcode::RestoreRegs {
+                        start_reg: 1, // Skip r0
+                        count: 0,     // Will be patched to next_reg - 1
+                        slot: self.save_regs_offset + 8, // Skip first 8 bytes (r0's slot)
+                    });
+                    func.emit(Opcode::Return);
                 }
                 self.has_returned = true;
                 result