Merge pull request #3 from Layr-Labs/sm-avsDiff

Only generate root based on diffs

internal/eigenState/avsOperators/avsOperators.go

@@ -4,8 +4,11 @@ import (
 	"database/sql"
 	"fmt"
 	"github.com/Layr-Labs/sidecar/internal/config"
-	"github.com/Layr-Labs/sidecar/internal/eigenState"
+	"github.com/Layr-Labs/sidecar/internal/eigenState/base"
+	"github.com/Layr-Labs/sidecar/internal/eigenState/stateManager"
+	"github.com/Layr-Labs/sidecar/internal/eigenState/types"
 	"github.com/Layr-Labs/sidecar/internal/storage"
+	"github.com/Layr-Labs/sidecar/internal/utils"
 	"github.com/wealdtech/go-merkletree/v2"
 	"github.com/wealdtech/go-merkletree/v2/keccak256"
 	orderedmap "github.com/wk8/go-ordered-map/v2"
@@ -41,53 +44,80 @@ type AvsOperatorChange struct {
 
 // EigenState model for AVS operators that implements IEigenStateModel
 type AvsOperators struct {
-	eigenState.BaseEigenState
-	StateTransitions eigenState.StateTransitions[AvsOperatorChange]
+	base.BaseEigenState
+	StateTransitions types.StateTransitions[AvsOperatorChange]
 	Db               *gorm.DB
 	Network          config.Network
 	Environment      config.Environment
 	logger           *zap.Logger
 	globalConfig     *config.Config
 }
 
+type RegisteredAvsOperatorDiff struct {
+	Operator    string
+	Avs         string
+	BlockNumber uint64
+	Registered  bool
+}
+
 // Create new instance of AvsOperators state model
 func NewAvsOperators(
-	esm *eigenState.EigenStateManager,
+	esm *stateManager.EigenStateManager,
 	grm *gorm.DB,
 	Network config.Network,
 	Environment config.Environment,
 	logger *zap.Logger,
 	globalConfig *config.Config,
 ) (*AvsOperators, error) {
 	s := &AvsOperators{
-		BaseEigenState: eigenState.BaseEigenState{},
-		Db:             grm,
-		Network:        Network,
-		Environment:    Environment,
-		logger:         logger,
-		globalConfig:   globalConfig,
+		BaseEigenState: base.BaseEigenState{
+			Logger: logger,
+		},
+		Db:           grm,
+		Network:      Network,
+		Environment:  Environment,
+		logger:       logger,
+		globalConfig: globalConfig,
 	}
-	esm.RegisterState(s)
+	esm.RegisterState(s, 0)
 	return s, nil
 }
 
+func (a *AvsOperators) GetModelName() string {
+	return "AvsOperators"
+}
+
 // Get the state transitions for the AvsOperators state model
 //
 // Each state transition is function indexed by a block number.
 // BlockNumber 0 is the catchall state
 //
 // Returns the map and a reverse sorted list of block numbers that can be traversed when
 // processing a log to determine which state change to apply.
-func (a *AvsOperators) GetStateTransitions() (eigenState.StateTransitions[AvsOperatorChange], []uint64) {
-	stateChanges := make(eigenState.StateTransitions[AvsOperatorChange])
+func (a *AvsOperators) GetStateTransitions() (types.StateTransitions[AvsOperatorChange], []uint64) {
+	stateChanges := make(types.StateTransitions[AvsOperatorChange])
 
 	// TODO(seanmcgary): make this not a closure so this function doesnt get big an messy...
 	stateChanges[0] = func(log *storage.TransactionLog) (*AvsOperatorChange, error) {
-		// TODO(seanmcgary): actually parse the log
+		arguments, err := a.ParseLogArguments(log)
+		if err != nil {
+			return nil, err
+		}
+
+		outputData, err := a.ParseLogOutput(log)
+		if err != nil {
+			return nil, err
+		}
+
+		registered := false
+		if val, ok := outputData["status"]; ok {
+			registered = uint64(val.(float64)) == 1
+		}
+
 		change := &AvsOperatorChange{
-			Operator:         "operator",
-			Avs:              "avs",
-			Registered:       true,
+			Operator:         arguments[0].Value.(string),
+			Avs:              arguments[1].Value.(string),
+			Registered:       registered,
 			TransactionHash:  log.TransactionHash,
 			TransactionIndex: log.TransactionIndex,
 			LogIndex:         log.LogIndex,
@@ -203,6 +233,7 @@ func (a *AvsOperators) WriteFinalState(blockNumber uint64) error {
 				and aoc.operator = nc.operator
 				and aoc.log_index = nc.log_index
 				and aoc.transaction_index = nc.transaction_index
+				and aoc.block_number = nc.block_number
 			)
 		),
 		unregistrations as (
@@ -241,85 +272,132 @@ func (a *AvsOperators) WriteFinalState(blockNumber uint64) error {
 	return nil
 }
 
+func (a *AvsOperators) getDifferenceInStates(blockNumber uint64) ([]RegisteredAvsOperatorDiff, error) {
+	query := `
+		with new_states as (
+			select
+				avs,
+				operator,
+				block_number,
+				true as registered
+			from registered_avs_operators
+			where block_number = @currentBlock
+		),
+		previous_states as (
+			select
+				avs,
+				operator,
+				block_number,
+				true as registered
+			from registered_avs_operators
+			where block_number = @previousBlock
+		),
+		unregistered as (
+			(select avs, operator, registered from previous_states)
+			except
+			(select avs, operator, registered from new_states)
+		),
+		new_registered as (
+			(select avs, operator, registered from new_states)
+			except
+			(select avs, operator, registered from previous_states)
+		)
+		select avs, operator, false as registered from unregistered
+		union all
+		select avs, operator, true as registered from new_registered;
+	`
+	results := make([]RegisteredAvsOperatorDiff, 0)
+	res := a.Db.Model(&RegisteredAvsOperatorDiff{}).
+		Raw(query,
+			sql.Named("currentBlock", blockNumber),
+			sql.Named("previousBlock", blockNumber-1),
+		).
+		Scan(&results)
+
+	if res.Error != nil {
+		a.logger.Sugar().Errorw("Failed to fetch registered_avs_operators", zap.Error(res.Error))
+		return nil, res.Error
+	}
+	return results, nil
+}
+
 // Generates a state root for the given block number.
 //
 // 1. Select all registered_avs_operators for the given block number ordered by avs and operator asc
 // 2. Create an ordered map, with AVSs at the top level that point to an ordered map of operators and block numbers
 // 3. Create a merkle tree for each AVS, with the operator:block_number pairs as leaves
 // 4. Create a merkle tree for all AVS trees
 // 5. Return the root of the full tree
-func (a *AvsOperators) GenerateStateRoot(blockNumber uint64) (eigenState.StateRoot, error) {
-	query := `
-		select
-			avs,
-			operator,
-			block_number
-		from registered_avs_operators
-		where
-			block_number = @blockNumber
-		order by avs asc, operator asc
-	`
-	results := make([]RegisteredAvsOperators, 0)
-	res := a.Db.Model(&results).Raw(query, sql.Named("blockNumber", blockNumber))
+func (a *AvsOperators) GenerateStateRoot(blockNumber uint64) (types.StateRoot, error) {
+	results, err := a.getDifferenceInStates(blockNumber)
+	if err != nil {
+		return "", err
+	}
 
-	if res.Error != nil {
-		a.logger.Sugar().Errorw("Failed to fetch registered_avs_operators", zap.Error(res.Error))
-		return "", res.Error
+	fullTree, err := a.merkelizeState(blockNumber, results)
+	if err != nil {
+		return "", err
 	}
+	return types.StateRoot(utils.ConvertBytesToString(fullTree.Root())), nil
+}
 
-	// Avs -> operator:block_number
-	om := orderedmap.New[string, *orderedmap.OrderedMap[string, uint64]]()
+func (a *AvsOperators) merkelizeState(blockNumber uint64, avsOperators []RegisteredAvsOperatorDiff) (*merkletree.MerkleTree, error) {
+	// Avs -> operator:registered
+	om := orderedmap.New[string, *orderedmap.OrderedMap[string, bool]]()
 
-	for _, result := range results {
+	for _, result := range avsOperators {
 		existingAvs, found := om.Get(result.Avs)
 		if !found {
-			existingAvs = orderedmap.New[string, uint64]()
+			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 "", fmt.Errorf("avs not in order")
+				return nil, fmt.Errorf("avs not in order")
 			}
 		}
-		existingAvs.Set(result.Operator, result.BlockNumber)
+		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 "", fmt.Errorf("operator not in order")
+			return nil, fmt.Errorf("operator not in order")
 		}
 	}
 
-	avsLeaves := make([][]byte, 0)
+	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
-			block := operator.Value
-			operatorLeafs = append(operatorLeafs, []byte(fmt.Sprintf("%s:%d", operatorAddr, block)))
+			registered := operator.Value
+			operatorLeafs = append(operatorLeafs, encodeOperatorLeaf(operatorAddr, registered))
 		}
 
 		avsTree, err := merkletree.NewTree(
 			merkletree.WithData(operatorLeafs),
 			merkletree.WithHashType(keccak256.New()),
 		)
 		if err != nil {
-			return "", err
+			return nil, err
 		}
 
-		avsBytes := []byte(avs.Key)
-		root := avsTree.Root()
-		avsLeaves = append(avsLeaves, append(avsBytes, root[:]...))
+		avsLeaves = append(avsLeaves, encodeAvsLeaf(avs.Key, avsTree.Root()))
 	}
 
-	fullTree, err := merkletree.NewTree(
+	return merkletree.NewTree(
 		merkletree.WithData(avsLeaves),
 		merkletree.WithHashType(keccak256.New()),
 	)
-	if err != nil {
-		return "", err
-	}
-	return eigenState.StateRoot(fullTree.Root()), nil
+}
+
+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[:]...)
 }