Merge pull request #47 from Layr-Labs/sm-keyOrdering

Merkleization refactoring

cmd/debugger/main.go

@@ -71,16 +71,16 @@ func main() {
 
 	sm := stateManager.NewEigenStateManager(l, grm)
 
-	if _, err := avsOperators.NewAvsOperators(sm, grm, cfg.Network, cfg.Environment, l, cfg); err != nil {
+	if _, err := avsOperators.NewAvsOperators(sm, grm, l, cfg); err != nil {
 		l.Sugar().Fatalw("Failed to create AvsOperatorsModel", zap.Error(err))
 	}
-	if _, err := operatorShares.NewOperatorSharesModel(sm, grm, cfg.Network, cfg.Environment, l, cfg); err != nil {
+	if _, err := operatorShares.NewOperatorSharesModel(sm, grm, l, cfg); err != nil {
 		l.Sugar().Fatalw("Failed to create OperatorSharesModel", zap.Error(err))
 	}
-	if _, err := stakerDelegations.NewStakerDelegationsModel(sm, grm, cfg.Network, cfg.Environment, l, cfg); err != nil {
+	if _, err := stakerDelegations.NewStakerDelegationsModel(sm, grm, l, cfg); err != nil {
 		l.Sugar().Fatalw("Failed to create StakerDelegationsModel", zap.Error(err))
 	}
-	if _, err := stakerShares.NewStakerSharesModel(sm, grm, cfg.Network, cfg.Environment, l, cfg); err != nil {
+	if _, err := stakerShares.NewStakerSharesModel(sm, grm, l, cfg); err != nil {
 		l.Sugar().Fatalw("Failed to create StakerSharesModel", zap.Error(err))
 	}
 

cmd/sidecar/main.go

@@ -75,22 +75,22 @@ func main() {
 
 	sm := stateManager.NewEigenStateManager(l, grm)
 
-	if _, err := avsOperators.NewAvsOperators(sm, grm, cfg.Network, cfg.Environment, l, cfg); err != nil {
+	if _, err := avsOperators.NewAvsOperators(sm, grm, l, cfg); err != nil {
 		l.Sugar().Fatalw("Failed to create AvsOperatorsModel", zap.Error(err))
 	}
-	if _, err := operatorShares.NewOperatorSharesModel(sm, grm, cfg.Network, cfg.Environment, l, cfg); err != nil {
+	if _, err := operatorShares.NewOperatorSharesModel(sm, grm, l, cfg); err != nil {
 		l.Sugar().Fatalw("Failed to create OperatorSharesModel", zap.Error(err))
 	}
-	if _, err := stakerDelegations.NewStakerDelegationsModel(sm, grm, cfg.Network, cfg.Environment, l, cfg); err != nil {
+	if _, err := stakerDelegations.NewStakerDelegationsModel(sm, grm, l, cfg); err != nil {
 		l.Sugar().Fatalw("Failed to create StakerDelegationsModel", zap.Error(err))
 	}
-	if _, err := stakerShares.NewStakerSharesModel(sm, grm, cfg.Network, cfg.Environment, l, cfg); err != nil {
+	if _, err := stakerShares.NewStakerSharesModel(sm, grm, l, cfg); err != nil {
 		l.Sugar().Fatalw("Failed to create StakerSharesModel", zap.Error(err))
 	}
-	if _, err := submittedDistributionRoots.NewSubmittedDistributionRootsModel(sm, grm, cfg.Network, cfg.Environment, l, cfg); err != nil {
+	if _, err := submittedDistributionRoots.NewSubmittedDistributionRootsModel(sm, grm, l, cfg); err != nil {
 		l.Sugar().Fatalw("Failed to create SubmittedDistributionRootsModel", zap.Error(err))
 	}
-	if _, err := rewardSubmissions.NewRewardSubmissionsModel(sm, grm, cfg.Network, cfg.Environment, l, cfg); err != nil {
+	if _, err := rewardSubmissions.NewRewardSubmissionsModel(sm, grm, l, cfg); err != nil {
 		l.Sugar().Fatalw("Failed to create RewardSubmissionsModel", zap.Error(err))
 	}
 

internal/clients/etherscan/client.go

@@ -1,10 +1,11 @@
 package etherscan
 
 import (
+	"crypto/rand"
 	"encoding/json"
 	"fmt"
 	"io"
-	"math/rand/v2"
+	"math/big"
 	"net/http"
 	"net/url"
 	"regexp"
@@ -138,14 +139,18 @@ func (ec *EtherscanClient) makeRequestWithBackoff(values url.Values) (*Etherscan
 
 func (ec *EtherscanClient) selectApiKey() string {
 	maxNumber := len(ec.ApiKeys)
-	randInt := rand.IntN(maxNumber)
+	randInt, err := rand.Int(rand.Reader, big.NewInt(int64(maxNumber)))
+	if err != nil {
+		ec.logger.Sugar().Errorw("Failed to generate random number", zap.Error(err))
+		return ec.ApiKeys[0]
+	}
 
-	if randInt > maxNumber {
+	if randInt.Uint64() > uint64(maxNumber) {
 		ec.logger.Sugar().Warnw("Random number is greater than the number of api keys", "randInt", randInt, "maxNumber", maxNumber)
-		randInt = 0
+		randInt = big.NewInt(int64(0))
 	}
 
-	return ec.ApiKeys[randInt]
+	return ec.ApiKeys[randInt.Int64()]
 }
 
 func (ec *EtherscanClient) buildBaseUrlParams(module string, action string) url.Values {

internal/eigenState/avsOperators/avsOperators.go

@@ -14,9 +14,6 @@ import (
 	"github.com/Layr-Labs/go-sidecar/internal/eigenState/types"
 	"github.com/Layr-Labs/go-sidecar/internal/storage"
 	"github.com/Layr-Labs/go-sidecar/internal/utils"
-	"github.com/wealdtech/go-merkletree/v2"
-	"github.com/wealdtech/go-merkletree/v2/keccak256"
-	orderedmap "github.com/wk8/go-ordered-map/v2"
 	"go.uber.org/zap"
 	"golang.org/x/xerrors"
 	"gorm.io/gorm"
@@ -47,47 +44,38 @@ type RegisteredAvsOperatorDiff struct {
 	Registered  bool
 }
 
-// SlotId represents a unique identifier for a slot.
-type SlotId string
-
-func NewSlotId(avs string, operator string) SlotId {
-	return SlotId(fmt.Sprintf("%s_%s", avs, operator))
+func NewSlotID(avs string, operator string) types.SlotID {
+	return types.SlotID(fmt.Sprintf("%s_%s", avs, operator))
 }
 
 // EigenState model for AVS operators that implements IEigenStateModel.
 type AvsOperatorsModel struct {
 	base.BaseEigenState
 	StateTransitions types.StateTransitions[AccumulatedStateChange]
-	Db               *gorm.DB
-	Network          config.Network
-	Environment      config.Environment
+	DB               *gorm.DB
 	logger           *zap.Logger
 	globalConfig     *config.Config
 
 	// Accumulates state changes for SlotIds, grouped by block number
-	stateAccumulator map[uint64]map[SlotId]*AccumulatedStateChange
+	stateAccumulator map[uint64]map[types.SlotID]*AccumulatedStateChange
 }
 
-// Create new instance of AvsOperatorsModel state model.
+// NewAvsOperators creates a new AvsOperatorsModel.
 func NewAvsOperators(
 	esm *stateManager.EigenStateManager,
 	grm *gorm.DB,
-	Network config.Network,
-	Environment config.Environment,
 	logger *zap.Logger,
 	globalConfig *config.Config,
 ) (*AvsOperatorsModel, error) {
 	s := &AvsOperatorsModel{
 		BaseEigenState: base.BaseEigenState{
 			Logger: logger,
 		},
-		Db:           grm,
-		Network:      Network,
-		Environment:  Environment,
+		DB:           grm,
 		logger:       logger,
 		globalConfig: globalConfig,
 
-		stateAccumulator: make(map[uint64]map[SlotId]*AccumulatedStateChange),
+		stateAccumulator: make(map[uint64]map[types.SlotID]*AccumulatedStateChange),
 	}
 	esm.RegisterState(s, 0)
 	return s, nil
@@ -132,21 +120,21 @@ func (a *AvsOperatorsModel) GetStateTransitions() (types.StateTransitions[Accumu
 			registered = uint64(val.(float64)) == 1
 		}
 
-		slotId := NewSlotId(avs, operator)
-		record, ok := a.stateAccumulator[log.BlockNumber][slotId]
+		slotID := NewSlotID(avs, operator)
+		record, ok := a.stateAccumulator[log.BlockNumber][slotID]
 		if !ok {
 			record = &AccumulatedStateChange{
 				Avs:         avs,
 				Operator:    operator,
 				BlockNumber: log.BlockNumber,
 			}
-			a.stateAccumulator[log.BlockNumber][slotId] = record
+			a.stateAccumulator[log.BlockNumber][slotID] = record
 		}
 		if registered == false && ok {
 			// In this situation, we've encountered a register and unregister in the same block
 			// which functionally results in no state change at all so we want to remove the record
 			// from the accumulated state.
-			delete(a.stateAccumulator[log.BlockNumber], slotId)
+			delete(a.stateAccumulator[log.BlockNumber], slotID)
 			return nil, nil
 		}
 		record.Registered = registered
@@ -184,7 +172,7 @@ func (a *AvsOperatorsModel) IsInterestingLog(log *storage.TransactionLog) bool {
 }
 
 func (a *AvsOperatorsModel) InitBlockProcessing(blockNumber uint64) error {
-	a.stateAccumulator[blockNumber] = make(map[SlotId]*AccumulatedStateChange)
+	a.stateAccumulator[blockNumber] = make(map[types.SlotID]*AccumulatedStateChange)
 	return nil
 }
 
@@ -222,7 +210,7 @@ func (a *AvsOperatorsModel) clonePreviousBlocksToNewBlock(blockNumber uint64) er
 			from registered_avs_operators
 			where block_number = @previousBlock
 	`
-	res := a.Db.Exec(query,
+	res := a.DB.Exec(query,
 		sql.Named("currentBlock", blockNumber),
 		sql.Named("previousBlock", blockNumber-1),
 	)
@@ -274,7 +262,7 @@ func (a *AvsOperatorsModel) CommitFinalState(blockNumber uint64) error {
 	}
 
 	for _, record := range recordsToDelete {
-		res := a.Db.Delete(&RegisteredAvsOperators{}, "avs = ? and operator = ? and block_number = ?", record.Avs, record.Operator, record.BlockNumber)
+		res := a.DB.Delete(&RegisteredAvsOperators{}, "avs = ? and operator = ? and block_number = ?", record.Avs, record.Operator, record.BlockNumber)
 		if res.Error != nil {
 			a.logger.Sugar().Errorw("Failed to delete record",
 				zap.Error(res.Error),
@@ -286,7 +274,7 @@ func (a *AvsOperatorsModel) CommitFinalState(blockNumber uint64) error {
 		}
 	}
 	if len(recordsToInsert) > 0 {
-		res := a.Db.Model(&RegisteredAvsOperators{}).Clauses(clause.Returning{}).Create(&recordsToInsert)
+		res := a.DB.Model(&RegisteredAvsOperators{}).Clauses(clause.Returning{}).Create(&recordsToInsert)
 		if res.Error != nil {
 			a.logger.Sugar().Errorw("Failed to insert records", zap.Error(res.Error))
 			return res.Error
@@ -307,91 +295,47 @@ func (a *AvsOperatorsModel) GenerateStateRoot(blockNumber uint64) (types.StateRo
 		return "", err
 	}
 
-	combinedResults := make([]RegisteredAvsOperatorDiff, 0)
+	combinedResults := make([]*RegisteredAvsOperatorDiff, 0)
 	for _, record := range inserts {
-		combinedResults = append(combinedResults, RegisteredAvsOperatorDiff{
+		combinedResults = append(combinedResults, &RegisteredAvsOperatorDiff{
 			Avs:         record.Avs,
 			Operator:    record.Operator,
 			BlockNumber: record.BlockNumber,
 			Registered:  true,
 		})
 	}
 	for _, record := range deletes {
-		combinedResults = append(combinedResults, RegisteredAvsOperatorDiff{
+		combinedResults = append(combinedResults, &RegisteredAvsOperatorDiff{
 			Avs:         record.Avs,
 			Operator:    record.Operator,
 			BlockNumber: record.BlockNumber,
 			Registered:  false,
 		})
 	}
 
-	fullTree, err := a.merkelizeState(blockNumber, combinedResults)
+	inputs := a.sortValuesForMerkleTree(combinedResults)
+
+	fullTree, err := a.MerkleizeState(blockNumber, inputs)
 	if err != nil {
 		return "", err
 	}
 	return types.StateRoot(utils.ConvertBytesToString(fullTree.Root())), nil
 }
 
-func (a *AvsOperatorsModel) merkelizeState(blockNumber uint64, avsOperators []RegisteredAvsOperatorDiff) (*merkletree.MerkleTree, error) {
-	// Avs -> operator:registered
-	om := orderedmap.New[string, *orderedmap.OrderedMap[string, bool]]()
-
-	for _, result := range avsOperators {
-		existingAvs, found := om.Get(result.Avs)
-		if !found {
-			existingAvs = orderedmap.New[string, bool]()
-			om.Set(result.Avs, existingAvs)
-
-			prev := om.GetPair(result.Avs).Prev()
-			if prev != nil && strings.Compare(prev.Key, result.Avs) >= 0 {
-				om.Delete(result.Avs)
-				return nil, fmt.Errorf("avs not in order")
-			}
-		}
-		existingAvs.Set(result.Operator, result.Registered)
-
-		prev := existingAvs.GetPair(result.Operator).Prev()
-		if prev != nil && strings.Compare(prev.Key, result.Operator) >= 0 {
-			existingAvs.Delete(result.Operator)
-			return nil, fmt.Errorf("operator not in order")
-		}
-	}
-
-	avsLeaves := a.InitializeMerkleTreeBaseStateWithBlock(blockNumber)
-
-	for avs := om.Oldest(); avs != nil; avs = avs.Next() {
-		operatorLeafs := make([][]byte, 0)
-		for operator := avs.Value.Oldest(); operator != nil; operator = operator.Next() {
-			operatorAddr := operator.Key
-			registered := operator.Value
-			operatorLeafs = append(operatorLeafs, encodeOperatorLeaf(operatorAddr, registered))
-		}
-
-		avsTree, err := merkletree.NewTree(
-			merkletree.WithData(operatorLeafs),
-			merkletree.WithHashType(keccak256.New()),
-		)
-		if err != nil {
-			return nil, err
-		}
-
-		avsLeaves = append(avsLeaves, encodeAvsLeaf(avs.Key, avsTree.Root()))
+func (a *AvsOperatorsModel) sortValuesForMerkleTree(diffs []*RegisteredAvsOperatorDiff) []*base.MerkleTreeInput {
+	inputs := make([]*base.MerkleTreeInput, 0)
+	for _, diff := range diffs {
+		inputs = append(inputs, &base.MerkleTreeInput{
+			SlotID: NewSlotID(diff.Avs, diff.Operator),
+			Value:  []byte(fmt.Sprintf("%t", diff.Registered)),
+		})
 	}
-
-	return merkletree.NewTree(
-		merkletree.WithData(avsLeaves),
-		merkletree.WithHashType(keccak256.New()),
-	)
+	slices.SortFunc(inputs, func(i, j *base.MerkleTreeInput) int {
+		return strings.Compare(string(i.SlotID), string(j.SlotID))
+	})
+	return inputs
 }
 
 func (a *AvsOperatorsModel) DeleteState(startBlockNumber uint64, endBlockNumber uint64) error {
-	return a.BaseEigenState.DeleteState("registered_avs_operators", startBlockNumber, endBlockNumber, a.Db)
-}
-
-func encodeOperatorLeaf(operator string, registered bool) []byte {
-	return []byte(fmt.Sprintf("%s:%t", operator, registered))
-}
-
-func encodeAvsLeaf(avs string, avsOperatorRoot []byte) []byte {
-	return append([]byte(avs), avsOperatorRoot...)
+	return a.BaseEigenState.DeleteState("registered_avs_operators", startBlockNumber, endBlockNumber, a.DB)
 }

internal/eigenState/avsOperators/avsOperators_test.go

@@ -37,8 +37,8 @@ func setup() (
 }
 
 func teardown(model *AvsOperatorsModel) {
-	model.Db.Exec("delete from avs_operator_changes")
-	model.Db.Exec("delete from registered_avs_operators")
+	model.DB.Exec("delete from avs_operator_changes")
+	model.DB.Exec("delete from registered_avs_operators")
 }
 
 func Test_AvsOperatorState(t *testing.T) {
@@ -50,7 +50,7 @@ func Test_AvsOperatorState(t *testing.T) {
 
 	t.Run("Should create a new AvsOperatorState", func(t *testing.T) {
 		esm := stateManager.NewEigenStateManager(l, grm)
-		avsOperatorState, err := NewAvsOperators(esm, grm, cfg.Network, cfg.Environment, l, cfg)
+		avsOperatorState, err := NewAvsOperators(esm, grm, l, cfg)
 		assert.Nil(t, err)
 		assert.NotNil(t, avsOperatorState)
 	})
@@ -71,7 +71,7 @@ func Test_AvsOperatorState(t *testing.T) {
 			DeletedAt:        time.Time{},
 		}
 
-		avsOperatorState, err := NewAvsOperators(esm, grm, cfg.Network, cfg.Environment, l, cfg)
+		avsOperatorState, err := NewAvsOperators(esm, grm, l, cfg)
 
 		assert.Equal(t, true, avsOperatorState.IsInterestingLog(&log))
 
@@ -102,7 +102,7 @@ func Test_AvsOperatorState(t *testing.T) {
 			DeletedAt:        time.Time{},
 		}
 
-		avsOperatorState, err := NewAvsOperators(esm, grm, cfg.Network, cfg.Environment, l, cfg)
+		avsOperatorState, err := NewAvsOperators(esm, grm, l, cfg)
 		assert.Nil(t, err)
 
 		assert.Equal(t, true, avsOperatorState.IsInterestingLog(&log))
@@ -118,7 +118,7 @@ func Test_AvsOperatorState(t *testing.T) {
 		assert.Nil(t, err)
 
 		states := []RegisteredAvsOperators{}
-		statesRes := avsOperatorState.Db.
+		statesRes := avsOperatorState.DB.
 			Model(&RegisteredAvsOperators{}).
 			Raw("select * from registered_avs_operators where block_number = @blockNumber", sql.Named("blockNumber", blockNumber)).
 			Scan(&states)
@@ -170,7 +170,7 @@ func Test_AvsOperatorState(t *testing.T) {
 			},
 		}
 
-		avsOperatorState, err := NewAvsOperators(esm, grm, cfg.Network, cfg.Environment, l, cfg)
+		avsOperatorState, err := NewAvsOperators(esm, grm, l, cfg)
 		assert.Nil(t, err)
 
 		for _, log := range logs {
@@ -187,7 +187,7 @@ func Test_AvsOperatorState(t *testing.T) {
 			assert.Nil(t, err)
 
 			states := []RegisteredAvsOperators{}
-			statesRes := avsOperatorState.Db.
+			statesRes := avsOperatorState.DB.
 				Model(&RegisteredAvsOperators{}).
 				Raw("select * from registered_avs_operators where block_number = @blockNumber", sql.Named("blockNumber", log.BlockNumber)).
 				Scan(&states)