Merge pull request #218 from cosmosquad-labs/improve-amm-pool

feat: add more functionalities to the amm package

x/liquidity/amm/match_bench_test.go

@@ -27,7 +27,7 @@ func BenchmarkFindMatchPrice(b *testing.B) {
 			var poolOrderSources []amm.OrderSource
 			for i := 0; i < 1000; i++ {
 				rx, ry := squad.RandomInt(r, minReserveAmt, maxReserveAmt), squad.RandomInt(r, minReserveAmt, maxReserveAmt)
-				pool := amm.NewBasicPool(rx, ry, sdk.ZeroInt())
+				pool := amm.NewBasicPool(rx, ry, sdk.Int{})
 				poolOrderSources = append(poolOrderSources, amm.NewMockPoolOrderSource(pool, "denom1", "denom2"))
 			}
 			os := amm.MergeOrderSources(append(poolOrderSources, ob)...)

x/liquidity/amm/orderbook.go

@@ -3,6 +3,7 @@ package amm
 import (
 	"fmt"
 	"sort"
+	"strings"
 
 	sdk "github.com/cosmos/cosmos-sdk/types"
 )
@@ -35,12 +36,14 @@ func (ob *OrderBook) Add(orders ...Order) {
 
 // HighestBuyPrice returns the highest buy price in the order book.
 func (ob *OrderBook) HighestBuyPrice() (sdk.Dec, bool) {
-	return ob.buys.highestPrice()
+	price, _, found := ob.buys.highestPrice()
+	return price, found
 }
 
 // LowestSellPrice returns the lowest sell price in the order book.
 func (ob *OrderBook) LowestSellPrice() (sdk.Dec, bool) {
-	return ob.sells.lowestPrice()
+	price, _, found := ob.sells.lowestPrice()
+	return price, found
 }
 
 // BuyAmountOver returns the amount of buy orders in the order book
@@ -67,6 +70,85 @@ func (ob *OrderBook) SellOrdersUnder(price sdk.Dec) []Order {
 	return ob.sells.ordersUnder(price)
 }
 
+func (ob *OrderBook) HighestPrice() (sdk.Dec, bool) {
+	highestBuyPrice, _, foundBuy := ob.buys.highestPrice()
+	highestSellPrice, _, foundSell := ob.sells.highestPrice()
+	switch {
+	case foundBuy && foundSell:
+		return sdk.MaxDec(highestBuyPrice, highestSellPrice), true
+	case foundBuy:
+		return highestBuyPrice, true
+	case foundSell:
+		return highestSellPrice, true
+	default:
+		return sdk.Dec{}, false
+	}
+}
+
+func (ob *OrderBook) LowestPrice() (sdk.Dec, bool) {
+	lowestBuyPrice, _, foundBuy := ob.buys.lowestPrice()
+	lowestSellPrice, _, foundSell := ob.sells.lowestPrice()
+	switch {
+	case foundBuy && foundSell:
+		return sdk.MinDec(lowestBuyPrice, lowestSellPrice), true
+	case foundBuy:
+		return lowestBuyPrice, true
+	case foundSell:
+		return lowestSellPrice, true
+	default:
+		return sdk.Dec{}, false
+	}
+}
+
+func (ob *OrderBook) stringRepresentation(prices []sdk.Dec) string {
+	if len(prices) == 0 {
+		return "<nil>"
+	}
+	sort.Slice(prices, func(i, j int) bool {
+		return prices[i].GT(prices[j])
+	})
+	var b strings.Builder
+	b.WriteString("+--------buy---------+------------price-------------+--------sell--------+\n")
+	for _, price := range prices {
+		buyAmt, sellAmt := sdk.ZeroInt(), sdk.ZeroInt()
+		if i, exact := ob.buys.findPrice(price); exact {
+			buyAmt = TotalOpenAmount(ob.buys[i].orders)
+		}
+		if i, exact := ob.sells.findPrice(price); exact {
+			sellAmt = TotalOpenAmount(ob.sells[i].orders)
+		}
+		_, _ = fmt.Fprintf(&b, "| %18s | %28s | %-18s |\n", buyAmt, price.String(), sellAmt)
+	}
+	b.WriteString("+--------------------+------------------------------+--------------------+")
+	return b.String()
+}
+
+// FullString returns a full string representation of the order book.
+// FullString includes all possible price ticks from the order book's
+// highest price to the lowest price.
+func (ob *OrderBook) FullString(tickPrec int) string {
+	var prices []sdk.Dec
+	highest, found := ob.HighestPrice()
+	if !found {
+		return "<nil>"
+	}
+	lowest, _ := ob.LowestPrice()
+	for ; lowest.LTE(highest); lowest = UpTick(lowest, tickPrec) {
+		prices = append(prices, lowest)
+	}
+	return ob.stringRepresentation(prices)
+}
+
+// String returns a compact string representation of the order book.
+// String includes a tick only when there is at least one order on it.
+func (ob *OrderBook) String() string {
+	var prices []sdk.Dec
+	for _, tick := range append(ob.buys, ob.sells...) {
+		prices = append(prices, tick.price)
+	}
+	return ob.stringRepresentation(prices)
+}
+
 // orderBookTicks represents a list of orderBookTick.
 // This type is used for both buy/sell sides of OrderBook.
 type orderBookTicks []*orderBookTick
@@ -96,28 +178,28 @@ func (ticks *orderBookTicks) add(order Order) {
 	}
 }
 
-func (ticks orderBookTicks) highestPrice() (sdk.Dec, bool) {
+func (ticks orderBookTicks) highestPrice() (sdk.Dec, int, bool) {
 	if len(ticks) == 0 {
-		return sdk.Dec{}, false
+		return sdk.Dec{}, 0, false
 	}
-	for _, tick := range ticks {
+	for i, tick := range ticks {
 		if TotalOpenAmount(tick.orders).IsPositive() {
-			return tick.price, true
+			return tick.price, i, true
 		}
 	}
-	return sdk.Dec{}, false
+	return sdk.Dec{}, 0, false
 }
 
-func (ticks orderBookTicks) lowestPrice() (sdk.Dec, bool) {
+func (ticks orderBookTicks) lowestPrice() (sdk.Dec, int, bool) {
 	if len(ticks) == 0 {
-		return sdk.Dec{}, false
+		return sdk.Dec{}, 0, false
 	}
 	for i := len(ticks) - 1; i >= 0; i-- {
 		if TotalOpenAmount(ticks[i].orders).IsPositive() {
-			return ticks[i].price, true
+			return ticks[i].price, i, true
 		}
 	}
-	return sdk.Dec{}, false
+	return sdk.Dec{}, 0, false
 }
 
 func (ticks orderBookTicks) amountOver(price sdk.Dec) sdk.Int {

x/liquidity/amm/pool.go

@@ -13,6 +13,8 @@ var (
 // It also satisfies OrderView interface.
 type Pool interface {
 	OrderView
+	Balances() (rx, ry sdk.Int)
+	PoolCoinSupply() sdk.Int
 	Price() sdk.Dec
 	IsDepleted() bool
 	Deposit(x, y sdk.Int) (ax, ay, pc sdk.Int)
@@ -21,26 +23,35 @@ type Pool interface {
 
 // BasicPool is the basic pool type.
 type BasicPool struct {
-	rx, ry sdk.Dec
-	ps     sdk.Dec
+	rx, ry sdk.Int
+	ps     sdk.Int
 }
 
 // NewBasicPool returns a new BasicPool.
-// Pass sdk.ZeroInt() to ps when ps is not going to be used.
 func NewBasicPool(rx, ry, ps sdk.Int) *BasicPool {
 	return &BasicPool{
-		rx: rx.ToDec(),
-		ry: ry.ToDec(),
-		ps: ps.ToDec(),
+		rx: rx,
+		ry: ry,
+		ps: ps,
 	}
 }
 
+// Balances returns the balances of the pool.
+func (pool *BasicPool) Balances() (rx, ry sdk.Int) {
+	return pool.rx, pool.ry
+}
+
+// PoolCoinSupply returns the pool coin supply.
+func (pool *BasicPool) PoolCoinSupply() sdk.Int {
+	return pool.ps
+}
+
 // Price returns the pool price.
 func (pool *BasicPool) Price() sdk.Dec {
 	if pool.rx.IsZero() || pool.ry.IsZero() {
 		panic("pool price is not defined for a depleted pool")
 	}
-	return pool.rx.Quo(pool.ry)
+	return pool.rx.ToDec().Quo(pool.ry.ToDec())
 }
 
 // IsDepleted returns whether the pool is depleted or not.
@@ -55,33 +66,36 @@ func (pool *BasicPool) Deposit(x, y sdk.Int) (ax, ay, pc sdk.Int) {
 	// Note that we take as many coins as possible(by ceiling numbers)
 	// from depositor and mint as little coins as possible.
 
+	rx, ry := pool.rx.ToDec(), pool.ry.ToDec()
+	ps := pool.ps.ToDec()
+
 	// pc = floor(ps * min(x / rx, y / ry))
-	pc = pool.ps.MulTruncate(sdk.MinDec(
-		x.ToDec().QuoTruncate(pool.rx),
-		y.ToDec().QuoTruncate(pool.ry),
+	pc = ps.MulTruncate(sdk.MinDec(
+		x.ToDec().QuoTruncate(rx),
+		y.ToDec().QuoTruncate(ry),
 	)).TruncateInt()
 
-	mintProportion := pc.ToDec().Quo(pool.ps)             // pc / ps
-	ax = pool.rx.Mul(mintProportion).Ceil().TruncateInt() // ceil(rx * mintProportion)
-	ay = pool.ry.Mul(mintProportion).Ceil().TruncateInt() // ceil(ry * mintProportion)
+	mintProportion := pc.ToDec().Quo(ps)             // pc / ps
+	ax = rx.Mul(mintProportion).Ceil().TruncateInt() // ceil(rx * mintProportion)
+	ay = ry.Mul(mintProportion).Ceil().TruncateInt() // ceil(ry * mintProportion)
 	return
 }
 
 // Withdraw returns withdrawn x and y coin amount when someone withdraws
 // pc pool coin.
 // Withdraw also takes care of the fee rate.
 func (pool *BasicPool) Withdraw(pc sdk.Int, feeRate sdk.Dec) (x, y sdk.Int) {
-	if pc.ToDec().Equal(pool.ps) {
+	if pc.Equal(pool.ps) {
 		// Redeeming the last pool coin - give all remaining rx and ry.
-		x = pool.rx.TruncateInt()
-		y = pool.ry.TruncateInt()
+		x = pool.rx
+		y = pool.ry
 		return
 	}
 
-	proportion := pc.ToDec().QuoTruncate(pool.ps)                             // pc / ps
-	multiplier := sdk.OneDec().Sub(feeRate)                                   // 1 - feeRate
-	x = pool.rx.MulTruncate(proportion).MulTruncate(multiplier).TruncateInt() // floor(rx * proportion * multiplier)
-	y = pool.ry.MulTruncate(proportion).MulTruncate(multiplier).TruncateInt() // floor(ry * proportion * multiplier)
+	proportion := pc.ToDec().QuoTruncate(pool.ps.ToDec())                             // pc / ps
+	multiplier := sdk.OneDec().Sub(feeRate)                                           // 1 - feeRate
+	x = pool.rx.ToDec().MulTruncate(proportion).MulTruncate(multiplier).TruncateInt() // floor(rx * proportion * multiplier)
+	y = pool.ry.ToDec().MulTruncate(proportion).MulTruncate(multiplier).TruncateInt() // floor(ry * proportion * multiplier)
 	return
 }
 
@@ -105,7 +119,7 @@ func (pool *BasicPool) BuyAmountOver(price sdk.Dec) sdk.Int {
 	if price.GTE(pool.Price()) {
 		return sdk.ZeroInt()
 	}
-	return pool.rx.QuoTruncate(price).Sub(pool.ry).TruncateInt()
+	return pool.rx.ToDec().QuoTruncate(price).Sub(pool.ry.ToDec()).TruncateInt()
 }
 
 // SellAmountUnder returns the amount of sell orders for price less or equal
@@ -114,7 +128,45 @@ func (pool *BasicPool) SellAmountUnder(price sdk.Dec) sdk.Int {
 	if price.LTE(pool.Price()) {
 		return sdk.ZeroInt()
 	}
-	return pool.ry.Sub(pool.rx.QuoRoundUp(price)).TruncateInt()
+	return pool.ry.ToDec().Sub(pool.rx.ToDec().QuoRoundUp(price)).TruncateInt()
+}
+
+// PoolsOrderBook returns an order book with orders made by pools.
+// The order book has at most (numTicks*2+1) ticks visible, which includes
+// basePrice, numTicks ticks over basePrice and numTicks ticks under basePrice.
+// PoolsOrderBook assumes that basePrice is on ticks.
+func PoolsOrderBook(pools []Pool, basePrice sdk.Dec, numTicks, tickPrec int) *OrderBook {
+	prec := TickPrecision(tickPrec)
+	i := prec.TickToIndex(basePrice)
+	highestTick := prec.TickFromIndex(i + numTicks)
+	lowestTick := prec.TickFromIndex(i - numTicks)
+	ob := NewOrderBook()
+	for _, pool := range pools {
+		poolPrice := pool.Price()
+		if poolPrice.GT(lowestTick) { // Buy orders
+			startTick := sdk.MinDec(prec.DownTick(poolPrice), highestTick)
+			accAmt := sdk.ZeroInt()
+			for tick := startTick; tick.GTE(lowestTick); tick = prec.DownTick(tick) {
+				amt := pool.BuyAmountOver(tick).Sub(accAmt)
+				if amt.IsPositive() {
+					ob.Add(NewBaseOrder(Buy, tick, amt, sdk.Coin{}, "denom"))
+					accAmt = accAmt.Add(amt)
+				}
+			}
+		}
+		if poolPrice.LT(highestTick) { // Sell orders
+			startTick := sdk.MaxDec(prec.UpTick(poolPrice), lowestTick)
+			accAmt := sdk.ZeroInt()
+			for tick := startTick; tick.LTE(highestTick); tick = prec.UpTick(tick) {
+				amt := pool.SellAmountUnder(tick).Sub(accAmt)
+				if amt.IsPositive() {
+					ob.Add(NewBaseOrder(Sell, tick, amt, sdk.Coin{}, "denom"))
+					accAmt = accAmt.Add(amt)
+				}
+			}
+		}
+	}
+	return ob
 }
 
 // MockPoolOrderSource demonstrates how to implement a pool OrderSource.

x/liquidity/amm/pool_test.go

@@ -1,6 +1,7 @@
 package amm_test
 
 import (
+	"fmt"
 	"math/rand"
 	"testing"
 
@@ -16,7 +17,7 @@ func TestBasicPool(t *testing.T) {
 	r := rand.New(rand.NewSource(0))
 	for i := 0; i < 1000; i++ {
 		rx, ry := sdk.NewInt(1+r.Int63n(100000000)), sdk.NewInt(1+r.Int63n(100000000))
-		pool := amm.NewBasicPool(rx, ry, sdk.ZeroInt())
+		pool := amm.NewBasicPool(rx, ry, sdk.Int{})
 
 		highest, found := pool.HighestBuyPrice()
 		require.True(t, found)
@@ -326,7 +327,7 @@ func TestBasicPool_Withdraw(t *testing.T) {
 }
 
 func TestBasicPool_Amount(t *testing.T) {
-	pool := amm.NewBasicPool(sdk.NewInt(1000000), sdk.NewInt(1000000), sdk.ZeroInt())
+	pool := amm.NewBasicPool(sdk.NewInt(1000000), sdk.NewInt(1000000), sdk.Int{})
 	require.True(t, squad.DecApproxEqual(
 		squad.ParseDec("1000000"),
 		pool.BuyAmountOver(defTickPrec.LowestTick()).ToDec().Mul(defTickPrec.LowestTick()),
@@ -337,8 +338,33 @@ func TestBasicPool_Amount(t *testing.T) {
 	)
 }
 
+func ExamplePoolsOrderBook() {
+	pools := []amm.Pool{
+		amm.NewBasicPool(sdk.NewInt(1000000), sdk.NewInt(1000000), sdk.Int{}),
+	}
+	ob := amm.PoolsOrderBook(pools, squad.ParseDec("1.0"), 6, int(defTickPrec))
+	fmt.Println(ob.FullString(int(defTickPrec)))
+
+	// Output:
+	// +--------buy---------+------------price-------------+--------sell--------+
+	// |                  0 |         1.006000000000000000 | 989                |
+	// |                  0 |         1.005000000000000000 | 991                |
+	// |                  0 |         1.004000000000000000 | 993                |
+	// |                  0 |         1.003000000000000000 | 995                |
+	// |                  0 |         1.002000000000000000 | 997                |
+	// |                  0 |         1.001000000000000000 | 999                |
+	// |                  0 |         1.000000000000000000 | 0                  |
+	// |                100 |         0.999900000000000000 | 0                  |
+	// |                100 |         0.999800000000000000 | 0                  |
+	// |                100 |         0.999700000000000000 | 0                  |
+	// |                100 |         0.999600000000000000 | 0                  |
+	// |                100 |         0.999500000000000000 | 0                  |
+	// |                100 |         0.999400000000000000 | 0                  |
+	// +--------------------+------------------------------+--------------------+
+}
+
 func TestMockPoolOrderSource_Orders(t *testing.T) {
-	pool := amm.NewBasicPool(sdk.NewInt(1000000), sdk.NewInt(1000000), sdk.ZeroInt())
+	pool := amm.NewBasicPool(sdk.NewInt(1000000), sdk.NewInt(1000000), sdk.Int{})
 	os := amm.NewMockPoolOrderSource(pool, "denom1", "denom2")
 	buyOrders := os.BuyOrdersOver(defTickPrec.LowestTick())
 	require.Len(t, buyOrders, 1)

x/liquidity/simulation/operations.go

@@ -374,7 +374,7 @@ func SimulateMsgLimitOrder(ak types.AccountKeeper, bk types.BankKeeper, k keeper
 			minPrice, maxPrice = minMaxPrice(k, ctx, *pair.LastPrice)
 		} else {
 			rx, ry := k.GetPoolBalances(ctx, pool)
-			ammPool := amm.NewBasicPool(rx.Amount, ry.Amount, sdk.ZeroInt())
+			ammPool := amm.NewBasicPool(rx.Amount, ry.Amount, sdk.Int{})
 			minPrice, maxPrice = minMaxPrice(k, ctx, ammPool.Price())
 		}
 		price := amm.PriceToDownTick(squad.RandomDec(r, minPrice, maxPrice), int(params.TickPrecision))