feat(ProtoRev): Supporting Two Pool Routes

CHANGELOG.md

@@ -70,6 +70,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 * [#5923] (https://github.com/osmosis-labs/osmosis/pull/5923) CL: Lower gas for initializing ticks
 * [#5927] (https://github.com/osmosis-labs/osmosis/pull/5927) Add gas metering to x/tokenfactory trackBeforeSend hook
 * [#5890](https://github.com/osmosis-labs/osmosis/pull/5890) feat: CreateCLPool & LinkCFMMtoCL pool into one gov-prop
+* [#5953] (https://github.com/osmosis-labs/osmosis/pull/5953) Supporting two pool routes in ProtoRev
 
 ### Minor improvements & Bug Fixes
 

x/protorev/keeper/keeper_test.go

@@ -142,6 +142,7 @@ func (s *KeeperTestSuite) SetupTest() {
 		sdk.NewCoin("gamm/pool/1", sdk.NewInt(9000000000000000000)),
 		sdk.NewCoin(apptesting.DefaultTransmuterDenomA, sdk.NewInt(9000000000000000000)),
 		sdk.NewCoin(apptesting.DefaultTransmuterDenomB, sdk.NewInt(9000000000000000000)),
+		sdk.NewCoin("stake", sdk.NewInt(9000000000000000000)),
 	)
 	s.fundAllAccountsWith()
 	s.Commit()
@@ -902,6 +903,91 @@ func (s *KeeperTestSuite) setUpPools() {
 	// Pool 50
 	s.PrepareCosmWasmPool()
 
+	// Create a duplicate pool for testing
+	// Pool 51
+	s.createGAMMPool(
+		[]balancer.PoolAsset{
+			{
+				Token:  sdk.NewCoin("Atom", sdk.NewInt(10000)),
+				Weight: sdk.NewInt(1),
+			},
+			{
+				Token:  sdk.NewCoin(types.OsmosisDenomination, sdk.NewInt(10000)),
+				Weight: sdk.NewInt(1),
+			},
+		},
+		sdk.NewDecWithPrec(2, 3),
+		sdk.NewDecWithPrec(0, 2),
+	)
+
+	// Create a duplicate pool for testing
+	// Pool 52
+	s.createGAMMPool(
+		[]balancer.PoolAsset{
+			{
+				Token:  sdk.NewCoin("usdc", sdk.NewInt(10000)),
+				Weight: sdk.NewInt(1),
+			},
+			{
+				Token:  sdk.NewCoin(types.OsmosisDenomination, sdk.NewInt(10000)),
+				Weight: sdk.NewInt(1),
+			},
+		},
+		sdk.NewDecWithPrec(2, 3),
+		sdk.NewDecWithPrec(0, 2),
+	)
+
+	// Create a duplicate pool for testing
+	// Pool 53
+	s.createGAMMPool(
+		[]balancer.PoolAsset{
+			{
+				Token:  sdk.NewCoin("stake", sdk.NewInt(10000)),
+				Weight: sdk.NewInt(1),
+			},
+			{
+				Token:  sdk.NewCoin(types.OsmosisDenomination, sdk.NewInt(10000)),
+				Weight: sdk.NewInt(1),
+			},
+		},
+		sdk.NewDecWithPrec(2, 3),
+		sdk.NewDecWithPrec(0, 2),
+	)
+
+	// Create a duplicate pool for testing
+	// Pool 54
+	s.createGAMMPool(
+		[]balancer.PoolAsset{
+			{
+				Token:  sdk.NewCoin("stake", sdk.NewInt(100000000)),
+				Weight: sdk.NewInt(1),
+			},
+			{
+				Token:  sdk.NewCoin(types.OsmosisDenomination, sdk.NewInt(1000000000)),
+				Weight: sdk.NewInt(1),
+			},
+		},
+		sdk.NewDecWithPrec(2, 3),
+		sdk.NewDecWithPrec(0, 2),
+	)
+
+	// Create a duplicate pool for testing
+	// Pool 55
+	s.createGAMMPool(
+		[]balancer.PoolAsset{
+			{
+				Token:  sdk.NewCoin("bitcoin", sdk.NewInt(100)),
+				Weight: sdk.NewInt(1),
+			},
+			{
+				Token:  sdk.NewCoin("Atom", sdk.NewInt(100)),
+				Weight: sdk.NewInt(1),
+			},
+		},
+		sdk.NewDecWithPrec(2, 3),
+		sdk.NewDecWithPrec(0, 2),
+	)
+
 	// Set all of the pool info into the stores
 	err := s.App.ProtoRevKeeper.UpdatePools(s.Ctx)
 	s.Require().NoError(err)

x/protorev/keeper/posthandler_test.go

@@ -340,7 +340,7 @@ func (s *KeeperTestSuite) TestAnteHandle() {
 			params: param{
 				trades: []types.Trade{
 					{
-						Pool:     51,
+						Pool:     56,
 						TokenOut: "ibc/0CD3A0285E1341859B5E86B6AB7682F023D03E97607CCC1DC95706411D866DF7",
 						TokenIn:  "uosmo",
 					},

x/protorev/keeper/protorev.go

@@ -118,7 +118,7 @@ func (k Keeper) DeleteBaseDenoms(ctx sdk.Context) {
 	k.DeleteAllEntriesForKeyPrefix(ctx, types.KeyPrefixBaseDenoms)
 }
 
-// GetPoolForDenomPair returns the id of the highest liquidty pool between the base denom and the denom to match
+// GetPoolForDenomPair returns the id of the highest liquidity pool between the base denom and the denom to match
 func (k Keeper) GetPoolForDenomPair(ctx sdk.Context, baseDenom, denomToMatch string) (uint64, error) {
 	store := prefix.NewStore(ctx.KVStore(k.storeKey), types.KeyPrefixDenomPairToPool)
 	key := types.GetKeyPrefixDenomPairToPool(baseDenom, denomToMatch)

x/protorev/keeper/routes.go

@@ -102,6 +102,10 @@ func (k Keeper) BuildHighestLiquidityRoutes(ctx sdk.Context, tokenIn, tokenOut s
 		if newRoute, err := k.BuildHighestLiquidityRoute(ctx, baseDenom, tokenIn, tokenOut, poolId); err == nil {
 			routes = append(routes, newRoute)
 		}
+
+		if newRoute, err := k.BuildTwoPoolRoute(ctx, baseDenom, tokenIn, tokenOut, poolId); err == nil {
+			routes = append(routes, newRoute)
+		}
 	}
 
 	return routes, nil
@@ -149,6 +153,51 @@ func (k Keeper) BuildHighestLiquidityRoute(ctx sdk.Context, swapDenom types.Base
 	}, nil
 }
 
+// BuildTwoPoolRoute will attempt to create a two pool route that will rebalance pools that are paired with the base denom.
+// This is useful for pools that contain the same assets but are imbalanced.
+func (k Keeper) BuildTwoPoolRoute(ctx sdk.Context, baseDenom types.BaseDenom, tokenIn, tokenOut string, poolId uint64) (RouteMetaData, error) {
+	if baseDenom.Denom != tokenOut {
+		return RouteMetaData{}, fmt.Errorf("the token out denom must be the same as the base denom: got %s, expected %s", tokenOut, baseDenom.Denom)
+	}
+
+	// Get the highest liquidity pool for the tokenIn and base denom
+	highestLiquidityPool, err := k.GetPoolForDenomPair(ctx, baseDenom.Denom, tokenIn)
+	if err != nil {
+		return RouteMetaData{}, err
+	}
+
+	// If the swapped on pool is already the pool with the highest liquidity, we cannot build a two pool route (it would be a two hop route with the same pool)
+	if highestLiquidityPool == poolId {
+		return RouteMetaData{}, fmt.Errorf("the pool id must be different from the highest liquidity pool: id %d", poolId)
+	}
+
+	// Create the first swap for the MultiHopSwap Route
+	entryHop := poolmanagertypes.SwapAmountInRoute{
+		PoolId:        highestLiquidityPool,
+		TokenOutDenom: tokenIn,
+	}
+
+	// Creating the second swap in the arb
+	exitHop := poolmanagertypes.SwapAmountInRoute{
+		PoolId:        poolId,
+		TokenOutDenom: baseDenom.Denom,
+	}
+
+	newRoute := poolmanagertypes.SwapAmountInRoutes{entryHop, exitHop}
+
+	// Check that the route is valid and update the number of pool points that this route will consume when simulating and executing trades
+	routePoolPoints, err := k.CalculateRoutePoolPoints(ctx, newRoute)
+	if err != nil {
+		return RouteMetaData{}, err
+	}
+
+	return RouteMetaData{
+		Route:      newRoute,
+		PoolPoints: routePoolPoints,
+		StepSize:   baseDenom.StepSize,
+	}, nil
+}
+
 // CalculateRoutePoolPoints calculates the number of pool points that will be consumed by a route when simulating and executing trades. This
 // is only added to the global pool point counter if the route simulated is minimally profitable i.e. it will make a profit.
 func (k Keeper) CalculateRoutePoolPoints(ctx sdk.Context, route poolmanagertypes.SwapAmountInRoutes) (uint64, error) {

x/protorev/keeper/routes_test.go

@@ -44,13 +44,17 @@ func (s *KeeperTestSuite) TestBuildRoutes() {
 			description: "Route exists for swap in Bitcoin and swap out Atom",
 			inputDenom:  "bitcoin",
 			outputDenom: "Atom",
-			poolID:      4,
+			poolID:      55,
 			expectedRoutes: [][]TestRoute{
 				{
 					{PoolId: 25, InputDenom: types.OsmosisDenomination, OutputDenom: "Atom"},
-					{PoolId: 4, InputDenom: "Atom", OutputDenom: "bitcoin"},
+					{PoolId: 55, InputDenom: "Atom", OutputDenom: "bitcoin"},
 					{PoolId: 10, InputDenom: "bitcoin", OutputDenom: types.OsmosisDenomination},
 				},
+				{
+					{PoolId: 4, InputDenom: "Atom", OutputDenom: "bitcoin"},
+					{PoolId: 55, InputDenom: "bitcoin", OutputDenom: "Atom"},
+				},
 			},
 		},
 		{
@@ -91,6 +95,18 @@ func (s *KeeperTestSuite) TestBuildRoutes() {
 				},
 			},
 		},
+		{
+			description: "Two Pool Route exists for (osmo, atom)",
+			inputDenom:  "Atom",
+			outputDenom: types.OsmosisDenomination,
+			poolID:      51,
+			expectedRoutes: [][]TestRoute{
+				{
+					{PoolId: 25, InputDenom: types.OsmosisDenomination, OutputDenom: "Atom"},
+					{PoolId: 51, InputDenom: "Atom", OutputDenom: types.OsmosisDenomination},
+				},
+			},
+		},
 	}
 
 	for _, tc := range cases {
@@ -225,6 +241,77 @@ func (s *KeeperTestSuite) TestBuildHighestLiquidityRoute() {
 	}
 }
 
+// TestBuildTwoPoolRoute tests the BuildTwoPoolRoute function
+func (s *KeeperTestSuite) TestBuildTwoPoolRoute() {
+	cases := []struct {
+		description   string
+		swapDenom     types.BaseDenom
+		tokenIn       string
+		tokenOut      string
+		poolId        uint64
+		expectedRoute []TestRoute
+		hasRoute      bool
+	}{
+		{
+			description: "two pool route can be created",
+			swapDenom: types.BaseDenom{
+				Denom:    types.OsmosisDenomination,
+				StepSize: sdk.NewInt(1_000_000),
+			},
+			tokenIn:  "stake",
+			tokenOut: types.OsmosisDenomination,
+			poolId:   53,
+			expectedRoute: []TestRoute{
+				{PoolId: 54, InputDenom: types.OsmosisDenomination, OutputDenom: "stake"},
+				{PoolId: 53, InputDenom: "stake", OutputDenom: types.OsmosisDenomination},
+			},
+			hasRoute: true,
+		},
+		{
+			description: "two pool route where swap is on the highest liquidity pool",
+			swapDenom: types.BaseDenom{
+				Denom:    types.OsmosisDenomination,
+				StepSize: sdk.NewInt(1_000_000),
+			},
+			tokenIn:       "stake",
+			tokenOut:      types.OsmosisDenomination,
+			poolId:        54,
+			expectedRoute: []TestRoute{},
+			hasRoute:      false,
+		},
+		{
+			description: "two pool route where swap in is the base denom",
+			swapDenom: types.BaseDenom{
+				Denom:    types.OsmosisDenomination,
+				StepSize: sdk.NewInt(1_000_000),
+			},
+			tokenIn:       types.OsmosisDenomination,
+			tokenOut:      "stake",
+			poolId:        53,
+			expectedRoute: []TestRoute{},
+			hasRoute:      false,
+		},
+	}
+
+	for _, tc := range cases {
+		s.Run(tc.description, func() {
+			routeMetaData, err := s.App.ProtoRevKeeper.BuildTwoPoolRoute(s.Ctx, tc.swapDenom, tc.tokenIn, tc.tokenOut, tc.poolId)
+
+			if tc.hasRoute {
+				s.Require().NoError(err)
+				s.Require().Equal(len(tc.expectedRoute), len(routeMetaData.Route.PoolIds()))
+
+				for index, trade := range tc.expectedRoute {
+					s.Require().Equal(trade.PoolId, routeMetaData.Route.PoolIds()[index])
+				}
+			} else {
+				s.Require().Equal(len(tc.expectedRoute), len(routeMetaData.Route.PoolIds()))
+				s.Require().Error(err)
+			}
+		})
+	}
+}
+
 // TestBuildHotRoutes tests the BuildHotRoutes function
 func (s *KeeperTestSuite) TestBuildHotRoutes() {
 	cases := []struct {