Skip to content

Commit

Permalink
tests: treasury
Browse files Browse the repository at this point in the history
  • Loading branch information
thounyy committed Nov 1, 2024
1 parent 12ba504 commit 56197cb
Show file tree
Hide file tree
Showing 2 changed files with 808 additions and 52 deletions.
108 changes: 56 additions & 52 deletions packages/actions/sources/treasury.move
Original file line number Diff line number Diff line change
Expand Up @@ -34,19 +34,23 @@ use account_actions::{
#[error]
const ETreasuryDoesntExist: vector<u8> = b"No Treasury with this name";
#[error]
const ETreasuryAlreadyExists: vector<u8> = b"A treasury already exists with this name";
const EAlreadyExists: vector<u8> = b"A treasury already exists with this name";
#[error]
const ETypesAmountsNotSameLength: vector<u8> = b"Types and amounts vectors not same length";
const ENotSameLength: vector<u8> = b"Recipients and amounts vectors not same length";
#[error]
const ETreasuryNotEmpty: vector<u8> = b"Treasury must be emptied before closing";
const ENotEmpty: vector<u8> = b"Treasury must be emptied before closing";
#[error]
const EInsufficientFunds: vector<u8> = b"Insufficient funds for this coin type in treasury";
#[error]
const ECoinTypeDoesntExist: vector<u8> = b"Coin type doesn't exist in treasury";

// === Structs ===

/// Dynamic Field key for the Treasury
public struct TreasuryKey has copy, drop, store { name: String }
/// Dynamic field holding a budget with different coin types, key is name
public struct Treasury has store {
// heterogeneous array of Balances, String -> Balance<C>
// heterogeneous array of Balances, TypeName -> Balance<CoinType>
bag: Bag
}

Expand All @@ -58,16 +62,14 @@ public struct TransferProposal() has copy, drop;
public struct PayProposal() has copy, drop;

/// [ACTION] action to be used with specific proposals making good use of the returned coins, similar to owned::withdraw
public struct SpendAction has store {
// coin types to withdraw
coin_type: String,
public struct SpendAction<phantom CoinType> has store {
// amount to withdraw
amount: u64,
}

// === View Functions ===

public fun treasury_exists<Config, Outcome>(
public fun has_treasury<Config, Outcome>(
account: &Account<Config, Outcome>,
name: String
): bool {
Expand All @@ -78,19 +80,16 @@ public fun borrow_treasury<Config, Outcome>(
account: &Account<Config, Outcome>,
name: String
): &Treasury {
assert!(has_treasury(account, name), ETreasuryDoesntExist);
account.borrow_managed_asset(TreasuryKey { name }, version::current())
}

public fun coin_type_string<C: drop>(): String {
type_name::get<C>().into_string().to_string()
public fun coin_type_exists<CoinType: drop>(treasury: &Treasury): bool {
treasury.bag.contains(type_name::get<CoinType>())
}

public fun coin_type_exists(treasury: &Treasury, coin_type: String): bool {
treasury.bag.contains(coin_type)
}

public fun coin_type_value<C: drop>(treasury: &Treasury, coin_type: String): u64 {
treasury.bag.borrow<String, Balance<C>>(coin_type).value()
public fun coin_type_value<CoinType: drop>(treasury: &Treasury): u64 {
treasury.bag.borrow<TypeName, Balance<CoinType>>(type_name::get<CoinType>()).value()
}

// === [MEMBER] Public Functions ===
Expand All @@ -103,41 +102,40 @@ public fun open<Config, Outcome>(
ctx: &mut TxContext
) {
auth.verify(account.addr());
assert!(!treasury_exists(account, name), ETreasuryAlreadyExists);
assert!(!has_treasury(account, name), EAlreadyExists);

account.add_managed_asset(TreasuryKey { name }, Treasury { bag: bag::new(ctx) }, version::current());
}

/// Deposits coins owned by the account into a treasury
public fun deposit_owned<Config, Outcome, C: drop>(
public fun deposit_owned<Config, Outcome, CoinType: drop>(
auth: Auth,
account: &mut Account<Config, Outcome>,
name: String,
receiving: Receiving<Coin<C>>,
receiving: Receiving<Coin<CoinType>>,
) {
let coin = account.receive(receiving, version::current());
deposit<Config, Outcome, C>(auth, account, name, coin);
deposit<Config, Outcome, CoinType>(auth, account, name, coin);
}

/// Deposits coins owned by a member into a treasury
public fun deposit<Config, Outcome, C: drop>(
public fun deposit<Config, Outcome, CoinType: drop>(
auth: Auth,
account: &mut Account<Config, Outcome>,
name: String,
coin: Coin<C>,
coin: Coin<CoinType>,
) {
auth.verify_with_role<Deposit>(account.addr(), name);
assert!(treasury_exists(account, name), ETreasuryDoesntExist);
assert!(has_treasury(account, name), ETreasuryDoesntExist);

let treasury: &mut Treasury =
account.borrow_managed_asset_mut(TreasuryKey { name }, version::current());
let coin_type = coin_type_string<C>();

if (treasury.coin_type_exists(coin_type)) {
let balance: &mut Balance<C> = treasury.bag.borrow_mut(coin_type);
if (treasury.coin_type_exists<CoinType>()) {
let balance: &mut Balance<CoinType> = treasury.bag.borrow_mut(type_name::get<CoinType>());
balance.join(coin.into_balance());
} else {
treasury.bag.add(coin_type, coin.into_balance());
treasury.bag.add(type_name::get<CoinType>(), coin.into_balance());
};
}

Expand All @@ -151,14 +149,14 @@ public fun close<Config, Outcome>(

let Treasury { bag } =
account.remove_managed_asset(TreasuryKey { name }, version::current());
assert!(bag.is_empty(), ETreasuryNotEmpty);
assert!(bag.is_empty(), ENotEmpty);
bag.destroy_empty();
}

// === [PROPOSAL] Public Functions ===

// step 1: propose to send managed coins
public fun propose_transfer<Config, Outcome>(
public fun propose_transfer<Config, Outcome, CoinType: drop>(
auth: Auth,
account: &mut Account<Config, Outcome>,
outcome: Outcome,
Expand All @@ -167,12 +165,15 @@ public fun propose_transfer<Config, Outcome>(
execution_time: u64,
expiration_epoch: u64,
treasury_name: String,
coin_types: vector<String>,
amounts: vector<u64>,
mut recipients: vector<address>,
recipients: vector<address>,
ctx: &mut TxContext
) {
assert!(coin_types.length() == amounts.length(), ETypesAmountsNotSameLength);
assert!(amounts.length() == recipients.length(), ENotSameLength);
let treasury = borrow_treasury(account, treasury_name);
assert!(treasury.coin_type_exists<CoinType>(), ECoinTypeDoesntExist);
let sum = amounts.fold!(0, |sum, amount| sum + amount);
if (treasury.coin_type_value<CoinType>() < sum) assert!(sum <= treasury.coin_type_value<CoinType>(), EInsufficientFunds);

let mut proposal = account.create_proposal(
auth,
Expand All @@ -187,9 +188,9 @@ public fun propose_transfer<Config, Outcome>(
ctx
);

coin_types.zip_do!(amounts, |coin_type, amount| {
new_spend(&mut proposal, coin_type, amount, TransferProposal());
acc_transfer::new_transfer(&mut proposal, recipients.remove(0), TransferProposal());
recipients.zip_do!(amounts, |recipient, amount| {
new_spend<Outcome, CoinType, TransferProposal>(&mut proposal, amount, TransferProposal());
acc_transfer::new_transfer(&mut proposal, recipient, TransferProposal());
});

account.add_proposal(proposal, version::current(), TransferProposal());
Expand All @@ -199,12 +200,12 @@ public fun propose_transfer<Config, Outcome>(
// step 3: execute the proposal and return the action (account::execute_proposal)

// step 4: loop over transfer
public fun execute_transfer<Config, Outcome, C: drop>(
public fun execute_transfer<Config, Outcome, CoinType: drop>(
executable: &mut Executable,
account: &mut Account<Config, Outcome>,
ctx: &mut TxContext
) {
let coin: Coin<C> = do_spend(executable, account, version::current(), TransferProposal(), ctx);
let coin: Coin<CoinType> = do_spend(executable, account, version::current(), TransferProposal(), ctx);
acc_transfer::do_transfer(executable, account, coin, version::current(), TransferProposal());
}

Expand All @@ -214,7 +215,7 @@ public fun complete_transfer(executable: Executable) {
}

// step 1(bis): same but from a treasury
public fun propose_vesting<Config, Outcome>(
public fun propose_vesting<Config, Outcome, CoinType: drop>(
auth: Auth,
account: &mut Account<Config, Outcome>,
outcome: Outcome,
Expand All @@ -223,13 +224,16 @@ public fun propose_vesting<Config, Outcome>(
execution_time: u64,
expiration_epoch: u64,
treasury_name: String,
coin_type: String,
coin_amount: u64,
start_timestamp: u64,
end_timestamp: u64,
recipient: address,
ctx: &mut TxContext
) {
let treasury = borrow_treasury(account, treasury_name);
assert!(treasury.coin_type_exists<CoinType>(), ECoinTypeDoesntExist);
assert!(treasury.coin_type_value<CoinType>() >= coin_amount, EInsufficientFunds);

let mut proposal = account.create_proposal(
auth,
outcome,
Expand All @@ -243,7 +247,7 @@ public fun propose_vesting<Config, Outcome>(
ctx
);

new_spend(&mut proposal, coin_type, coin_amount, PayProposal());
new_spend<Outcome, CoinType, PayProposal>(&mut proposal, coin_amount, PayProposal());
vesting::new_vesting(&mut proposal, start_timestamp, end_timestamp, recipient, PayProposal());
account.add_proposal(proposal, version::current(), PayProposal());
}
Expand All @@ -252,46 +256,46 @@ public fun propose_vesting<Config, Outcome>(
// step 3: execute the proposal and return the action (account::execute_proposal)

// step 4: loop over it in PTB, sends last object from the Send action
public fun execute_vesting<Config, Outcome, C: drop>(
public fun execute_vesting<Config, Outcome, CoinType: drop>(
mut executable: Executable,
account: &mut Account<Config, Outcome>,
ctx: &mut TxContext
) {
let coin: Coin<C> = do_spend(&mut executable, account, version::current(), PayProposal(), ctx);
let coin: Coin<CoinType> = do_spend(&mut executable, account, version::current(), PayProposal(), ctx);
vesting::do_vesting(&mut executable, account, coin, version::current(), PayProposal(), ctx);
executable.destroy(version::current(), PayProposal());
}

// === [ACTION] Public Functions ===

public fun new_spend<Outcome, W: drop>(
public fun new_spend<Outcome, CoinType: drop, W: drop>(
proposal: &mut Proposal<Outcome>,
coin_type: String,
amount: u64,
witness: W,
) {
proposal.add_action(SpendAction { coin_type, amount }, witness);
proposal.add_action(SpendAction<CoinType> { amount }, witness);
}

public fun do_spend<Config, Outcome, C: drop, W: copy + drop>(
public fun do_spend<Config, Outcome, CoinType: drop, W: copy + drop>(
executable: &mut Executable,
account: &mut Account<Config, Outcome>,
version: TypeName,
witness: W,
ctx: &mut TxContext
): Coin<C> {
): Coin<CoinType> {
let name = executable.issuer().role_name();
let SpendAction { coin_type, amount } = executable.action(account.addr(), version, witness);
let SpendAction<CoinType> { amount } = executable.action(account.addr(), version, witness);

let treasury: &mut Treasury = account.borrow_managed_asset_mut(TreasuryKey { name }, version);
let balance: &mut Balance<C> = treasury.bag.borrow_mut(coin_type);
let balance: &mut Balance<CoinType> = treasury.bag.borrow_mut(type_name::get<CoinType>());
let coin = coin::take(balance, amount, ctx);

if (balance.value() == 0) treasury.bag.remove<String, Balance<C>>(coin_type).destroy_zero();
if (balance.value() == 0)
treasury.bag.remove<TypeName, Balance<CoinType>>(type_name::get<CoinType>()).destroy_zero();

coin
}

public fun delete_spend_action<Outcome>(expired: &mut Expired<Outcome>) {
let SpendAction { .. } = expired.remove_expired_action();
public fun delete_spend_action<Outcome, CoinType>(expired: &mut Expired<Outcome>) {
let SpendAction<CoinType> { .. } = expired.remove_expired_action();
}
Loading

0 comments on commit 56197cb

Please sign in to comment.