From 11e5f3195d3ee6d3d85187323cc79ecc2300160f Mon Sep 17 00:00:00 2001
From: Dmytro Kozhevin <dmytro@stellar.org>
Date: Thu, 23 Jan 2025 14:58:35 -0500
Subject: [PATCH] Tighten the budget estimations in e2e tests. (#1512)

### What

Tighten the budget estimations in e2e tests.

There was a small bug in the module cache construction - we didn't
consider the proper version of the code entry when building the module
cache to pass to the enforcing mode. This fix allows us to bring the
discrepancy between the recording and enforcing modes back to 0.02.

### Why

Test fixes

### Known limitations

N/A
---
 soroban-env-host/src/test/e2e_tests.rs | 51 +++++++++++++++++++-------
 1 file changed, 38 insertions(+), 13 deletions(-)

diff --git a/soroban-env-host/src/test/e2e_tests.rs b/soroban-env-host/src/test/e2e_tests.rs
index 2b6f9ab78..785647337 100644
--- a/soroban-env-host/src/test/e2e_tests.rs
+++ b/soroban-env-host/src/test/e2e_tests.rs
@@ -1,7 +1,9 @@
 use crate::builtin_contracts::testutils::AccountContractSigner;
+use crate::crypto::sha256_hash_from_bytes_raw;
 use crate::e2e_invoke::RecordingInvocationAuthMode;
 use crate::e2e_testutils::{account_entry, bytes_sc_val, upload_wasm_host_fn};
 use crate::testutils::simple_account_sign_fn;
+use crate::vm::VersionedContractCodeCostInputs;
 use crate::{
     budget::{AsBudget, Budget},
     builtin_contracts::testutils::TestSigner,
@@ -17,10 +19,10 @@ use crate::{
     },
     testutils::MockSnapshotSource,
     xdr::{
-        AccountId, ContractDataDurability, ContractDataEntry, ContractEvent, ContractExecutable,
-        ContractIdPreimage, ContractIdPreimageFromAddress, CreateContractArgs, DiagnosticEvent,
-        ExtensionPoint, HashIdPreimage, HashIdPreimageSorobanAuthorization, HostFunction,
-        InvokeContractArgs, LedgerEntry, LedgerEntryData, LedgerFootprint, LedgerKey,
+        AccountId, ContractCodeEntryExt, ContractDataDurability, ContractDataEntry, ContractEvent,
+        ContractExecutable, ContractIdPreimage, ContractIdPreimageFromAddress, CreateContractArgs,
+        DiagnosticEvent, ExtensionPoint, Hash, HashIdPreimage, HashIdPreimageSorobanAuthorization,
+        HostFunction, InvokeContractArgs, LedgerEntry, LedgerEntryData, LedgerFootprint, LedgerKey,
         LedgerKeyContractCode, LedgerKeyContractData, Limits, ReadXdr, ScAddress,
         ScContractInstance, ScErrorCode, ScErrorType, ScMap, ScNonceKey, ScVal, ScVec,
         SorobanAuthorizationEntry, SorobanCredentials, SorobanResources, TtlEntry, Uint256,
@@ -50,7 +52,7 @@ use std::rc::Rc;
 // We don't anticipate this divergence to be too high though: specifically,
 // we expect the estimated instructions to be within a range of
 // [1 - RECORDING_MODE_INSTRUCTIONS_RANGE, 1 + RECORDING_MODE_INSTRUCTIONS_RANGE] * real_instructions
-const RECORDING_MODE_INSTRUCTIONS_RANGE: f64 = 0.3;
+const RECORDING_MODE_INSTRUCTIONS_RANGE: f64 = 0.02;
 
 fn wasm_entry_size(wasm: &[u8]) -> u32 {
     wasm_entry(wasm).to_xdr(Limits::none()).unwrap().len() as u32
@@ -318,13 +320,7 @@ fn invoke_host_function_helper(
         })
         .collect();
 
-    let ctx = E2eTestCompilationContext::new()?;
-    let cache = ModuleCache::new(&ctx)?;
-    for (e, _) in ledger_entries_with_ttl.iter() {
-        if let LedgerEntryData::ContractCode(cd) = &e.data {
-            cache.parse_and_cache_module_simple(&ctx, ledger_info.protocol_version, &cd.code)?;
-        }
-    }
+    let module_cache = build_module_cache_for_entries(ledger_info, ledger_entries_with_ttl)?;
 
     let budget = Budget::default();
     budget
@@ -344,7 +340,7 @@ fn invoke_host_function_helper(
         prng_seed.to_vec(),
         &mut diagnostic_events,
         None,
-        Some(cache),
+        Some(module_cache),
     )?;
     Ok(InvokeHostFunctionHelperResult {
         invoke_result: res
@@ -361,6 +357,35 @@ fn invoke_host_function_helper(
     })
 }
 
+fn build_module_cache_for_entries(
+    ledger_info: &LedgerInfo,
+    ledger_entries_with_ttl: Vec<(LedgerEntry, Option<u32>)>,
+) -> Result<ModuleCache, HostError> {
+    let ctx = E2eTestCompilationContext::new()?;
+    let cache = ModuleCache::new(&ctx)?;
+    for (e, _) in ledger_entries_with_ttl.iter() {
+        if let LedgerEntryData::ContractCode(cd) = &e.data {
+            let contract_id = Hash(sha256_hash_from_bytes_raw(&cd.code, ctx.as_budget())?);
+            let code_cost_inputs = match &cd.ext {
+                ContractCodeEntryExt::V0 => VersionedContractCodeCostInputs::V0 {
+                    wasm_bytes: cd.code.len(),
+                },
+                ContractCodeEntryExt::V1(v1) => {
+                    VersionedContractCodeCostInputs::V1(v1.cost_inputs.clone())
+                }
+            };
+            cache.parse_and_cache_module(
+                &ctx,
+                ledger_info.protocol_version,
+                &contract_id,
+                &cd.code,
+                code_cost_inputs,
+            )?;
+        }
+    }
+    Ok(cache)
+}
+
 fn invoke_host_function_recording_helper(
     enable_diagnostics: bool,
     host_fn: &HostFunction,