Refactor extend {code, instance} TTL (#1373)

### What To avoid code duplication. Previous [attempt](#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 <[email protected]>
stellar · Mar 26, 2024 · c03cc14 · c03cc14
1 parent a54b40f
commit c03cc14
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 24 deletions.
diff --git 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
@@ -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,