feat: base mempool setup (#2)

mempool/mempool.go

@@ -0,0 +1,89 @@
+package mempool
+
+import (
+	"context"
+	"crypto/sha256"
+	"encoding/base64"
+	"fmt"
+
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	sdkmempool "github.com/cosmos/cosmos-sdk/types/mempool"
+)
+
+var _ sdkmempool.Mempool = (*AuctionMempool)(nil)
+
+// AuctionMempool defines an auction mempool. It can be seen as an extension of
+// an SDK PriorityNonceMempool, i.e. a mempool that supports <sender, nonce>
+// two-dimensional priority ordering, with the additional support of prioritizing
+// and indexing auction bids.
+type AuctionMempool struct {
+	// globalIndex defines the index of all transactions in the mempool. It uses
+	// the SDK's builtin PriorityNonceMempool. Once a bid if selected for top-of-block,
+	// all subsequent transactions in the mempool will be selected from this index.
+	globalIndex sdkmempool.PriorityNonceMempool
+
+	// txIndex defines an index of all transactions in the mempool by hash.
+	txIndex map[string]*WrappedTx
+
+	// txEncoder defines the sdk.Tx encoder that allows us to encode transactions
+	// and construct their hashes.
+	txEncoder sdk.TxEncoder
+
+	// auctionIndex *heap.Heap[PriorityTx]
+}
+
+func NewAuctionMempool(txEncoder sdk.TxEncoder, opts ...sdkmempool.PriorityNonceMempoolOption) *AuctionMempool {
+	return &AuctionMempool{
+		globalIndex: *sdkmempool.NewPriorityMempool(opts...),
+		txIndex:     make(map[string]*WrappedTx),
+		txEncoder:   txEncoder,
+	}
+}
+
+func (am *AuctionMempool) Insert(ctx context.Context, tx sdk.Tx) error {
+	bz, err := am.txEncoder(tx)
+	if err != nil {
+		return fmt.Errorf("failed to encode tx: %w", err)
+	}
+
+	hash := sha256.Sum256(bz)
+	hashStr := base64.StdEncoding.EncodeToString(hash[:])
+	if _, ok := am.txIndex[hashStr]; ok {
+		return fmt.Errorf("tx already exists: %s", hashStr)
+	}
+
+	wrappedTx := &WrappedTx{
+		Tx:   tx,
+		hash: hash,
+	}
+
+	msg, err := GetMsgAuctionBidFromTx(tx)
+	if err != nil {
+		return err
+	}
+
+	if msg != nil {
+		// TODO: Insert into auctionIndex and update wrappedTx to reflect the index
+		// pointer.
+	}
+
+	if err := am.globalIndex.Insert(ctx, wrappedTx); err != nil {
+		return fmt.Errorf("failed to insert tx into global index: %w", err)
+	}
+
+	am.txIndex[hashStr] = wrappedTx
+
+	return nil
+}
+
+func (am *AuctionMempool) Remove(tx sdk.Tx) error {
+	panic("not implemented")
+}
+
+func (am *AuctionMempool) Select(ctx context.Context, txs [][]byte) sdkmempool.Iterator {
+	return am.globalIndex.Select(ctx, txs)
+}
+
+func (am *AuctionMempool) CountTx() int {
+	return am.globalIndex.CountTx()
+}

mempool/tx.go

@@ -0,0 +1,54 @@
+package mempool
+
+import (
+	"errors"
+
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	auctiontypes "github.com/skip-mev/pob/x/auction/types"
+)
+
+type (
+	// WrappedTx defines a wrapper around an sdk.Tx with additional metadata.
+	WrappedTx struct {
+		sdk.Tx
+
+		hash [32]byte
+	}
+
+	// WrappedBidTx defines a wrapper around an sdk.Tx that contains a single
+	// MsgAuctionBid message with additional metadata.
+	WrappedBidTx struct {
+		sdk.Tx
+
+		hash [32]byte
+		bid  sdk.Coins
+	}
+)
+
+// GetMsgAuctionBidFromTx attempts to retrieve a MsgAuctionBid from an sdk.Tx if
+// one exists. If a MsgAuctionBid does exist and other messages are also present,
+// an error is returned. If no MsgAuctionBid is present, <nil, nil> is returned.
+func GetMsgAuctionBidFromTx(tx sdk.Tx) (*auctiontypes.MsgAuctionBid, error) {
+	auctionBidMsgs := make([]*auctiontypes.MsgAuctionBid, 0)
+	for _, msg := range tx.GetMsgs() {
+		t, ok := msg.(*auctiontypes.MsgAuctionBid)
+		if ok {
+			auctionBidMsgs = append(auctionBidMsgs, t)
+		}
+	}
+
+	switch {
+	case len(auctionBidMsgs) == 0:
+		// a normal transaction without a MsgAuctionBid message
+		return nil, nil
+
+	case len(auctionBidMsgs) == 1 && len(tx.GetMsgs()) == 1:
+		// a single MsgAuctionBid message transaction
+		return auctionBidMsgs[0], nil
+
+	default:
+		// A transaction with at at least one MsgAuctionBid message and some other
+		// message.
+		return nil, errors.New("invalid MsgAuctionBid transaction")
+	}
+}

mempool/tx_test.go

@@ -0,0 +1,48 @@
+package mempool_test
+
+import (
+	"testing"
+
+	banktypes "github.com/cosmos/cosmos-sdk/x/bank/types"
+	pobcodec "github.com/skip-mev/pob/codec"
+	"github.com/skip-mev/pob/mempool"
+	auctiontypes "github.com/skip-mev/pob/x/auction/types"
+	"github.com/stretchr/testify/require"
+)
+
+func TestGetMsgAuctionBidFromTx_Valid(t *testing.T) {
+	encCfg := pobcodec.CreateEncodingConfig()
+
+	txBuilder := encCfg.TxConfig.NewTxBuilder()
+	txBuilder.SetMsgs(&auctiontypes.MsgAuctionBid{})
+
+	msg, err := mempool.GetMsgAuctionBidFromTx(txBuilder.GetTx())
+	require.NoError(t, err)
+	require.NotNil(t, msg)
+}
+
+func TestGetMsgAuctionBidFromTx_MultiMsgBid(t *testing.T) {
+	encCfg := pobcodec.CreateEncodingConfig()
+
+	txBuilder := encCfg.TxConfig.NewTxBuilder()
+	txBuilder.SetMsgs(
+		&auctiontypes.MsgAuctionBid{},
+		&auctiontypes.MsgAuctionBid{},
+		&banktypes.MsgSend{},
+	)
+
+	msg, err := mempool.GetMsgAuctionBidFromTx(txBuilder.GetTx())
+	require.Error(t, err)
+	require.Nil(t, msg)
+}
+
+func TestGetMsgAuctionBidFromTx_NoBid(t *testing.T) {
+	encCfg := pobcodec.CreateEncodingConfig()
+
+	txBuilder := encCfg.TxConfig.NewTxBuilder()
+	txBuilder.SetMsgs(&banktypes.MsgSend{})
+
+	msg, err := mempool.GetMsgAuctionBidFromTx(txBuilder.GetTx())
+	require.NoError(t, err)
+	require.Nil(t, msg)
+}