Change sandbox restore to be a bit more judicious about when to restore

When we are loading from a snapshot, we can get away with just a
restore and never going through a wasmtime load.

Signed-off-by: Lucy Menon <168595099+syntactically@users.noreply.github.com>

Author: syntactically

src/hyperlight_wasm/src/sandbox/wasm_sandbox.rs

@@ -41,6 +41,7 @@ pub struct WasmSandbox {
     // Snapshot of state of an initial WasmSandbox (runtime loaded, but no guest module code loaded).
     // Used for LoadedWasmSandbox to be able restore state back to WasmSandbox
     snapshot: Option<Arc<Snapshot>>,
+    needs_restore: bool,
 }
 
 const MAPPED_BINARY_VA: u64 = 0x1_0000_0000u64;
@@ -56,6 +57,7 @@ impl WasmSandbox {
         Ok(WasmSandbox {
             inner: Some(inner),
             snapshot: Some(snapshot),
+            needs_restore: false,
         })
     }
 
@@ -64,24 +66,33 @@ impl WasmSandbox {
     /// the snapshot has already been created in that case.
     /// Expects a snapshot of the state where wasm runtime is loaded, but no guest module code is loaded.
     pub(super) fn new_from_loaded(
-        mut loaded: MultiUseSandbox,
+        loaded: MultiUseSandbox,
         snapshot: Arc<Snapshot>,
     ) -> Result<Self> {
-        loaded.restore(snapshot.clone())?;
         metrics::gauge!(METRIC_ACTIVE_WASM_SANDBOXES).increment(1);
         metrics::counter!(METRIC_TOTAL_WASM_SANDBOXES).increment(1);
         Ok(WasmSandbox {
             inner: Some(loaded),
             snapshot: Some(snapshot),
+            needs_restore: true,
         })
     }
 
+    fn restore_if_needed(&mut self) -> Result<()> {
+        if self.needs_restore {
+            self.inner.as_mut().ok_or(new_error!("WasmSandbox is none"))?.restore(self.snapshot.as_ref().ok_or(new_error!("Snapshot is none"))?.clone())?;
+            self.needs_restore = false;
+        }
+        Ok(())
+    }
+
     /// Load a Wasm module at the given path into the sandbox and return a `LoadedWasmSandbox`
     /// able to execute code in the loaded Wasm Module.
     ///
     /// Before you can call guest functions in the sandbox, you must call
     /// this function and use the returned value to call guest functions.
     pub fn load_module(mut self, file: impl AsRef<Path>) -> Result<LoadedWasmSandbox> {
+        self.restore_if_needed()?;
         let inner = self
             .inner
             .as_mut()
@@ -97,6 +108,17 @@ impl WasmSandbox {
         self.finalize_module_load()
     }
 
+    /// Load a Wasm module by restoring a Hyperlight snapshot taken
+    /// from a `LoadedWasmSandbox`.
+    pub fn load_from_snapshot(mut self, snapshot: Arc<Snapshot>) -> Result<LoadedWasmSandbox> {
+        let mut sb = self.inner.take().unwrap();
+        sb.restore(snapshot)?;
+        LoadedWasmSandbox::new(
+            sb,
+            self.snapshot.take().unwrap(),
+        )
+    }
+
     /// Load a Wasm module that is currently present in a buffer in
     /// host memory, by mapping the host memory directly into the
     /// sandbox.
@@ -114,6 +136,7 @@ impl WasmSandbox {
         base: *mut libc::c_void,
         len: usize,
     ) -> Result<LoadedWasmSandbox> {
+        self.restore_if_needed()?;
         let inner = self
             .inner
             .as_mut()