From 5e4332ee35507515644d62f94f5828678cef7c5c Mon Sep 17 00:00:00 2001
From: Andrew Fitzgerald <apfitzge@gmail.com>
Date: Mon, 22 Jan 2024 09:22:49 -0800
Subject: [PATCH] Feature Impl: cost model uses number of requested write locks
 (#34820)

---
 cost-model/src/cost_model.rs         | 45 ++++++++++++++++++++++++----
 sdk/program/src/message/sanitized.rs |  2 ++
 sdk/src/feature_set.rs               |  5 ++++
 3 files changed, 47 insertions(+), 5 deletions(-)

diff --git a/cost-model/src/cost_model.rs b/cost-model/src/cost_model.rs
index ba01ed9fe993a5..1e15735426737f 100644
--- a/cost-model/src/cost_model.rs
+++ b/cost-model/src/cost_model.rs
@@ -18,7 +18,7 @@ use {
     solana_sdk::{
         borsh1::try_from_slice_unchecked,
         compute_budget::{self, ComputeBudgetInstruction},
-        feature_set::{include_loaded_accounts_data_size_in_fee_calculation, FeatureSet},
+        feature_set::{self, include_loaded_accounts_data_size_in_fee_calculation, FeatureSet},
         fee::FeeStructure,
         instruction::CompiledInstruction,
         program_utils::limited_deserialize,
@@ -44,7 +44,7 @@ impl CostModel {
             let mut tx_cost = UsageCostDetails::new_with_default_capacity();
 
             tx_cost.signature_cost = Self::get_signature_cost(transaction);
-            Self::get_write_lock_cost(&mut tx_cost, transaction);
+            Self::get_write_lock_cost(&mut tx_cost, transaction, feature_set);
             Self::get_transaction_cost(&mut tx_cost, transaction, feature_set);
             tx_cost.account_data_size = Self::calculate_account_data_size(transaction);
 
@@ -73,10 +73,19 @@ impl CostModel {
             .collect()
     }
 
-    fn get_write_lock_cost(tx_cost: &mut UsageCostDetails, transaction: &SanitizedTransaction) {
+    fn get_write_lock_cost(
+        tx_cost: &mut UsageCostDetails,
+        transaction: &SanitizedTransaction,
+        feature_set: &FeatureSet,
+    ) {
         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);
+        let num_write_locks =
+            if feature_set.is_active(&feature_set::cost_model_requested_write_lock_cost::id()) {
+                transaction.message().num_write_locks()
+            } else {
+                tx_cost.writable_accounts.len() as u64
+            };
+        tx_cost.write_lock_cost = WRITE_LOCK_UNITS.saturating_mul(num_write_locks);
     }
 
     fn get_transaction_cost(
@@ -329,6 +338,32 @@ mod tests {
         assert_eq!(0, tx_cost.data_bytes_cost);
     }
 
+    #[test]
+    fn test_cost_model_demoted_write_lock() {
+        let (mint_keypair, start_hash) = test_setup();
+
+        // Cannot write-lock the system program, it will be demoted when taking locks.
+        // However, the cost should be calculated as if it were taken.
+        let simple_transaction = SanitizedTransaction::from_transaction_for_tests(
+            system_transaction::transfer(&mint_keypair, &system_program::id(), 2, start_hash),
+        );
+
+        // Feature not enabled - write lock is demoted and does not count towards cost
+        {
+            let tx_cost = CostModel::calculate_cost(&simple_transaction, &FeatureSet::default());
+            assert_eq!(WRITE_LOCK_UNITS, tx_cost.write_lock_cost());
+            assert_eq!(1, tx_cost.writable_accounts().len());
+        }
+
+        // Feature enabled - write lock is demoted but still counts towards cost
+        {
+            let tx_cost =
+                CostModel::calculate_cost(&simple_transaction, &FeatureSet::all_enabled());
+            assert_eq!(2 * WRITE_LOCK_UNITS, tx_cost.write_lock_cost());
+            assert_eq!(1, tx_cost.writable_accounts().len());
+        }
+    }
+
     #[test]
     fn test_cost_model_compute_budget_transaction() {
         let (mint_keypair, start_hash) = test_setup();
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()
diff --git a/sdk/src/feature_set.rs b/sdk/src/feature_set.rs
index f2e9c63ff1b2c9..2941c94ae81cb3 100644
--- a/sdk/src/feature_set.rs
+++ b/sdk/src/feature_set.rs
@@ -768,6 +768,10 @@ pub mod enable_zk_proof_from_account {
     solana_sdk::declare_id!("zkiTNuzBKxrCLMKehzuQeKZyLtX2yvFcEKMML8nExU8");
 }
 
+pub mod cost_model_requested_write_lock_cost {
+    solana_sdk::declare_id!("wLckV1a64ngtcKPRGU4S4grVTestXjmNjxBjaKZrAcn");
+}
+
 lazy_static! {
     /// Map of feature identifiers to user-visible description
     pub static ref FEATURE_NAMES: HashMap<Pubkey, &'static str> = [
@@ -955,6 +959,7 @@ lazy_static! {
         (deprecate_executable_meta_update_in_bpf_loader::id(), "deprecate executable meta flag update in bpf loader #34194"),
         (enable_zk_proof_from_account::id(), "Enable zk token proof program to read proof from accounts instead of instruction data #34750"),
         (curve25519_restrict_msm_length::id(), "restrict curve25519 multiscalar multiplication vector lengths #34763"),
+        (cost_model_requested_write_lock_cost::id(), "cost model uses number of requested write locks #34819"),
         /*************** ADD NEW FEATURES HERE ***************/
     ]
     .iter()