From 95926a60d40db2e3d3f01718315ac03dcd4e00e5 Mon Sep 17 00:00:00 2001 From: D'Angelo Rodriguez <70290504+dangelo352@users.noreply.github.com> Date: Sun, 21 Jun 2026 22:21:14 -0400 Subject: [PATCH] feat: weight allocations by apy and headroom --- contracts/allocation_logic/src/benchmarks.rs | 43 ++++---- contracts/allocation_logic/src/lib.rs | 103 ++++++++++++++++-- contracts/allocation_logic/src/tests.rs | 102 +++++++++++++++++ ...ION_LOGIC_POOL_REGISTRY_AND_RISK_LEVELS.md | 24 +++- 4 files changed, 238 insertions(+), 34 deletions(-) diff --git a/contracts/allocation_logic/src/benchmarks.rs b/contracts/allocation_logic/src/benchmarks.rs index 3ade348b..b293df14 100644 --- a/contracts/allocation_logic/src/benchmarks.rs +++ b/contracts/allocation_logic/src/benchmarks.rs @@ -98,18 +98,20 @@ fn benchmark_allocate() { let e = Env::default(); let (contract_id, admin) = setup_test_env(&e); - // Register a pool first - e.as_contract(&contract_id, || { - AllocationStrategiesContract::register_pool( - e.clone(), - admin.clone(), - 1, - RiskLevel::Low, - 500, - 10000_0000000, - ) - .unwrap(); - }); + // Register weighted low-risk pools for allocation cost coverage. + for (pool_id, apy, capacity) in [(1, 500, 10000_0000000), (2, 1200, 5000_0000000)] { + e.as_contract(&contract_id, || { + AllocationStrategiesContract::register_pool( + e.clone(), + admin.clone(), + pool_id, + RiskLevel::Low, + apy, + capacity, + ) + .unwrap(); + }); + } let caller = Address::generate(&e); let mut metrics = BenchmarkMetrics::new("allocate"); @@ -119,7 +121,7 @@ fn benchmark_allocate() { let _ = AllocationStrategiesContract::allocate( e.clone(), caller.clone(), - 1, + String::from_str(&e, "bench_allocate"), 1000_0000000, Strategy::Safe, ); @@ -147,21 +149,23 @@ fn benchmark_get_allocation() { 10000_0000000, ) .unwrap(); - AllocationStrategiesContract::allocate( + let _ = AllocationStrategiesContract::allocate( e.clone(), caller.clone(), - 1, + String::from_str(&e, "bench_get_allocation"), 1000_0000000, Strategy::Safe, - ) - .unwrap(); + ); }); let mut metrics = BenchmarkMetrics::new("get_allocation"); e.as_contract(&contract_id, || { let start = e.ledger().sequence(); - AllocationStrategiesContract::get_allocation(e.clone(), 1); + AllocationStrategiesContract::get_allocation( + e.clone(), + String::from_str(&e, "bench_get_allocation"), + ); let end = e.ledger().sequence(); metrics.record_gas(start, end); }); @@ -228,11 +232,12 @@ fn benchmark_batch_allocate() { let _ = AllocationStrategiesContract::allocate( e.clone(), caller.clone(), - i, + String::from_str(&e, "bench_batch_allocate"), 1000_0000000, Strategy::Safe, ); }); + let _ = i; } let end = e.ledger().sequence(); metrics.record_gas(start, end); diff --git a/contracts/allocation_logic/src/lib.rs b/contracts/allocation_logic/src/lib.rs index 32378b49..6710fb56 100644 --- a/contracts/allocation_logic/src/lib.rs +++ b/contracts/allocation_logic/src/lib.rs @@ -1140,22 +1140,65 @@ impl AllocationStrategiesContract { return Err(Error::PoolCapacityExceeded); } - // Base allocation with deterministic remainder. - let base = amount / pool_count as i128; - let remainder = amount % pool_count as i128; - let mut allocations = Vec::::new(env); let mut leftover_capacity = Vec::::new(env); + let mut weights = Vec::::new(env); + let mut remainders = Vec::::new(env); + let mut total_weight: i128 = 0; + + for pool in pools.iter() { + let available_capacity = pool + .max_capacity + .checked_sub(pool.total_liquidity) + .ok_or(Error::ArithmeticOverflow)?; + + let apy_weight = (pool.apy as i128) + .checked_mul(available_capacity) + .ok_or(Error::ArithmeticOverflow)?; + weights.push_back(apy_weight); + total_weight = total_weight + .checked_add(apy_weight) + .ok_or(Error::ArithmeticOverflow)?; + } + + // If every eligible pool has zero APY, fall back to headroom-only weights + // so capacity can still be used deterministically. + if total_weight == 0 { + weights = Vec::::new(env); + for pool in pools.iter() { + let available_capacity = pool + .max_capacity + .checked_sub(pool.total_liquidity) + .ok_or(Error::ArithmeticOverflow)?; + weights.push_back(available_capacity); + total_weight = total_weight + .checked_add(available_capacity) + .ok_or(Error::ArithmeticOverflow)?; + } + } + + if total_weight <= 0 { + return Err(Error::PoolCapacityExceeded); + } let mut allocated_total: i128 = 0; - for (i, pool) in pools.iter().enumerate() { + for i in 0..pool_count { + let pool = pools.get(i).unwrap(); let available_capacity = pool .max_capacity .checked_sub(pool.total_liquidity) .ok_or(Error::ArithmeticOverflow)?; + let weight = weights.get(i).unwrap(); + let weighted_amount = amount + .checked_mul(weight) + .ok_or(Error::ArithmeticOverflow)?; - // target_i sums to `amount` across all pools (before capping). - let target = base + if (i as i128) < remainder { 1 } else { 0 }; + let target = weighted_amount + .checked_div(total_weight) + .ok_or(Error::ArithmeticOverflow)?; + let remainder = weighted_amount + .checked_rem(total_weight) + .ok_or(Error::ArithmeticOverflow)?; let alloc = if target > available_capacity { available_capacity @@ -1169,19 +1212,58 @@ impl AllocationStrategiesContract { allocations.push_back(alloc); leftover_capacity.push_back(available_capacity - alloc); + remainders.push_back(remainder); } - // Redistribute any shortfall caused by per-pool caps, filling pools in deterministic registry order. + // Redistribute integer-division dust to the largest fractional remainders first. let mut shortfall = amount .checked_sub(allocated_total) .ok_or(Error::ArithmeticOverflow)?; + if shortfall > 0 { + for _ in 0..pool_count { + if shortfall == 0 { + break; + } + + let mut best_index: Option = None; + let mut best_remainder: i128 = -1; + for candidate in 0..pool_count { + let left = leftover_capacity.get(candidate).unwrap(); + let remainder = remainders.get(candidate).unwrap(); + if left > 0 && remainder > best_remainder { + best_index = Some(candidate); + best_remainder = remainder; + } + } + + let Some(best_index) = best_index else { + break; + }; + if best_remainder <= 0 { + break; + } + + let prev = allocations.get(best_index).unwrap(); + let left = leftover_capacity.get(best_index).unwrap(); + let new_alloc = prev.checked_add(1).ok_or(Error::ArithmeticOverflow)?; + let new_left = left.checked_sub(1).ok_or(Error::ArithmeticOverflow)?; + allocations.set(best_index, new_alloc); + leftover_capacity.set(best_index, new_left); + remainders.set(best_index, -1); + shortfall = shortfall + .checked_sub(1) + .ok_or(Error::ArithmeticOverflow)?; + } + } + + // Any larger shortfall is caused by capped pools; fill remaining capacity + // in deterministic registry order after the weighted targets are exhausted. if shortfall > 0 { for i in 0..pool_count { if shortfall == 0 { break; } - let pool = pools.get(i).unwrap(); let prev = allocations.get(i).unwrap(); let left = leftover_capacity.get(i).unwrap(); @@ -1194,9 +1276,6 @@ impl AllocationStrategiesContract { shortfall = shortfall .checked_sub(extra) .ok_or(Error::ArithmeticOverflow)?; - - // We don't write to the allocation_map yet; we finalize below. - let _ = pool; } } diff --git a/contracts/allocation_logic/src/tests.rs b/contracts/allocation_logic/src/tests.rs index 515b211b..032fd6fe 100644 --- a/contracts/allocation_logic/src/tests.rs +++ b/contracts/allocation_logic/src/tests.rs @@ -90,6 +90,108 @@ fn setup_test_pools(_env: &Env, client: &AllocationStrategiesContractClient, adm client.register_pool(admin, &5, &RiskLevel::High, &2500, &500_000_000); } +fn allocation_amount(summary: &crate::AllocationSummary, pool_id: u32) -> i128 { + for allocation in summary.allocations.iter() { + if allocation.pool_id == pool_id { + return allocation.amount; + } + } + 0 +} + +#[test] +fn test_allocate_weights_safe_pools_by_apy_and_headroom() { + let env = Env::default(); + env.mock_all_auths(); + let (admin, core_id, client) = create_contract(&env); + let user = Address::generate(&env); + + client.register_pool(&admin, &0, &RiskLevel::Low, &100, &1_000); + client.register_pool(&admin, &1, &RiskLevel::Low, &300, &1_000); + create_mock_commitment(&env, &core_id, "weighted_safe", 400, "active"); + + let summary = client.allocate( + &user, + &String::from_str(&env, "weighted_safe"), + &400, + &Strategy::Safe, + ); + + assert_eq!(summary.total_allocated, 400); + assert_eq!(allocation_amount(&summary, 0), 100); + assert_eq!(allocation_amount(&summary, 1), 300); +} + +#[test] +fn test_allocate_weighted_distribution_respects_headroom() { + let env = Env::default(); + env.mock_all_auths(); + let (admin, core_id, client) = create_contract(&env); + let user = Address::generate(&env); + + client.register_pool(&admin, &0, &RiskLevel::Low, &1000, &100); + client.register_pool(&admin, &1, &RiskLevel::Low, &100, &1_000); + create_mock_commitment(&env, &core_id, "weighted_headroom", 200, "active"); + + let summary = client.allocate( + &user, + &String::from_str(&env, "weighted_headroom"), + &200, + &Strategy::Safe, + ); + + assert_eq!(summary.total_allocated, 200); + assert_eq!(allocation_amount(&summary, 0), 100); + assert_eq!(allocation_amount(&summary, 1), 100); + assert_eq!(client.get_pool(&0).total_liquidity, 100); +} + +#[test] +fn test_allocate_zero_apy_falls_back_to_headroom_weights() { + let env = Env::default(); + env.mock_all_auths(); + let (admin, core_id, client) = create_contract(&env); + let user = Address::generate(&env); + + client.register_pool(&admin, &0, &RiskLevel::Low, &0, &100); + client.register_pool(&admin, &1, &RiskLevel::Low, &0, &300); + create_mock_commitment(&env, &core_id, "zero_apy", 200, "active"); + + let summary = client.allocate( + &user, + &String::from_str(&env, "zero_apy"), + &200, + &Strategy::Safe, + ); + + assert_eq!(summary.total_allocated, 200); + assert_eq!(allocation_amount(&summary, 0), 50); + assert_eq!(allocation_amount(&summary, 1), 150); +} + +#[test] +fn test_allocate_weighted_dust_goes_to_largest_remainder() { + let env = Env::default(); + env.mock_all_auths(); + let (admin, core_id, client) = create_contract(&env); + let user = Address::generate(&env); + + client.register_pool(&admin, &0, &RiskLevel::Low, &100, &1_000); + client.register_pool(&admin, &1, &RiskLevel::Low, &200, &1_000); + create_mock_commitment(&env, &core_id, "weighted_dust", 10, "active"); + + let summary = client.allocate( + &user, + &String::from_str(&env, "weighted_dust"), + &10, + &Strategy::Safe, + ); + + assert_eq!(summary.total_allocated, 10); + assert_eq!(allocation_amount(&summary, 0), 3); + assert_eq!(allocation_amount(&summary, 1), 7); +} + #[test] fn test_migrate_rejects_wrong_from_version_without_mutating_state() { let env = Env::default(); diff --git a/docs/ALLOCATION_LOGIC_POOL_REGISTRY_AND_RISK_LEVELS.md b/docs/ALLOCATION_LOGIC_POOL_REGISTRY_AND_RISK_LEVELS.md index 7156dc0f..899d3768 100644 --- a/docs/ALLOCATION_LOGIC_POOL_REGISTRY_AND_RISK_LEVELS.md +++ b/docs/ALLOCATION_LOGIC_POOL_REGISTRY_AND_RISK_LEVELS.md @@ -68,13 +68,27 @@ Validation and invariants: - `Balanced`: `Low` + `Medium` + `High` - `Aggressive`: `Medium` + `High` -## Allocation Rounding, Determinism, and Capacity Failure +## Weighted Allocation, Rounding, and Capacity Failure All allocations use integer arithmetic and are deterministic given the same on-chain state. -### Deterministic Remainder Handling +### APY and Headroom Weighting + +Within each eligible risk bucket, allocation is weighted by both yield and available capacity: + +`pool_weight = pool.apy * (pool.max_capacity - pool.total_liquidity)` + +Each pool receives: + +`floor(bucket_amount * pool_weight / total_bucket_weight)` + +If all pools in a bucket have `apy == 0`, the contract falls back to headroom-only weights so otherwise usable capacity is not stranded: -When an amount is split across pools (or across risk-level buckets and then pools), integer division can produce a remainder. The contract assigns remainder units deterministically in the same order pools are iterated (registry order after filtering for strategy and `active`). +`pool_weight = pool.max_capacity - pool.total_liquidity` + +Checked arithmetic is used for the weight and allocation products. If the eligible pools cannot fully satisfy the requested amount, allocation fails with `Error::PoolCapacityExceeded`. + +### Strategy Risk Buckets Balanced strategy risk split: @@ -89,6 +103,10 @@ Aggressive strategy risk split: Within each bucket, the bucket amount is distributed across pools with the same deterministic remainder behavior. +### Deterministic Remainder Handling + +When weighted integer division creates dust, remainder units are assigned to pools with the largest fractional remainder first, with registry order as the deterministic tie-breaker. If pool caps still leave a shortfall after weighted targets are exhausted, remaining capacity is filled in registry order. + ### Capacity Enforcement and Failure Mode Each pool has an enforced capacity: