Minimize stream storage load

StreamSwapProtocol · Dec 2, 2024 · cb4a3f9 · cb4a3f9
1 parent 2ceb3da
commit cb4a3f9
Show file tree

Hide file tree

Showing 11 changed files with 314 additions and 254 deletions.
diff --git a/contracts/stream/src/circuit_ops.rs b/contracts/stream/src/circuit_ops.rs
@@ -1,6 +1,6 @@
 use crate::helpers::build_u128_bank_send_msg;
 use crate::pool::pool_refund;
-use crate::state::{CONTROLLER_PARAMS, POSITIONS, STREAM_STATE};
+use crate::state::{CONTROLLER_PARAMS, POSITIONS, POST_STREAM, STREAM_INFO, STREAM_STATE};
 use crate::stream::{sync_stream, sync_stream_status};
 use crate::ContractError;
 use cosmwasm_std::{attr, BankMsg, CosmosMsg, DepsMut, Env, MessageInfo, Response, Timestamp};
@@ -97,15 +97,24 @@ pub fn execute_cancel_stream(
 
     // Refund all out tokens to stream creator(treasury)
     let mut refund_coins = vec![stream.out_asset.clone()];
-    let pool_refund_coins = pool_refund(&deps, stream.pool_config, stream.out_asset.denom.clone())?;
 
-    refund_coins.extend(pool_refund_coins);
+    // refund pool creation if any
+    let post_stream_ops = POST_STREAM.may_load(deps.storage)?;
+    if let Some(post_stream_ops) = post_stream_ops {
+        let pool_refund_coins = pool_refund(
+            &deps,
+            post_stream_ops.pool_config,
+            stream.out_asset.denom.clone(),
+        )?;
+        refund_coins.extend(pool_refund_coins);
+    }
 
+    let stream_info = STREAM_INFO.load(deps.storage)?;
     let funds_msgs: Vec<CosmosMsg> = refund_coins
         .iter()
         .map(|coin| {
             CosmosMsg::Bank(BankMsg::Send {
-                to_address: stream.treasury.to_string(),
+                to_address: stream_info.treasury.to_string(),
                 amount: vec![coin.clone()],
             })
         })
@@ -123,8 +132,9 @@ pub fn execute_cancel_stream_with_threshold(
     info: MessageInfo,
 ) -> Result<Response, ContractError> {
     let mut stream = STREAM_STATE.load(deps.storage)?;
+    let stream_info = STREAM_INFO.load(deps.storage)?;
     // Only stream creator can cancel the stream with threshold not reached
-    if info.sender != stream.stream_admin {
+    if info.sender != stream_info.stream_admin {
         return Err(ContractError::Unauthorized {});
     }
     sync_stream_status(&mut stream, env.block.time);
@@ -154,15 +164,23 @@ pub fn execute_cancel_stream_with_threshold(
 
     // Refund all out tokens to stream creator(treasury)
     let mut refund_coins = vec![stream.out_asset.clone()];
-    let pool_refund_coins = pool_refund(&deps, stream.pool_config, stream.out_asset.denom.clone())?;
 
-    refund_coins.extend(pool_refund_coins);
+    // refund pool creation if any
+    let post_stream_ops = POST_STREAM.may_load(deps.storage)?;
+    if let Some(post_stream_ops) = post_stream_ops {
+        let pool_refund_coins = pool_refund(
+            &deps,
+            post_stream_ops.pool_config,
+            stream.out_asset.denom.clone(),
+        )?;
+        refund_coins.extend(pool_refund_coins);
+    }
 
     let funds_msgs: Vec<CosmosMsg> = refund_coins
         .iter()
         .map(|coin| {
             CosmosMsg::Bank(BankMsg::Send {
-                to_address: stream.treasury.to_string(),
+                to_address: stream_info.treasury.to_string(),
                 amount: vec![coin.clone()],
             })
         })
@@ -179,8 +197,9 @@ pub fn execute_stream_admin_cancel(
     info: MessageInfo,
 ) -> Result<Response, ContractError> {
     let mut stream = STREAM_STATE.load(deps.storage)?;
+    let stream_info = STREAM_INFO.load(deps.storage)?;
     // Only stream admin can cancel the stream with this method
-    if info.sender != stream.stream_admin {
+    if info.sender != stream_info.stream_admin {
         return Err(ContractError::Unauthorized {});
     }
 
@@ -198,14 +217,23 @@ pub fn execute_stream_admin_cancel(
 
     // Refund all out tokens to stream creator(treasury)
     let mut refund_coins = vec![stream.out_asset.clone()];
-    let pool_refund_coins = pool_refund(&deps, stream.pool_config, stream.out_asset.denom.clone())?;
 
-    refund_coins.extend(pool_refund_coins);
+    // refund pool creation if any
+    let post_stream_ops = POST_STREAM.may_load(deps.storage)?;
+    if let Some(post_stream_ops) = post_stream_ops {
+        let pool_refund_coins = pool_refund(
+            &deps,
+            post_stream_ops.pool_config,
+            stream.out_asset.denom.clone(),
+        )?;
+        refund_coins.extend(pool_refund_coins);
+    }
+
     let funds_msgs: Vec<CosmosMsg> = refund_coins
         .iter()
         .map(|coin| {
             CosmosMsg::Bank(BankMsg::Send {
-                to_address: stream.treasury.to_string(),
+                to_address: stream_info.treasury.to_string(),
                 amount: vec![coin.clone()],
             })
         })