From b423df22006320d31ab7d070edbce79f4f6440fe Mon Sep 17 00:00:00 2001
From: Islam El-Ashi <islam.elashi@dfinity.org>
Date: Thu, 5 Dec 2024 15:45:09 +0100
Subject: [PATCH 1/4] feat: support loading stable memory from a file.

---
 canbench-bin/src/lib.rs                  | 38 +++++++++++++++++++++---
 canbench-bin/src/main.rs                 | 12 ++++++++
 canbench-bin/tests/tests.rs              | 29 ++++++++++++++++++
 tests/Cargo.toml                         |  8 +++++
 tests/stable_memory/canbench.yml         |  9 ++++++
 tests/stable_memory/src/main.rs          | 16 ++++++++++
 tests/stable_memory/stable_memory.bin    |  1 +
 tests/stable_memory_invalid/canbench.yml |  9 ++++++
 tests/stable_memory_invalid/src/main.rs  |  1 +
 9 files changed, 119 insertions(+), 4 deletions(-)
 create mode 100644 tests/stable_memory/canbench.yml
 create mode 100644 tests/stable_memory/src/main.rs
 create mode 100644 tests/stable_memory/stable_memory.bin
 create mode 100644 tests/stable_memory_invalid/canbench.yml
 create mode 100644 tests/stable_memory_invalid/src/main.rs

diff --git a/canbench-bin/src/lib.rs b/canbench-bin/src/lib.rs
index ccfd09f7..77868c34 100644
--- a/canbench-bin/src/lib.rs
+++ b/canbench-bin/src/lib.rs
@@ -1,13 +1,14 @@
 //! A module for running benchmarks.
+mod print_benchmark;
+mod results_file;
 use canbench_rs::BenchResult;
 use candid::Principal;
 use flate2::read::GzDecoder;
+use pocket_ic::common::rest::BlobCompression;
 use pocket_ic::{PocketIc, PocketIcBuilder, WasmResult};
-use std::{collections::BTreeMap, env, fs::File, io::Read, path::PathBuf, process::Command};
-mod print_benchmark;
-mod results_file;
 use print_benchmark::print_benchmark;
 use results_file::VersionError;
+use std::{collections::BTreeMap, env, fs::File, io::Read, path::PathBuf, process::Command};
 use wasmparser::Parser as WasmParser;
 
 // The prefix benchmarks are expected to have in their name.
@@ -29,6 +30,7 @@ pub fn run_benchmarks(
     verbose: bool,
     integrity_check: bool,
     runtime_path: &PathBuf,
+    stable_memory_path: Option<PathBuf>,
 ) {
     maybe_download_pocket_ic(runtime_path, verbose, integrity_check);
 
@@ -47,7 +49,12 @@ pub fn run_benchmarks(
     let benchmark_fns = extract_benchmark_fns(canister_wasm_path);
 
     // Initialize PocketIC
-    let (pocket_ic, canister_id) = init_pocket_ic(runtime_path, canister_wasm_path, init_args);
+    let (pocket_ic, canister_id) = init_pocket_ic(
+        runtime_path,
+        canister_wasm_path,
+        stable_memory_path,
+        init_args,
+    );
 
     // Run the benchmarks
     let mut results = BTreeMap::new();
@@ -246,6 +253,7 @@ fn set_env_var_if_unset(key: &str, target_value: &str) {
 fn init_pocket_ic(
     path: &PathBuf,
     canister_wasm_path: &PathBuf,
+    stable_memory_path: Option<PathBuf>,
     init_args: Vec<u8>,
 ) -> (PocketIc, Principal) {
     // PocketIC is used for running the benchmark.
@@ -264,6 +272,28 @@ fn init_pocket_ic(
         init_args,
         None,
     );
+
+    // Load the canister's stable memory if a stable memory file is specified.
+    if let Some(stable_memory_path) = stable_memory_path {
+        let stable_memory_bytes = match std::fs::read(&stable_memory_path) {
+            Ok(bytes) => bytes,
+            Err(err) => {
+                eprintln!(
+                    "Error reading stable memory file {}",
+                    &stable_memory_path.display()
+                );
+                eprintln!("Error: {}", err);
+                std::process::exit(1);
+            }
+        };
+
+        pocket_ic.set_stable_memory(
+            canister_id,
+            stable_memory_bytes,
+            BlobCompression::NoCompression,
+        );
+    }
+
     (pocket_ic, canister_id)
 }
 
diff --git a/canbench-bin/src/main.rs b/canbench-bin/src/main.rs
index 26b562bf..5efedbc0 100644
--- a/canbench-bin/src/main.rs
+++ b/canbench-bin/src/main.rs
@@ -43,6 +43,12 @@ struct InitArgs {
     hex: String,
 }
 
+#[derive(Debug, Deserialize)]
+struct StableMemory {
+    // File path to load stable memory from.
+    file: String,
+}
+
 #[derive(Debug, Deserialize)]
 struct Config {
     // If provided, instructs canbench to build the canister
@@ -57,6 +63,9 @@ struct Config {
 
     // If provided, the init arguments to pass to the canister
     init_args: Option<InitArgs>,
+
+    // The stable memory to load into the canister.
+    stable_memory: Option<StableMemory>,
 }
 
 // Path to the canbench directory where we keep internal data.
@@ -113,6 +122,8 @@ fn main() {
         );
     }
 
+    let stable_memory_path = cfg.stable_memory.map(|sm| PathBuf::from(sm.file));
+
     let init_args = cfg
         .init_args
         .map(|args| hex::decode(args.hex).expect("invalid init_args hex value"))
@@ -128,5 +139,6 @@ fn main() {
         !args.less_verbose,
         !args.no_runtime_integrity_check,
         &args.runtime_path.unwrap_or_else(default_runtime_path),
+        stable_memory_path,
     );
 }
diff --git a/canbench-bin/tests/tests.rs b/canbench-bin/tests/tests.rs
index 9da81f4f..378e6ec1 100644
--- a/canbench-bin/tests/tests.rs
+++ b/canbench-bin/tests/tests.rs
@@ -387,3 +387,32 @@ Benchmark: write_stable_memory (new)
             );
         });
 }
+
+#[test]
+fn loads_stable_memory_file() {
+    BenchTest::canister("stable_memory").run(|output| {
+        assert_success!(
+            output,
+            "
+---------------------------------------------------
+
+Benchmark: read_from_stable_memory (new)
+  total:
+    instructions: 589 (new)
+    heap_increase: 0 pages (new)
+    stable_memory_increase: 0 pages (new)
+
+---------------------------------------------------
+"
+        );
+    });
+}
+
+#[test]
+fn stable_memory_file_not_exit_prints_error() {
+    BenchTest::canister("stable_memory_invalid").run(|output| {
+        assert_err!(output, "
+Error reading stable memory file stable_memory_does_not_exist.bin
+Error: No such file or directory");
+    });
+}
diff --git a/tests/Cargo.toml b/tests/Cargo.toml
index ace3a5ab..4b867b7e 100644
--- a/tests/Cargo.toml
+++ b/tests/Cargo.toml
@@ -17,6 +17,14 @@ path = "gzipped_wasm/src/main.rs"
 name = "init_arg"
 path = "init_arg/src/main.rs"
 
+[[bin]]
+name = "stable_memory"
+path = "stable_memory/src/main.rs"
+
+[[bin]]
+name = "stable_memory_invalid"
+path = "stable_memory_invalid/src/main.rs"
+
 [dependencies]
 canbench-rs = { path = "../canbench-rs" }
 candid.workspace = true
diff --git a/tests/stable_memory/canbench.yml b/tests/stable_memory/canbench.yml
new file mode 100644
index 00000000..e951e01f
--- /dev/null
+++ b/tests/stable_memory/canbench.yml
@@ -0,0 +1,9 @@
+build_cmd:
+  cargo build --release --target wasm32-unknown-unknown
+
+wasm_path:
+  ../../target/wasm32-unknown-unknown/release/stable_memory.wasm
+
+stable_memory:
+  file:
+    stable_memory.bin
diff --git a/tests/stable_memory/src/main.rs b/tests/stable_memory/src/main.rs
new file mode 100644
index 00000000..d8325f79
--- /dev/null
+++ b/tests/stable_memory/src/main.rs
@@ -0,0 +1,16 @@
+use canbench_rs::bench;
+
+#[bench]
+fn read_from_stable_memory() {
+    let mut buf = [0; 10];
+    ic_cdk::api::stable::stable_read(0, &mut buf);
+
+    // There should be one page in stable memory.
+    assert_eq!(ic_cdk::api::stable::stable_size(), 1);
+
+    // The `stable_memory.bin` specified in canbench.yml only has the first give bytes set.
+    // The rest should be zero.
+    assert_eq!(&buf, &[0x41, 0x42, 0x43, 0x44, 0x45, 0, 0, 0, 0, 0]);
+}
+
+fn main() {}
diff --git a/tests/stable_memory/stable_memory.bin b/tests/stable_memory/stable_memory.bin
new file mode 100644
index 00000000..402476b0
--- /dev/null
+++ b/tests/stable_memory/stable_memory.bin
@@ -0,0 +1 @@
+ABCDE
\ No newline at end of file
diff --git a/tests/stable_memory_invalid/canbench.yml b/tests/stable_memory_invalid/canbench.yml
new file mode 100644
index 00000000..6c9726c2
--- /dev/null
+++ b/tests/stable_memory_invalid/canbench.yml
@@ -0,0 +1,9 @@
+build_cmd:
+  cargo build --release --target wasm32-unknown-unknown
+
+wasm_path:
+  ../../target/wasm32-unknown-unknown/release/stable_memory_invalid.wasm
+
+stable_memory:
+  file:
+    stable_memory_does_not_exist.bin
diff --git a/tests/stable_memory_invalid/src/main.rs b/tests/stable_memory_invalid/src/main.rs
new file mode 100644
index 00000000..f328e4d9
--- /dev/null
+++ b/tests/stable_memory_invalid/src/main.rs
@@ -0,0 +1 @@
+fn main() {}

From 85425bc36061567b575228513449215ed87149bb Mon Sep 17 00:00:00 2001
From: Islam El-Ashi <islam.elashi@dfinity.org>
Date: Thu, 5 Dec 2024 15:51:44 +0100
Subject: [PATCH 2/4] docs

---
 canbench-rs/src/lib.rs | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/canbench-rs/src/lib.rs b/canbench-rs/src/lib.rs
index fce410a7..7d7b0d3f 100644
--- a/canbench-rs/src/lib.rs
+++ b/canbench-rs/src/lib.rs
@@ -42,6 +42,18 @@
 //! init_args:
 //!   hex: 4449444c0001710568656c6c6f
 //! ```
+//!
+//! #### Stable Memory
+//!
+//! A stable memory file can be specified to be loaded in the canister's stable memory _after_
+//! initialization.
+//!
+//! ```yml
+//! stable_memory:
+//!   file:
+//!     stable_memory.bin
+//! ```
+//!
 //! ### 4. Start benching! 🏋🏽
 //!
 //! Let's say we have a canister that exposes a `query` computing the fibonacci sequence of a given number.

From b256ffb6328422f1ac2837324e1154c1ae2085d5 Mon Sep 17 00:00:00 2001
From: Islam El-Ashi <islam.elashi@dfinity.org>
Date: Thu, 5 Dec 2024 15:52:52 +0100
Subject: [PATCH 3/4] cargo fmt

---
 canbench-bin/tests/tests.rs | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/canbench-bin/tests/tests.rs b/canbench-bin/tests/tests.rs
index 378e6ec1..41f75449 100644
--- a/canbench-bin/tests/tests.rs
+++ b/canbench-bin/tests/tests.rs
@@ -411,8 +411,11 @@ Benchmark: read_from_stable_memory (new)
 #[test]
 fn stable_memory_file_not_exit_prints_error() {
     BenchTest::canister("stable_memory_invalid").run(|output| {
-        assert_err!(output, "
+        assert_err!(
+            output,
+            "
 Error reading stable memory file stable_memory_does_not_exist.bin
-Error: No such file or directory");
+Error: No such file or directory"
+        );
     });
 }

From 7275abbbf18d9fa2ef0170545ac6819c5d2a7d10 Mon Sep 17 00:00:00 2001
From: Islam El-Ashi <islam.elashi@dfinity.org>
Date: Fri, 6 Dec 2024 10:23:06 +0100
Subject: [PATCH 4/4] .

---
 canbench-bin/tests/tests.rs | 17 +++--------------
 canbench-rs/src/lib.rs      |  6 ++++--
 2 files changed, 7 insertions(+), 16 deletions(-)

diff --git a/canbench-bin/tests/tests.rs b/canbench-bin/tests/tests.rs
index 41f75449..b0bf0325 100644
--- a/canbench-bin/tests/tests.rs
+++ b/canbench-bin/tests/tests.rs
@@ -391,20 +391,9 @@ Benchmark: write_stable_memory (new)
 #[test]
 fn loads_stable_memory_file() {
     BenchTest::canister("stable_memory").run(|output| {
-        assert_success!(
-            output,
-            "
----------------------------------------------------
-
-Benchmark: read_from_stable_memory (new)
-  total:
-    instructions: 589 (new)
-    heap_increase: 0 pages (new)
-    stable_memory_increase: 0 pages (new)
-
----------------------------------------------------
-"
-        );
+        // There are assertions in the code of that canister itself, so
+        // all is needed is to assert that the run succeeded.
+        assert_eq!(output.status.code(), Some(0), "output: {:?}", output);
     });
 }
 
diff --git a/canbench-rs/src/lib.rs b/canbench-rs/src/lib.rs
index 7d7b0d3f..3c97e5af 100644
--- a/canbench-rs/src/lib.rs
+++ b/canbench-rs/src/lib.rs
@@ -45,8 +45,7 @@
 //!
 //! #### Stable Memory
 //!
-//! A stable memory file can be specified to be loaded in the canister's stable memory _after_
-//! initialization.
+//! A file can be specified to be loaded in the canister's stable memory _after_ initialization.
 //!
 //! ```yml
 //! stable_memory:
@@ -54,6 +53,9 @@
 //!     stable_memory.bin
 //! ```
 //!
+//! <div class="warning">Contents of the stable memory file are loaded <i>after</i> the call to the canister's init method.
+//! Therefore, changes made to stable memory in the init method would be overwritten.</div>
+//!
 //! ### 4. Start benching! 🏋🏽
 //!
 //! Let's say we have a canister that exposes a `query` computing the fibonacci sequence of a given number.