From c03cc14b44d7591df24283fe65c92e032e589b46 Mon Sep 17 00:00:00 2001
From: Jay Geng <jay@stellar.org>
Date: Tue, 26 Mar 2024 19:19:13 -0400
Subject: [PATCH] Refactor extend {code, instance} TTL (#1373)

### What

To avoid code duplication.
Previous [attempt](https://github.com/stellar/rs-soroban-env/pull/1368)
broke some code observation so it had to be reverted. This PR fixes it.

### 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 <siddharth@stellar.org>
---
 soroban-env-host/src/host.rs             | 33 ++++++++++++------------
 soroban-env-host/src/host/data_helper.rs | 28 +++++++++++++++-----
 2 files changed, 37 insertions(+), 24 deletions(-)

diff --git a/soroban-env-host/src/host.rs b/soroban-env-host/src/host.rs
index e2068ba72..d89447827 100644
--- a/soroban-env-host/src/host.rs
+++ b/soroban-env-host/src/host.rs
@@ -2073,11 +2073,14 @@ impl VmCallerEnv for Host {
         extend_to: U32Val,
     ) -> Result<Void, HostError> {
         let contract_id = self.get_current_contract_id_internal()?;
-        self.extend_contract_instance_and_code_ttl_from_contract_id(
+        let key = self.contract_instance_ledger_key(&contract_id)?;
+        self.extend_contract_instance_ttl_from_contract_id(
             &contract_id,
+            key.clone(),
             threshold.into(),
             extend_to.into(),
         )?;
+        self.extend_contract_code_ttl_from_contract_id(key, threshold.into(), extend_to.into())?;
         Ok(Val::VOID)
     }
 
@@ -2090,9 +2093,13 @@ impl VmCallerEnv for Host {
     ) -> Result<Void, Self::Error> {
         let contract_id = self.contract_id_from_address(contract)?;
         let key = self.contract_instance_ledger_key(&contract_id)?;
-        self.try_borrow_storage_mut()?
-            .extend_ttl(self, key, threshold.into(), extend_to.into())
-            .map_err(|e| self.decorate_contract_instance_storage_error(e, &contract_id))?;
+
+        self.extend_contract_instance_ttl_from_contract_id(
+            &contract_id,
+            key,
+            threshold.into(),
+            extend_to.into(),
+        )?;
 
         Ok(Val::VOID)
     }
@@ -2105,11 +2112,14 @@ impl VmCallerEnv for Host {
         extend_to: U32Val,
     ) -> Result<Void, Self::Error> {
         let contract_id = self.contract_id_from_address(contract)?;
-        self.extend_contract_instance_and_code_ttl_from_contract_id(
+        let key = self.contract_instance_ledger_key(&contract_id)?;
+        self.extend_contract_instance_ttl_from_contract_id(
             &contract_id,
+            key.clone(),
             threshold.into(),
             extend_to.into(),
         )?;
+        self.extend_contract_code_ttl_from_contract_id(key, threshold.into(), extend_to.into())?;
         Ok(Val::VOID)
     }
 
@@ -2123,18 +2133,7 @@ impl VmCallerEnv for Host {
         let contract_id = self.contract_id_from_address(contract)?;
         let key = self.contract_instance_ledger_key(&contract_id)?;
 
-        match self
-            .retrieve_contract_instance_from_storage(&key)?
-            .executable
-        {
-            ContractExecutable::Wasm(wasm_hash) => {
-                let key = self.contract_code_ledger_key(&wasm_hash)?;
-                self.try_borrow_storage_mut()?
-                    .extend_ttl(self, key, threshold.into(), extend_to.into())
-                    .map_err(|e| self.decorate_contract_code_storage_error(e, &wasm_hash))?;
-            }
-            ContractExecutable::StellarAsset => {}
-        }
+        self.extend_contract_code_ttl_from_contract_id(key, threshold.into(), extend_to.into())?;
 
         Ok(Val::VOID)
     }
diff --git a/soroban-env-host/src/host/data_helper.rs b/soroban-env-host/src/host/data_helper.rs
index cada88a83..5bb2e22c3 100644
--- a/soroban-env-host/src/host/data_helper.rs
+++ b/soroban-env-host/src/host/data_helper.rs
@@ -253,18 +253,14 @@ impl Host {
         Ok(())
     }
 
-    pub(crate) fn extend_contract_instance_and_code_ttl_from_contract_id(
+    pub(crate) fn extend_contract_code_ttl_from_contract_id(
         &self,
-        contract_id: &Hash,
+        instance_key: Rc<LedgerKey>,
         threshold: u32,
         extend_to: u32,
     ) -> Result<(), HostError> {
-        let key = self.contract_instance_ledger_key(&contract_id)?;
-        self.try_borrow_storage_mut()?
-            .extend_ttl(self, key.metered_clone(self)?, threshold, extend_to)
-            .map_err(|e| self.decorate_contract_instance_storage_error(e, &contract_id))?;
         match self
-            .retrieve_contract_instance_from_storage(&key)?
+            .retrieve_contract_instance_from_storage(&instance_key)?
             .executable
         {
             ContractExecutable::Wasm(wasm_hash) => {
@@ -278,6 +274,24 @@ impl Host {
         Ok(())
     }
 
+    pub(crate) fn extend_contract_instance_ttl_from_contract_id(
+        &self,
+        contract_id: &Hash,
+        instance_key: Rc<LedgerKey>,
+        threshold: u32,
+        extend_to: u32,
+    ) -> Result<(), HostError> {
+        self.try_borrow_storage_mut()?
+            .extend_ttl(
+                self,
+                instance_key.metered_clone(self)?,
+                threshold,
+                extend_to,
+            )
+            .map_err(|e| self.decorate_contract_instance_storage_error(e, &contract_id))?;
+        Ok(())
+    }
+
     // metering: covered by components
     pub(crate) fn get_full_contract_id_preimage(
         &self,