refactor: extract fee calculation and distribution logic into a separate function

Author: armandocodecr

contracts/escrow/src/core/dispute.rs

@@ -4,6 +4,7 @@ use soroban_sdk::{Address, Env, Map, String};
 use crate::core::escrow::EscrowManager;
 use crate::core::validators::milestone::validate_and_convert_milestone_index;
 use crate::error::ContractError;
+use crate::modules::fee::distribution::calculate_and_distribute_fees;
 use crate::modules::{
     fee::{FeeCalculator, FeeCalculatorTrait},
     math::{BasicArithmetic, BasicMath},
@@ -59,29 +60,17 @@ impl DisputeManager {
 
         let fee_result =
             FeeCalculator::calculate_standard_fees(total, escrow.platform_fee)?;
-        let total_fees = BasicMath::safe_add(fee_result.trustless_work_fee, fee_result.platform_fee)?;
-
-        if fee_result.trustless_work_fee > 0 {
-            token_client.transfer(&contract_address, &trustless_work_address, &fee_result.trustless_work_fee);
-        }
-        if fee_result.platform_fee > 0 {
-            token_client.transfer(
-                &contract_address,
-                &escrow.roles.platform,
-                &fee_result.platform_fee,
-            );
-        }
-
-        for (addr, amount) in distributions.iter() {
-            if amount > 0 {
-                let fee_share =
-                    BasicMath::safe_div(BasicMath::safe_mul(amount, total_fees)?, total)?;
-                let net_amount = BasicMath::safe_sub(amount, fee_share)?;
-                if net_amount > 0 {
-                    token_client.transfer(&contract_address, &addr, &net_amount);
-                }
-            }
-        }
+            
+        calculate_and_distribute_fees(
+            e,
+            &token_client,
+            &contract_address,
+            &trustless_work_address,
+            &escrow.roles.platform,
+            &fee_result,
+            &distributions,
+            total,
+        )?;
 
         e.storage().persistent().set(&DataKey::Escrow, &escrow);
         e.storage()
@@ -131,28 +120,17 @@ impl DisputeManager {
 
         let fee_result =
             FeeCalculator::calculate_standard_fees(total, escrow.platform_fee)?;
-        let total_fees = BasicMath::safe_add(fee_result.trustless_work_fee, fee_result.platform_fee)?;
-
-        if fee_result.trustless_work_fee > 0 {
-            token_client.transfer(&contract_address, &trustless_work_address, &fee_result.trustless_work_fee);
-        }
-        if fee_result.platform_fee > 0 {
-            token_client.transfer(
-                &contract_address,
-                &escrow.roles.platform,
-                &fee_result.platform_fee,
-            );
-        }
-
-        for (addr, amount) in distributions.iter() {
-            if amount > 0 {
-                let fee_share = BasicMath::safe_div(BasicMath::safe_mul(amount, total_fees)?, total)?;
-                let net_amount = BasicMath::safe_sub(amount, fee_share)?;
-                if net_amount > 0 {
-                    token_client.transfer(&contract_address, &addr, &net_amount);
-                }
-            }
-        }
+        
+        calculate_and_distribute_fees(
+            e,
+            &token_client,
+            &contract_address,
+            &trustless_work_address,
+            &escrow.roles.platform,
+            &fee_result,
+            &distributions,
+            total,
+        )?;
 
         let mut updated_milestones = escrow.milestones.clone();
         let mut new_flags = milestone.flags.clone();

contracts/escrow/src/lib.rs

@@ -29,6 +29,7 @@ mod modules {
 
     pub mod fee {
         pub mod calculator;
+        pub mod distribution;
 
         pub use calculator::*;
     }

contracts/escrow/src/modules/fee/distribution.rs

@@ -0,0 +1,60 @@
+use soroban_sdk::token::Client as TokenClient;
+use soroban_sdk::{Address, Env, Map, Vec};
+
+use crate::error::ContractError;
+use super::StandardFeeResult;
+use crate::modules::math::{BasicArithmetic, BasicMath};
+
+pub fn calculate_and_distribute_fees(
+    e: &Env,
+    token_client: &TokenClient,
+    contract_address: &Address,
+    trustless_work_address: &Address,
+    platform_address: &Address,
+    fee_result: &StandardFeeResult,
+    distributions: &Map<Address, i128>,
+    total: i128,
+) -> Result<(), ContractError> {
+    let mut actual_trustless_fees = 0i128;
+    let mut actual_platform_fees = 0i128;
+    let mut net_distributions: Vec<(Address, i128)> = Vec::new(e);
+
+    for (addr, amount) in distributions.iter() {
+        if amount > 0 {
+            let recipient_trustless_fee = BasicMath::safe_div(
+                BasicMath::safe_mul(amount, fee_result.trustless_work_fee)?,
+                total,
+            )?;
+            let recipient_platform_fee = BasicMath::safe_div(
+                BasicMath::safe_mul(amount, fee_result.platform_fee)?,
+                total,
+            )?;
+
+            let total_recipient_fee =
+                BasicMath::safe_add(recipient_trustless_fee, recipient_platform_fee)?;
+            let net_amount = BasicMath::safe_sub(amount, total_recipient_fee)?;
+
+            actual_trustless_fees =
+                BasicMath::safe_add(actual_trustless_fees, recipient_trustless_fee)?;
+            actual_platform_fees =
+                BasicMath::safe_add(actual_platform_fees, recipient_platform_fee)?;
+
+            if net_amount > 0 {
+                net_distributions.push_back((addr.clone(), net_amount));
+            }
+        }
+    }
+
+    if actual_trustless_fees > 0 {
+        token_client.transfer(contract_address, trustless_work_address, &actual_trustless_fees);
+    }
+    if actual_platform_fees > 0 {
+        token_client.transfer(contract_address, platform_address, &actual_platform_fees);
+    }
+
+    for (addr, net_amount) in net_distributions.iter() {
+        token_client.transfer(contract_address, &addr, &net_amount);
+    }
+
+    Ok(())
+}