fix(ProtoRev): Parameterizing Pool Type Information (#5948)

* init

* testing update

* changelog

* comments

* nit

* refactor

* clean up

* nit

* comment

* cr

* cl are not symetric

* init

* smh more symetric changes

* simpler safe swap

* map update

* more stuff

* testing fixes

* more testing

* nits

* benchmark + CLI clean up

* test fix

* no map no problems

* update with merge

* better docs

CHANGELOG.md

@@ -74,6 +74,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 * [#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
 * [#5964](https://github.com/osmosis-labs/osmosis/pull/5964) fix e2e test concurrency bugs
+* [#5948] (https://github.com/osmosis-labs/osmosis/pull/5948) Parameterizing Pool Type Information in Protorev
 
 ### Minor improvements & Bug Fixes
 

app/upgrades/v17/upgrades.go

@@ -123,12 +123,7 @@ func CreateUpgradeHandler(
 		keepers.DistrKeeper.SetFeePool(ctx, feePool)
 
 		// Reset the pool weights upon upgrade. This will add support for CW pools on ProtoRev.
-		keepers.ProtoRevKeeper.SetPoolWeights(ctx, types.PoolWeights{
-			BalancerWeight:     1,
-			StableWeight:       4,
-			ConcentratedWeight: 300,
-			CosmwasmWeight:     300,
-		})
+		keepers.ProtoRevKeeper.SetInfoByPoolType(ctx, types.DefaultPoolTypeInfo)
 
 		return migrations, nil
 	}

proto/osmosis/protorev/v1beta1/genesis.proto

@@ -27,6 +27,9 @@ message GenesisState {
   ];
   // The pool weights that are being used to calculate the weight (compute cost)
   // of each route.
+  //
+  // DEPRECATED: This field is deprecated and will be removed in the next
+  // release. It is replaced by the `info_by_pool_type` field.
   PoolWeights pool_weights = 4 [
     (gogoproto.nullable) = false,
     (gogoproto.moretags) = "yaml:\"pool_weights\""
@@ -63,4 +66,10 @@ message GenesisState {
     (gogoproto.nullable) = false,
     (gogoproto.moretags) = "yaml:\"profits\""
   ];
+  // Information that is used to estimate execution time / gas
+  // consumption of a swap on a given pool type.
+  InfoByPoolType info_by_pool_type = 13 [
+    (gogoproto.nullable) = false,
+    (gogoproto.moretags) = "yaml:\"info_by_pool_type\""
+  ];
 }

proto/osmosis/protorev/v1beta1/protorev.proto

@@ -78,6 +78,9 @@ message RouteStatistics {
 // significantly between the different pool types. Each weight roughly
 // corresponds to the amount of time (in ms) it takes to execute a swap on that
 // pool type.
+//
+// DEPRECATED: This field is deprecated and will be removed in the next
+// release. It is replaced by the `info_by_pool_type` field.
 message PoolWeights {
   // The weight of a stableswap pool
   uint64 stable_weight = 1 [ (gogoproto.moretags) = "yaml:\"stable_weight\"" ];
@@ -92,6 +95,74 @@ message PoolWeights {
       [ (gogoproto.moretags) = "yaml:\"cosmwasm_weight\"" ];
 }
 
+// InfoByPoolType contains information pertaining to how expensive (in terms of
+// gas and time) it is to execute a swap on a given pool type. This distinction
+// is made and necessary because the execution time ranges significantly between
+// the different pool types.
+message InfoByPoolType {
+  // The stable pool info
+  StablePoolInfo stable = 1 [
+    (gogoproto.moretags) = "yaml:\"stable\"",
+    (gogoproto.nullable) = false
+  ];
+  // The balancer pool info
+  BalancerPoolInfo balancer = 2 [
+    (gogoproto.moretags) = "yaml:\"balancer\"",
+    (gogoproto.nullable) = false
+  ];
+  // The concentrated pool info
+  ConcentratedPoolInfo concentrated = 3 [
+    (gogoproto.moretags) = "yaml:\"concentrated\"",
+    (gogoproto.nullable) = false
+  ];
+  // The cosmwasm pool info
+  CosmwasmPoolInfo cosmwasm = 4 [
+    (gogoproto.moretags) = "yaml:\"cosmwasm\"",
+    (gogoproto.nullable) = false
+  ];
+}
+
+// StablePoolInfo contains meta data pertaining to a stableswap pool type.
+message StablePoolInfo {
+  // The weight of a stableswap pool
+  uint64 weight = 1 [ (gogoproto.moretags) = "yaml:\"weight\"" ];
+}
+
+// BalancerPoolInfo contains meta data pertaining to a balancer pool type.
+message BalancerPoolInfo {
+  // The weight of a balancer pool
+  uint64 weight = 1 [ (gogoproto.moretags) = "yaml:\"weight\"" ];
+}
+
+// ConcentratedPoolInfo contains meta data pertaining to a concentrated pool
+// type.
+message ConcentratedPoolInfo {
+  // The weight of a concentrated pool
+  uint64 weight = 1 [ (gogoproto.moretags) = "yaml:\"weight\"" ];
+  // The maximum number of ticks we can move when rebalancing
+  uint64 max_ticks_crossed = 2
+      [ (gogoproto.moretags) = "yaml:\"max_ticks_crossed\"" ];
+}
+
+// CosmwasmPoolInfo contains meta data pertaining to a cosmwasm pool type.
+message CosmwasmPoolInfo {
+  // The weight of a cosmwasm pool (by contract address)
+  repeated WeightMap weight_maps = 1 [
+    (gogoproto.moretags) = "yaml:\"weight_maps\"",
+    (gogoproto.nullable) = false
+  ];
+}
+
+// WeightMap maps a contract address to a weight. The weight of an address
+// corresponds to the amount of ms required to execute a swap on that contract.
+message WeightMap {
+  // The weight of a cosmwasm pool (by contract address)
+  uint64 weight = 1 [ (gogoproto.moretags) = "yaml:\"weight\"" ];
+  // The contract address
+  string contract_address = 2
+      [ (gogoproto.moretags) = "yaml:\"contract_address\"" ];
+}
+
 // BaseDenom represents a single base denom that the module uses for its
 // arbitrage trades. It contains the denom name alongside the step size of the
 // binary search that is used to find the optimal swap amount

proto/osmosis/protorev/v1beta1/query.proto

@@ -71,11 +71,11 @@ service Query {
     option (google.api.http).get = "/osmosis/protorev/developer_account";
   }
 
-  // GetProtoRevPoolWeights queries the weights of each pool type currently
-  // being used by the module
-  rpc GetProtoRevPoolWeights(QueryGetProtoRevPoolWeightsRequest)
-      returns (QueryGetProtoRevPoolWeightsResponse) {
-    option (google.api.http).get = "/osmosis/protorev/pool_weights";
+  // GetProtoRevInfoByPoolType queries pool type information that is currently
+  // being utilized by the module
+  rpc GetProtoRevInfoByPoolType(QueryGetProtoRevInfoByPoolTypeRequest)
+      returns (QueryGetProtoRevInfoByPoolTypeResponse) {
+    option (google.api.http).get = "/osmosis/protorev/info_by_pool_type";
   }
 
   // GetProtoRevMaxPoolPointsPerTx queries the maximum number of pool points
@@ -242,16 +242,17 @@ message QueryGetProtoRevDeveloperAccountResponse {
       [ (gogoproto.moretags) = "yaml:\"developer_account\"" ];
 }
 
-// QueryGetProtoRevPoolWeightsRequest is request type for the
-// Query/GetProtoRevPoolWeights RPC method.
-message QueryGetProtoRevPoolWeightsRequest {}
-
-// QueryGetProtoRevPoolWeightsResponse is response type for the
-// Query/GetProtoRevPoolWeights RPC method.
-message QueryGetProtoRevPoolWeightsResponse {
-  // pool_weights is a list of all of the pool weights
-  PoolWeights pool_weights = 1 [
-    (gogoproto.moretags) = "yaml:\"pool_weights\"",
+// QueryGetProtoRevInfoByPoolTypeRequest is request type for the
+// Query/GetProtoRevInfoByPoolType RPC method.
+message QueryGetProtoRevInfoByPoolTypeRequest {}
+
+// QueryGetProtoRevInfoByPoolTypeResponse is response type for the
+// Query/GetProtoRevInfoByPoolType RPC method.
+message QueryGetProtoRevInfoByPoolTypeResponse {
+  // InfoByPoolType contains all information pertaining to how different
+  // pool types are handled by the module.
+  InfoByPoolType info_by_pool_type = 1 [
+    (gogoproto.moretags) = "yaml:\"info_by_pool_type\"",
     (gogoproto.nullable) = false
   ];
 }

proto/osmosis/protorev/v1beta1/tx.proto

@@ -39,10 +39,11 @@ service Msg {
         "/osmosis/protorev/set_max_pool_points_per_block";
   };
 
-  // SetPoolWeights sets the weights of each pool type in the store. Can only be
-  // called by the admin account.
-  rpc SetPoolWeights(MsgSetPoolWeights) returns (MsgSetPoolWeightsResponse) {
-    option (google.api.http).post = "/osmosis/protorev/set_pool_weights";
+  // SetInfoByPoolType sets the pool type information needed to make smart
+  // assumptions about swapping on different pool types
+  rpc SetInfoByPoolType(MsgSetInfoByPoolType)
+      returns (MsgSetInfoByPoolTypeResponse) {
+    option (google.api.http).post = "/osmosis/protorev/set_info_by_pool_type";
   };
 
   // SetBaseDenoms sets the base denoms that will be used to create cyclic
@@ -90,24 +91,24 @@ message MsgSetDeveloperAccount {
 // type.
 message MsgSetDeveloperAccountResponse {}
 
-// MsgSetPoolWeights defines the Msg/SetPoolWeights request type.
-message MsgSetPoolWeights {
-  option (amino.name) = "osmosis/MsgSetPoolWeights";
+// MsgSetInfoByPoolType defines the Msg/SetInfoByPoolType request type.
+message MsgSetInfoByPoolType {
+  option (amino.name) = "osmosis/MsgSetInfoByPoolType";
 
   // admin is the account that is authorized to set the pool weights.
   string admin = 1 [
     (gogoproto.moretags) = "yaml:\"admin\"",
     (cosmos_proto.scalar) = "cosmos.AddressString"
   ];
-  // pool_weights is the list of pool weights to set.
-  PoolWeights pool_weights = 2 [
-    (gogoproto.moretags) = "yaml:\"pool_weights\"",
+  // info_by_pool_type contains information about the pool types.
+  InfoByPoolType info_by_pool_type = 2 [
+    (gogoproto.moretags) = "yaml:\"info_by_pool_type\"",
     (gogoproto.nullable) = false
   ];
 }
 
-// MsgSetPoolWeightsResponse defines the Msg/SetPoolWeights response type.
-message MsgSetPoolWeightsResponse {}
+// MsgSetInfoByPoolTypeResponse defines the Msg/SetInfoByPoolType response type.
+message MsgSetInfoByPoolTypeResponse {}
 
 // MsgSetMaxPoolPointsPerTx defines the Msg/SetMaxPoolPointsPerTx request type.
 message MsgSetMaxPoolPointsPerTx {

tests/e2e/configurer/chain/queries.go

@@ -118,20 +118,20 @@ func (n *NodeConfig) QueryProtoRevDeveloperAccount() (sdk.AccAddress, error) {
 	return account, nil
 }
 
-// QueryProtoRevPoolWeights gets the pool point weights of the module.
-func (n *NodeConfig) QueryProtoRevPoolWeights() (protorevtypes.PoolWeights, error) {
-	path := "/osmosis/protorev/pool_weights"
+// QueryProtoRevInfoByPoolType gets information on how the module handles different pool types.
+func (n *NodeConfig) QueryProtoRevInfoByPoolType() (*protorevtypes.InfoByPoolType, error) {
+	path := "/osmosis/protorev/info_by_pool_type"
 
 	bz, err := n.QueryGRPCGateway(path)
 	if err != nil {
-		return protorevtypes.PoolWeights{}, err
+		return nil, err
 	}
 
 	// nolint: staticcheck
-	var response protorevtypes.QueryGetProtoRevPoolWeightsResponse
+	var response protorevtypes.QueryGetProtoRevInfoByPoolTypeResponse
 	err = util.Cdc.UnmarshalJSON(bz, &response)
 	require.NoError(n.t, err) // this error should not happen
-	return response.PoolWeights, nil
+	return &response.InfoByPoolType, nil
 }
 
 // QueryProtoRevMaxPoolPointsPerTx gets the max pool points per tx of the module.

tests/e2e/e2e_test.go

@@ -222,10 +222,10 @@ func (s *IntegrationTestSuite) ProtoRev() {
 	s.Require().Error(err)
 
 	// The module should have pool weights by default.
-	poolWeights, err := chainANode.QueryProtoRevPoolWeights()
-	s.T().Logf("checking that the protorev module has pool weights on init: %v", poolWeights)
+	info, err := chainANode.QueryProtoRevInfoByPoolType()
+	s.T().Logf("checking that the protorev module has pool info on init: %v", info)
 	s.Require().NoError(err)
-	s.Require().NotNil(poolWeights)
+	s.Require().NotNil(info)
 
 	// The module should have max pool points per tx by default.
 	maxPoolPointsPerTx, err := chainANode.QueryProtoRevMaxPoolPointsPerTx()

x/protorev/client/cli/query.go

@@ -28,7 +28,7 @@ func NewCmdQuery() *cobra.Command {
 	osmocli.AddQueryCmd(cmd, types.NewQueryClient, NewQueryMaxPoolPointsPerBlockCmd)
 	osmocli.AddQueryCmd(cmd, types.NewQueryClient, NewQueryBaseDenomsCmd)
 	osmocli.AddQueryCmd(cmd, types.NewQueryClient, NewQueryEnabledCmd)
-	osmocli.AddQueryCmd(cmd, types.NewQueryClient, NewQueryPoolWeightsCmd)
+	osmocli.AddQueryCmd(cmd, types.NewQueryClient, NewQueryInfoByPoolTypeCmd)
 	osmocli.AddQueryCmd(cmd, types.NewQueryClient, NewQueryPoolCmd)
 
 	return cmd
@@ -141,12 +141,12 @@ func NewQueryEnabledCmd() (*osmocli.QueryDescriptor, *types.QueryGetProtoRevEnab
 	}, &types.QueryGetProtoRevEnabledRequest{}
 }
 
-// NewQueryPoolWeightsCmd returns the command to query the pool weights of protorev
-func NewQueryPoolWeightsCmd() (*osmocli.QueryDescriptor, *types.QueryGetProtoRevPoolWeightsRequest) {
+// NewQueryInfoByPoolTypeCmd returns the command to query the pool type info of protorev
+func NewQueryInfoByPoolTypeCmd() (*osmocli.QueryDescriptor, *types.QueryGetProtoRevInfoByPoolTypeRequest) {
 	return &osmocli.QueryDescriptor{
-		Use:   "pool-weights",
-		Short: "Query the pool weights used to determine how computationally expensive a route is",
-	}, &types.QueryGetProtoRevPoolWeightsRequest{}
+		Use:   "info-by-pool-type",
+		Short: "Query the pool info used to determine how computationally expensive a route is",
+	}, &types.QueryGetProtoRevInfoByPoolTypeRequest{}
 }
 
 // NewQueryPoolCmd returns the command to query the pool id for a given denom pair stored via the highest liquidity method in ProtoRev

x/protorev/client/cli/tx.go

@@ -29,7 +29,7 @@ func NewCmdTx() *cobra.Command {
 	osmocli.AddTxCmd(txCmd, CmdSetMaxPoolPointsPerBlock)
 	txCmd.AddCommand(
 		CmdSetDeveloperHotRoutes().BuildCommandCustomFn(),
-		CmdSetPoolWeights().BuildCommandCustomFn(),
+		CmdSetInfoByPoolType().BuildCommandCustomFn(),
 		CmdSetBaseDenoms().BuildCommandCustomFn(),
 		CmdSetProtoRevAdminAccountProposal(),
 		CmdSetProtoRevEnabledProposal(),
@@ -141,22 +141,35 @@ func CmdSetMaxPoolPointsPerBlock() (*osmocli.TxCliDesc, *types.MsgSetMaxPoolPoin
 	}, &types.MsgSetMaxPoolPointsPerBlock{}
 }
 
-// CmdSetPoolWeights implements the command to set the pool weights used to estimate execution costs
-func CmdSetPoolWeights() *osmocli.TxCliDesc {
+// CmdSetInfoByPoolType implements the command to set the pool information used throughout the module
+func CmdSetInfoByPoolType() *osmocli.TxCliDesc {
 	desc := osmocli.TxCliDesc{
-		Use:   "set-pool-weights [path/to/routes.json]",
-		Short: "set the protorev pool weights",
-		Long: `Must provide a json file with all the pool weights that will be set. 
+		Use:   "set-info-by-pool-type [path/to/pool_info.json]",
+		Short: "set the protorev pool type info",
+		Long: `Must provide a json file with all the pool info that will be set. This does NOT set info for a single pool type.
+		All information must be provided across all pool types in the json file.
 		Sample json file:
 		{
-			"stable_weight" : 1,
-			"balancer_weight" : 1,
-			"concentrated_weight" : 1
+			"stable" : {
+				"weight" : 1,
+			},
+			"concentrated" : {
+				"weight" : 1,
+				"max_ticks_crossed": 10,
+			},
+			"balancer" : {
+				"weight" : 1,
+			},
+			"cosmwasm" : {
+				"weight_maps" : [
+					{"contract_address" : "cosmos123...", "weight" : 1}
+				],
+			},
 		}
 		`,
-		Example:          fmt.Sprintf(`$ %s tx protorev set-pool-weights weights.json --from mykey`, version.AppName),
+		Example:          fmt.Sprintf(`$ %s tx protorev set-info-by-pool-type pool_info.json --from mykey`, version.AppName),
 		NumArgs:          1,
-		ParseAndBuildMsg: BuildSetPoolWeightsMsg,
+		ParseAndBuildMsg: BuildSetInfoByPoolTypeMsg,
 	}
 
 	return &desc