library.cairo

// AMM LIBRARY
// The purpose of the library is to hold any abstracted logic from the AMM. This allows cleaner test suites.

// MIT LICENCE

%lang starknet

from starkware.cairo.common.cairo_builtins import HashBuiltin
from starkware.cairo.common.math import assert_not_equal
from starkware.cairo.common.uint256 import (
    Uint256,
    uint256_add,
    uint256_sub,
    uint256_mul,
    uint256_unsigned_div_rem,
    uint256_le,
    uint256_eq,
)

namespace AMM {
    // input arguments are:
    // 1) exact amount of token to sell
    // 2) currency reserve amount
    // 3) token reserve amount
    func get_sell_price{syscall_ptr: felt*, pedersen_ptr: HashBuiltin*, range_check_ptr}(
        token_amount: Uint256,
        currency_reserves: Uint256,
        token_reserves: Uint256,
        lp_fee_thousands_: Uint256,
    ) -> (price: Uint256) {
        alloc_locals;

        // LP fee is used to withold currency as reward to LP providers
        let (lp_fee) = uint256_sub(Uint256(1000, 0), lp_fee_thousands_);

        // Apply LP fee to token amount
        let (token_amount_w_fee, mul_overflow) = uint256_mul(token_amount, lp_fee);
        with_attr error_message("LP fee overflow") {
            assert mul_overflow = Uint256(0, 0);
        }

        // Calculate price
        let (numerator, mul_overflow) = uint256_mul(token_amount_w_fee, currency_reserves);
        with_attr error_message("Price numerator overflow") {
            assert mul_overflow = Uint256(0, 0);
        }
        let (denominator, mul_overflow) = uint256_mul(token_reserves, Uint256(1000, 0));
        with_attr error_message("LP fee buffer denominator overflow") {
            assert mul_overflow = Uint256(0, 0);
        }
        let (denominator_fee, is_overflow) = uint256_add(denominator, token_amount_w_fee);
        with_attr error_message("Price denominator overflow") {
            assert is_overflow = 0;
        }

        let (price, _) = uint256_unsigned_div_rem(numerator, denominator_fee);

        // Sanity fallback check to see if price is 0
        with_attr error_message("Price cannot be 0") {
            let (is_zero) = uint256_le(price, Uint256(0, 0));
            assert_not_equal(is_zero, 1);
        }

        // Rounding errors favour the contract
        return (price,);
    }

    // input arguments are:
    // 1) exact amount of token to buy
    // 2) currency reserve amount
    // 3) token reserve amount
    func get_buy_price{syscall_ptr: felt*, pedersen_ptr: HashBuiltin*, range_check_ptr}(
        token_amount: Uint256,
        currency_reserves: Uint256,
        token_reserves: Uint256,
        lp_fee_thousands_: Uint256,
    ) -> (price: Uint256) {
        alloc_locals;

        // LP fee is used to withold currency as reward to LP providers
        let (lp_fee) = uint256_sub(Uint256(1000, 0), lp_fee_thousands_);

        // Calculate price
        let (numerator, mul_overflow) = uint256_mul(currency_reserves, token_amount);
        with_attr error_message("Values too large") {
            assert mul_overflow = Uint256(0, 0);
        }

        let (denominator) = uint256_sub(token_reserves, token_amount);
        with_attr error_message("Can't buy 100% of the AMM") {
            let (is_z) = uint256_eq(denominator, Uint256(0, 0));
            assert is_z = 0;
        }
        // Add LP fee
        let (numerator, mul_overflow) = uint256_mul(numerator, Uint256(1000, 0));
        with_attr error_message("Numerator overflow") {
            assert mul_overflow = Uint256(0, 0);
        }
        let (denominator, mul_overflow) = uint256_mul(denominator, lp_fee);
        with_attr error_message("Denominator overflow") {
            assert mul_overflow = Uint256(0, 0);
        }

        // Calculate price
        let (price, remainder) = uint256_unsigned_div_rem(numerator, denominator);

        // Return value if no remainder
        let (is_z) = uint256_eq(remainder, Uint256(0, 0));
        if (is_z == (1)) {
            return (price,);
        }

        // Round up when there is a remainder, to favour LP providers
        let (price, _) = uint256_add(price, Uint256(1, 0));
        return (price,);
    }
}