[x/TWAP] Expose a geometric TWAP API

x/twap/api.go

@@ -51,12 +51,40 @@ func (k Keeper) GetArithmeticTwap(
 	quoteAssetDenom string,
 	startTime time.Time,
 	endTime time.Time,
+) (sdk.Dec, error) {
+	arithmeticStrategy := &arithmetic{k}
+	return k.getTwap(ctx, poolId, baseAssetDenom, quoteAssetDenom, startTime, endTime, arithmeticStrategy)
+}
+
+// GetArithmeticTwapToNow returns arithmetic twap from start time until the current block time for quote and base
+// assets in a given pool.
+func (k Keeper) GetArithmeticTwapToNow(
+	ctx sdk.Context,
+	poolId uint64,
+	baseAssetDenom string,
+	quoteAssetDenom string,
+	startTime time.Time,
+) (sdk.Dec, error) {
+	arithmeticStrategy := &arithmetic{k}
+	return k.getTwapToNow(ctx, poolId, baseAssetDenom, quoteAssetDenom, startTime, arithmeticStrategy)
+}
+
+// getTwap computes and returns twap from the start time until the end time. The type
+// of twap returned depends on the strategy given and can be either arithmetic or geometric.
+func (k Keeper) getTwap(
+	ctx sdk.Context,
+	poolId uint64,
+	baseAssetDenom string,
+	quoteAssetDenom string,
+	startTime time.Time,
+	endTime time.Time,
+	strategy twapStrategy,
 ) (sdk.Dec, error) {
 	if startTime.After(endTime) {
 		return sdk.Dec{}, types.StartTimeAfterEndTimeError{StartTime: startTime, EndTime: endTime}
 	}
 	if endTime.Equal(ctx.BlockTime()) {
-		return k.GetArithmeticTwapToNow(ctx, poolId, baseAssetDenom, quoteAssetDenom, startTime)
+		return k.getTwapToNow(ctx, poolId, baseAssetDenom, quoteAssetDenom, startTime, strategy)
 	} else if endTime.After(ctx.BlockTime()) {
 		return sdk.Dec{}, types.EndTimeInFutureError{EndTime: endTime, BlockTime: ctx.BlockTime()}
 	}
@@ -68,17 +96,19 @@ func (k Keeper) GetArithmeticTwap(
 	if err != nil {
 		return sdk.Dec{}, err
 	}
-	return computeTwap(startRecord, endRecord, quoteAssetDenom, arithmeticTwapType)
+
+	return strategy.computeTwap(startRecord, endRecord, quoteAssetDenom)
 }
 
-// GetArithmeticTwapToNow returns GetArithmeticTwap on the input, with endTime being fixed to ctx.BlockTime()
-// This function does not mutate records.
-func (k Keeper) GetArithmeticTwapToNow(
+// getTwapToNow computes and returns twap from the start time until the current block time. The type
+// of twap returned depends on the strategy given and can be either arithmetic or geometric.
+func (k Keeper) getTwapToNow(
 	ctx sdk.Context,
 	poolId uint64,
 	baseAssetDenom string,
 	quoteAssetDenom string,
 	startTime time.Time,
+	strategy twapStrategy,
 ) (sdk.Dec, error) {
 	if startTime.After(ctx.BlockTime()) {
 		return sdk.Dec{}, types.StartTimeAfterEndTimeError{StartTime: startTime, EndTime: ctx.BlockTime()}
@@ -92,14 +122,12 @@ func (k Keeper) GetArithmeticTwapToNow(
 	if err != nil {
 		return sdk.Dec{}, err
 	}
-	return computeTwap(startRecord, endRecord, quoteAssetDenom, arithmeticTwapType)
+
+	return strategy.computeTwap(startRecord, endRecord, quoteAssetDenom)
 }
 
 // GetBeginBlockAccumulatorRecord returns a TwapRecord struct corresponding to the state of pool `poolId`
 // as of the beginning of the block this is called on.
-// This uses the state of the beginning of the block, as if there were swaps since the block has started,
-// these swaps have had no time to be arbitraged back.
-// This accumulator can be stored, to compute wider ranged twaps.
 func (k Keeper) GetBeginBlockAccumulatorRecord(ctx sdk.Context, poolId uint64, asset0Denom string, asset1Denom string) (types.TwapRecord, error) {
 	return k.getMostRecentRecord(ctx, poolId, asset0Denom, asset1Denom)
 }

x/twap/strategy.go

@@ -0,0 +1,24 @@
+package twap
+
+import (
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/osmosis-labs/osmosis/v13/x/twap/types"
+)
+
+// twapStrategy is an interface for computing TWAPs.
+// We have two strategies implementing the interface - arithmetic and geometric.
+// We expose a common TWAP API to reduce duplication and avoid complexity.
+type twapStrategy interface {
+	// computeTwap calculates the TWAP with specific startRecord and endRecord.
+	computeTwap(startRecord types.TwapRecord, endRecord types.TwapRecord, quoteAsset string) (sdk.Dec, error)
+}
+
+type arithmetic struct {
+	keeper Keeper
+}
+
+var _ twapStrategy = &arithmetic{}
+
+func (s *arithmetic) computeTwap(startRecord types.TwapRecord, endRecord types.TwapRecord, quoteAsset string) (sdk.Dec, error) {
+	return computeTwap(startRecord, endRecord, quoteAsset, arithmeticTwapType)
+}