DeFi Example - Flash Loans

module flash_lender::example {
    use sui::balance::{Self, Balance};
    use sui::coin::{Self, Coin};

    /// Shared object offering flash loans
    public struct FlashLender<phantom T> has key {
        id: UID,
        to_lend: Balance<T>,  // Available funds
        fee: u64,             // Fee per loan
    }

    /// Hot potato receipt - MUST be repaid
    public struct Receipt<phantom T> {
        flash_lender_id: ID,
        repay_amount: u64,
    }

    /// Admin capability
    public struct AdminCap has key, store {
        id: UID,
        flash_lender_id: ID,
    }

    const ELoanTooLarge: u64 = 0;
    const EInvalidRepaymentAmount: u64 = 1;
    const ERepayToWrongLender: u64 = 2;
    const EAdminOnly: u64 = 3;

    /// Create a flash lender
    public fun new<T>(
        to_lend: Balance<T>,
        fee: u64,
        ctx: &mut TxContext
    ): AdminCap {
        let id = object::new(ctx);
        let flash_lender_id = id.uid_to_inner();
        let flash_lender = FlashLender { id, to_lend, fee };

        // Share the lender so anyone can borrow
        transfer::share_object(flash_lender);

        // Return admin capability to creator
        AdminCap { id: object::new(ctx), flash_lender_id }
    }

    /// Borrow funds - returns loan and receipt
    public fun loan<T>(
        self: &mut FlashLender<T>,
        amount: u64,
        ctx: &mut TxContext,
    ): (Coin<T>, Receipt<T>) {
        assert!(self.to_lend.value() >= amount, ELoanTooLarge);

        let loan = coin::take(&mut self.to_lend, amount, ctx);
        let repay_amount = amount + self.fee;
        let receipt = Receipt {
            flash_lender_id: object::id(self),
            repay_amount,
        };

        (loan, receipt)
    }

    /// Repay loan - consumes receipt
    public fun repay<T>(
        self: &mut FlashLender<T>,
        payment: Coin<T>,
        receipt: Receipt<T>
    ) {
        let Receipt { flash_lender_id, repay_amount } = receipt;

        assert!(object::id(self) == flash_lender_id, ERepayToWrongLender);
        assert!(payment.value() == repay_amount, EInvalidRepaymentAmount);

        coin::put(&mut self.to_lend, payment)
    }

    /// Admin can withdraw funds
    public fun withdraw<T>(
        self: &mut FlashLender<T>,
        admin: &AdminCap,
        amount: u64,
        ctx: &mut TxContext,
    ): Coin<T> {
        assert!(object::borrow_id(self) == &admin.flash_lender_id, EAdminOnly);
        coin::take(&mut self.to_lend, amount, ctx)
    }
}

public struct Receipt<phantom T> {
    flash_lender_id: ID,
    repay_amount: u64,
    // No abilities! Must be consumed
}

transfer::share_object(flash_lender);

public struct AdminCap has key, store {
    id: UID,
    flash_lender_id: ID,  // Tied to specific lender
}

public fun arbitrage(
    lender: &mut FlashLender<SUI>,
    dex1: &mut Pool,
    dex2: &mut Pool,
    ctx: &mut TxContext,
) {
    // 1. Borrow funds
    let (loan, receipt) = lender.loan(1000000, ctx);

    // 2. Trade on DEX 1
    let token = dex1.swap_sui_to_token(loan);

    // 3. Trade on DEX 2
    let mut profit = dex2.swap_token_to_sui(token);

    // 4. Separate repayment from profit
    let repayment = profit.split(receipt.repay_amount(), ctx);

    // 5. Repay loan
    lender.repay(repayment, receipt);

    // 6. Keep profit
    transfer::public_transfer(profit, ctx.sender());
}

#[test]
fun test_flash_loan() {
    use sui::test_scenario as ts;
    use sui::coin;
    use sui::sui::SUI;

    let mut ts = ts::begin(@0x0);

    // Admin creates flash lender with 100 SUI, fee of 1
    {
        ts.next_tx(ADMIN);
        let coin = coin::mint_for_testing<SUI>(100, ts.ctx());
        let bal = coin.into_balance();
        let cap = new(bal, 1, ts.ctx());
        transfer::public_transfer(cap, ADMIN);
    };

    // Alice borrows and repays
    {
        ts.next_tx(ALICE);
        let mut lender: FlashLender<SUI> = ts.take_shared();
        let (loan, receipt) = lender.loan(10, ts.ctx());

        // Simulate profit
        let mut profit = coin::mint_for_testing<SUI>(1, ts.ctx());
        profit.join(loan);

        lender.repay(profit, receipt);
        ts::return_shared(lender);
    };

    // Admin withdraws fees
    {
        ts.next_tx(ADMIN);
        let cap = ts.take_from_sender();
        let mut lender: FlashLender<SUI> = ts.take_shared();

        // Pool grew from 100 to 101 (fee collected)
        assert!(lender.max_loan() == 101, 0);

        let coin = lender.withdraw(&cap, 1, ts.ctx());
        transfer::public_transfer(coin, ADMIN);

        ts::return_shared(lender);
        ts.return_to_sender(cap);
    };

    ts.end();
}

public struct Pool<phantom X, phantom Y> has key {
    id: UID,
    reserve_x: Balance<X>,
    reserve_y: Balance<Y>,
    lp_supply: Supply<LP<X, Y>>,
}

public fun swap<X, Y>(
    pool: &mut Pool<X, Y>,
    input: Coin<X>,
    ctx: &mut TxContext,
): Coin<Y> {
    let input_amount = input.value();
    let output_amount = calculate_output(
        input_amount,
        pool.reserve_x.value(),
        pool.reserve_y.value(),
    );

    coin::put(&mut pool.reserve_x, input);
    coin::take(&mut pool.reserve_y, output_amount, ctx)
}

// Receipt MUST be consumed
public struct Receipt<phantom T> {
    // No abilities
}

assert!(payment.value() == repay_amount, EInvalidRepayment);
assert!(lender_id == receipt.lender_id, EWrongLender);

public fun admin_only(cap: &AdminCap, ...) { /* ... */ }
public fun anyone(...) { /* ... */ }

event::emit(LoanTaken {
    borrower: ctx.sender(),
    amount,
    fee,
});