Refactor estimate swap (#11)

* Add type conversions

* Refactor EstimatePrice to use same logic as SwapMsg

app/wasm/bindings/query.go

@@ -42,6 +42,14 @@ type EstimatePrice struct {
 	Amount   SwapAmount `json:"amount"`
 }
 
+func (e *EstimatePrice) ToSwapMsg() *SwapMsg {
+	return &SwapMsg{
+		First:  e.First,
+		Route:  e.Route,
+		Amount: e.Amount.Unlimited(),
+	}
+}
+
 type FullDenomResponse struct {
 	Denom string `json:"denom"`
 }

app/wasm/bindings/types.go

@@ -1,6 +1,10 @@
 package wasmbindings
 
-import sdk "github.com/cosmos/cosmos-sdk/types"
+import (
+	"math"
+
+	sdk "github.com/cosmos/cosmos-sdk/types"
+)
 
 type Swap struct {
 	PoolId   uint64 `json:"pool_id"`
@@ -18,11 +22,47 @@ type SwapAmount struct {
 	Out *sdk.Int `json:"out,omitempty"`
 }
 
+// This returns SwapAmountWithLimit with the largest possible limits (that will never be hit)
+func (s SwapAmount) Unlimited() SwapAmountWithLimit {
+	if s.In != nil {
+		return SwapAmountWithLimit{
+			ExactIn: &ExactIn{
+				Input:     *s.In,
+				MinOutput: sdk.NewInt(1),
+			},
+		}
+	}
+	if s.Out != nil {
+		return SwapAmountWithLimit{
+			ExactOut: &ExactOut{
+				Output:   *s.Out,
+				MaxInput: sdk.NewInt(math.MaxInt64),
+			},
+		}
+	}
+	panic("Must define In or Out")
+}
+
 type SwapAmountWithLimit struct {
 	ExactIn  *ExactIn  `json:"exact_in,omitempty"`
 	ExactOut *ExactOut `json:"exact_out,omitempty"`
 }
 
+// This returns the amount without min/max to use as simpler argument
+func (s SwapAmountWithLimit) RemoveLimit() SwapAmount {
+	if s.ExactIn != nil {
+		return SwapAmount{
+			In: &s.ExactIn.Input,
+		}
+	}
+	if s.ExactOut != nil {
+		return SwapAmount{
+			Out: &s.ExactOut.Output,
+		}
+	}
+	panic("Must define ExactIn or ExactOut")
+}
+
 type ExactIn struct {
 	Input     sdk.Int `json:"input"`
 	MinOutput sdk.Int `json:"min_output"`

app/wasm/message_plugin.go

@@ -78,8 +78,14 @@ func (m *MintTokenMessenger) mintTokens(ctx sdk.Context, contractAddr sdk.AccAdd
 // TODO: this is very close to QueryPlugin.EstimatePrice, maybe we can pull out common code into one function
 // that these both use? at least the routes / token In/Out calculation
 func (m *MintTokenMessenger) swapTokens(ctx sdk.Context, contractAddr sdk.AccAddress, swap *wasmbindings.SwapMsg) ([]sdk.Event, [][]byte, error) {
+	_, err := performSwap(m.gammKeeper, ctx, contractAddr, swap)
+	return nil, nil, sdkerrors.Wrap(err, "gamm estimate price exact amount out")
+}
+
+// This can be used both for the real swap as well as with EstimatePrice query
+func performSwap(keeper *gammkeeper.Keeper, ctx sdk.Context, contractAddr sdk.AccAddress, swap *wasmbindings.SwapMsg) (*wasmbindings.SwapAmount, error) {
 	if len(swap.Route) != 0 {
-		return nil, nil, wasmvmtypes.UnsupportedRequest{Kind: "TODO: multi-hop swaps"}
+		return nil, wasmvmtypes.UnsupportedRequest{Kind: "TODO: multi-hop swaps"}
 	}
 	if swap.Amount.ExactIn != nil {
 		routes := []gammtypes.SwapAmountInRoute{{
@@ -91,11 +97,11 @@ func (m *MintTokenMessenger) swapTokens(ctx sdk.Context, contractAddr sdk.AccAdd
 			Amount: swap.Amount.ExactIn.Input,
 		}
 		tokenOutMinAmount := swap.Amount.ExactIn.MinOutput
-		_, err := m.gammKeeper.MultihopSwapExactAmountIn(ctx, contractAddr, routes, tokenIn, tokenOutMinAmount)
+		estimatedAmount, err := keeper.MultihopSwapExactAmountIn(ctx, contractAddr, routes, tokenIn, tokenOutMinAmount)
 		if err != nil {
-			return nil, nil, sdkerrors.Wrap(err, "gamm estimate price exact amount in")
+			return nil, sdkerrors.Wrap(err, "gamm estimate price exact amount in")
 		}
-		return nil, nil, nil
+		return &wasmbindings.SwapAmount{Out: &estimatedAmount}, nil
 	} else if swap.Amount.ExactOut != nil {
 		routes := []gammtypes.SwapAmountOutRoute{{
 			PoolId:       swap.First.PoolId,
@@ -106,13 +112,13 @@ func (m *MintTokenMessenger) swapTokens(ctx sdk.Context, contractAddr sdk.AccAdd
 			Denom:  swap.First.DenomOut,
 			Amount: swap.Amount.ExactOut.Output,
 		}
-		_, err := m.gammKeeper.MultihopSwapExactAmountOut(ctx, contractAddr, routes, tokenInMaxAmount, tokenOut)
+		estimatedAmount, err := keeper.MultihopSwapExactAmountOut(ctx, contractAddr, routes, tokenInMaxAmount, tokenOut)
 		if err != nil {
-			return nil, nil, sdkerrors.Wrap(err, "gamm estimate price exact amount out")
+			return nil, sdkerrors.Wrap(err, "gamm estimate price exact amount out")
 		}
-		return nil, nil, nil
+		return &wasmbindings.SwapAmount{In: &estimatedAmount}, nil
 	} else {
-		return nil, nil, wasmvmtypes.UnsupportedRequest{Kind: "must support either Swap.ExactIn or Swap.ExactOut"}
+		return nil, wasmvmtypes.UnsupportedRequest{Kind: "must support either Swap.ExactIn or Swap.ExactOut"}
 	}
 }
 

app/wasm/queries.go

@@ -1,16 +1,12 @@
 package wasm
 
 import (
-	wasmvmtypes "github.com/CosmWasm/wasmvm/types"
-	"math"
-
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
 
-	bindings "github.com/osmosis-labs/osmosis/v7/app/wasm/bindings"
+	wasmbindings "github.com/osmosis-labs/osmosis/v7/app/wasm/bindings"
 	"github.com/osmosis-labs/osmosis/v7/app/wasm/types"
 	gammkeeper "github.com/osmosis-labs/osmosis/v7/x/gamm/keeper"
-	gammtypes "github.com/osmosis-labs/osmosis/v7/x/gamm/types"
 )
 
 type QueryPlugin struct {
@@ -41,7 +37,7 @@ func (qp QueryPlugin) GetPoolState(ctx sdk.Context, poolId uint64) (*types.PoolS
 	return &poolState, nil
 }
 
-func (qp QueryPlugin) GetSpotPrice(ctx sdk.Context, spotPrice *bindings.SpotPrice) (*sdk.Dec, error) {
+func (qp QueryPlugin) GetSpotPrice(ctx sdk.Context, spotPrice *wasmbindings.SpotPrice) (*sdk.Dec, error) {
 	poolId := spotPrice.Swap.PoolId
 	denomIn := spotPrice.Swap.DenomIn
 	denomOut := spotPrice.Swap.DenomOut
@@ -59,57 +55,12 @@ func (qp QueryPlugin) GetSpotPrice(ctx sdk.Context, spotPrice *bindings.SpotPric
 	return &price, nil
 }
 
-// TODO: this is very close to MintTokenMessenger.swapTokens, maybe we can pull out common code into one function
-// that these both use? at least the routes / token In/Out calculation
-func (qp QueryPlugin) EstimatePrice(ctx sdk.Context, estimatePrice *bindings.EstimatePrice) (*bindings.SwapAmount, error) {
-	sender := estimatePrice.Contract
-	poolId := estimatePrice.First.PoolId
-	denomIn := estimatePrice.First.DenomIn
-	denomOut := estimatePrice.First.DenomOut
-	route := estimatePrice.Route
-	if len(route) != 0 {
-		return nil, wasmvmtypes.UnsupportedRequest{Kind: "TODO: multi-hop swap price estimation"}
-	}
-
-	senderAddress, err := sdk.AccAddressFromBech32(sender)
+func (qp QueryPlugin) EstimatePrice(ctx sdk.Context, estimatePrice *wasmbindings.EstimatePrice) (*wasmbindings.SwapAmount, error) {
+	contractAddr, err := sdk.AccAddressFromBech32(estimatePrice.Contract)
 	if err != nil {
 		return nil, sdkerrors.Wrap(err, "gamm estimate price sender address")
 	}
 
-	if estimatePrice.Amount.In != nil {
-		tokenIn := sdk.NewCoin(denomIn, *estimatePrice.Amount.In)
-
-		// Populate route
-		var steps gammtypes.SwapAmountInRoutes
-		firstStep := gammtypes.SwapAmountInRoute{
-			PoolId:        poolId,
-			TokenOutDenom: denomOut,
-		}
-		steps = append(steps, firstStep)
-
-		tokenOutMinAmount := sdk.NewInt(1)
-		estimatedAmount, err := qp.gammKeeper.MultihopSwapExactAmountIn(ctx, senderAddress, steps, tokenIn, tokenOutMinAmount)
-		if err != nil {
-			return nil, sdkerrors.Wrap(err, "gamm estimate price exact amount in")
-		}
-		return &bindings.SwapAmount{Out: &estimatedAmount}, nil
-	} else if estimatePrice.Amount.Out != nil {
-		tokenOut := sdk.NewCoin(denomOut, *estimatePrice.Amount.Out)
-
-		// Populate route
-		var steps gammtypes.SwapAmountOutRoutes
-		firstStep := gammtypes.SwapAmountOutRoute{
-			PoolId:       poolId,
-			TokenInDenom: denomIn,
-		}
-		steps = append(steps, firstStep)
-
-		tokenInMaxAmount := sdk.NewInt(math.MaxInt64)
-		estimatedAmount, err := qp.gammKeeper.MultihopSwapExactAmountOut(ctx, senderAddress, steps, tokenInMaxAmount, tokenOut)
-		if err != nil {
-			return nil, sdkerrors.Wrap(err, "gamm estimate price exact amount out")
-		}
-		return &bindings.SwapAmount{In: &estimatedAmount}, nil
-	}
-	return nil, wasmvmtypes.UnsupportedRequest{Kind: "must support either EstimatePrice.Amount.In or EstimatePrice.Amount.Out"}
+	estimate, err := performSwap(qp.gammKeeper, ctx, contractAddr, estimatePrice.ToSwapMsg())
+	return estimate, err
 }