Accept custom snapshot sources in Env::from_ledger_snapshot

soroban-sdk/src/env.rs

@@ -259,7 +259,7 @@ struct EnvTestState {
     config: EnvTestConfig,
     generators: Rc<RefCell<Generators>>,
     auth_snapshot: Rc<RefCell<AuthSnapshot>>,
-    snapshot: Option<Rc<LedgerSnapshot>>,
+    snapshot: Option<Rc<RefCell<LedgerSnapshot>>>,
 }
 
 /// Config for changing the default behavior of the Env when used in tests.
@@ -479,7 +479,8 @@ use crate::{
     testutils::{
         budget::Budget, default_ledger_info, Address as _, AuthSnapshot, AuthorizedInvocation,
         ContractFunctionSet, EventsSnapshot, Generators, Ledger as _, MockAuth, MockAuthContract,
-        Register, Snapshot, StellarAssetContract, StellarAssetIssuer,
+        Register, Snapshot, SnapshotSourceCache, SnapshotSourceInput, StellarAssetContract,
+        StellarAssetIssuer,
     },
     Bytes, BytesN, ConstructorArgs,
 };
@@ -542,7 +543,7 @@ impl Env {
         recording_footprint: Rc<dyn internal::storage::SnapshotSource>,
         generators: Option<Rc<RefCell<Generators>>>,
         ledger_info: Option<internal::LedgerInfo>,
-        snapshot: Option<Rc<LedgerSnapshot>>,
+        snapshot: Option<Rc<RefCell<LedgerSnapshot>>>,
     ) -> Env {
         // Store in the Env the name of the test it is for, and a number so that within a test
         // where one or more Env's have been created they can be uniquely identified relative to
@@ -577,6 +578,40 @@ impl Env {
             1
         };
 
+        // Apply fallback to read from test_snapshots_before if snapshot is None
+        let snapshot = snapshot.or_else(|| {
+            // Try to read from test_snapshots_before file for the current test.
+            if let Some(test_name) = test_name.as_ref() {
+                // Construct path similar to to_test_ledger_snapshot_before_file.
+                let test_name_path = test_name
+                    .split("::")
+                    .map(|p| std::path::Path::new(p).to_path_buf())
+                    .reduce(|p0, p1| p0.join(p1))
+                    .expect("test name to not be empty");
+                let dir = std::path::Path::new("test_snapshots_before");
+                let p = dir
+                    .join(&test_name_path)
+                    .with_extension(format!("{number}.json"));
+                if let Ok(snapshot) = LedgerSnapshot::read_file(&p) {
+                    eprintln!(
+                        "Reading test snapshot before file for test {test_name:?} from {p:?}."
+                    );
+                    return Some(Rc::new(RefCell::new(snapshot)));
+                }
+            }
+            None
+        });
+
+        // Default to an empty snapshot if none exists.
+        let snapshot = snapshot.unwrap_or_else(|| Rc::new(RefCell::new(LedgerSnapshot::default())));
+
+        // Wrap the recording footprint into a layer that'll record the initial state of anything
+        // loaded into the snapshot.
+        let recording_footprint = Rc::new(SnapshotSourceCache::new(
+            recording_footprint,
+            snapshot.clone(),
+        ));
+
         let storage = internal::storage::Storage::with_recording_footprint(recording_footprint);
         let budget = internal::budget::Budget::default();
         let env_impl = internal::EnvImpl::with_storage_and_budget(storage, budget.clone());
@@ -616,7 +651,7 @@ impl Env {
                 number,
                 config,
                 generators: generators.unwrap_or_default(),
-                snapshot,
+                snapshot: Some(snapshot),
                 auth_snapshot,
             },
         };
@@ -1561,7 +1596,7 @@ impl Env {
             Rc::new(s.ledger.clone()),
             Some(Rc::new(RefCell::new(s.generators))),
             Some(s.ledger.ledger_info()),
-            Some(Rc::new(s.ledger.clone())),
+            Some(Rc::new(RefCell::new(s.ledger.clone()))),
         )
     }
 
@@ -1597,16 +1632,24 @@ impl Env {
         self.to_snapshot().write_file(p).unwrap();
     }
 
-    /// Creates a new Env loaded with the [`LedgerSnapshot`].
+    /// Creates a new Env loaded with the snapshot source.
     ///
-    /// The ledger info and state in the snapshot are loaded into the Env.
-    pub fn from_ledger_snapshot(s: LedgerSnapshot) -> Env {
+    /// The ledger info and state from the snapshot source are loaded into the Env.
+    pub fn from_ledger_snapshot(input: impl Into<SnapshotSourceInput>) -> Env {
+        let SnapshotSourceInput {
+            source,
+            ledger_info,
+            snapshot,
+        } = input.into();
+
+        let snapshot = snapshot.map(|s| Rc::new(RefCell::new((*s).clone())));
+
         Env::new_for_testutils(
             EnvTestConfig::default(), // TODO: Allow setting the config.
-            Rc::new(s.clone()),
+            source,
             None,
-            Some(s.ledger_info()),
-            Some(Rc::new(s.clone())),
+            ledger_info,
+            snapshot,
         )
     }
 
@@ -1621,13 +1664,21 @@ impl Env {
 
     /// Create a snapshot from the Env's current state.
     pub fn to_ledger_snapshot(&self) -> LedgerSnapshot {
-        let snapshot = self.test_state.snapshot.clone().unwrap_or_default();
-        let mut snapshot = (*snapshot).clone();
+        let mut snapshot = self.to_ledger_snapshot_before();
         snapshot.set_ledger_info(self.ledger().get());
         snapshot.update_entries(&self.host().get_stored_entries().unwrap());
         snapshot
     }
 
+    /// Create a snapshot from all data loaded by the Env prior to any changes.
+    pub fn to_ledger_snapshot_before(&self) -> LedgerSnapshot {
+        self.test_state
+            .snapshot
+            .as_ref()
+            .map(|s| s.borrow().clone())
+            .unwrap_or_else(LedgerSnapshot::default)
+    }
+
     /// Create a snapshot file from the Env's current state.
     ///
     /// ### Panics
@@ -1677,6 +1728,7 @@ impl Drop for Env {
         // because it is only when there are no other references that the host
         // is being dropped.
         if self.env_impl.can_finish() && self.test_state.config.capture_snapshot_at_drop {
+            self.to_test_ledger_snapshot_before_file();
             self.to_test_snapshot_file();
         }
     }
@@ -1685,6 +1737,53 @@ impl Drop for Env {
 #[doc(hidden)]
 #[cfg(any(test, feature = "testutils"))]
 impl Env {
+    /// Create a snapshot file for the currently executing test containing the ledger entries
+    /// loaded but not modified.
+    ///
+    /// Writes the file to the `test_snapshots_before/{test-name}.N.json` path where
+    /// `N` is incremented for each unique `Env` in the test.
+    ///
+    /// Use to record the beginning state of a test.
+    ///
+    /// No file will be created if the environment has no meaningful data such
+    /// as stored entries or events.
+    ///
+    /// ### Panics
+    ///
+    /// If there is any error writing the file.
+    pub(crate) fn to_test_ledger_snapshot_before_file(&self) {
+        let snapshot = self.to_ledger_snapshot_before();
+
+        // Don't write a snapshot that has no data in it.
+        if snapshot.entries().into_iter().count() == 0 {
+            return;
+        }
+
+        // Determine path to write test snapshots to.
+        let Some(test_name) = &self.test_state.test_name else {
+            // If there's no test name, we're not in a test context, so don't write snapshots.
+            return;
+        };
+        let number = self.test_state.number;
+        // Break up the test name into directories, using :: as the separator.
+        // The :: module separator cannot be written into the filename because
+        // some operating systems (e.g. Windows) do not allow the : character in
+        // filenames.
+        let test_name_path = test_name
+            .split("::")
+            .map(|p| std::path::Path::new(p).to_path_buf())
+            .reduce(|p0, p1| p0.join(p1))
+            .expect("test name to not be empty");
+        let dir = std::path::Path::new("test_snapshots_before");
+        let p = dir
+            .join(&test_name_path)
+            .with_extension(format!("{number}.json"));
+
+        // Write test snapshots to file.
+        eprintln!("Writing test snapshot before file for test {test_name:?} to {p:?}.");
+        snapshot.write_file(p).unwrap();
+    }
+
     /// Create a snapshot file for the currently executing test.
     ///
     /// Writes the file to the `test_snapshots/{test-name}.N.json` path where

soroban-sdk/src/testutils.rs

@@ -20,11 +20,16 @@ pub mod storage;
 
 pub mod cost_estimate;
 
-use crate::{xdr, ConstructorArgs, Env, Val, Vec};
+use crate::{
+    xdr::{self, LedgerEntry, LedgerKey},
+    ConstructorArgs, Env, Val, Vec,
+};
 use soroban_ledger_snapshot::LedgerSnapshot;
 
 pub use crate::env::EnvTestConfig;
 
+pub use crate::env::internal::{storage::SnapshotSource, HostError};
+
 pub trait Register {
     fn register<'i, I, A>(self, env: &Env, id: I, args: A) -> crate::Address
     where
@@ -605,3 +610,78 @@ impl StellarAssetContract {
         self.asset.clone()
     }
 }
+
+/// A snapshot source that wraps another snapshot source and tracks entries
+/// in a cached snapshot as they are accessed.
+///
+/// The cached snapshot maintains a pre-change version of the ledger state,
+/// where entries are added as they are loaded from the underlying snapshot source.
+pub struct SnapshotSourceCache {
+    /// The underlying snapshot source to delegate to
+    source: Rc<dyn SnapshotSource>,
+    /// The snapshot to read from first
+    cache: Rc<std::cell::RefCell<LedgerSnapshot>>,
+}
+
+impl SnapshotSourceCache {
+    /// Create a new SnapshotSourceCache that wraps the given snapshot source
+    /// and uses the provided shared snapshot.
+    pub fn new(
+        source: Rc<dyn SnapshotSource>,
+        cache: Rc<std::cell::RefCell<LedgerSnapshot>>,
+    ) -> Self {
+        Self { source, cache }
+    }
+
+    /// Get a reference to the cache
+    pub fn cache(&self) -> std::cell::Ref<'_, LedgerSnapshot> {
+        self.cache.borrow()
+    }
+}
+
+impl SnapshotSource for SnapshotSourceCache {
+    fn get(
+        &self,
+        key: &Rc<LedgerKey>,
+    ) -> Result<Option<(Rc<LedgerEntry>, Option<u32>)>, HostError> {
+        // First, check if the entry exists in the cache
+        {
+            let cache = self.cache.borrow();
+            for (entry_key, (entry, live_until_ledger)) in cache.entries() {
+                if entry_key.as_ref() == key.as_ref() {
+                    return Ok(Some((Rc::new((**entry).clone()), *live_until_ledger)));
+                }
+            }
+        }
+
+        // If not in cache, try to get from the underlying source
+        let result = self.source.get(key)?;
+
+        // If we got an entry, add it to our cache using update_entries
+        if let Some((entry, live_until_ledger)) = &result {
+            let mut cache = self.cache.borrow_mut();
+            // Use update_entries to add or update the entry in the cache
+            let entry_tuple = (key.clone(), Some((entry.clone(), *live_until_ledger)));
+            cache.update_entries(std::iter::once(&entry_tuple));
+        }
+
+        Ok(result)
+    }
+}
+
+pub struct SnapshotSourceInput {
+    pub source: Rc<dyn SnapshotSource>,
+    pub ledger_info: Option<LedgerInfo>,
+    pub snapshot: Option<Rc<LedgerSnapshot>>,
+}
+
+impl From<LedgerSnapshot> for SnapshotSourceInput {
+    fn from(s: LedgerSnapshot) -> Self {
+        let s = Rc::new(s);
+        Self {
+            source: s.clone(),
+            ledger_info: Some(s.ledger_info()),
+            snapshot: Some(s),
+        }
+    }
+}

soroban-sdk/test_snapshots_before/tests/contract_snapshot/test.2.json

@@ -0,0 +1,96 @@
+{
+  "protocol_version": 23,
+  "sequence_number": 0,
+  "timestamp": 0,
+  "network_id": "0000000000000000000000000000000000000000000000000000000000000000",
+  "base_reserve": 0,
+  "min_persistent_entry_ttl": 4096,
+  "min_temp_entry_ttl": 16,
+  "max_entry_ttl": 6312000,
+  "ledger_entries": [
+    [
+      {
+        "contract_data": {
+          "contract": "CAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAD2KM",
+          "key": {
+            "i32": 2
+          },
+          "durability": "persistent"
+        }
+      },
+      [
+        {
+          "last_modified_ledger_seq": 0,
+          "data": {
+            "contract_data": {
+              "ext": "v0",
+              "contract": "CAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAD2KM",
+              "key": {
+                "i32": 2
+              },
+              "durability": "persistent",
+              "val": {
+                "i32": 4
+              }
+            }
+          },
+          "ext": "v0"
+        },
+        4095
+      ]
+    ],
+    [
+      {
+        "contract_data": {
+          "contract": "CAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAD2KM",
+          "key": "ledger_key_contract_instance",
+          "durability": "persistent"
+        }
+      },
+      [
+        {
+          "last_modified_ledger_seq": 0,
+          "data": {
+            "contract_data": {
+              "ext": "v0",
+              "contract": "CAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAD2KM",
+              "key": "ledger_key_contract_instance",
+              "durability": "persistent",
+              "val": {
+                "contract_instance": {
+                  "executable": {
+                    "wasm": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
+                  },
+                  "storage": null
+                }
+              }
+            }
+          },
+          "ext": "v0"
+        },
+        4095
+      ]
+    ],
+    [
+      {
+        "contract_code": {
+          "hash": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
+        }
+      },
+      [
+        {
+          "last_modified_ledger_seq": 0,
+          "data": {
+            "contract_code": {
+              "ext": "v0",
+              "hash": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855",
+              "code": ""
+            }
+          },
+          "ext": "v0"
+        },
+        4095
+      ]
+    ]
+  ]
+}