feat: weight allocations by apy and headroom

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);

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::<i128>::new(env);
         let mut leftover_capacity = Vec::<i128>::new(env);
+        let mut weights = Vec::<i128>::new(env);
+        let mut remainders = Vec::<i128>::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::<i128>::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<u32> = 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;
             }
         }
 

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();

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: