7. Liquidity Pool

Full working code: The complete source code for this example is available in example-contracts/pool/

This example closely mirrors the AMM contract’s logic but supports only one token pair. This design allows the pool to implement a token interface, where the pool’s token represents shares in the liquidity pool, enabling interaction through standard token methods (demonstrated at the test’s end). In a production environment, a middleware application would manage pool creation for users and maintain a mapping of token pairs to pool contract IDs.

WIT Interface

package root:component;

world root {
    include kontor:built-in/built-in;
    use kontor:built-in/context.{view-context, proc-context, signer};
    use kontor:built-in/error.{error};
    use kontor:built-in/foreign.{contract-address};
    use kontor:built-in/numbers.{integer, decimal};

  	record deposit-result {
  		lp-shares: integer,
  		deposit-a: integer,
  		deposit-b: integer,
  	}

  	record withdraw-result {
  		amount-a: integer,
  		amount-b: integer,
  	}

    export init: func(ctx: borrow<proc-context>);

    export re-init: func(ctx: borrow<proc-context>, token-a: contract-address, amount-a: integer, token-b: contract-address, amount-b: integer, fee: integer) -> result<integer, error>;

  	export fee: func(ctx: borrow<view-context>) -> integer;

    export balance: func(ctx: borrow<view-context>, acc: string) -> option<integer>;
    export transfer: func(ctx: borrow<proc-context>, to: string, n: integer) -> result<_, error>;

    export token-balance: func(ctx: borrow<view-context>, token: contract-address) -> result<integer, error>;
    export quote-deposit: func(ctx: borrow<view-context>, amount-a: integer, amount-b: integer) -> result<deposit-result, error>;
    export deposit: func(ctx: borrow<proc-context>, amount-a: integer, amount-b: integer) -> result<deposit-result, error>;
    export quote-withdraw: func(ctx: borrow<view-context>, shares: integer) -> result<withdraw-result, error>;
    export withdraw: func(ctx: borrow<proc-context>, shares: integer) -> result<withdraw-result, error>;

    export swap: func(ctx: borrow<proc-context>, token-in: contract-address, amount-in: integer, min-out: integer) -> result<integer, error>;
    export quote-swap: func(ctx: borrow<view-context>, token-in: contract-address, amount-in: integer) -> result<integer, error>;
}

Rust Implementation

use stdlib::*;

contract!(name = "pool");

interface!(name = "token_dyn", path = "../token/contract/wit");

#[derive(Clone, StorageRoot)]
struct PoolStorage {
    pub token_a: ContractAddress,
    pub token_b: ContractAddress,
    pub fee_bps: Integer,

    pub lp_total_supply: Integer,
    pub lp_ledger: Map<String, Integer>,

    pub custodian: String,
}

impl PoolStorage {
    pub fn new(
        ctx: &ProcContext,
        token_a: ContractAddress,
        amount_a: Integer,
        token_b: ContractAddress,
        amount_b: Integer,
        fee_bps: Integer,
    ) -> Result<Self, Error> {
        validate_amount(amount_a)?;
        validate_amount(amount_b)?;

        let lp_shares = (amount_a * amount_b).sqrt()?;
        let custodian = ctx.contract_signer().to_string();
        let pool = PoolStorage {
            token_a: token_a.clone(),
            token_b: token_b.clone(),
            fee_bps,
            lp_total_supply: lp_shares,
            lp_ledger: Map::new(&[(ctx.signer().to_string(), lp_shares)]),
            custodian: custodian.clone(),
        };

        token_dyn::transfer(&token_a, ctx.signer(), &custodian, amount_a)?;
        token_dyn::transfer(&token_b, ctx.signer(), &custodian, amount_b)?;

        Ok(pool)
    }
}

fn token_out(ctx: &impl ReadContext, token_in: &ContractAddress) -> Result<ContractAddress, Error> {
    let token_a = ctx.model().token_a();
    let token_b = ctx.model().token_b();
    if token_in == &token_a {
        Ok(token_b)
    } else if token_in == &token_b {
        Ok(token_a)
    } else {
        Err(Error::Message(format!("token {} not in pair", token_in)))
    }
}

fn validate_amount(amount: Integer) -> Result<(), Error> {
    // 0 < amount < sqrt(MAX_INT)
    if amount <= Integer::default() || amount > "340_282_366_920_938_463_463_374_607_431".into() {
        return Err(Error::Message("bad amount".to_string()));
    }
    Ok(())
}

fn calc_swap_result(
    amount_in: Integer,
    bal_in: Integer,
    bal_out: Integer,
    fee_bps: Integer,
) -> Result<Integer, Error> {
    validate_amount(amount_in)?;
    validate_amount(bal_in)?;
    validate_amount(bal_out)?;

    // input amount less fee, round down
    let bps_in_100pct = 10000.into();
    let in_less_fee = amount_in * (bps_in_100pct - fee_bps) / bps_in_100pct;

    let new_bal_in = bal_in + in_less_fee;
    validate_amount(new_bal_in)?;

    // calculate output amount from delta in output-token balance, round down
    let k = bal_in * bal_out;
    Ok((bal_out * new_bal_in - k) / new_bal_in)
}

fn quote_swap(
    ctx: &impl ReadContext,
    token_in: &ContractAddress,
    amount_in: Integer,
) -> Result<Integer, Error> {
    let custodian = ctx.model().custodian();
    let bal_in = token_dyn::balance(token_in, &custodian).unwrap_or_default();
    let bal_out = token_dyn::balance(&token_out(ctx, token_in)?, &custodian).unwrap_or_default();
    calc_swap_result(amount_in, bal_in, bal_out, ctx.model().fee_bps())
}

fn quote_deposit(
    ctx: &impl ReadContext,
    amount_a: Integer,
    amount_b: Integer,
) -> Result<DepositResult, Error> {
    validate_amount(amount_a)?;
    validate_amount(amount_b)?;

    let token_a = ctx.model().token_a();
    let token_b = ctx.model().token_b();
    let lp_supply = ctx.model().lp_total_supply();

    let custodian = ctx.model().custodian();
    let bal_a = token_dyn::balance(&token_a, &custodian).unwrap_or_default();
    let bal_b = token_dyn::balance(&token_b, &custodian).unwrap_or_default();

    let lp_shares = if amount_a * bal_b < amount_b * bal_a {
        amount_a * lp_supply / bal_a
    } else {
        amount_b * lp_supply / bal_b
    };

    let supply_minus_one = lp_supply - 1.into();
    Ok(DepositResult {
        deposit_a: (lp_shares * bal_a + supply_minus_one) / lp_supply, // round up
        deposit_b: (lp_shares * bal_b + supply_minus_one) / lp_supply, // round up
        lp_shares,
    })
}

fn quote_withdraw(ctx: &impl ReadContext, shares: Integer) -> Result<WithdrawResult, Error> {
    validate_amount(shares)?;

    let token_a = ctx.model().token_a();
    let token_b = ctx.model().token_b();
    let lp_supply = ctx.model().lp_total_supply();

    let custodian = ctx.model().custodian();
    let bal_a = token_dyn::balance(&token_a, &custodian).unwrap_or_default();
    let bal_b = token_dyn::balance(&token_b, &custodian).unwrap_or_default();

    Ok(WithdrawResult {
        amount_a: shares * bal_a / lp_supply,
        amount_b: shares * bal_b / lp_supply,
    })
}

impl Guest for Pool {
    // Dummy implementation for testing purposes.
    fn init(ctx: &ProcContext) {
        PoolStorage {
            token_a: ContractAddress {
                name: "".to_string(),
                height: 0,
                tx_index: 0,
            },
            token_b: ContractAddress {
                name: "".to_string(),
                height: 0,
                tx_index: 0,
            },
            lp_ledger: Map::default(),
            lp_total_supply: 0.into(),
            fee_bps: 0.into(),
            custodian: "".to_string(),
        }
        .init(ctx);
    }

    // This represents the production init function.
    // Only for local testing purposes.
    fn re_init(
        ctx: &ProcContext,
        token_a: ContractAddress,
        amount_a: Integer,
        token_b: ContractAddress,
        amount_b: Integer,
        fee: Integer,
    ) -> Result<Integer, Error> {
        PoolStorage::new(ctx, token_a, amount_a, token_b, amount_b, fee)?.init(ctx);
        Ok(ctx.model().lp_total_supply())
    }

    fn fee(ctx: &ViewContext) -> Integer {
        ctx.model().fee_bps()
    }

    fn balance(ctx: &ViewContext, acc: String) -> Option<Integer> {
        ctx.model().lp_ledger().get(acc)
    }

    fn transfer(ctx: &ProcContext, to: String, n: Integer) -> Result<(), Error> {
        let from = ctx.signer().to_string();
        let ledger = ctx.model().lp_ledger();

        let from_balance = ledger.get(&from).unwrap_or_default();
        let to_balance = ledger.get(&to).unwrap_or_default();
        if from_balance < n {
            return Err(Error::Message("insufficient funds".to_string()));
        }

        ledger.set(from, from_balance - n);
        ledger.set(to, to_balance + n);
        Ok(())
    }

    fn token_balance(ctx: &ViewContext, token: ContractAddress) -> Result<Integer, Error> {
        token_out(ctx, &token)?;
        Ok(token_dyn::balance(&token, &ctx.model().custodian()).unwrap_or_default())
    }

    fn quote_deposit(
        ctx: &ViewContext,
        amount_a: Integer,
        amount_b: Integer,
    ) -> Result<DepositResult, Error> {
        quote_deposit(ctx, amount_a, amount_b)
    }

    fn deposit(
        ctx: &ProcContext,
        amount_a: Integer,
        amount_b: Integer,
    ) -> Result<DepositResult, Error> {
        let res = quote_deposit(ctx, amount_a, amount_b)?;
        let ledger = ctx.model().lp_ledger();
        let custodian = ctx.model().custodian();

        let user = ctx.signer().to_string();
        let bal = ledger.get(&user).unwrap_or_default();
        ledger.set(user, bal + res.lp_shares);
        ctx.model().set_lp_total_supply(ctx.model().lp_total_supply() + res.lp_shares);

        token_dyn::transfer(
            &ctx.model().token_a(),
            ctx.signer(),
            &custodian,
            res.deposit_a,
        )?;
        token_dyn::transfer(
            &ctx.model().token_b(),
            ctx.signer(),
            &custodian,
            res.deposit_b,
        )?;

        Ok(res)
    }

    fn quote_withdraw(ctx: &ViewContext, shares: Integer) -> Result<WithdrawResult, Error> {
        quote_withdraw(ctx, shares)
    }

    fn withdraw(ctx: &ProcContext, shares: Integer) -> Result<WithdrawResult, Error> {
        let res = quote_withdraw(ctx, shares)?;

        let ledger = ctx.model().lp_ledger();
        let user = ctx.signer().to_string();

        let total = ctx.model().lp_total_supply();
        let bal = ledger.get(&user).unwrap_or_default();

        if total < shares {
            return Err(Error::Message("insufficient total supply".to_string()));
        }
        if bal < shares {
            return Err(Error::Message("insufficient share balance".to_string()));
        }

        ledger.set(user.clone(), bal - shares);
        ctx.model().set_lp_total_supply(total - shares);

        token_dyn::transfer(
            &ctx.model().token_a(),
            ctx.contract_signer(),
            &user,
            res.amount_a,
        )?;
        token_dyn::transfer(
            &ctx.model().token_b(),
            ctx.contract_signer(),
            &user,
            res.amount_b,
        )?;

        Ok(res)
    }

    fn quote_swap(
        ctx: &ViewContext,
        token_in: ContractAddress,
        amount_in: Integer,
    ) -> Result<Integer, Error> {
        quote_swap(ctx, &token_in, amount_in)
    }

    fn swap(
        ctx: &ProcContext,
        token_in: ContractAddress,
        amount_in: Integer,
        min_out: Integer,
    ) -> Result<Integer, Error> {
        let token_out = token_out(ctx, &token_in)?;
        let amount_out = quote_swap(ctx, &token_in, amount_in)?;

        if amount_out < min_out {
            return Err(Error::Message(format!(
                "amount out ({}) below minimum",
                amount_out
            )));
        }

        token_dyn::transfer(
            &token_in,
            ctx.signer(),
            &ctx.model().custodian(),
            amount_in,
        )?;
        token_dyn::transfer(
            &token_out,
            ctx.contract_signer(),
            &ctx.signer().to_string(),
            amount_out,
        )?;

        Ok(amount_out)
    }
}

Testing

#[cfg(test)]
mod tests {
    use testlib::*;

    interface!(name = "pool");
    interface!(name = "token_a", path = "../token/contract/wit");
    interface!(name = "token_b", path = "../token/contract/wit");

    #[testlib::test]
    async fn test_contract() -> Result<()> {
        let admin = runtime.identity().await?;
        let minter = runtime.identity().await?;

        let pool = runtime.publish(&admin, "pool").await?;
        let token_a = runtime.publish(&admin, "token").await?;
        let token_b = runtime.publish(&admin, "token").await?;

        token_a::mint(runtime, &token_a, &minter, 1000.into()).await??;
        token_b::mint(runtime, &token_b, &minter, 1000.into()).await??;

        token_a::transfer(runtime, &token_a, &minter, &admin, 100.into()).await??;
        token_b::transfer(runtime, &token_b, &minter, &admin, 500.into()).await??;

        let res = pool::re_init(
            runtime,
            &pool,
            &admin,
            token_a.clone(),
            100.into(),
            token_b.clone(),
            500.into(),
            0.into(),
        )
        .await?;
        assert_eq!(res, Ok(223.into()));

        let bal_a = pool::token_balance(runtime, &pool, token_a.clone()).await?;
        assert_eq!(bal_a, Ok(100.into()));
        let bal_b = pool::token_balance(runtime, &pool, token_b.clone()).await?;
        assert_eq!(bal_b, Ok(500.into()));
        let k1 = bal_a.unwrap() * bal_b.unwrap();

        let res = pool::quote_swap(runtime, &pool, token_a.clone(), 10.into()).await?;
        assert_eq!(res, Ok(45.into()));

        let res = pool::quote_swap(runtime, &pool, token_a.clone(), 100.into()).await?;
        assert_eq!(res, Ok(250.into()));

        let res = pool::quote_swap(runtime, &pool, token_a.clone(), 1000.into()).await?;
        assert_eq!(res, Ok(454.into()));

        let res = pool::swap(runtime, &pool, &minter, token_a.clone(), 10.into(), 46.into()).await?;
        assert!(res.is_err()); // below minimum

        let res = pool::swap(runtime, &pool, &minter, token_a.clone(), 10.into(), 45.into()).await?;
        assert_eq!(res, Ok(45.into()));

        let bal_a = pool::token_balance(runtime, &pool, token_a.clone()).await?;
        let bal_b = pool::token_balance(runtime, &pool, token_b.clone()).await?;
        let k2 = bal_a.unwrap() * bal_b.unwrap();
        assert!(k2 >= k1);

        let res = pool::quote_swap(runtime, &pool, token_b.clone(), 45.into()).await?;
        assert_eq!(res, Ok(9.into()));
        let res = pool::swap(runtime, &pool, &minter, token_b.clone(), 45.into(), 0.into()).await?;
        assert_eq!(res, Ok(9.into()));

        let bal_a = pool::token_balance(runtime, &pool, token_a.clone()).await?;
        let bal_b = pool::token_balance(runtime, &pool, token_b.clone()).await?;
        let k3 = bal_a.unwrap() * bal_b.unwrap();
        assert!(k3 >= k2);

        // use token interface to transfer shares
        let res = token::balance(runtime, &pool, &admin).await?;
        assert_eq!(res, Some(223.into()));
        let res = token::balance(runtime, &pool, &minter).await?;
        assert_eq!(res, None);

        token::transfer(runtime, &pool, &admin, &minter, 23.into()).await??;

        let res = token::balance(runtime, &pool, &admin).await?;
        assert_eq!(res, Some(200.into()));
        let res = token::balance(runtime, &pool, &minter).await?;
        assert_eq!(res, Some(23.into()));

        Ok(())
    }
}

Basics

Reference

Examples

WIT Interface

Rust Implementation

Testing

Basics

Reference

Examples

​WIT Interface

​Rust Implementation

​Testing

WIT Interface

Rust Implementation

Testing