Rounding improvement for budget model with iterations > 1 (#1450)

### What

This is a simple improvement to the rounding of budget model evaluation,
when the model is linear and iteration > 1.

Current: 1. evaluates the linear equation, 2. scale the result back,
then 3. multiply by iteration.
This change: 1. evaluates the linear equation, 2. multiply by iteration
3. scale the result back.

The latter results in less rounding and better precision when the
iteration count is large.
Also added a test to illustrated expected behavior.

### Why

[TODO: Why this change is being made. Include any context required to
understand the why.]

### Known limitations

[TODO or N/A]

---------

Co-authored-by: Siddharth Suresh <[email protected]>

soroban-env-host/src/budget/dimension.rs

@@ -165,7 +165,7 @@ impl BudgetDimension {
         is_shadow: IsShadowMode,
     ) -> Result<u64, HostError> {
         let cm = self.get_cost_model(ty)?;
-        let amount = cm.evaluate(input)?.saturating_mul(iterations);
+        let amount = cm.evaluate(iterations, input)?;
 
         #[cfg(all(not(target_family = "wasm"), feature = "tracy"))]
         if _is_cpu.0 {

soroban-env-host/src/budget/model.rs

@@ -6,13 +6,16 @@ use core::fmt::{Debug, Display};
 
 /// We provide a "cost model" object that evaluates a linear expression:
 ///
-///    f(x) = a + b * Option<x>
+///    f(x) = I * (a + b * Option<x>)
 ///
-/// Where a, b are "fixed" parameters at construction time (extracted from an
-/// on-chain cost schedule, so technically not _totally_ fixed) and Option<x>
-/// is some abstract input variable -- say, event counts or object sizes --
-/// provided at runtime. If the input cannot be defined, i.e., the cost is
-/// constant, input-independent, then pass in `None` as the input.
+/// I is a simple scale factor that represents number of iterations the linear
+/// model inside the parentheses is repeated.
+///
+/// a, b are "fixed" parameters at construction time (extracted from an on-chain
+/// cost schedule, so technically not _totally_ fixed) and Option<x> is some
+/// abstract input variable -- say, event counts or object sizes -- provided at
+/// runtime. If the input cannot be defined, i.e., the cost is constant,
+/// input-independent, then pass in `None` as the input.
 ///
 /// The same `CostModel` type, i.e. `CostType` (applied to different parameters
 /// and variables) is used for calculating memory as well as CPU time.
@@ -22,10 +25,10 @@ use core::fmt::{Debug, Display};
 /// sufficiently by a linear model and 2. they together encompass the vast
 /// majority of available operations done by the `env` -- the host and the VM.
 ///
-/// The parameters for a `CostModel` are calibrated empirically. See this crate's
-/// benchmarks for more details.
+/// The parameters for a `CostModel` are calibrated empirically. See this
+/// crate's benchmarks for more details.
 pub trait HostCostModel {
-    fn evaluate(&self, input: Option<u64>) -> Result<u64, HostError>;
+    fn evaluate(&self, iterations: u64, input: Option<u64>) -> Result<u64, HostError>;
 
     #[cfg(any(test, feature = "testutils", feature = "bench"))]
     fn reset(&mut self);
@@ -110,14 +113,17 @@ impl TryFrom<ContractCostParamEntry> for MeteredCostComponent {
 }
 
 impl HostCostModel for MeteredCostComponent {
-    fn evaluate(&self, input: Option<u64>) -> Result<u64, HostError> {
-        let const_term = self.const_term;
+    fn evaluate(&self, iterations: u64, input: Option<u64>) -> Result<u64, HostError> {
+        let const_term = self.const_term.saturating_mul(iterations);
         match input {
             Some(input) => {
                 let mut res = const_term;
                 if !self.lin_term.is_zero() {
-                    let lin_cost = self.lin_term.saturating_mul(input).unscale();
-                    res = res.saturating_add(lin_cost)
+                    let lin_cost = self
+                        .lin_term
+                        .saturating_mul(input)
+                        .saturating_mul(iterations);
+                    res = res.saturating_add(lin_cost.unscale())
                 }
                 Ok(res)
             }
@@ -131,3 +137,25 @@ impl HostCostModel for MeteredCostComponent {
         self.lin_term = ScaledU64(0);
     }
 }
+
+mod test {
+    #[allow(unused)]
+    use super::{HostCostModel, MeteredCostComponent, ScaledU64};
+
+    #[test]
+    fn test_model_evaluation_with_rounding() {
+        let test_model = MeteredCostComponent {
+            const_term: 3,
+            lin_term: ScaledU64(5),
+        };
+        // low iteration + low input
+        // the constant part is 3, the linear part is 5 >> 7 == 0, total is 3
+        assert_eq!(3, test_model.evaluate(1, Some(1)).unwrap());
+        // low iteration + high input
+        // the contant part is 3, the linear part is (26 * 5) >> 7 == 1, total is 4
+        assert_eq!(4, test_model.evaluate(1, Some(26)).unwrap());
+        // high iteration + low input
+        // the constant part is 26 * 3 == 78, the linear part is (26 * 5) >> 7 == 1, total is 79
+        assert_eq!(79, test_model.evaluate(26, Some(1)).unwrap());
+    }
+}

soroban-env-host/src/host/declared_size.rs

@@ -411,8 +411,20 @@ mod test {
         expect!["16"].assert_eq(size_of::<HostError>().to_string().as_str());
         #[cfg(target_arch = "x86_64")]
         expect!["512"].assert_eq(size_of::<Context>().to_string().as_str());
+
+        #[rustversion::before(1.79)]
+        #[cfg(target_arch = "aarch64")]
+        fn check_aarch64_size_that_changed_at_rust_1_79() {
+            expect!["496"].assert_eq(size_of::<Context>().to_string().as_str());
+        }
+        #[rustversion::since(1.79)]
         #[cfg(target_arch = "aarch64")]
-        expect!["496"].assert_eq(size_of::<Context>().to_string().as_str());
+        fn check_aarch64_size_that_changed_at_rust_1_79() {
+            expect!["488"].assert_eq(size_of::<Context>().to_string().as_str());
+        }
+        #[cfg(target_arch = "aarch64")]
+        check_aarch64_size_that_changed_at_rust_1_79();
+
         expect!["16"].assert_eq(size_of::<Address>().to_string().as_str());
 
         expect!["1"].assert_eq(size_of::<AccessType>().to_string().as_str());