TWAP proto, types package, and osmoutils (#2175)

* Extract osmoutils, proto & types from main TWAP PR

* Delete merge conflict

* Add TestGetAllUniqueDenomPairs

* Resolve more comments

* Fix incorrect pool interface (but was a no-op)

* Apply suggestions from code review

Co-authored-by: Aleksandr Bezobchuk <[email protected]>

* apply review comment

* Make NewTwapRecord return an error

Co-authored-by: Aleksandr Bezobchuk <[email protected]>

osmoutils/slice_helper.go

@@ -7,6 +7,7 @@ import (
 )
 
 // SortSlice sorts a slice of type T elements that implement constraints.Ordered.
+// Mutates input slice s
 func SortSlice[T constraints.Ordered](s []T) {
 	sort.Slice(s, func(i, j int) bool {
 		return s[i] < s[j]
@@ -22,3 +23,14 @@ func Filter[T interface{}](filter func(T) bool, s []T) []T {
 	}
 	return filteredSlice
 }
+
+// ReverseSlice returns a reversed copy of the input slice.
+// Does not mutate argument.
+func ReverseSlice[T any](s []T) []T {
+	newSlice := make([]T, len(s))
+	maxIndex := len(s) - 1
+	for i := 0; i < len(s); i++ {
+		newSlice[maxIndex-i] = s[i]
+	}
+	return newSlice
+}

osmoutils/store_helper.go

@@ -0,0 +1,97 @@
+package osmoutils
+
+import (
+	"errors"
+
+	"github.com/cosmos/cosmos-sdk/store"
+	"github.com/gogo/protobuf/proto"
+)
+
+func GatherAllKeysFromStore(storeObj store.KVStore) []string {
+	iterator := storeObj.Iterator(nil, nil)
+	defer iterator.Close()
+
+	keys := []string{}
+	for ; iterator.Valid(); iterator.Next() {
+		keys = append(keys, string(iterator.Key()))
+	}
+	return keys
+}
+
+func GatherValuesFromStore[T any](storeObj store.KVStore, keyStart []byte, keyEnd []byte, parseValue func([]byte) (T, error)) ([]T, error) {
+	iterator := storeObj.Iterator(keyStart, keyEnd)
+	defer iterator.Close()
+
+	values := []T{}
+	for ; iterator.Valid(); iterator.Next() {
+		val, err := parseValue(iterator.Value())
+		if err != nil {
+			return nil, err
+		}
+		values = append(values, val)
+	}
+	return values, nil
+}
+
+func GetValuesUntilDerivedStop[T any](storeObj store.KVStore, keyStart []byte, stopFn func([]byte) bool, parseValue func([]byte) (T, error)) ([]T, error) {
+	// SDK iterator is broken for nil end time, and non-nil start time
+	// https://github.com/cosmos/cosmos-sdk/issues/12661
+	// hence we use []byte{0xff}
+	keyEnd := []byte{0xff}
+	return GetIterValuesWithStop(storeObj, keyStart, keyEnd, false, stopFn, parseValue)
+}
+
+func GetIterValuesWithStop[T any](
+	storeObj store.KVStore,
+	keyStart []byte,
+	keyEnd []byte,
+	reverse bool,
+	stopFn func([]byte) bool,
+	parseValue func([]byte) (T, error)) ([]T, error) {
+	var iter store.Iterator
+	if reverse {
+		iter = storeObj.ReverseIterator(keyStart, keyEnd)
+	} else {
+		iter = storeObj.Iterator(keyStart, keyEnd)
+	}
+	defer iter.Close()
+
+	values := []T{}
+	for ; iter.Valid(); iter.Next() {
+		if stopFn(iter.Key()) {
+			break
+		}
+		val, err := parseValue(iter.Value())
+		if err != nil {
+			return nil, err
+		}
+		values = append(values, val)
+	}
+	return values, nil
+}
+
+func GetFirstValueAfterPrefix[T any](storeObj store.KVStore, keyStart []byte, parseValue func([]byte) (T, error)) (T, error) {
+	// SDK iterator is broken for nil end time, and non-nil start time
+	// https://github.com/cosmos/cosmos-sdk/issues/12661
+	// hence we use []byte{0xff}
+	iterator := storeObj.Iterator(keyStart, []byte{0xff})
+	defer iterator.Close()
+
+	if !iterator.Valid() {
+		var blankValue T
+		return blankValue, errors.New("No values in iterator")
+	}
+
+	return parseValue(iterator.Value())
+}
+
+// MustSet runs store.Set(key, proto.Marshal(value))
+// but panics on any error.
+func MustSet(storeObj store.KVStore, key []byte, value proto.Message) {
+	bz, err := proto.Marshal(value)
+	if err != nil {
+		panic(err)
+	}
+
+	storeObj.Set(key, bz)
+}

osmoutils/time_helper.go

@@ -0,0 +1,20 @@
+package osmoutils
+
+import (
+	"time"
+
+	sdk "github.com/cosmos/cosmos-sdk/types"
+)
+
+func FormatTimeString(t time.Time) string {
+	return t.UTC().Round(0).Format(sdk.SortableTimeFormat)
+}
+
+// Parses a string encoded using FormatTimeString back into a time.Time
+func ParseTimeString(s string) (time.Time, error) {
+	t, err := time.Parse(sdk.SortableTimeFormat, s)
+	if err != nil {
+		return t, err
+	}
+	return t.UTC().Round(0), nil
+}

proto/osmosis/gamm/twap/v1beta1/twap_record.proto

@@ -0,0 +1,63 @@
+syntax = "proto3";
+package osmosis.gamm.twap.v1beta1;
+
+import "gogoproto/gogo.proto";
+import "google/protobuf/any.proto";
+import "cosmos_proto/cosmos.proto";
+import "cosmos/base/v1beta1/coin.proto";
+import "google/protobuf/timestamp.proto";
+
+option go_package = "github.com/osmosis-labs/osmosis/v10/x/gamm/twap/types";
+
+// A TWAP record should be indexed in state by pool_id, (asset pair), timestamp
+// The asset pair assets should be lexicographically sorted.
+// Technically (pool_id, asset_0_denom, asset_1_denom, height) do not need to
+// appear in the struct however we view this as the wrong performance tradeoff
+// given SDK today. Would rather we optimize for readability and correctness,
+// than an optimal state storage format. The system bottleneck is elsewhere for
+// now.
+message TwapRecord {
+  uint64 pool_id = 1;
+  // Lexicographically smaller denom of the pair
+  string asset0_denom = 2;
+  // Lexicographically larger denom of the pair
+  string asset1_denom = 3;
+  // height this record corresponds to, for debugging purposes
+  int64 height = 4 [
+    (gogoproto.moretags) = "yaml:\"record_height\"",
+    (gogoproto.jsontag) = "record_height"
+  ];
+  // This field should only exist until we have a global registry in the state
+  // machine, mapping prior block heights within {TIME RANGE} to times.
+  google.protobuf.Timestamp time = 5 [
+    (gogoproto.nullable) = false,
+    (gogoproto.stdtime) = true,
+    (gogoproto.moretags) = "yaml:\"record_time\""
+  ];
+
+  // We store the last spot prices in the struct, so that we can interpolate
+  // accumulator values for times between when accumulator records are stored.
+  string p0_last_spot_price = 6 [
+    (gogoproto.customtype) = "github.com/cosmos/cosmos-sdk/types.Dec",
+    (gogoproto.nullable) = false
+  ];
+  string p1_last_spot_price = 7 [
+    (gogoproto.customtype) = "github.com/cosmos/cosmos-sdk/types.Dec",
+    (gogoproto.nullable) = false
+  ];
+
+  string p0_arithmetic_twap_accumulator = 8 [
+    (gogoproto.customtype) = "github.com/cosmos/cosmos-sdk/types.Dec",
+    (gogoproto.nullable) = false
+  ];
+  string p1_arithmetic_twap_accumulator = 9 [
+    (gogoproto.customtype) = "github.com/cosmos/cosmos-sdk/types.Dec",
+    (gogoproto.nullable) = false
+  ];
+  // string geometric_twap_accumulator = 7 [(gogoproto.customtype) =
+  // "github.com/cosmos/cosmos-sdk/types.Dec",
+  // (gogoproto.nullable) = false];
+}
+
+// GenesisState defines the gamm module's genesis state.
+message GenesisState { repeated TwapRecord twaps = 1; }

x/gamm/twap/types/expected_interfaces.go

@@ -0,0 +1,16 @@
+package types
+
+import sdk "github.com/cosmos/cosmos-sdk/types"
+
+// AmmInterface is the functionality needed from a given pool ID, in order to maintain records and serve TWAPs.
+type AmmInterface interface {
+	GetPoolDenoms(ctx sdk.Context, poolId uint64) (denoms []string, err error)
+	// CalculateSpotPrice returns the spot price of the quote asset in terms of the base asset,
+	// using the specified pool.
+	// E.g. if pool 1 traded 2 atom for 3 osmo, the quote asset was atom, and the base asset was osmo,
+	// this would return 1.5. (Meaning that 1 atom costs 1.5 osmo)
+	CalculateSpotPrice(ctx sdk.Context,
+		poolID uint64,
+		baseAssetDenom string,
+		quoteAssetDenom string) (price sdk.Dec, err error)
+}

x/gamm/twap/types/keys.go

@@ -1,5 +1,17 @@
 package types
 
+import (
+	"errors"
+	fmt "fmt"
+	"strings"
+	time "time"
+
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/gogo/protobuf/proto"
+
+	"github.com/osmosis-labs/osmosis/v10/osmoutils"
+)
+
 const (
 	ModuleName = "twap"
 
@@ -8,6 +20,78 @@ const (
 	RouterKey         = ModuleName
 
 	QuerierRoute = ModuleName
+	// Contract: Coin denoms cannot contain this character
+	KeySeparator = "|"
 )
 
-var AlteredPoolIdsPrefix = []byte{0}
+var (
+	mostRecentTWAPsNoSeparator         = "recent_twap"
+	historicalTWAPTimeIndexNoSeparator = "historical_time_index"
+	historicalTWAPPoolIndexNoSeparator = "historical_pool_index"
+
+	mostRecentTWAPsPrefix         = mostRecentTWAPsNoSeparator + KeySeparator
+	historicalTWAPTimeIndexPrefix = historicalTWAPTimeIndexNoSeparator + KeySeparator
+	historicalTWAPPoolIndexPrefix = historicalTWAPPoolIndexNoSeparator + KeySeparator
+)
+
+// TODO: make utility command to automatically interlace separators
+
+func FormatMostRecentTWAPKey(poolId uint64, denom1, denom2 string) []byte {
+	return []byte(fmt.Sprintf("%s%d%s%s%s%s", mostRecentTWAPsPrefix, poolId, KeySeparator, denom1, KeySeparator, denom2))
+}
+
+// TODO: Replace historical management with ORM, we currently accept 2x write amplification right now.
+func FormatHistoricalTimeIndexTWAPKey(accumulatorWriteTime time.Time, poolId uint64, denom1, denom2 string) []byte {
+	timeS := osmoutils.FormatTimeString(accumulatorWriteTime)
+	return []byte(fmt.Sprintf("%s%s%s%d%s%s%s%s", historicalTWAPTimeIndexPrefix, timeS, KeySeparator, poolId, KeySeparator, denom1, KeySeparator, denom2))
+}
+
+func FormatHistoricalPoolIndexTWAPKey(poolId uint64, accumulatorWriteTime time.Time, denom1, denom2 string) []byte {
+	timeS := osmoutils.FormatTimeString(accumulatorWriteTime)
+	return []byte(fmt.Sprintf("%s%d%s%s%s%s%s%s", historicalTWAPPoolIndexPrefix, poolId, KeySeparator, timeS, KeySeparator, denom1, KeySeparator, denom2))
+}
+
+func FormatHistoricalPoolIndexTimePrefix(poolId uint64, accumulatorWriteTime time.Time) []byte {
+	timeS := osmoutils.FormatTimeString(accumulatorWriteTime)
+	return []byte(fmt.Sprintf("%s%d%s%s%s", historicalTWAPPoolIndexPrefix, poolId, KeySeparator, timeS, KeySeparator))
+}
+
+func ParseTimeFromHistoricalTimeIndexKey(key []byte) time.Time {
+	keyS := string(key)
+	s := strings.Split(keyS, KeySeparator)
+	if len(s) != 5 || s[0] != historicalTWAPTimeIndexNoSeparator {
+		panic("Called ParseTimeFromHistoricalTimeIndexKey on incorrectly formatted key")
+	}
+	t, err := osmoutils.ParseTimeString(s[1])
+	if err != nil {
+		panic("incorrectly formatted time string in key")
+	}
+	return t
+}
+
+func ParseTimeFromHistoricalPoolIndexKey(key []byte) (time.Time, error) {
+	keyS := string(key)
+	s := strings.Split(keyS, KeySeparator)
+	if len(s) != 5 || s[0] != historicalTWAPPoolIndexNoSeparator {
+		return time.Time{}, fmt.Errorf("Called ParseTimeFromHistoricalPoolIndexKey on incorrectly formatted key: %v", s)
+	}
+	t, err := osmoutils.ParseTimeString(s[2])
+	if err != nil {
+		return time.Time{}, fmt.Errorf("incorrectly formatted time string in key %s : %v", keyS, err)
+	}
+	return t, nil
+}
+
+func GetAllMostRecentTwapsForPool(store sdk.KVStore, poolId uint64) ([]TwapRecord, error) {
+	startPrefix := fmt.Sprintf("%s%s%d%s", mostRecentTWAPsPrefix, KeySeparator, poolId, KeySeparator)
+	endPrefix := fmt.Sprintf("%s%s%d%s", mostRecentTWAPsPrefix, KeySeparator, poolId+1, KeySeparator)
+	return osmoutils.GatherValuesFromStore(store, []byte(startPrefix), []byte(endPrefix), ParseTwapFromBz)
+}
+
+func ParseTwapFromBz(bz []byte) (twap TwapRecord, err error) {
+	if len(bz) == 0 {
+		return TwapRecord{}, errors.New("twap not found")
+	}
+	err = proto.Unmarshal(bz, &twap)
+	return twap, err
+}