a demo of BAMM

BAMM.py

"""
This is a demo of BAMM. Frax's recent article introduced BAMM, which is quite intriguing. 
As the article didn't provide more details, I wanted to create a simple model for better understanding and simulation. 

This includes some of my personal speculations on implementation details.
If there are any inaccuracies, feel free to point them out!

original article link: <https://flywheeldefi.com/article/bamm-revolutionary-primative>
"""
from typing import List, Tuple


class CPAMM:
    def __init__(self):
        self.token_reserve = 0
        self.eth_reserve = 0
        self.total_supply = 0

    def mint(self, token_amount: float, eth_amount: float) -> float:
        if self.total_supply == 0:
            lp_token = token_amount * eth_amount
        else:
            lp_token = (
                min(token_amount / self.token_reserve, eth_amount / self.eth_reserve)
                * self.total_supply
            )

        self.total_supply += lp_token
        self.token_reserve += token_amount
        self.eth_reserve += eth_amount

        return lp_token

    def burn(self, liquidity_amount: float) -> Tuple[float, float]:
        token_amount = (liquidity_amount * self.token_reserve) / self.total_supply
        eth_amount = (liquidity_amount * self.eth_reserve) / self.total_supply

        self.token_reserve -= token_amount
        self.eth_reserve -= eth_amount

        return token_amount, eth_amount

    def swap(self, i: int, amount_in: float) -> float:
        k = self.get_k()
        # frax in eth out
        if i == 0:
            self.token_reserve += amount_in
            new_eth_reserve = k / self.token_reserve
            amount_out = self.eth_reserve - new_eth_reserve
            self.eth_reserve = new_eth_reserve
        else:
            self.eth_reserve += amount_in
            new_token_reserve = k / self.eth_reserve
            amount_out = self.eth_reserve - new_token_reserve
            self.token_reserve = new_token_reserve

        return amount_out

    def quote_amount_in(self, i: int, amount_out: float) -> float:
        k = self.get_k()
        if i == 0:
            return k / (self.eth_reserve - amount_out) - self.token_reserve
        else:
            return k / (self.token_reserve - amount_out) - self.eth_reserve

    def quote_amount_out(self, i: int, amount_in: float) -> float:
        k = self.get_k()
        if i == 0:
            return self.eth_reserve - k / (self.token_reserve + amount_in)
        else:
            return self.token_reserve - k / (self.eth_reserve + amount_in)

    def get_k(self):
        return self.token_reserve * self.eth_reserve


class LoanUser:
    def __init__(self, initial_collateral, debt, shares):
        self.initial_collateral = initial_collateral
        self.initial_debt = debt  # equals lp_token amount
        self.shares = shares


class BAMM:
    def __init__(self, amm: CPAMM):
        self.AMM = amm
        self.liquidation_discount = 0.999
        self.total_debt = 0
        self.total_collateral = 0
        self.total_shares = 0
        self.users = {}  # this demo only support one user for now

    def loan(self, user_address: str, collateral_amount: float, debt: float) -> float:
        # debt equals lp_token, we could get the amount of frax and eth add into pool
        frax_add_in_pool = debt / self.AMM.total_supply * self.AMM.token_reserve
        eth_add_in_pool = debt / self.AMM.total_supply * self.AMM.eth_reserve

        # We do not deposit the FRAX token into the AMM pool.
        # Instead, we keep a virtual record, and later transfer the LP tokens to the user.
        # Since the LP token price is always equal to 1, we can consider that the user has
        # borrowed the FRAX equivalent to the amount of LP tokens received.
        lp_token = self.AMM.mint(frax_add_in_pool, eth_add_in_pool)
        user_shares = (
            collateral_amount / self.total_collateral
            if self.total_collateral > 0
            else collateral_amount
        )

        self.users[user_address] = LoanUser(collateral_amount, lp_token, user_shares)
        self.total_debt += debt
        self.total_collateral += collateral_amount

        return lp_token

    def repay(self, user_address: str) -> Tuple[float, float]:
        user_state: LoanUser = self.users[user_address]
        (frax_amount, eth_amount) = self.AMM.burn(user_state.initial_debt)
        eth_amount += (
            (self.total_collateral - self.AMM.eth_reserve)
            * user_state.shares
            / self.total_shares
            if self.total_shares > 0
            else 0
        )

        self.total_debt -= user_state.initial_debt
        self.total_collateral -= eth_amount

        self.users.pop(user_address)

        return frax_amount, eth_amount

    def healthy(self, user_address: str) -> float:
        # set lowest healthy as 105%
        user_state: LoanUser = self.users[user_address]
        amm_p = 1e-6 / self.AMM.quote_amount_out(0, 1e-6)
        return user_state.initial_collateral * amm_p / user_state.initial_debt - 1.05

    def soft_liquidate(self, user_address: str, pump: bool, amount) -> float:
        assert self.healthy(user_address) <= 0

        # liquidator buy collateral
        if pump:
            amount_out = self.AMM.quote_amount_in(0, amount)
            # transfer frax token from liquidator with liquidation_discount
            # frax_token.transferFrom(liquidator, address(this), amount * liquidation_discount)

            # and swap to AMM pool
            self.AMM.swap(0, amount)
            self.total_collateral -= amount_out
        else:
            # liquidator sell collateral
            amount_out = self.AMM.quote_amount_in(1, amount)
            # transfer ETH from liquidator with liquidation_discount
            # eth_in_amount = address(this).balance - beforeBalance
            # assert eth_in_amount >= amount * liquidation_discount

            # and swap to AMM pool
            self.AMM.swap(1, amount)
            self.total_collateral += amount * self.liquidation_discount

        return amount_out
 
 
 # In this example we have that:
#   sfrxETH=$1000
#   FRAX= $1
#   sfrxETH-FRAX LvP- $25000 (5000 frax * 5 sfrxETH)

amm = CPAMM()
bamm = BAMM(amm)

# initial liquidity
amm.mint(5000, 5)
print("initial total_supply", amm.total_supply)

user_address = "loan_user_address"
print("user put 1 sfrxETH in the BAMM and take out a loan of 250 FRAX")
initial_collateral = 1
initial_debt = 250
lp_token = bamm.loan(user_address, initial_collateral, initial_debt)
print("get lp_token amount:", lp_token)

print("user healthy before price drop: %.4f" % bamm.healthy(user_address))

print("price drop down...")
amm.swap(1, 5)
print("Now ETH price is", 1e-6 / amm.quote_amount_out(0, 1e-6))
print("user healthy after price drop: %.4f, can be soft-liquidate." % bamm.healthy(user_address))

print("do soft-liquidation while price drop")
for i in range(10):
    # ETH out price drop, arbitragures sell eth to BAMM
    amm.swap(1, 0.1)

    # liquidator do soft-liquidation to get profits, 
    # buy ETH from BAMM
    amount_out = bamm.soft_liquidate(user_address, pump=True, amount=0.1)
    assert amount_out > 0.1


repay_frax, repay_eth = bamm.repay(user_address)
print("user repay debt and get frax %.4f, eth %.4f" % (repay_frax, repay_eth))


# ========== output: ===========

initial total_supply 25000
user put 1 sfrxETH in the BAMM and take out a loan of 250 FRAX
get lp_token amount: 250.0
user healthy before price drop: 2.9500
price drop down...
Now ETH price is 252.49374469460878
user healthy after price drop: -0.0400, can be soft-liquidate.
do soft-liquidation while price drop
user repay debt and get frax 22.8598, eth 0.1094