From 496cd23422e67ba14a2afb713199e5c15b4d9849 Mon Sep 17 00:00:00 2001
From: Andrew Fitzgerald <apfitzge@gmail.com>
Date: Wed, 17 Jan 2024 15:46:38 -0800
Subject: [PATCH] transaction cost based on requested write-locks

---
 cost-model/src/cost_model.rs         | 2 +-
 sdk/program/src/message/sanitized.rs | 2 ++
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/cost-model/src/cost_model.rs b/cost-model/src/cost_model.rs
index 8f8ee9ab752f24..3efc5d94148293 100644
--- a/cost-model/src/cost_model.rs
+++ b/cost-model/src/cost_model.rs
@@ -76,7 +76,7 @@ impl CostModel {
     fn get_write_lock_cost(tx_cost: &mut UsageCostDetails, transaction: &SanitizedTransaction) {
         tx_cost.writable_accounts = Self::get_writable_accounts(transaction);
         tx_cost.write_lock_cost =
-            WRITE_LOCK_UNITS.saturating_mul(tx_cost.writable_accounts.len() as u64);
+            WRITE_LOCK_UNITS.saturating_mul(transaction.message().num_write_locks());
     }
 
     fn get_transaction_cost(
diff --git a/sdk/program/src/message/sanitized.rs b/sdk/program/src/message/sanitized.rs
index 640159a7ad2dea..098a781ea4dbf7 100644
--- a/sdk/program/src/message/sanitized.rs
+++ b/sdk/program/src/message/sanitized.rs
@@ -360,6 +360,8 @@ impl SanitizedMessage {
         num_signatures
     }
 
+    /// Returns the number of requested write-locks in this message.
+    /// This does not consider if write-locks are demoted.
     pub fn num_write_locks(&self) -> u64 {
         self.account_keys()
             .len()