Update MultiUseSandbox to use dirty page tracking in restore and devolve methods

Signed-off-by: Simon Davies <[email protected]>

Author: simongdavies

src/hyperlight_host/src/sandbox/initialized_multi_use.rs

@@ -193,7 +193,8 @@ impl MultiUseSandbox {
     #[instrument(err(Debug), skip_all, parent = Span::current(), level = "Trace")]
     pub(crate) fn restore_state(&mut self) -> Result<()> {
         let mem_mgr = self.mem_mgr.unwrap_mgr_mut();
-        mem_mgr.restore_state_from_last_snapshot()
+        let dirty_pages = self.vm.get_and_clear_dirty_pages()?;
+        mem_mgr.restore_state_from_last_snapshot(&dirty_pages)
     }
 
     pub(crate) fn call_guest_function_by_name_no_reset(
@@ -275,9 +276,10 @@ impl DevolvableSandbox<MultiUseSandbox, MultiUseSandbox, Noop<MultiUseSandbox, M
     /// The devolve can be used to return the MultiUseSandbox to the state before the code was loaded. Thus avoiding initialisation overhead
     #[instrument(err(Debug), skip_all, parent = Span::current(), level = "Trace")]
     fn devolve(mut self, _tsn: Noop<MultiUseSandbox, MultiUseSandbox>) -> Result<MultiUseSandbox> {
+        let dirty_pages = self.vm.get_and_clear_dirty_pages()?;
         self.mem_mgr
             .unwrap_mgr_mut()
-            .pop_and_restore_state_from_snapshot()?;
+            .pop_and_restore_state_from_snapshot(&dirty_pages)?;
         Ok(self)
     }
 }
@@ -308,7 +310,10 @@ where
         let mut ctx = self.new_call_context();
         transition_func.call(&mut ctx)?;
         let mut sbox = ctx.finish_no_reset();
-        sbox.mem_mgr.unwrap_mgr_mut().push_state()?;
+        let vm_dirty_pages = sbox.vm.get_and_clear_dirty_pages()?;
+        sbox.mem_mgr
+            .unwrap_mgr_mut()
+            .push_state(Some(&vm_dirty_pages))?;
         Ok(sbox)
     }
 }