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

A few minor nits/documentation fixes/formatting fixes/cleanups

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

Author: syntactically

src/hyperlight_wasm/src/sandbox/loaded_wasm_sandbox.rs

@@ -112,11 +112,14 @@ impl LoadedWasmSandbox {
         }
     }
 
-    /// Unload the wasm module and return a `WasmSandbox` that can be used to load another module.
+    /// Unload the wasm module and return a `WasmSandbox` that can be
+    /// used to load another module.
     ///
-    /// This method internally calls [`restore()`](Self::restore) to reset the sandbox to its
-    /// pre-module state, which also clears any poisoned state. This means `unload_module()`
-    /// can be called on a poisoned sandbox to recover it.
+    /// This method defers calling [`restore()`](Self::restore) to
+    /// reset the sandbox to its pre-module state until a new module
+    /// is loaded. However, the sandbox will always be restored when a
+    /// new module is loaded, so a poisoned sandbox can be recovered
+    /// by unloading and reloading a module.
     pub fn unload_module(mut self) -> Result<WasmSandbox> {
         let sandbox = self
             .inner

src/hyperlight_wasm/src/sandbox/wasm_sandbox.rs

@@ -80,7 +80,15 @@ impl WasmSandbox {
 
     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.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(())
@@ -111,12 +119,13 @@ impl WasmSandbox {
     /// 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();
+        let sb = self
+            .inner
+            .as_mut()
+            .ok_or_else(|| new_error!("WasmSandbox is None"))?;
         sb.restore(snapshot)?;
-        LoadedWasmSandbox::new(
-            sb,
-            self.snapshot.take().unwrap(),
-        )
+
+        self.finalize_module_load()
     }
 
     /// Load a Wasm module that is currently present in a buffer in
@@ -165,6 +174,7 @@ impl WasmSandbox {
     /// 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_from_buffer(mut self, buffer: &[u8]) -> Result<LoadedWasmSandbox> {
+        self.restore_if_needed()?;
         let inner = self
             .inner
             .as_mut()
@@ -496,6 +506,46 @@ mod tests {
         }
     }
 
+    #[test]
+    fn test_load_from_snapshot() {
+        let mut sandbox = SandboxBuilder::new().build().unwrap();
+        sandbox
+            .register(
+                "GetTimeSinceBootMicrosecond",
+                get_time_since_boot_microsecond,
+            )
+            .unwrap();
+        let sb = sandbox.load_runtime().unwrap();
+
+        let helloworld_wasm = get_test_file_path("HelloWorld.aot").unwrap();
+        let runwasm_wasm = get_test_file_path("RunWasm.aot").unwrap();
+
+        // load one module, and make sure that a function in it
+        // can be called
+        let mut lb1 = sb.load_module(helloworld_wasm).unwrap();
+        let result: i32 = lb1
+            .call_guest_function("HelloWorld", "Message from Rust Test".to_string())
+            .unwrap();
+        assert_eq!(result, 0);
+        let snapshot = lb1.snapshot().unwrap();
+
+        // load another module, and make sure that a function in
+        // it can be called
+        let sb = lb1.unload_module().unwrap();
+        let mut lb2 = sb.load_module(runwasm_wasm).unwrap();
+        let result: i32 = lb2.call_guest_function("CalcFib", 10i32).unwrap();
+        assert_eq!(result, 55);
+
+        // reload the first module via snapshot, and make sure the
+        // original function can be called again
+        let sb = lb2.unload_module().unwrap();
+        let mut lb3 = sb.load_from_snapshot(snapshot).unwrap();
+        let result: i32 = lb3
+            .call_guest_function("HelloWorld", "Message from Rust Test".to_string())
+            .unwrap();
+        assert_eq!(result, 0);
+    }
+
     #[test]
     fn test_load_module_buffer() {
         let sandboxes = get_test_wasm_sandboxes().unwrap();