Skip to content

Commit 626ae37

Browse files
ritik4everritik4ever
authored
Merge pull request #1331 from Olasunkanmi975/feat/portfolio-onchain-rebalance-transfers
feat(contracts): replace internal rebalance bookkeeping with on-chain asset transfers
2 parents 165c1df + 97b9be7 commit 626ae37

3 files changed

Lines changed: 271 additions & 43 deletions

File tree

‎contracts/src/lib.rs‎

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -229,6 +229,9 @@ impl PortfolioRebalancer {
229229
.unwrap_or(portfolio.user.clone());
230230
steward.require_auth();
231231

232+
let token_client = TokenClient::new(&env, &asset);
233+
token_client.transfer(&steward, &env.current_contract_address(), &amount);
234+
232235
let current_balance = portfolio.current_balances.get(asset.clone()).unwrap_or(0);
233236
portfolio
234237
.current_balances
@@ -264,6 +267,9 @@ impl PortfolioRebalancer {
264267
return Err(Error::InsufficientBalance);
265268
}
266269

270+
let token_client = TokenClient::new(&env, &asset);
271+
token_client.transfer(&env.current_contract_address(), &portfolio.user, &amount);
272+
267273
let new_balance = current_balance - amount;
268274
if new_balance == 0 {
269275
portfolio.current_balances.remove(asset.clone());
@@ -1094,6 +1100,7 @@ impl PortfolioRebalancer {
10941100
} else {
10951101
0
10961102
};
1103+
let fee_recipient = fee_config.fee_recipient;
10971104

10981105
let mut has_actual_balances = false;
10991106
for (_, _) in actual_balances.iter() {
@@ -1155,13 +1162,29 @@ impl PortfolioRebalancer {
11551162
}
11561163
}
11571164

1165+
let contract_address = env.current_contract_address();
11581166
for (asset, amount) in trades.iter() {
1167+
let abs_amount = amount.abs();
11591168
let fee_amount = if effective_fee_bps > 0 {
1160-
(amount.abs() * effective_fee_bps as i128) / 10000
1169+
(abs_amount * effective_fee_bps as i128) / 10000
11611170
} else {
11621171
0
11631172
};
11641173
let effective_amount = amount - fee_amount;
1174+
1175+
let token_client = TokenClient::new(env, &asset);
1176+
if amount > 0 {
1177+
token_client.transfer(&steward, &contract_address, &abs_amount);
1178+
if fee_amount > 0 {
1179+
token_client.transfer(&contract_address, &fee_recipient, &fee_amount);
1180+
}
1181+
} else if amount < 0 {
1182+
token_client.transfer(&contract_address, &steward, &abs_amount);
1183+
if fee_amount > 0 {
1184+
token_client.transfer(&contract_address, &fee_recipient, &fee_amount);
1185+
}
1186+
}
1187+
11651188
let current = portfolio.current_balances.get(asset.clone()).unwrap_or(0);
11661189
portfolio
11671190
.current_balances

‎contracts/src/test.rs‎

Lines changed: 59 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,13 @@ fn allocation_decimals(
1818
asset_decimals
1919
}
2020

21+
fn create_token_and_mint(env: &Env, admin: &Address, to: &Address, amount: i128) -> Address {
22+
let token_id = env.register_stellar_asset_contract(admin.clone());
23+
let token = TokenClient::new(env, &token_id);
24+
token.mint(to, &amount);
25+
token_id
26+
}
27+
2128
fn create_portfolio_with_defaults(
2229
env: &Env,
2330
client: &PortfolioRebalancerClient,
@@ -228,15 +235,20 @@ fn test_deposit_valid() {
228235
let user = Address::generate(&env);
229236
client.initialize(&admin, &reflector_id);
230237

238+
let asset = create_token_and_mint(&env, &admin, &user, 2000);
239+
231240
let mut allocations = Map::new(&env);
232-
let asset = Address::generate(&env);
233241
allocations.set(asset.clone(), 10000);
234242
let pid = create_portfolio_with_defaults(&env, &client, &user, &allocations, 5, 50);
235243

236244
client.deposit(&pid, &asset, &1000, &String::from_str(&env, ""));
237245

238246
let portfolio = client.get_portfolio(&pid);
239-
assert_eq!(portfolio.current_balances.get(asset).unwrap(), 1000);
247+
assert_eq!(portfolio.current_balances.get(asset.clone()).unwrap(), 1000);
248+
249+
let token = TokenClient::new(&env, &asset);
250+
assert_eq!(token.balance(&contract_id), 1000);
251+
assert_eq!(token.balance(&user), 1000);
240252
}
241253

242254
#[test]
@@ -279,9 +291,10 @@ fn test_check_rebalance_needed_no_drift() {
279291
let user = Address::generate(&env);
280292
client.initialize(&admin, &reflector_id);
281293

294+
let asset1 = create_token_and_mint(&env, &admin, &user, 200);
295+
let asset2 = create_token_and_mint(&env, &admin, &user, 200);
296+
282297
let mut allocations = Map::new(&env);
283-
let asset1 = Address::generate(&env);
284-
let asset2 = Address::generate(&env);
285298
allocations.set(asset1.clone(), 5000);
286299
allocations.set(asset2.clone(), 5000);
287300

@@ -309,9 +322,10 @@ fn test_check_rebalance_needed_with_drift() {
309322
let user = Address::generate(&env);
310323
client.initialize(&admin, &reflector_id);
311324

325+
let asset1 = create_token_and_mint(&env, &admin, &user, 300);
326+
let asset2 = create_token_and_mint(&env, &admin, &user, 200);
327+
312328
let mut allocations = Map::new(&env);
313-
let asset1 = Address::generate(&env);
314-
let asset2 = Address::generate(&env);
315329
allocations.set(asset1.clone(), 5000);
316330
allocations.set(asset2.clone(), 5000);
317331

@@ -469,9 +483,10 @@ fn test_rebalance_applies_non_zero_fee_to_trade_amount() {
469483
let user = Address::generate(&env);
470484
client.initialize(&admin, &reflector_id);
471485

486+
let asset1 = create_token_and_mint(&env, &admin, &user, 10_000_000);
487+
let asset2 = create_token_and_mint(&env, &admin, &user, 5_000_000);
488+
472489
let mut allocations = Map::new(&env);
473-
let asset1 = Address::generate(&env);
474-
let asset2 = Address::generate(&env);
475490
allocations.set(asset1.clone(), 5000);
476491
allocations.set(asset2.clone(), 5000);
477492

@@ -482,7 +497,7 @@ fn test_rebalance_applies_non_zero_fee_to_trade_amount() {
482497
let config = FeeConfig {
483498
platform_name: String::from_str(&env, "Acme Vault"),
484499
fee_bps: 50,
485-
fee_recipient: recipient,
500+
fee_recipient: recipient.clone(),
486501
enabled: true,
487502
};
488503
client.set_fee_config(&config);
@@ -494,8 +509,15 @@ fn test_rebalance_applies_non_zero_fee_to_trade_amount() {
494509
client.execute_rebalance(&pid, &Map::new(&env));
495510

496511
let portfolio = client.get_portfolio(&pid);
497-
assert_eq!(portfolio.current_balances.get(asset1).unwrap(), 4_975_000);
498-
assert_eq!(portfolio.current_balances.get(asset2).unwrap(), 4_975_000);
512+
assert_eq!(portfolio.current_balances.get(asset1.clone()).unwrap(), 4_975_000);
513+
assert_eq!(portfolio.current_balances.get(asset2.clone()).unwrap(), 4_975_000);
514+
515+
let token1 = TokenClient::new(&env, &asset1);
516+
let token2 = TokenClient::new(&env, &asset2);
517+
assert_eq!(token1.balance(&contract_id), 4_975_000);
518+
assert_eq!(token2.balance(&contract_id), 4_975_000);
519+
assert_eq!(token1.balance(&recipient), 25_000);
520+
assert_eq!(token2.balance(&recipient), 25_000);
499521
}
500522

501523
#[test]
@@ -540,8 +562,9 @@ fn test_emergency_stop() {
540562
let user = Address::generate(&env);
541563
client.initialize(&admin, &reflector_id);
542564

565+
let asset = create_token_and_mint(&env, &admin, &user, 200);
566+
543567
let mut allocations = Map::new(&env);
544-
let asset = Address::generate(&env);
545568
allocations.set(asset.clone(), 10000);
546569
let pid = create_portfolio_with_defaults(&env, &client, &user, &allocations, 5, 50);
547570

@@ -1453,7 +1476,7 @@ fn test_emergency_stop_admin_pause_and_reactivate() {
14531476
client.initialize(&admin, &reflector_id);
14541477

14551478
let mut allocations = Map::new(&env);
1456-
let asset = Address::generate(&env);
1479+
let asset = create_token_and_mint(&env, &admin, &user, 100);
14571480
allocations.set(asset.clone(), 10000);
14581481
let pid = create_portfolio_with_defaults(&env, &client, &user, &allocations, 5, 50);
14591482

@@ -1618,12 +1641,12 @@ fn test_transfer_stewardship_steward_can_deposit() {
16181641
let user = Address::generate(&env);
16191642
client.initialize(&admin, &reflector_id);
16201643

1644+
let new_steward = Address::generate(&env);
16211645
let mut allocations = Map::new(&env);
1622-
let asset = Address::generate(&env);
1646+
let asset = create_token_and_mint(&env, &admin, &new_steward, 500);
16231647
allocations.set(asset.clone(), 10000);
16241648
let pid = create_portfolio_with_defaults(&env, &client, &user, &allocations, 5, 50);
16251649

1626-
let new_steward = Address::generate(&env);
16271650
client.transfer_stewardship(&pid, &new_steward);
16281651

16291652
client
@@ -1655,8 +1678,8 @@ fn test_preview_rebalance_reports_trades_and_thresholds() {
16551678
client.initialize(&admin, &reflector_id);
16561679

16571680
let mut allocations = Map::new(&env);
1658-
let asset1 = Address::generate(&env);
1659-
let asset2 = Address::generate(&env);
1681+
let asset1 = create_token_and_mint(&env, &admin, &user, 20_000_000);
1682+
let asset2 = create_token_and_mint(&env, &admin, &user, 10_000_000);
16601683
allocations.set(asset1.clone(), 5000);
16611684
allocations.set(asset2.clone(), 5000);
16621685
let pid = create_portfolio_with_defaults(&env, &client, &user, &allocations, 5, 50);
@@ -1685,7 +1708,7 @@ fn test_preview_rebalance_does_not_mutate_portfolio() {
16851708
client.initialize(&admin, &reflector_id);
16861709

16871710
let mut allocations = Map::new(&env);
1688-
let asset = Address::generate(&env);
1711+
let asset = create_token_and_mint(&env, &admin, &user, 1000);
16891712
allocations.set(asset.clone(), 10000);
16901713
let pid = create_portfolio_with_defaults(&env, &client, &user, &allocations, 5, 50);
16911714

@@ -1917,7 +1940,7 @@ fn test_deposit_rejects_paused_portfolio() {
19171940
client.initialize(&admin, &reflector_id);
19181941

19191942
let mut allocations = Map::new(&env);
1920-
let asset = Address::generate(&env);
1943+
let asset = create_token_and_mint(&env, &admin, &user, 100);
19211944
allocations.set(asset.clone(), 10000);
19221945
let pid = create_portfolio_with_defaults(&env, &client, &user, &allocations, 5, 50);
19231946
client.pause_portfolio(&pid, &PauseReason::VolatilityCircuitBreaker);
@@ -1981,7 +2004,7 @@ fn test_check_invariants_inactive_portfolio() {
19812004
client.initialize(&admin, &reflector_id);
19822005

19832006
let mut allocations = Map::new(&env);
1984-
let asset = Address::generate(&env);
2007+
let asset = create_token_and_mint(&env, &admin, &user, 100);
19852008
allocations.set(asset.clone(), 10000);
19862009
let pid = create_portfolio_with_defaults(&env, &client, &user, &allocations, 5, 50);
19872010
client.deposit(&pid, &asset, &100, &String::from_str(&env, ""));
@@ -2003,7 +2026,7 @@ fn test_withdraw_success() {
20032026
client.initialize(&admin, &reflector_id);
20042027

20052028
let mut allocations = Map::new(&env);
2006-
let asset = Address::generate(&env);
2029+
let asset = create_token_and_mint(&env, &admin, &user, 1000);
20072030
allocations.set(asset.clone(), 10000);
20082031
let pid = create_portfolio_with_defaults(&env, &client, &user, &allocations, 5, 50);
20092032
client.deposit(&pid, &asset, &1000, &String::from_str(&env, ""));
@@ -2026,7 +2049,7 @@ fn test_withdraw_insufficient_balance() {
20262049
client.initialize(&admin, &reflector_id);
20272050

20282051
let mut allocations = Map::new(&env);
2029-
let asset = Address::generate(&env);
2052+
let asset = create_token_and_mint(&env, &admin, &user, 100);
20302053
allocations.set(asset.clone(), 10000);
20312054
let pid = create_portfolio_with_defaults(&env, &client, &user, &allocations, 5, 50);
20322055
client.deposit(&pid, &asset, &100, &String::from_str(&env, ""));
@@ -2047,7 +2070,7 @@ fn test_withdraw_full_exit_deactivates_portfolio() {
20472070
client.initialize(&admin, &reflector_id);
20482071

20492072
let mut allocations = Map::new(&env);
2050-
let asset = Address::generate(&env);
2073+
let asset = create_token_and_mint(&env, &admin, &user, 100);
20512074
allocations.set(asset.clone(), 10000);
20522075
let pid = create_portfolio_with_defaults(&env, &client, &user, &allocations, 5, 50);
20532076
client.deposit(&pid, &asset, &100, &String::from_str(&env, ""));
@@ -2075,7 +2098,7 @@ fn test_admin_force_rebalance_bypasses_cooldown() {
20752098
client.initialize(&admin, &reflector_id);
20762099

20772100
let mut allocations = Map::new(&env);
2078-
let asset = Address::generate(&env);
2101+
let asset = create_token_and_mint(&env, &admin, &user, 100);
20792102
allocations.set(asset.clone(), 10000);
20802103
let pid = create_portfolio_with_defaults(&env, &client, &user, &allocations, 5, 50);
20812104
client.deposit(&pid, &asset, &100, &String::from_str(&env, ""));
@@ -2311,7 +2334,7 @@ fn benchmark_deposit_gas() {
23112334
let _ = client.initialize(&admin, &reflector_id);
23122335

23132336
let mut allocations = Map::new(&env);
2314-
let asset = Address::generate(&env);
2337+
let asset = create_token_and_mint(&env, &admin, &user, 100);
23152338
allocations.set(asset.clone(), 10000);
23162339
let pid = create_portfolio_with_defaults(&env, &client, &user, &allocations, 5, 50);
23172340

@@ -2453,8 +2476,8 @@ fn test_get_portfolio_value_usd_basic() {
24532476
client.initialize(&admin, &reflector_id);
24542477

24552478
let mut allocations = Map::new(&env);
2456-
let asset1 = Address::generate(&env);
2457-
let asset2 = Address::generate(&env);
2479+
let asset1 = create_token_and_mint(&env, &admin, &user, 100);
2480+
let asset2 = create_token_and_mint(&env, &admin, &user, 100);
24582481
allocations.set(asset1.clone(), 5000);
24592482
allocations.set(asset2.clone(), 5000);
24602483
let pid = create_portfolio_with_defaults(&env, &client, &user, &allocations, 5, 50);
@@ -2480,8 +2503,8 @@ fn test_get_portfolio_value_usd_drift() {
24802503
client.initialize(&admin, &reflector_id);
24812504

24822505
let mut allocations = Map::new(&env);
2483-
let asset1 = Address::generate(&env);
2484-
let asset2 = Address::generate(&env);
2506+
let asset1 = create_token_and_mint(&env, &admin, &user, 200);
2507+
let asset2 = create_token_and_mint(&env, &admin, &user, 100);
24852508
allocations.set(asset1.clone(), 5000);
24862509
allocations.set(asset2.clone(), 5000);
24872510
let pid = create_portfolio_with_defaults(&env, &client, &user, &allocations, 5, 50);
@@ -2634,8 +2657,8 @@ fn test_get_drift_preview_balanced_no_needs_rebalance() {
26342657
client.initialize(&admin, &reflector_id);
26352658

26362659
let mut allocations = Map::new(&env);
2637-
let asset1 = Address::generate(&env);
2638-
let asset2 = Address::generate(&env);
2660+
let asset1 = create_token_and_mint(&env, &admin, &user, 100);
2661+
let asset2 = create_token_and_mint(&env, &admin, &user, 100);
26392662
allocations.set(asset1.clone(), 5000u32);
26402663
allocations.set(asset2.clone(), 5000u32);
26412664

@@ -2672,8 +2695,8 @@ fn test_get_drift_preview_imbalanced_needs_rebalance() {
26722695
client.initialize(&admin, &reflector_id);
26732696

26742697
let mut allocations = Map::new(&env);
2675-
let asset1 = Address::generate(&env);
2676-
let asset2 = Address::generate(&env);
2698+
let asset1 = create_token_and_mint(&env, &admin, &user, 900);
2699+
let asset2 = create_token_and_mint(&env, &admin, &user, 100);
26772700
allocations.set(asset1.clone(), 5000u32);
26782701
allocations.set(asset2.clone(), 5000u32);
26792702

@@ -2733,7 +2756,7 @@ fn test_manual_nav_snapshot() {
27332756
client.initialize(&admin, &reflector_id);
27342757

27352758
let mut allocations = Map::new(&env);
2736-
let asset = Address::generate(&env);
2759+
let asset = create_token_and_mint(&env, &admin, &user, 100_0000000);
27372760
allocations.set(asset.clone(), 10000);
27382761
let pid = create_portfolio_with_defaults(&env, &client, &user, &allocations, 5, 50);
27392762

@@ -2811,7 +2834,7 @@ fn test_auto_nav_snapshot_on_rebalance() {
28112834
client.initialize(&admin, &reflector_id);
28122835

28132836
let mut allocations = Map::new(&env);
2814-
let asset = Address::generate(&env);
2837+
let asset = create_token_and_mint(&env, &admin, &user, 100_0000000);
28152838
allocations.set(asset.clone(), 10000);
28162839
let pid = create_portfolio_with_defaults(&env, &client, &user, &allocations, 5, 50);
28172840

@@ -3497,7 +3520,7 @@ fn test_benchmark_nav_operations() {
34973520
client.initialize(&admin, &reflector_id);
34983521

34993522
let mut allocations = Map::new(&env);
3500-
let asset = Address::generate(&env);
3523+
let asset = create_token_and_mint(&env, &admin, &user, 100_0000000);
35013524
allocations.set(asset.clone(), 10000);
35023525
let pid = create_portfolio_with_defaults(&env, &client, &user, &allocations, 5, 50);
35033526
client.deposit(&pid, &asset, &100_0000000, &String::from_str(&env, "init"));

0 commit comments

Comments
 (0)