Autoswap accounting fixes #NTRN-385

app/app.go

@@ -226,7 +226,9 @@ const (
 )
 
 var (
-	Upgrades = []upgrades.Upgrade{v500.Upgrade}
+	Upgrades = []upgrades.Upgrade{
+		v500.Upgrade,
+	}
 
 	// DefaultNodeHome default home directories for the application daemon
 	DefaultNodeHome string

app/upgrades/v5.0.0/upgrades.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 
+	"cosmossdk.io/math"
 	upgradetypes "cosmossdk.io/x/upgrade/types"
 	adminmoduletypes "github.com/cosmos/admin-module/v2/x/adminmodule/types"
 	"github.com/cosmos/cosmos-sdk/codec"
@@ -44,6 +45,11 @@ func CreateUpgradeHandler(
 			}
 		}
 
+		err = upgradePools(ctx, *keepers.DexKeeper)
+		if err != nil {
+			return nil, err
+		}
+
 		err = setMarketMapParams(ctx, keepers.MarketmapKeeper)
 		if err != nil {
 			return nil, err
@@ -77,6 +83,48 @@ func upgradeDexPause(ctx sdk.Context, k dexkeeper.Keeper) error {
 	return nil
 }
 
+func upgradePools(ctx sdk.Context, k dexkeeper.Keeper) error {
+	// Due to an issue with autoswap logic any pools with multiple shareholders must be withdrawn to ensure correct accounting
+	ctx.Logger().Info("Migrating Pools...")
+
+	allSharesholders := k.GetAllPoolShareholders(ctx)
+
+	for poolID, shareholders := range allSharesholders {
+		if len(shareholders) > 1 {
+			pool, found := k.GetPoolByID(ctx, poolID)
+			if !found {
+				return fmt.Errorf("cannot find pool with ID %d", poolID)
+			}
+			for _, shareholder := range shareholders {
+				addr := sdk.MustAccAddressFromBech32(shareholder.Address)
+				pairID := pool.LowerTick0.Key.TradePairId.MustPairID()
+				tick := pool.CenterTickIndexToken1()
+				fee := pool.Fee()
+				nShares := shareholder.Shares
+
+				reserve0Removed, reserve1Removed, sharesBurned, err := k.WithdrawCore(ctx, pairID, addr, addr, []math.Int{nShares}, []int64{tick}, []uint64{fee})
+				if err != nil {
+					return fmt.Errorf("user %s failed to withdraw from pool %d", addr, poolID)
+				}
+
+				ctx.Logger().Info(
+					"Withdrew user from pool",
+					"User", addr.String(),
+					"Pool", poolID,
+					"SharesBurned", sharesBurned.String(),
+					"Reserve0Withdrawn", reserve0Removed.String(),
+					"Reserve1Withdrawn", reserve1Removed.String(),
+				)
+
+			}
+		}
+	}
+
+	ctx.Logger().Info("Finished migrating Pools...")
+
+	return nil
+}
+
 func upgradeIbcRateLimitSetContract(ctx sdk.Context, k ibcratelimitkeeper.Keeper) error {
 	// Set the dex to paused
 	ctx.Logger().Info("Setting ibc rate limiting contract...")

app/upgrades/v5.0.0/upgrades_test.go

@@ -6,6 +6,8 @@ import (
 
 	"cosmossdk.io/math"
 	upgradetypes "cosmossdk.io/x/upgrade/types"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	bankkeeper "github.com/cosmos/cosmos-sdk/x/bank/keeper"
 	"github.com/stretchr/testify/require"
 	"github.com/stretchr/testify/suite"
 
@@ -85,6 +87,112 @@ func (suite *UpgradeTestSuite) TestUpgradeDexPause() {
 	suite.ErrorIs(err, dextypes.ErrDexPaused)
 }
 
+func (suite *UpgradeTestSuite) TestPoolMigrationSingleShareHolder() {
+	var (
+		app           = suite.GetNeutronZoneApp(suite.ChainA)
+		ctx           = suite.ChainA.GetContext().WithChainID("neutron-1")
+		alice         = []byte("alice")
+		pairID        = &dextypes.PairID{Token0: "TokenA", Token1: "TokenB"}
+		depositAmount = math.NewInt(10_000)
+	)
+
+	// create a pool with 1 shareholder
+	FundAccount(app.BankKeeper, ctx, alice, sdk.NewCoins(sdk.NewCoin("TokenA", depositAmount)))
+	shares, err := suite.makeDeposit(ctx, app.DexKeeper, alice, pairID, depositAmount, math.ZeroInt(), 0, 1)
+	suite.NoError(err)
+
+	// run upgrade
+	upgrade := upgradetypes.Plan{
+		Name:   v500.UpgradeName,
+		Info:   "some text here",
+		Height: 100,
+	}
+	suite.NoError(app.UpgradeKeeper.ApplyUpgrade(ctx, upgrade))
+
+	// assert pool and shareholder balance are unchanged
+	poolID, err := dextypes.ParsePoolIDFromDenom(shares[0].Denom)
+	suite.NoError(err)
+
+	pool, _ := app.DexKeeper.GetPoolByID(ctx, poolID)
+
+	suite.True(pool.LowerTick0.ReservesMakerDenom.Equal(depositAmount), "Pool value changed")
+	aliceBalance := app.BankKeeper.GetAllBalances(ctx, alice)
+	suite.True(aliceBalance.Equal(shares))
+}
+
+func (suite *UpgradeTestSuite) TestPoolMigrationMultiShareHolder() {
+	var (
+		app            = suite.GetNeutronZoneApp(suite.ChainA)
+		ctx            = suite.ChainA.GetContext().WithChainID("neutron-1")
+		alice          = []byte("alice")
+		bob            = []byte("bob")
+		pairID         = &dextypes.PairID{Token0: "TokenA", Token1: "TokenB"}
+		depositAmount  = math.NewInt(10_000)
+		initialBalance = sdk.NewCoins(sdk.NewCoin("TokenA", depositAmount))
+	)
+	FundAccount(app.BankKeeper, ctx, alice, initialBalance)
+	FundAccount(app.BankKeeper, ctx, bob, initialBalance)
+
+	// create a pool with 2 shareholders
+	shares, err := suite.makeDeposit(ctx, app.DexKeeper, alice, pairID, depositAmount, math.ZeroInt(), 0, 1)
+	suite.NoError(err)
+	aliceBalance := app.BankKeeper.GetAllBalances(ctx, alice)
+	suite.True(aliceBalance.Equal(shares))
+
+	shares, err = suite.makeDeposit(ctx, app.DexKeeper, bob, pairID, depositAmount, math.ZeroInt(), 0, 1)
+	suite.NoError(err)
+	bobBalance := app.BankKeeper.GetAllBalances(ctx, bob)
+	suite.True(bobBalance.Equal(shares))
+
+	// run upgrade
+	upgrade := upgradetypes.Plan{
+		Name:   v500.UpgradeName,
+		Info:   "some text here",
+		Height: 100,
+	}
+	suite.NoError(app.UpgradeKeeper.ApplyUpgrade(ctx, upgrade))
+
+	// assert that all users have withdrawn from the pool
+	poolID, err := dextypes.ParsePoolIDFromDenom(shares[0].Denom)
+	suite.NoError(err)
+
+	pool, _ := app.DexKeeper.GetPoolByID(ctx, poolID)
+	suite.True(pool.LowerTick0.ReservesMakerDenom.Equal(math.ZeroInt()), "Pool not withdrawn")
+
+	// AND funds are returned to the users
+	aliceBalance = app.BankKeeper.GetAllBalances(ctx, alice)
+	suite.True(aliceBalance.Equal(initialBalance))
+
+	bobBalance = app.BankKeeper.GetAllBalances(ctx, bob)
+	suite.True(bobBalance.Equal(initialBalance))
+}
+
+func FundAccount(bankKeeper bankkeeper.Keeper, ctx sdk.Context, addr sdk.AccAddress, amounts sdk.Coins) {
+	if err := bankKeeper.MintCoins(ctx, dextypes.ModuleName, amounts); err != nil {
+		panic(err)
+	}
+
+	if err := bankKeeper.SendCoinsFromModuleToAccount(ctx, dextypes.ModuleName, addr, amounts); err != nil {
+		panic(err)
+	}
+}
+
+func (suite *UpgradeTestSuite) makeDeposit(
+	ctx sdk.Context,
+	k dexkeeper.Keeper,
+	addr sdk.AccAddress,
+	pairID *dextypes.PairID,
+	amount0, amount1 math.Int,
+	tick int64,
+	fee uint64,
+) (sharesIssued sdk.Coins, err error) {
+	deposit0, deposit1, sharesIssued, _, err := k.DepositCore(ctx, pairID, addr, addr, []math.Int{amount0}, []math.Int{amount1}, []int64{tick}, []uint64{fee}, []*dextypes.DepositOptions{{}})
+	suite.True(deposit0[0].Equal(amount0))
+	suite.True(deposit1[0].Equal(amount1))
+
+	return sharesIssued, err
+}
+
 func (suite *UpgradeTestSuite) TestUpgradeSetRateLimitContractMainnet() {
 	var (
 		app = suite.GetNeutronZoneApp(suite.ChainA)

x/dex/keeper/core_helper_test.go

@@ -10,6 +10,7 @@ import (
 
 	neutronapp "github.com/neutron-org/neutron/v5/app"
 	"github.com/neutron-org/neutron/v5/testutil"
+	math_utils "github.com/neutron-org/neutron/v5/utils/math"
 	"github.com/neutron-org/neutron/v5/x/dex/types"
 )
 
@@ -65,7 +66,8 @@ func (s *CoreHelpersTestSuite) setLPAtFee1Pool(tickIndex int64, amountA, amountB
 
 	existingShares := s.app.BankKeeper.GetSupply(s.ctx, pool.GetPoolDenom()).Amount
 
-	totalShares := pool.CalcSharesMinted(amountAInt, amountBInt, existingShares)
+	depositAmountAsToken0 := types.CalcAmountAsToken0(amountAInt, amountBInt, pool.MustCalcPrice1To0Center())
+	totalShares := pool.CalcSharesMinted(depositAmountAsToken0, existingShares, math_utils.ZeroPrecDec())
 
 	err = s.app.DexKeeper.MintShares(s.ctx, s.alice, sdk.NewCoins(totalShares))
 	s.Require().NoError(err)

x/dex/keeper/grpc_query_pool_test.go

@@ -25,7 +25,7 @@ func TestPoolQuerySingle(t *testing.T) {
 			desc: "First",
 			request: &types.QueryPoolRequest{
 				PairId:    "TokenA<>TokenB",
-				TickIndex: msgs[0].CenterTickIndex(),
+				TickIndex: msgs[0].CenterTickIndexToken1(),
 				Fee:       msgs[0].Fee(),
 			},
 			response: &types.QueryPoolResponse{Pool: msgs[0]},
@@ -34,7 +34,7 @@ func TestPoolQuerySingle(t *testing.T) {
 			desc: "Second",
 			request: &types.QueryPoolRequest{
 				PairId:    "TokenA<>TokenB",
-				TickIndex: msgs[1].CenterTickIndex(),
+				TickIndex: msgs[1].CenterTickIndexToken1(),
 				Fee:       msgs[1].Fee(),
 			},
 			response: &types.QueryPoolResponse{Pool: msgs[1]},