Emulate account entry extensions in simulation. (#1363)

### What

Emulate account entry extensions in simulation.

This also adds a general mechanism for modifying any ledger entries
retrieved from the snapshot.

### Why

Improve simulation accuracy.

Fixes #1357

### Known limitations

N/A

soroban-simulation/src/simulation.rs

@@ -3,7 +3,9 @@ use crate::resources::{
     compute_adjusted_transaction_resources, compute_resource_fee, simulate_extend_ttl_op_resources,
     simulate_invoke_host_function_op_resources, simulate_restore_op_resources,
 };
-use crate::snapshot_source::SnapshotSourceWithArchive;
+use crate::snapshot_source::{
+    SimulationSnapshotSource, SimulationSnapshotSourceWithArchive, SnapshotSourceWithArchive,
+};
 use anyhow::Result;
 use soroban_env_host::{
     e2e_invoke::invoke_host_function_in_recording_mode,
@@ -125,6 +127,7 @@ pub fn simulate_invoke_host_function_op(
     base_prng_seed: [u8; 32],
     enable_diagnostics: bool,
 ) -> Result<InvokeHostFunctionSimulationResult> {
+    let snapshot_source = Rc::new(SimulationSnapshotSource::new_from_rc(snapshot_source));
     let budget = network_config.create_budget()?;
     let mut diagnostic_events = vec![];
     let recording_result = invoke_host_function_in_recording_mode(
@@ -219,9 +222,10 @@ pub fn simulate_extend_ttl_op(
     keys_to_extend: &[LedgerKey],
     extend_to: u32,
 ) -> Result<ExtendTtlOpSimulationResult> {
+    let snapshot_source = SimulationSnapshotSource::new(snapshot_source);
     let (mut resources, rent_changes) = simulate_extend_ttl_op_resources(
         keys_to_extend,
-        snapshot_source,
+        &snapshot_source,
         ledger_info.sequence_number,
         extend_to,
     )?;
@@ -266,8 +270,9 @@ pub fn simulate_restore_op(
     ledger_info: &LedgerInfo,
     keys_to_restore: &[LedgerKey],
 ) -> Result<RestoreOpSimulationResult> {
+    let snapshot_source = SimulationSnapshotSourceWithArchive::new(snapshot_source);
     let (mut resources, rent_changes) =
-        simulate_restore_op_resources(keys_to_restore, snapshot_source, ledger_info)?;
+        simulate_restore_op_resources(keys_to_restore, &snapshot_source, ledger_info)?;
     let operation = OperationBody::RestoreFootprint(RestoreFootprintOp {
         ext: ExtensionPoint::V0,
     });

soroban-simulation/src/snapshot_source.rs

@@ -3,14 +3,19 @@ use crate::simulation::{
     simulate_restore_op, RestoreOpSimulationResult, SimulationAdjustmentConfig,
 };
 use anyhow::{anyhow, Result};
+use soroban_env_host::xdr::{
+    AccountEntry, AccountEntryExt, AccountEntryExtensionV1, AccountEntryExtensionV1Ext,
+    AccountEntryExtensionV2, AccountEntryExtensionV2Ext, AccountEntryExtensionV3, ExtensionPoint,
+    LedgerEntryData, Liabilities, SponsorshipDescriptor, TimePoint,
+};
 use soroban_env_host::{
     ledger_info::get_key_durability,
     storage::{EntryWithLiveUntil, SnapshotSource},
     xdr::{ContractDataDurability, LedgerKey, ScErrorCode, ScErrorType},
     HostError, LedgerInfo,
 };
 use std::cell::RefCell;
-use std::collections::BTreeSet;
+use std::collections::{BTreeMap, BTreeSet};
 use std::rc::Rc;
 
 /// Read-only ledger snapshot accessor that also has access
@@ -115,3 +120,171 @@ impl<T: SnapshotSourceWithArchive> SnapshotSource for AutoRestoringSnapshotSourc
         }
     }
 }
+
+#[derive(Default)]
+struct LedgerEntryUpdater {
+    updated_entries_cache: BTreeMap<Rc<LedgerKey>, Option<EntryWithLiveUntil>>,
+}
+
+impl LedgerEntryUpdater {
+    fn maybe_update_entry(
+        &mut self,
+        key: &Rc<LedgerKey>,
+        entry: Option<EntryWithLiveUntil>,
+    ) -> Option<EntryWithLiveUntil> {
+        if let Some(e) = self.updated_entries_cache.get(key) {
+            return e.clone();
+        }
+        if let Some((entry, live_until)) = &entry {
+            match &entry.data {
+                LedgerEntryData::Account(_) => {
+                    let mut updated_entry = (**entry).clone();
+                    match &mut updated_entry.data {
+                        LedgerEntryData::Account(acc) => {
+                            update_account_entry(acc);
+                        }
+                        _ => (),
+                    }
+                    let entry_with_live_until = Some((Rc::new(updated_entry), *live_until));
+                    self.updated_entries_cache
+                        .insert(key.clone(), entry_with_live_until.clone());
+                    return entry_with_live_until;
+                }
+                _ => (),
+            }
+        }
+        entry
+    }
+}
+
+enum SnapshotSourceHolder<'a> {
+    Ref(&'a dyn SnapshotSource),
+    Rc(Rc<dyn SnapshotSource>),
+}
+
+impl<'a> SnapshotSource for SnapshotSourceHolder<'a> {
+    fn get(&self, key: &Rc<LedgerKey>) -> Result<Option<EntryWithLiveUntil>, HostError> {
+        match self {
+            SnapshotSourceHolder::Ref(r) => r.get(key),
+            SnapshotSourceHolder::Rc(r) => r.get(key),
+        }
+    }
+}
+
+// This is an internal wrapper for the snapshot sources used in the simulation.
+// The purpose of the wrapper is to emulate any ledger entry modifying logic
+// that Core might perform before the Soroban host operation is invoked. For example,
+// Core might create the account extensions, which would impact the read bytes
+// amount and simulated CPU instructions.
+pub(crate) struct SimulationSnapshotSource<'a> {
+    inner_snapshot: SnapshotSourceHolder<'a>,
+    entry_updater: RefCell<LedgerEntryUpdater>,
+}
+
+// This is the same as `SimulationSnapshotSource`, but for
+// `SnapshotSourceWithArchive` trait.
+// Note, that we don't implement both traits for `SimulationSnapshotSource` in order
+// to avoid confusion: we never want to accidentally create a `SnapshotSource` with
+// `SnapshotSourceWithArchive` inner snapshot, or vice versa.
+pub(crate) struct SimulationSnapshotSourceWithArchive<'a, T: SnapshotSourceWithArchive> {
+    inner_snapshot: &'a T,
+    entry_updater: RefCell<LedgerEntryUpdater>,
+}
+
+impl<'a> SimulationSnapshotSource<'a> {
+    pub(crate) fn new(snapshot: &'a dyn SnapshotSource) -> Self {
+        Self {
+            inner_snapshot: SnapshotSourceHolder::Ref(snapshot),
+            entry_updater: RefCell::new(Default::default()),
+        }
+    }
+
+    pub(crate) fn new_from_rc(snapshot: Rc<dyn SnapshotSource>) -> Self {
+        Self {
+            inner_snapshot: SnapshotSourceHolder::Rc(snapshot),
+            entry_updater: RefCell::new(Default::default()),
+        }
+    }
+}
+
+impl<'a, T: SnapshotSourceWithArchive> SimulationSnapshotSourceWithArchive<'a, T> {
+    pub(crate) fn new(snapshot: &'a T) -> Self {
+        Self {
+            inner_snapshot: &snapshot,
+            entry_updater: RefCell::new(Default::default()),
+        }
+    }
+}
+
+impl<'a> SnapshotSource for SimulationSnapshotSource<'a> {
+    fn get(&self, key: &Rc<LedgerKey>) -> Result<Option<EntryWithLiveUntil>, HostError> {
+        Ok(self
+            .entry_updater
+            .borrow_mut()
+            .maybe_update_entry(key, self.inner_snapshot.get(key)?))
+    }
+}
+
+impl<'a, T: SnapshotSourceWithArchive> SnapshotSourceWithArchive
+    for SimulationSnapshotSourceWithArchive<'a, T>
+{
+    fn get_including_archived(
+        &self,
+        key: &Rc<LedgerKey>,
+    ) -> Result<Option<EntryWithLiveUntil>, HostError> {
+        Ok(self
+            .entry_updater
+            .borrow_mut()
+            .maybe_update_entry(key, self.inner_snapshot.get_including_archived(key)?))
+    }
+}
+
+fn update_account_entry(account_entry: &mut AccountEntry) {
+    match &mut account_entry.ext {
+        AccountEntryExt::V0 => {
+            let mut ext = AccountEntryExtensionV1 {
+                liabilities: Liabilities {
+                    buying: 0,
+                    selling: 0,
+                },
+                ext: AccountEntryExtensionV1Ext::V0,
+            };
+            fill_account_ext_v2(&mut ext, account_entry.signers.len());
+            account_entry.ext = AccountEntryExt::V1(ext);
+        }
+        AccountEntryExt::V1(ext) => {
+            fill_account_ext_v2(ext, account_entry.signers.len());
+        }
+    }
+}
+
+fn fill_account_ext_v2(account_ext_v1: &mut AccountEntryExtensionV1, signers_count: usize) {
+    match &mut account_ext_v1.ext {
+        AccountEntryExtensionV1Ext::V0 => {
+            let mut ext = AccountEntryExtensionV2 {
+                num_sponsored: 0,
+                num_sponsoring: 0,
+                signer_sponsoring_i_ds: vec![SponsorshipDescriptor(None); signers_count]
+                    .try_into()
+                    .unwrap_or_default(),
+                ext: AccountEntryExtensionV2Ext::V0,
+            };
+            fill_account_ext_v3(&mut ext);
+            account_ext_v1.ext = AccountEntryExtensionV1Ext::V2(ext);
+        }
+        AccountEntryExtensionV1Ext::V2(ext) => fill_account_ext_v3(ext),
+    }
+}
+
+fn fill_account_ext_v3(account_ext_v2: &mut AccountEntryExtensionV2) {
+    match account_ext_v2.ext {
+        AccountEntryExtensionV2Ext::V0 => {
+            account_ext_v2.ext = AccountEntryExtensionV2Ext::V3(AccountEntryExtensionV3 {
+                ext: ExtensionPoint::V0,
+                seq_ledger: 0,
+                seq_time: TimePoint(0),
+            });
+        }
+        AccountEntryExtensionV2Ext::V3(_) => (),
+    }
+}

soroban-simulation/src/test/simulation.rs

@@ -481,15 +481,15 @@ fn test_simulate_invoke_contract_with_auth() {
                     .try_into()
                     .unwrap()
                 },
-                instructions: 38139496,
-                read_bytes: 7440,
+                instructions: 38210172,
+                read_bytes: 7492,
                 write_bytes: 76,
             },
-            resource_fee: 39290,
+            resource_fee: 39363,
         })
     );
-    assert_eq!(res.simulated_instructions, 38139496);
-    assert_eq!(res.simulated_memory, 19069620);
+    assert_eq!(res.simulated_instructions, 38210172);
+    assert_eq!(res.simulated_memory, 19104958);
     assert_eq!(
         res.modified_entries,
         vec![LedgerEntryDiff {