Make partial refund unit rounding partition-safe

contracts/chainmove-pool/src/lib.rs

@@ -24,6 +24,7 @@ pub enum ContractError {
     Overpayment = 11,
     RepayerMismatch = 12,
     NothingToRefund = 13,
+    RefundTooSmall = 14,
 }
 
 #[contracttype]
@@ -85,12 +86,21 @@ struct OperationReceipt {
     amount: i128,
 }
 
+#[contracttype]
+#[derive(Clone, Debug, Eq, PartialEq)]
+struct RefundBasis {
+    basis_invested: i128,
+    basis_units: u64,
+    cumulative_refunded: i128,
+}
+
 #[contracttype]
 #[derive(Clone)]
 enum DataKey {
     Pool(u64),
     InvestorPosition(u64, Address),
     Reference(String),
+    RefundBasis(u64, Address),
     LegacyPool(u64), // Legacy key format for migration testing
 }
 
@@ -235,6 +245,20 @@ impl ChainMovePoolContract {
         env.storage().persistent().set(&position_key, &position);
         env.storage().persistent().extend_ttl(&position_key, RENT_THRESHOLD, RENT_EXTEND_TO);
 
+        // Any new funding starts a new exact refund epoch from the combined
+        // principal/unit position. Subsequent partitioned refunds all resolve
+        // against this immutable basis instead of repeatedly rounding ratios.
+        let refund_basis_key = DataKey::RefundBasis(pool_id, investor.clone());
+        env.storage().persistent().set(
+            &refund_basis_key,
+            &RefundBasis {
+                basis_invested: position.invested,
+                basis_units: position.units,
+                cumulative_refunded: 0,
+            },
+        );
+        env.storage().persistent().extend_ttl(&refund_basis_key, RENT_THRESHOLD, RENT_EXTEND_TO);
+
         write_reference(
             &env,
             reference.clone(),
@@ -404,9 +428,20 @@ impl ChainMovePoolContract {
             return Err(ContractError::NothingToRefund);
         }
 
-        let refund_units = release_units(&position, amount)?;
+        let refund_basis_key = DataKey::RefundBasis(pool_id, investor.clone());
+        let mut refund_basis: RefundBasis = env
+            .storage()
+            .persistent()
+            .get(&refund_basis_key)
+            .unwrap_or(RefundBasis {
+                basis_invested: position.invested,
+                basis_units: position.units,
+                cumulative_refunded: 0,
+            });
+        let refund_units = release_units(&position, &refund_basis, amount)?;
         transfer_from_contract_to_participant(&env, &pool.asset, &investor, amount);
 
+        refund_basis.cumulative_refunded = checked_add_i128(refund_basis.cumulative_refunded, amount)?;
         position.refunded = checked_add_i128(position.refunded, amount)?;
         position.invested = checked_sub_i128(position.invested, amount)?;
         position.units = checked_sub_u64(position.units, refund_units)?;
@@ -417,6 +452,12 @@ impl ChainMovePoolContract {
         env.storage().persistent().extend_ttl(&pool_key, RENT_THRESHOLD, RENT_EXTEND_TO);
         env.storage().persistent().set(&position_key, &position);
         env.storage().persistent().extend_ttl(&position_key, RENT_THRESHOLD, RENT_EXTEND_TO);
+        if position.invested == 0 {
+            env.storage().persistent().remove(&refund_basis_key);
+        } else {
+            env.storage().persistent().set(&refund_basis_key, &refund_basis);
+            env.storage().persistent().extend_ttl(&refund_basis_key, RENT_THRESHOLD, RENT_EXTEND_TO);
+        }
 
         write_reference(
             &env,
@@ -640,20 +681,35 @@ fn allocate_units(pool: &Pool, amount: i128, new_total: i128) -> Result<u64, Con
     u64::try_from(unit_amount).map_err(|_| ContractError::ArithmeticOverflow)
 }
 
-fn release_units(position: &InvestorPosition, amount: i128) -> Result<u64, ContractError> {
+fn release_units(
+    position: &InvestorPosition,
+    basis: &RefundBasis,
+    amount: i128,
+) -> Result<u64, ContractError> {
     if amount == position.invested {
         return Ok(position.units);
     }
 
-    let unit_amount = amount
-        .checked_mul(position.units as i128)
-        .ok_or(ContractError::ArithmeticOverflow)?
-        .checked_div(position.invested)
+    let cumulative_refunded = checked_add_i128(basis.cumulative_refunded, amount)?;
+    let remaining_principal = checked_sub_i128(basis.basis_invested, cumulative_refunded)?;
+    let numerator = remaining_principal
+        .checked_mul(basis.basis_units as i128)
         .ok_or(ContractError::ArithmeticOverflow)?;
-    if unit_amount <= 0 {
-        return Err(ContractError::InvalidInput);
+    let quotient = numerator
+        .checked_div(basis.basis_invested)
+        .ok_or(ContractError::ArithmeticOverflow)?;
+    let remainder = numerator
+        .checked_rem(basis.basis_invested)
+        .ok_or(ContractError::ArithmeticOverflow)?;
+    let entitled_units = quotient + i128::from(remainder > 0);
+    let entitled_units = u64::try_from(entitled_units).map_err(|_| ContractError::ArithmeticOverflow)?;
+    let released = checked_sub_u64(position.units, entitled_units)?;
+    if released == 0 {
+        // The refund is below the current unit granularity. Reject it explicitly
+        // so principal cannot move while all corresponding units are retained.
+        return Err(ContractError::RefundTooSmall);
     }
-    u64::try_from(unit_amount).map_err(|_| ContractError::ArithmeticOverflow)
+    Ok(released)
 }
 
 fn transfer_from_participant_to_contract(

contracts/chainmove-pool/src/test.rs

@@ -387,6 +387,64 @@ fn refunds_return_custody_and_reduce_principal_units() {
     assert_eq!(token.balance(&fixture.investor), 8_500);
 }
 
+#[test]
+fn partitioned_refunds_match_an_equivalent_one_shot_refund() {
+    let partitioned = create_fixture();
+    fund(&partitioned, &partitioned.investor, 3_333, "fund-partitioned");
+    for (index, amount) in [1_111_i128, 1_111_i128].iter().enumerate() {
+        pool_client(&partitioned)
+            .try_refund_position(
+                &partitioned.owner,
+                &POOL_ID,
+                &partitioned.investor,
+                amount,
+                &String::from_str(&partitioned.env, if index == 0 { "refund-part-1" } else { "refund-part-2" }),
+            )
+            .unwrap()
+            .unwrap();
+    }
+    let partitioned_position = pool_client(&partitioned)
+        .try_read_investor_position(&partitioned.investor, &POOL_ID)
+        .unwrap()
+        .unwrap();
+
+    let one_shot = create_fixture();
+    fund(&one_shot, &one_shot.investor, 3_333, "fund-one-shot");
+    let one_shot_position = pool_client(&one_shot)
+        .try_refund_position(
+            &one_shot.owner,
+            &POOL_ID,
+            &one_shot.investor,
+            &2_222,
+            &String::from_str(&one_shot.env, "refund-one-shot"),
+        )
+        .unwrap()
+        .unwrap();
+
+    assert_eq!(partitioned_position.invested, one_shot_position.invested);
+    assert_eq!(partitioned_position.units, one_shot_position.units);
+}
+
+#[test]
+fn refund_below_unit_granularity_is_rejected_without_moving_principal() {
+    let fixture = create_fixture();
+    fund(&fixture, &fixture.investor, 1_000, "fund-dust-refund");
+    let result = pool_client(&fixture).try_refund_position(
+        &fixture.owner,
+        &POOL_ID,
+        &fixture.investor,
+        &1,
+        &String::from_str(&fixture.env, "refund-dust"),
+    );
+    assert_eq!(result.unwrap_err().unwrap(), ContractError::RefundTooSmall);
+    let position = pool_client(&fixture)
+        .try_read_investor_position(&fixture.investor, &POOL_ID)
+        .unwrap()
+        .unwrap();
+    assert_eq!(position.invested, 1_000);
+    assert_eq!(position.units, 10);
+}
+
 #[test]
 fn conservation_of_value_holds_across_funding_refund_and_repayment() {
     let fixture = create_fixture();