From 443dbf322dc7f26646877fb454d4ee05759e4fa0 Mon Sep 17 00:00:00 2001
From: Tao Zhu <tao@solana.com>
Date: Fri, 10 Mar 2023 18:06:43 +0000
Subject: [PATCH] add test for refaction heap size, fix size truncating by div
 op

---
 programs/bpf_loader/src/lib.rs | 34 +++++++++++++++++++++++++++++-----
 1 file changed, 29 insertions(+), 5 deletions(-)

diff --git a/programs/bpf_loader/src/lib.rs b/programs/bpf_loader/src/lib.rs
index bea7a4c6fb8e17..d1d0cc49696c84 100644
--- a/programs/bpf_loader/src/lib.rs
+++ b/programs/bpf_loader/src/lib.rs
@@ -295,6 +295,15 @@ fn check_loader_id(id: &Pubkey) -> bool {
         || bpf_loader_upgradeable::check_id(id)
 }
 
+fn calculate_heap_cost(heap_size: usize, heap_cost: u64) -> u64 {
+    const KILOBYTE: u64 = 1024;
+    ((heap_size as u64)
+        .saturating_add((32_u64 - 1).saturating_mul(KILOBYTE))
+        .saturating_div(32_u64.saturating_mul(KILOBYTE)))
+        .saturating_sub(1)
+        .saturating_mul(heap_cost)
+}
+
 /// Create the SBF virtual machine
 pub fn create_vm<'a, 'b>(
     program: &'a VerifiedExecutable<RequisiteVerifier, InvokeContext<'b>>,
@@ -304,11 +313,7 @@ pub fn create_vm<'a, 'b>(
 ) -> Result<EbpfVm<'a, RequisiteVerifier, InvokeContext<'b>>, EbpfError> {
     let compute_budget = invoke_context.get_compute_budget();
     let heap_size = compute_budget.heap_size.unwrap_or(HEAP_LENGTH);
-    let _ = invoke_context.consume_checked(
-        ((heap_size as u64).saturating_div(32_u64.saturating_mul(1024)))
-            .saturating_sub(1)
-            .saturating_mul(compute_budget.heap_cost),
-    );
+    let _ = invoke_context.consume_checked(calculate_heap_cost(heap_size, compute_budget.heap_cost));
     let heap =
         AlignedMemory::<HOST_ALIGN>::zero_filled(compute_budget.heap_size.unwrap_or(HEAP_LENGTH));
     let check_aligned = bpf_loader_deprecated::id()
@@ -3832,4 +3837,23 @@ mod tests {
             },
         );
     }
+
+    #[test]
+    fn test_calculate_heap_cost() {
+        let heap_cost = 8_u64;
+    
+        // heap allocations are in 32K block, `heap_cost` of CU is consumed per additional 32k 
+    
+        // assert less than 32K heap should cost zero unit
+        assert_eq!(0, calculate_heap_cost(31_usize * 1024, heap_cost));
+    
+        // assert exact 32K heap should be cost zero unit
+        assert_eq!(0, calculate_heap_cost(32_usize * 1024, heap_cost));
+    
+        // assert slightly more than 32K heap should cost 1 * heap_cost
+        assert_eq!(heap_cost, calculate_heap_cost(33_usize * 1024, heap_cost));
+    
+        // assert exact 64K heap should cost 1 * heap_cost
+        assert_eq!(heap_cost, calculate_heap_cost(64_usize * 1024, heap_cost));
+    }
 }