From 4d5abd654302f74358ddd85e9bc440d5e1edf90f Mon Sep 17 00:00:00 2001
From: Sean McGary <sean@seanmcgary.com>
Date: Thu, 5 Sep 2024 21:00:19 -0500
Subject: [PATCH 1/7] Change avsOperators model to only create a tree from
 diffs

---
 .../eigenState/avsOperators/avsOperators.go   | 150 +++++++++++++-----
 .../avsOperators/avsOperators_test.go         | 129 +++++++++++++--
 internal/eigenState/eigenstate.go             |  42 +++++
 .../operatorShares/operatorShares.go          |   2 -
 4 files changed, 265 insertions(+), 58 deletions(-)

diff --git a/internal/eigenState/avsOperators/avsOperators.go b/internal/eigenState/avsOperators/avsOperators.go
index 876a1ae8..6449b2cc 100644
--- a/internal/eigenState/avsOperators/avsOperators.go
+++ b/internal/eigenState/avsOperators/avsOperators.go
@@ -6,6 +6,7 @@ import (
 	"github.com/Layr-Labs/sidecar/internal/config"
 	"github.com/Layr-Labs/sidecar/internal/eigenState"
 	"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"
@@ -50,6 +51,13 @@ type AvsOperators struct {
 	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,
@@ -60,12 +68,14 @@ func NewAvsOperators(
 	globalConfig *config.Config,
 ) (*AvsOperators, error) {
 	s := &AvsOperators{
-		BaseEigenState: eigenState.BaseEigenState{},
-		Db:             grm,
-		Network:        Network,
-		Environment:    Environment,
-		logger:         logger,
-		globalConfig:   globalConfig,
+		BaseEigenState: eigenState.BaseEigenState{
+			Logger: logger,
+		},
+		Db:           grm,
+		Network:      Network,
+		Environment:  Environment,
+		logger:       logger,
+		globalConfig: globalConfig,
 	}
 	esm.RegisterState(s)
 	return s, nil
@@ -83,11 +93,25 @@ func (a *AvsOperators) GetStateTransitions() (eigenState.StateTransitions[AvsOpe
 
 	// 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 +227,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,6 +266,55 @@ 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
@@ -249,28 +323,23 @@ func (a *AvsOperators) WriteFinalState(blockNumber uint64) error {
 // 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))
+	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 eigenState.StateRoot(utils.ConvertBytesToString(fullTree.Root())), nil
+}
 
+func (a *AvsOperators) merkelizeState(blockNumber uint64, avsOperators []RegisteredAvsOperatorDiff) (*merkletree.MerkleTree, error) {
 	// Avs -> operator:block_number
 	om := orderedmap.New[string, *orderedmap.OrderedMap[string, uint64]]()
 
-	for _, result := range results {
+	for _, result := range avsOperators {
 		existingAvs, found := om.Get(result.Avs)
 		if !found {
 			existingAvs = orderedmap.New[string, uint64]()
@@ -279,7 +348,7 @@ func (a *AvsOperators) GenerateStateRoot(blockNumber uint64) (eigenState.StateRo
 			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)
@@ -287,18 +356,19 @@ func (a *AvsOperators) GenerateStateRoot(blockNumber uint64) (eigenState.StateRo
 		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.InitializeBaseStateWithBlock(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)))
+			operatorLeafs = append(operatorLeafs, encodeOperatorLeaf(operatorAddr, block))
 		}
 
 		avsTree, err := merkletree.NewTree(
@@ -306,20 +376,22 @@ func (a *AvsOperators) GenerateStateRoot(blockNumber uint64) (eigenState.StateRo
 			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, blockNumber uint64) []byte {
+	return []byte(fmt.Sprintf("%s:%d", operator, blockNumber))
+}
+
+func encodeAvsLeaf(avs string, avsOperatorRoot []byte) []byte {
+	return append([]byte(avs), avsOperatorRoot[:]...)
 }
diff --git a/internal/eigenState/avsOperators/avsOperators_test.go b/internal/eigenState/avsOperators/avsOperators_test.go
index 0de181d6..98c7a92a 100644
--- a/internal/eigenState/avsOperators/avsOperators_test.go
+++ b/internal/eigenState/avsOperators/avsOperators_test.go
@@ -2,7 +2,6 @@ package avsOperators
 
 import (
 	"database/sql"
-	"fmt"
 	"github.com/Layr-Labs/sidecar/internal/config"
 	"github.com/Layr-Labs/sidecar/internal/eigenState"
 	"github.com/Layr-Labs/sidecar/internal/logger"
@@ -32,6 +31,11 @@ func setup() (
 	return cfg, grm, l, eigenState, err
 }
 
+func teardown(model *AvsOperators) {
+	model.Db.Exec("truncate table avs_operator_changes cascade")
+	model.Db.Exec("truncate table registered_avs_operators cascade")
+}
+
 func Test_AvsOperatorState(t *testing.T) {
 	cfg, grm, l, esm, err := setup()
 
@@ -45,42 +49,45 @@ func Test_AvsOperatorState(t *testing.T) {
 		assert.NotNil(t, avsOperatorState)
 	})
 	t.Run("Should register AvsOperatorState", func(t *testing.T) {
+		blockNumber := uint64(200)
 		log := storage.TransactionLog{
 			TransactionHash:  "some hash",
 			TransactionIndex: 100,
-			BlockNumber:      200,
+			BlockNumber:      blockNumber,
 			BlockSequenceId:  300,
-			Address:          "some address",
-			Arguments:        "some arguments",
+			Address:          cfg.GetContractsMapForEnvAndNetwork().AvsDirectory,
+			Arguments:        `[{"Value": "0xdf25bdcdcdd9a3dd8c9069306c4dba8d90dd8e8e" }, { "Value": "0x870679e138bcdf293b7ff14dd44b70fc97e12fc0" }]`,
 			EventName:        "OperatorAVSRegistrationStatusUpdated",
 			LogIndex:         400,
-			OutputData:       "some output data",
+			OutputData:       `{ "status": 1 }`,
 			CreatedAt:        time.Time{},
 			UpdatedAt:        time.Time{},
 			DeletedAt:        time.Time{},
 		}
 
 		avsOperatorState, err := NewAvsOperators(esm, grm, cfg.Network, cfg.Environment, l, cfg)
-		fmt.Printf("avsOperatorState err: %+v\n", err)
+
+		assert.Equal(t, true, avsOperatorState.IsInterestingLog(&log))
 
 		res, err := avsOperatorState.HandleStateChange(&log)
 		assert.Nil(t, err)
-		t.Logf("res_typed: %+v\n", res)
+		assert.NotNil(t, res)
 
-		avsOperatorState.Db.Raw("truncate table avs_operator_changes cascade").Scan(&res)
-		avsOperatorState.Db.Raw("truncate table registered_avs_operators cascade").Scan(&res)
+		teardown(avsOperatorState)
 	})
 	t.Run("Should register AvsOperatorState and generate the table for the block", func(t *testing.T) {
+		blockNumber := uint64(200)
+
 		log := storage.TransactionLog{
 			TransactionHash:  "some hash",
 			TransactionIndex: 100,
-			BlockNumber:      200,
+			BlockNumber:      blockNumber,
 			BlockSequenceId:  300,
-			Address:          "some address",
-			Arguments:        "some arguments",
+			Address:          cfg.GetContractsMapForEnvAndNetwork().AvsDirectory,
+			Arguments:        `[{"Value": "0xdf25bdcdcdd9a3dd8c9069306c4dba8d90dd8e8e" }, { "Value": "0x870679e138bcdf293b7ff14dd44b70fc97e12fc0" }]`,
 			EventName:        "OperatorAVSRegistrationStatusUpdated",
 			LogIndex:         400,
-			OutputData:       "some output data",
+			OutputData:       `{ "status": 1 }`,
 			CreatedAt:        time.Time{},
 			UpdatedAt:        time.Time{},
 			DeletedAt:        time.Time{},
@@ -89,23 +96,111 @@ func Test_AvsOperatorState(t *testing.T) {
 		avsOperatorState, err := NewAvsOperators(esm, grm, cfg.Network, cfg.Environment, l, cfg)
 		assert.Nil(t, err)
 
+		assert.Equal(t, true, avsOperatorState.IsInterestingLog(&log))
+
 		stateChange, err := avsOperatorState.HandleStateChange(&log)
 		assert.Nil(t, err)
-		fmt.Printf("stateChange: %+v\n", stateChange)
+		assert.NotNil(t, stateChange)
 
-		err = avsOperatorState.WriteFinalState(200)
+		err = avsOperatorState.WriteFinalState(blockNumber)
 		assert.Nil(t, err)
 
 		states := []RegisteredAvsOperators{}
 		statesRes := avsOperatorState.Db.
 			Model(&RegisteredAvsOperators{}).
-			Raw("select * from registered_avs_operators where block_number = @blockNumber", sql.Named("blockNumber", 200)).
+			Raw("select * from registered_avs_operators where block_number = @blockNumber", sql.Named("blockNumber", blockNumber)).
 			Scan(&states)
 
 		if statesRes.Error != nil {
 			t.Fatalf("Failed to fetch registered_avs_operators: %v", statesRes.Error)
 		}
 		assert.Equal(t, 1, len(states))
-		fmt.Printf("states: %+v\n", states)
+
+		stateRoot, err := avsOperatorState.GenerateStateRoot(blockNumber)
+		assert.Nil(t, err)
+		assert.True(t, len(stateRoot) > 0)
+
+		teardown(avsOperatorState)
+	})
+	t.Run("Should correctly generate state across multiple blocks", func(t *testing.T) {
+		blocks := []uint64{
+			300,
+			301,
+		}
+
+		logs := []*storage.TransactionLog{
+			&storage.TransactionLog{
+				TransactionHash:  "some hash",
+				TransactionIndex: 100,
+				BlockNumber:      blocks[0],
+				BlockSequenceId:  300,
+				Address:          cfg.GetContractsMapForEnvAndNetwork().AvsDirectory,
+				Arguments:        `[{"Value": "0xdf25bdcdcdd9a3dd8c9069306c4dba8d90dd8e8e" }, { "Value": "0x870679e138bcdf293b7ff14dd44b70fc97e12fc0" }]`,
+				EventName:        "OperatorAVSRegistrationStatusUpdated",
+				LogIndex:         400,
+				OutputData:       `{ "status": 1 }`,
+				CreatedAt:        time.Time{},
+				UpdatedAt:        time.Time{},
+				DeletedAt:        time.Time{},
+			},
+			&storage.TransactionLog{
+				TransactionHash:  "some hash",
+				TransactionIndex: 100,
+				BlockNumber:      blocks[1],
+				BlockSequenceId:  300,
+				Address:          cfg.GetContractsMapForEnvAndNetwork().AvsDirectory,
+				Arguments:        `[{"Value": "0xdf25bdcdcdd9a3dd8c9069306c4dba8d90dd8e8e" }, { "Value": "0x870679e138bcdf293b7ff14dd44b70fc97e12fc0" }]`,
+				EventName:        "OperatorAVSRegistrationStatusUpdated",
+				LogIndex:         400,
+				OutputData:       `{ "status": 0 }`,
+				CreatedAt:        time.Time{},
+				UpdatedAt:        time.Time{},
+				DeletedAt:        time.Time{},
+			},
+		}
+
+		avsOperatorState, err := NewAvsOperators(esm, grm, cfg.Network, cfg.Environment, l, cfg)
+		assert.Nil(t, err)
+
+		for _, log := range logs {
+			assert.True(t, avsOperatorState.IsInterestingLog(log))
+
+			stateChange, err := avsOperatorState.HandleStateChange(log)
+			assert.Nil(t, err)
+			assert.NotNil(t, stateChange)
+
+			err = avsOperatorState.WriteFinalState(log.BlockNumber)
+			assert.Nil(t, err)
+
+			states := []RegisteredAvsOperators{}
+			statesRes := avsOperatorState.Db.
+				Model(&RegisteredAvsOperators{}).
+				Raw("select * from registered_avs_operators where block_number = @blockNumber", sql.Named("blockNumber", log.BlockNumber)).
+				Scan(&states)
+
+			if statesRes.Error != nil {
+				t.Fatalf("Failed to fetch registered_avs_operators: %v", statesRes.Error)
+			}
+
+			if log.BlockNumber == blocks[0] {
+				assert.Equal(t, 1, len(states))
+				diffs, err := avsOperatorState.getDifferenceInStates(log.BlockNumber)
+				assert.Nil(t, err)
+				assert.Equal(t, 1, len(diffs))
+				assert.Equal(t, true, diffs[0].Registered)
+			} else if log.BlockNumber == blocks[1] {
+				assert.Equal(t, 0, len(states))
+				diffs, err := avsOperatorState.getDifferenceInStates(log.BlockNumber)
+				assert.Nil(t, err)
+				assert.Equal(t, 1, len(diffs))
+				assert.Equal(t, false, diffs[0].Registered)
+			}
+
+			stateRoot, err := avsOperatorState.GenerateStateRoot(log.BlockNumber)
+			assert.Nil(t, err)
+			assert.True(t, len(stateRoot) > 0)
+		}
+
+		teardown(avsOperatorState)
 	})
 }
diff --git a/internal/eigenState/eigenstate.go b/internal/eigenState/eigenstate.go
index 1198ae3b..c23ca149 100644
--- a/internal/eigenState/eigenstate.go
+++ b/internal/eigenState/eigenstate.go
@@ -1,6 +1,9 @@
 package eigenState
 
 import (
+	"encoding/json"
+	"fmt"
+	"github.com/Layr-Labs/sidecar/internal/parser"
 	"github.com/Layr-Labs/sidecar/internal/storage"
 	"go.uber.org/zap"
 )
@@ -78,6 +81,45 @@ type IEigenStateModel interface {
 }
 
 type BaseEigenState struct {
+	Logger *zap.Logger
+}
+
+func (b *BaseEigenState) ParseLogArguments(log *storage.TransactionLog) ([]parser.Argument, error) {
+	arguments := make([]parser.Argument, 0)
+	err := json.Unmarshal([]byte(log.Arguments), &arguments)
+	if err != nil {
+		b.Logger.Sugar().Errorw("Failed to unmarshal arguments",
+			zap.Error(err),
+			zap.String("transactionHash", log.TransactionHash),
+			zap.Uint64("transactionIndex", log.TransactionIndex),
+		)
+		return nil, err
+	}
+	return arguments, nil
+}
+
+func (b *BaseEigenState) ParseLogOutput(log *storage.TransactionLog) (map[string]interface{}, error) {
+	outputData := make(map[string]interface{})
+	err := json.Unmarshal([]byte(log.OutputData), &outputData)
+	if err != nil {
+		b.Logger.Sugar().Errorw("Failed to unmarshal outputData",
+			zap.Error(err),
+			zap.String("transactionHash", log.TransactionHash),
+			zap.Uint64("transactionIndex", log.TransactionIndex),
+		)
+		return nil, err
+	}
+	return outputData, nil
+}
+
+// Include the block number as the first item in the tree.
+// This does two things:
+// 1. Ensures that the tree is always different for different blocks
+// 2. Allows us to have at least 1 value if there are no model changes for a block
+func (b *BaseEigenState) InitializeBaseStateWithBlock(blockNumber uint64) [][]byte {
+	return [][]byte{
+		[]byte(fmt.Sprintf("%d", blockNumber)),
+	}
 }
 
 // Map of block number to function that will transition the state to the next block
diff --git a/internal/eigenState/operatorShares/operatorShares.go b/internal/eigenState/operatorShares/operatorShares.go
index b3e98c0c..0b67440b 100644
--- a/internal/eigenState/operatorShares/operatorShares.go
+++ b/internal/eigenState/operatorShares/operatorShares.go
@@ -100,7 +100,6 @@ func (osm *OperatorSharesModel) GetStateTransitions() (eigenState.StateTransitio
 			)
 			return nil, err
 		}
-		fmt.Printf("Outputdata: %+v\n", outputData)
 		shares := big.Int{}
 		sharesInt, _ := shares.SetString(outputData["shares"].(string), 10)
 
@@ -117,7 +116,6 @@ func (osm *OperatorSharesModel) GetStateTransitions() (eigenState.StateTransitio
 			LogIndex:         log.LogIndex,
 			BlockNumber:      log.BlockNumber,
 		}
-		fmt.Printf("Change: %+v\n", change)
 		return change, nil
 	}
 

From e5046be635f9b38b4f33a6ce7ac38a584b7d3682 Mon Sep 17 00:00:00 2001
From: Sean McGary <sean@seanmcgary.com>
Date: Thu, 5 Sep 2024 21:02:03 -0500
Subject: [PATCH 2/7] Use registered flag, not block number

---
 internal/eigenState/avsOperators/avsOperators.go | 16 ++++++++--------
 1 file changed, 8 insertions(+), 8 deletions(-)

diff --git a/internal/eigenState/avsOperators/avsOperators.go b/internal/eigenState/avsOperators/avsOperators.go
index 6449b2cc..61d34a69 100644
--- a/internal/eigenState/avsOperators/avsOperators.go
+++ b/internal/eigenState/avsOperators/avsOperators.go
@@ -336,13 +336,13 @@ func (a *AvsOperators) GenerateStateRoot(blockNumber uint64) (eigenState.StateRo
 }
 
 func (a *AvsOperators) merkelizeState(blockNumber uint64, avsOperators []RegisteredAvsOperatorDiff) (*merkletree.MerkleTree, error) {
-	// Avs -> operator:block_number
-	om := orderedmap.New[string, *orderedmap.OrderedMap[string, uint64]]()
+	// 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, uint64]()
+			existingAvs = orderedmap.New[string, bool]()
 			om.Set(result.Avs, existingAvs)
 
 			prev := om.GetPair(result.Avs).Prev()
@@ -351,7 +351,7 @@ func (a *AvsOperators) merkelizeState(blockNumber uint64, avsOperators []Registe
 				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 {
@@ -367,8 +367,8 @@ func (a *AvsOperators) merkelizeState(blockNumber uint64, avsOperators []Registe
 		operatorLeafs := make([][]byte, 0)
 		for operator := avs.Value.Oldest(); operator != nil; operator = operator.Next() {
 			operatorAddr := operator.Key
-			block := operator.Value
-			operatorLeafs = append(operatorLeafs, encodeOperatorLeaf(operatorAddr, block))
+			registered := operator.Value
+			operatorLeafs = append(operatorLeafs, encodeOperatorLeaf(operatorAddr, registered))
 		}
 
 		avsTree, err := merkletree.NewTree(
@@ -388,8 +388,8 @@ func (a *AvsOperators) merkelizeState(blockNumber uint64, avsOperators []Registe
 	)
 }
 
-func encodeOperatorLeaf(operator string, blockNumber uint64) []byte {
-	return []byte(fmt.Sprintf("%s:%d", operator, blockNumber))
+func encodeOperatorLeaf(operator string, registered bool) []byte {
+	return []byte(fmt.Sprintf("%s:%t", operator, registered))
 }
 
 func encodeAvsLeaf(avs string, avsOperatorRoot []byte) []byte {

From a952fc2db3c4c75a3f3ebe8640ded38c06dabffe Mon Sep 17 00:00:00 2001
From: Sean McGary <sean@seanmcgary.com>
Date: Thu, 5 Sep 2024 21:06:33 -0500
Subject: [PATCH 3/7] Include blockNumber in operator shares root

---
 internal/eigenState/avsOperators/avsOperators.go     | 2 +-
 internal/eigenState/eigenstate.go                    | 2 +-
 internal/eigenState/operatorShares/operatorShares.go | 6 +++---
 3 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/internal/eigenState/avsOperators/avsOperators.go b/internal/eigenState/avsOperators/avsOperators.go
index 61d34a69..e91345dc 100644
--- a/internal/eigenState/avsOperators/avsOperators.go
+++ b/internal/eigenState/avsOperators/avsOperators.go
@@ -360,7 +360,7 @@ func (a *AvsOperators) merkelizeState(blockNumber uint64, avsOperators []Registe
 		}
 	}
 
-	avsLeaves := a.InitializeBaseStateWithBlock(blockNumber)
+	avsLeaves := a.InitializeMerkleTreeBaseStateWithBlock(blockNumber)
 
 	for avs := om.Oldest(); avs != nil; avs = avs.Next() {
 
diff --git a/internal/eigenState/eigenstate.go b/internal/eigenState/eigenstate.go
index c23ca149..79f80aa8 100644
--- a/internal/eigenState/eigenstate.go
+++ b/internal/eigenState/eigenstate.go
@@ -116,7 +116,7 @@ func (b *BaseEigenState) ParseLogOutput(log *storage.TransactionLog) (map[string
 // This does two things:
 // 1. Ensures that the tree is always different for different blocks
 // 2. Allows us to have at least 1 value if there are no model changes for a block
-func (b *BaseEigenState) InitializeBaseStateWithBlock(blockNumber uint64) [][]byte {
+func (b *BaseEigenState) InitializeMerkleTreeBaseStateWithBlock(blockNumber uint64) [][]byte {
 	return [][]byte{
 		[]byte(fmt.Sprintf("%d", blockNumber)),
 	}
diff --git a/internal/eigenState/operatorShares/operatorShares.go b/internal/eigenState/operatorShares/operatorShares.go
index 0b67440b..4ed96b80 100644
--- a/internal/eigenState/operatorShares/operatorShares.go
+++ b/internal/eigenState/operatorShares/operatorShares.go
@@ -285,14 +285,14 @@ func (osm *OperatorSharesModel) GenerateStateRoot(blockNumber uint64) (eigenStat
 		return "", err
 	}
 
-	fullTree, err := osm.merkelizeState(diffs)
+	fullTree, err := osm.merkelizeState(blockNumber, diffs)
 	if err != nil {
 		return "", err
 	}
 	return eigenState.StateRoot(utils.ConvertBytesToString(fullTree.Root())), nil
 }
 
-func (osm *OperatorSharesModel) merkelizeState(diffs []OperatorShares) (*merkletree.MerkleTree, error) {
+func (osm *OperatorSharesModel) merkelizeState(blockNumber uint64, diffs []OperatorShares) (*merkletree.MerkleTree, error) {
 	// Create a merkle tree with the structure:
 	// strategy: map[operators]: shares
 	om := orderedmap.New[string, *orderedmap.OrderedMap[string, string]]()
@@ -318,7 +318,7 @@ func (osm *OperatorSharesModel) merkelizeState(diffs []OperatorShares) (*merklet
 		}
 	}
 
-	leaves := make([][]byte, 0)
+	leaves := osm.InitializeMerkleTreeBaseStateWithBlock(blockNumber)
 	for strat := om.Oldest(); strat != nil; strat = strat.Next() {
 
 		operatorLeaves := make([][]byte, 0)

From 2761271c80ed0a25f4455275d84e1e0abca81148 Mon Sep 17 00:00:00 2001
From: Sean McGary <sean@seanmcgary.com>
Date: Thu, 5 Sep 2024 21:08:43 -0500
Subject: [PATCH 4/7] Better organization of files

---
 internal/eigenState/baseEigenState.go |  51 +++++++++++
 internal/eigenState/eigenstate.go     | 126 --------------------------
 internal/eigenState/stateManager.go   |  60 ++++++++++++
 internal/eigenState/types.go          |  26 ++++++
 4 files changed, 137 insertions(+), 126 deletions(-)
 create mode 100644 internal/eigenState/baseEigenState.go
 delete mode 100644 internal/eigenState/eigenstate.go
 create mode 100644 internal/eigenState/stateManager.go
 create mode 100644 internal/eigenState/types.go

diff --git a/internal/eigenState/baseEigenState.go b/internal/eigenState/baseEigenState.go
new file mode 100644
index 00000000..c584d3e5
--- /dev/null
+++ b/internal/eigenState/baseEigenState.go
@@ -0,0 +1,51 @@
+package eigenState
+
+import (
+	"encoding/json"
+	"fmt"
+	"github.com/Layr-Labs/sidecar/internal/parser"
+	"github.com/Layr-Labs/sidecar/internal/storage"
+	"go.uber.org/zap"
+)
+
+type BaseEigenState struct {
+	Logger *zap.Logger
+}
+
+func (b *BaseEigenState) ParseLogArguments(log *storage.TransactionLog) ([]parser.Argument, error) {
+	arguments := make([]parser.Argument, 0)
+	err := json.Unmarshal([]byte(log.Arguments), &arguments)
+	if err != nil {
+		b.Logger.Sugar().Errorw("Failed to unmarshal arguments",
+			zap.Error(err),
+			zap.String("transactionHash", log.TransactionHash),
+			zap.Uint64("transactionIndex", log.TransactionIndex),
+		)
+		return nil, err
+	}
+	return arguments, nil
+}
+
+func (b *BaseEigenState) ParseLogOutput(log *storage.TransactionLog) (map[string]interface{}, error) {
+	outputData := make(map[string]interface{})
+	err := json.Unmarshal([]byte(log.OutputData), &outputData)
+	if err != nil {
+		b.Logger.Sugar().Errorw("Failed to unmarshal outputData",
+			zap.Error(err),
+			zap.String("transactionHash", log.TransactionHash),
+			zap.Uint64("transactionIndex", log.TransactionIndex),
+		)
+		return nil, err
+	}
+	return outputData, nil
+}
+
+// Include the block number as the first item in the tree.
+// This does two things:
+// 1. Ensures that the tree is always different for different blocks
+// 2. Allows us to have at least 1 value if there are no model changes for a block
+func (b *BaseEigenState) InitializeMerkleTreeBaseStateWithBlock(blockNumber uint64) [][]byte {
+	return [][]byte{
+		[]byte(fmt.Sprintf("%d", blockNumber)),
+	}
+}
diff --git a/internal/eigenState/eigenstate.go b/internal/eigenState/eigenstate.go
deleted file mode 100644
index 79f80aa8..00000000
--- a/internal/eigenState/eigenstate.go
+++ /dev/null
@@ -1,126 +0,0 @@
-package eigenState
-
-import (
-	"encoding/json"
-	"fmt"
-	"github.com/Layr-Labs/sidecar/internal/parser"
-	"github.com/Layr-Labs/sidecar/internal/storage"
-	"go.uber.org/zap"
-)
-
-type EigenStateManager struct {
-	StateModels []IEigenStateModel
-	logger      *zap.Logger
-}
-
-func NewEigenStateManager(logger *zap.Logger) *EigenStateManager {
-	return &EigenStateManager{
-		StateModels: make([]IEigenStateModel, 0),
-		logger:      logger,
-	}
-}
-
-// Allows a model to register itself with the state manager
-func (e *EigenStateManager) RegisterState(state IEigenStateModel) {
-	e.StateModels = append(e.StateModels, state)
-}
-
-// Given a log, allow each state model to determine if/how to process it
-func (e *EigenStateManager) HandleLogStateChange(log *storage.TransactionLog) error {
-	for _, state := range e.StateModels {
-		if state.IsInterestingLog(log) {
-			_, err := state.HandleStateChange(log)
-			if err != nil {
-				return err
-			}
-		}
-	}
-	return nil
-}
-
-// With all transactions/logs processed for a block, commit the final state to the table
-func (e *EigenStateManager) CommitFinalState(blockNumber uint64) error {
-	for _, state := range e.StateModels {
-		err := state.WriteFinalState(blockNumber)
-		if err != nil {
-			return err
-		}
-	}
-	return nil
-}
-
-type StateRoot string
-
-func (e *EigenStateManager) GenerateStateRoot(blockNumber uint64) (StateRoot, error) {
-	roots := make([]StateRoot, len(e.StateModels))
-	for i, state := range e.StateModels {
-		root, err := state.GenerateStateRoot(blockNumber)
-		if err != nil {
-			return "", err
-		}
-		roots[i] = root
-	}
-	// TODO: generate this
-	return "", nil
-}
-
-type IEigenStateModel interface {
-	// Determine if the log is interesting to the state model
-	IsInterestingLog(log *storage.TransactionLog) bool
-
-	// Allow the state model to handle the state change
-	//
-	// Returns the saved value. Listed as an interface because go generics suck
-	HandleStateChange(log *storage.TransactionLog) (interface{}, error)
-
-	// Once all state changes are processed, calculate and write final state
-	WriteFinalState(blockNumber uint64) error
-
-	// Generate the state root for the model
-	GenerateStateRoot(blockNumber uint64) (StateRoot, error)
-}
-
-type BaseEigenState struct {
-	Logger *zap.Logger
-}
-
-func (b *BaseEigenState) ParseLogArguments(log *storage.TransactionLog) ([]parser.Argument, error) {
-	arguments := make([]parser.Argument, 0)
-	err := json.Unmarshal([]byte(log.Arguments), &arguments)
-	if err != nil {
-		b.Logger.Sugar().Errorw("Failed to unmarshal arguments",
-			zap.Error(err),
-			zap.String("transactionHash", log.TransactionHash),
-			zap.Uint64("transactionIndex", log.TransactionIndex),
-		)
-		return nil, err
-	}
-	return arguments, nil
-}
-
-func (b *BaseEigenState) ParseLogOutput(log *storage.TransactionLog) (map[string]interface{}, error) {
-	outputData := make(map[string]interface{})
-	err := json.Unmarshal([]byte(log.OutputData), &outputData)
-	if err != nil {
-		b.Logger.Sugar().Errorw("Failed to unmarshal outputData",
-			zap.Error(err),
-			zap.String("transactionHash", log.TransactionHash),
-			zap.Uint64("transactionIndex", log.TransactionIndex),
-		)
-		return nil, err
-	}
-	return outputData, nil
-}
-
-// Include the block number as the first item in the tree.
-// This does two things:
-// 1. Ensures that the tree is always different for different blocks
-// 2. Allows us to have at least 1 value if there are no model changes for a block
-func (b *BaseEigenState) InitializeMerkleTreeBaseStateWithBlock(blockNumber uint64) [][]byte {
-	return [][]byte{
-		[]byte(fmt.Sprintf("%d", blockNumber)),
-	}
-}
-
-// Map of block number to function that will transition the state to the next block
-type StateTransitions[T interface{}] map[uint64]func(log *storage.TransactionLog) (*T, error)
diff --git a/internal/eigenState/stateManager.go b/internal/eigenState/stateManager.go
new file mode 100644
index 00000000..4f83407a
--- /dev/null
+++ b/internal/eigenState/stateManager.go
@@ -0,0 +1,60 @@
+package eigenState
+
+import (
+	"github.com/Layr-Labs/sidecar/internal/storage"
+	"go.uber.org/zap"
+)
+
+type EigenStateManager struct {
+	StateModels []IEigenStateModel
+	logger      *zap.Logger
+}
+
+func NewEigenStateManager(logger *zap.Logger) *EigenStateManager {
+	return &EigenStateManager{
+		StateModels: make([]IEigenStateModel, 0),
+		logger:      logger,
+	}
+}
+
+// Allows a model to register itself with the state manager
+func (e *EigenStateManager) RegisterState(state IEigenStateModel) {
+	e.StateModels = append(e.StateModels, state)
+}
+
+// Given a log, allow each state model to determine if/how to process it
+func (e *EigenStateManager) HandleLogStateChange(log *storage.TransactionLog) error {
+	for _, state := range e.StateModels {
+		if state.IsInterestingLog(log) {
+			_, err := state.HandleStateChange(log)
+			if err != nil {
+				return err
+			}
+		}
+	}
+	return nil
+}
+
+// With all transactions/logs processed for a block, commit the final state to the table
+func (e *EigenStateManager) CommitFinalState(blockNumber uint64) error {
+	for _, state := range e.StateModels {
+		err := state.WriteFinalState(blockNumber)
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func (e *EigenStateManager) GenerateStateRoot(blockNumber uint64) (StateRoot, error) {
+	roots := make([]StateRoot, len(e.StateModels))
+	for i, state := range e.StateModels {
+		root, err := state.GenerateStateRoot(blockNumber)
+		if err != nil {
+			return "", err
+		}
+		roots[i] = root
+	}
+	// TODO: generate this
+	return "", nil
+}
diff --git a/internal/eigenState/types.go b/internal/eigenState/types.go
new file mode 100644
index 00000000..83bd893c
--- /dev/null
+++ b/internal/eigenState/types.go
@@ -0,0 +1,26 @@
+package eigenState
+
+import (
+	"github.com/Layr-Labs/sidecar/internal/storage"
+)
+
+type StateRoot string
+
+type IEigenStateModel interface {
+	// Determine if the log is interesting to the state model
+	IsInterestingLog(log *storage.TransactionLog) bool
+
+	// Allow the state model to handle the state change
+	//
+	// Returns the saved value. Listed as an interface because go generics suck
+	HandleStateChange(log *storage.TransactionLog) (interface{}, error)
+
+	// Once all state changes are processed, calculate and write final state
+	WriteFinalState(blockNumber uint64) error
+
+	// Generate the state root for the model
+	GenerateStateRoot(blockNumber uint64) (StateRoot, error)
+}
+
+// Map of block number to function that will transition the state to the next block
+type StateTransitions[T interface{}] map[uint64]func(log *storage.TransactionLog) (*T, error)

From 1bbab6c1053748ec81ecc0e9b53c6c59271759d0 Mon Sep 17 00:00:00 2001
From: Sean McGary <sean@seanmcgary.com>
Date: Thu, 5 Sep 2024 21:20:25 -0500
Subject: [PATCH 5/7] Give models a name and make them ordered

---
 internal/eigenState/avsOperators/avsOperators.go      |  6 +++++-
 internal/eigenState/avsOperators/avsOperators_test.go | 11 ++++++-----
 internal/eigenState/operatorShares/operatorShares.go  |  6 +++++-
 .../eigenState/operatorShares/operatorShares_test.go  | 10 +++++-----
 internal/eigenState/stateManager.go                   | 11 +++++++----
 internal/eigenState/types.go                          | 11 ++++++++++-
 6 files changed, 38 insertions(+), 17 deletions(-)

diff --git a/internal/eigenState/avsOperators/avsOperators.go b/internal/eigenState/avsOperators/avsOperators.go
index e91345dc..4b1e42c4 100644
--- a/internal/eigenState/avsOperators/avsOperators.go
+++ b/internal/eigenState/avsOperators/avsOperators.go
@@ -77,10 +77,14 @@ func NewAvsOperators(
 		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.
diff --git a/internal/eigenState/avsOperators/avsOperators_test.go b/internal/eigenState/avsOperators/avsOperators_test.go
index 98c7a92a..cb70999e 100644
--- a/internal/eigenState/avsOperators/avsOperators_test.go
+++ b/internal/eigenState/avsOperators/avsOperators_test.go
@@ -18,7 +18,6 @@ func setup() (
 	*config.Config,
 	*gorm.DB,
 	*zap.Logger,
-	*eigenState.EigenStateManager,
 	error,
 ) {
 	cfg := tests.GetConfig()
@@ -26,9 +25,7 @@ func setup() (
 
 	_, grm, err := tests.GetDatabaseConnection(cfg)
 
-	eigenState := eigenState.NewEigenStateManager(l)
-
-	return cfg, grm, l, eigenState, err
+	return cfg, grm, l, err
 }
 
 func teardown(model *AvsOperators) {
@@ -37,18 +34,20 @@ func teardown(model *AvsOperators) {
 }
 
 func Test_AvsOperatorState(t *testing.T) {
-	cfg, grm, l, esm, err := setup()
+	cfg, grm, l, err := setup()
 
 	if err != nil {
 		t.Fatal(err)
 	}
 
 	t.Run("Should create a new AvsOperatorState", func(t *testing.T) {
+		esm := eigenState.NewEigenStateManager(l)
 		avsOperatorState, err := NewAvsOperators(esm, grm, cfg.Network, cfg.Environment, l, cfg)
 		assert.Nil(t, err)
 		assert.NotNil(t, avsOperatorState)
 	})
 	t.Run("Should register AvsOperatorState", func(t *testing.T) {
+		esm := eigenState.NewEigenStateManager(l)
 		blockNumber := uint64(200)
 		log := storage.TransactionLog{
 			TransactionHash:  "some hash",
@@ -76,6 +75,7 @@ func Test_AvsOperatorState(t *testing.T) {
 		teardown(avsOperatorState)
 	})
 	t.Run("Should register AvsOperatorState and generate the table for the block", func(t *testing.T) {
+		esm := eigenState.NewEigenStateManager(l)
 		blockNumber := uint64(200)
 
 		log := storage.TransactionLog{
@@ -123,6 +123,7 @@ func Test_AvsOperatorState(t *testing.T) {
 		teardown(avsOperatorState)
 	})
 	t.Run("Should correctly generate state across multiple blocks", func(t *testing.T) {
+		esm := eigenState.NewEigenStateManager(l)
 		blocks := []uint64{
 			300,
 			301,
diff --git a/internal/eigenState/operatorShares/operatorShares.go b/internal/eigenState/operatorShares/operatorShares.go
index 4ed96b80..0d9c16eb 100644
--- a/internal/eigenState/operatorShares/operatorShares.go
+++ b/internal/eigenState/operatorShares/operatorShares.go
@@ -72,10 +72,14 @@ func NewOperatorSharesModel(
 		globalConfig:   globalConfig,
 	}
 
-	esm.RegisterState(model)
+	esm.RegisterState(model, 1)
 	return model, nil
 }
 
+func (osm *OperatorSharesModel) GetModelName() string {
+	return "OperatorSharesModel"
+}
+
 func (osm *OperatorSharesModel) GetStateTransitions() (eigenState.StateTransitions[OperatorShareChange], []uint64) {
 	stateChanges := make(eigenState.StateTransitions[OperatorShareChange])
 
diff --git a/internal/eigenState/operatorShares/operatorShares_test.go b/internal/eigenState/operatorShares/operatorShares_test.go
index c5894690..0571fbc8 100644
--- a/internal/eigenState/operatorShares/operatorShares_test.go
+++ b/internal/eigenState/operatorShares/operatorShares_test.go
@@ -20,7 +20,6 @@ func setup() (
 	*config.Config,
 	*gorm.DB,
 	*zap.Logger,
-	*eigenState.EigenStateManager,
 	error,
 ) {
 	cfg := tests.GetConfig()
@@ -28,9 +27,7 @@ func setup() (
 
 	_, grm, err := tests.GetDatabaseConnection(cfg)
 
-	eigenState := eigenState.NewEigenStateManager(l)
-
-	return cfg, grm, l, eigenState, err
+	return cfg, grm, l, err
 }
 
 func teardown(model *OperatorSharesModel) {
@@ -39,18 +36,20 @@ func teardown(model *OperatorSharesModel) {
 }
 
 func Test_OperatorSharesState(t *testing.T) {
-	cfg, grm, l, esm, err := setup()
+	cfg, grm, l, err := setup()
 
 	if err != nil {
 		t.Fatal(err)
 	}
 
 	t.Run("Should create a new OperatorSharesState", func(t *testing.T) {
+		esm := eigenState.NewEigenStateManager(l)
 		model, err := NewOperatorSharesModel(esm, grm, cfg.Network, cfg.Environment, l, cfg)
 		assert.Nil(t, err)
 		assert.NotNil(t, model)
 	})
 	t.Run("Should register OperatorSharesState", func(t *testing.T) {
+		esm := eigenState.NewEigenStateManager(l)
 		blockNumber := uint64(200)
 		log := storage.TransactionLog{
 			TransactionHash:  "some hash",
@@ -76,6 +75,7 @@ func Test_OperatorSharesState(t *testing.T) {
 		teardown(model)
 	})
 	t.Run("Should register AvsOperatorState and generate the table for the block", func(t *testing.T) {
+		esm := eigenState.NewEigenStateManager(l)
 		blockNumber := uint64(200)
 		log := storage.TransactionLog{
 			TransactionHash:  "some hash",
diff --git a/internal/eigenState/stateManager.go b/internal/eigenState/stateManager.go
index 4f83407a..4082fd23 100644
--- a/internal/eigenState/stateManager.go
+++ b/internal/eigenState/stateManager.go
@@ -6,20 +6,23 @@ import (
 )
 
 type EigenStateManager struct {
-	StateModels []IEigenStateModel
+	StateModels map[int]IEigenStateModel
 	logger      *zap.Logger
 }
 
 func NewEigenStateManager(logger *zap.Logger) *EigenStateManager {
 	return &EigenStateManager{
-		StateModels: make([]IEigenStateModel, 0),
+		StateModels: make(map[int]IEigenStateModel),
 		logger:      logger,
 	}
 }
 
 // Allows a model to register itself with the state manager
-func (e *EigenStateManager) RegisterState(state IEigenStateModel) {
-	e.StateModels = append(e.StateModels, state)
+func (e *EigenStateManager) RegisterState(model IEigenStateModel, index int) {
+	if m, ok := e.StateModels[index]; ok {
+		e.logger.Sugar().Fatalf("Registering model model at index %d which already exists and belongs to %s", index, m.GetModelName())
+	}
+	e.StateModels[index] = model
 }
 
 // Given a log, allow each state model to determine if/how to process it
diff --git a/internal/eigenState/types.go b/internal/eigenState/types.go
index 83bd893c..b0a984e9 100644
--- a/internal/eigenState/types.go
+++ b/internal/eigenState/types.go
@@ -7,20 +7,29 @@ import (
 type StateRoot string
 
 type IEigenStateModel interface {
-	// Determine if the log is interesting to the state model
+	// GetModelName
+	// Get the name of the model
+	GetModelName() string
+
+	// IsInterestingLog
+	//Determine if the log is interesting to the state model
 	IsInterestingLog(log *storage.TransactionLog) bool
 
+	// HandleStateChange
 	// Allow the state model to handle the state change
 	//
 	// Returns the saved value. Listed as an interface because go generics suck
 	HandleStateChange(log *storage.TransactionLog) (interface{}, error)
 
+	// WriteFinalState
 	// Once all state changes are processed, calculate and write final state
 	WriteFinalState(blockNumber uint64) error
 
+	// GenerateStateRoot
 	// Generate the state root for the model
 	GenerateStateRoot(blockNumber uint64) (StateRoot, error)
 }
 
+// StateTransitions
 // Map of block number to function that will transition the state to the next block
 type StateTransitions[T interface{}] map[uint64]func(log *storage.TransactionLog) (*T, error)

From 2f732081d1a13fd00bba02e8c81a7a6979673c33 Mon Sep 17 00:00:00 2001
From: Sean McGary <sean@seanmcgary.com>
Date: Thu, 5 Sep 2024 21:34:05 -0500
Subject: [PATCH 6/7] Actually generate the root from all models

---
 internal/eigenState/stateManager.go | 53 ++++++++++++++++++++++++-----
 1 file changed, 45 insertions(+), 8 deletions(-)

diff --git a/internal/eigenState/stateManager.go b/internal/eigenState/stateManager.go
index 4082fd23..1c0860b1 100644
--- a/internal/eigenState/stateManager.go
+++ b/internal/eigenState/stateManager.go
@@ -1,8 +1,13 @@
 package eigenState
 
 import (
+	"fmt"
 	"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"
 	"go.uber.org/zap"
+	"slices"
 )
 
 type EigenStateManager struct {
@@ -27,7 +32,8 @@ func (e *EigenStateManager) RegisterState(model IEigenStateModel, index int) {
 
 // Given a log, allow each state model to determine if/how to process it
 func (e *EigenStateManager) HandleLogStateChange(log *storage.TransactionLog) error {
-	for _, state := range e.StateModels {
+	for _, index := range e.getSortedModelIndexes() {
+		state := e.StateModels[index]
 		if state.IsInterestingLog(log) {
 			_, err := state.HandleStateChange(log)
 			if err != nil {
@@ -40,7 +46,8 @@ func (e *EigenStateManager) HandleLogStateChange(log *storage.TransactionLog) er
 
 // With all transactions/logs processed for a block, commit the final state to the table
 func (e *EigenStateManager) CommitFinalState(blockNumber uint64) error {
-	for _, state := range e.StateModels {
+	for _, index := range e.getSortedModelIndexes() {
+		state := e.StateModels[index]
 		err := state.WriteFinalState(blockNumber)
 		if err != nil {
 			return err
@@ -50,14 +57,44 @@ func (e *EigenStateManager) CommitFinalState(blockNumber uint64) error {
 }
 
 func (e *EigenStateManager) GenerateStateRoot(blockNumber uint64) (StateRoot, error) {
-	roots := make([]StateRoot, len(e.StateModels))
-	for i, state := range e.StateModels {
-		root, err := state.GenerateStateRoot(blockNumber)
+	sortedIndexes := e.getSortedModelIndexes()
+	roots := [][]byte{
+		[]byte(fmt.Sprintf("%d", blockNumber)),
+	}
+
+	for _, state := range sortedIndexes {
+		state := e.StateModels[state]
+		leaf, err := e.encodeModelLeaf(state, blockNumber)
 		if err != nil {
 			return "", err
 		}
-		roots[i] = root
+		roots = append(roots, leaf)
+	}
+
+	tree, err := merkletree.NewTree(
+		merkletree.WithData(roots),
+		merkletree.WithHashType(keccak256.New()),
+	)
+	if err != nil {
+		return "", err
+	}
+
+	return StateRoot(utils.ConvertBytesToString(tree.Root())), nil
+}
+
+func (e *EigenStateManager) encodeModelLeaf(model IEigenStateModel, blockNumber uint64) ([]byte, error) {
+	root, err := model.GenerateStateRoot(blockNumber)
+	if err != nil {
+		return nil, err
+	}
+	return append([]byte(model.GetModelName()), []byte(root)[:]...), nil
+}
+
+func (e *EigenStateManager) getSortedModelIndexes() []int {
+	indexes := make([]int, 0, len(e.StateModels))
+	for i := range e.StateModels {
+		indexes = append(indexes, i)
 	}
-	// TODO: generate this
-	return "", nil
+	slices.Sort(indexes)
+	return indexes
 }

From ee200716da409d6e6dd37c27f1d481a1226b3f82 Mon Sep 17 00:00:00 2001
From: Sean McGary <sean@seanmcgary.com>
Date: Thu, 5 Sep 2024 21:49:21 -0500
Subject: [PATCH 7/7] Fix dependency nonsense

---
 .../eigenState/avsOperators/avsOperators.go   | 20 +++---
 .../avsOperators/avsOperators_test.go         | 10 +--
 .../eigenState/{ => base}/baseEigenState.go   |  2 +-
 internal/eigenState/eigenstate_test.go        | 68 +++++++++++++++++++
 .../operatorShares/operatorShares.go          | 32 +++++----
 .../operatorShares/operatorShares_test.go     |  8 +--
 .../{ => stateManager}/stateManager.go        | 23 ++++---
 internal/eigenState/{ => types}/types.go      |  2 +-
 8 files changed, 120 insertions(+), 45 deletions(-)
 rename internal/eigenState/{ => base}/baseEigenState.go (98%)
 create mode 100644 internal/eigenState/eigenstate_test.go
 rename internal/eigenState/{ => stateManager}/stateManager.go (75%)
 rename internal/eigenState/{ => types}/types.go (98%)

diff --git a/internal/eigenState/avsOperators/avsOperators.go b/internal/eigenState/avsOperators/avsOperators.go
index 4b1e42c4..ed83c66d 100644
--- a/internal/eigenState/avsOperators/avsOperators.go
+++ b/internal/eigenState/avsOperators/avsOperators.go
@@ -4,7 +4,9 @@ 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"
@@ -42,8 +44,8 @@ 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
@@ -60,7 +62,7 @@ type RegisteredAvsOperatorDiff struct {
 
 // Create new instance of AvsOperators state model
 func NewAvsOperators(
-	esm *eigenState.EigenStateManager,
+	esm *stateManager.EigenStateManager,
 	grm *gorm.DB,
 	Network config.Network,
 	Environment config.Environment,
@@ -68,7 +70,7 @@ func NewAvsOperators(
 	globalConfig *config.Config,
 ) (*AvsOperators, error) {
 	s := &AvsOperators{
-		BaseEigenState: eigenState.BaseEigenState{
+		BaseEigenState: base.BaseEigenState{
 			Logger: logger,
 		},
 		Db:           grm,
@@ -92,8 +94,8 @@ func (a *AvsOperators) GetModelName() string {
 //
 // 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) {
@@ -326,7 +328,7 @@ func (a *AvsOperators) getDifferenceInStates(blockNumber uint64) ([]RegisteredAv
 // 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) {
+func (a *AvsOperators) GenerateStateRoot(blockNumber uint64) (types.StateRoot, error) {
 	results, err := a.getDifferenceInStates(blockNumber)
 	if err != nil {
 		return "", err
@@ -336,7 +338,7 @@ func (a *AvsOperators) GenerateStateRoot(blockNumber uint64) (eigenState.StateRo
 	if err != nil {
 		return "", err
 	}
-	return eigenState.StateRoot(utils.ConvertBytesToString(fullTree.Root())), nil
+	return types.StateRoot(utils.ConvertBytesToString(fullTree.Root())), nil
 }
 
 func (a *AvsOperators) merkelizeState(blockNumber uint64, avsOperators []RegisteredAvsOperatorDiff) (*merkletree.MerkleTree, error) {
diff --git a/internal/eigenState/avsOperators/avsOperators_test.go b/internal/eigenState/avsOperators/avsOperators_test.go
index cb70999e..14233343 100644
--- a/internal/eigenState/avsOperators/avsOperators_test.go
+++ b/internal/eigenState/avsOperators/avsOperators_test.go
@@ -3,7 +3,7 @@ package avsOperators
 import (
 	"database/sql"
 	"github.com/Layr-Labs/sidecar/internal/config"
-	"github.com/Layr-Labs/sidecar/internal/eigenState"
+	"github.com/Layr-Labs/sidecar/internal/eigenState/stateManager"
 	"github.com/Layr-Labs/sidecar/internal/logger"
 	"github.com/Layr-Labs/sidecar/internal/storage"
 	"github.com/Layr-Labs/sidecar/internal/tests"
@@ -41,13 +41,13 @@ func Test_AvsOperatorState(t *testing.T) {
 	}
 
 	t.Run("Should create a new AvsOperatorState", func(t *testing.T) {
-		esm := eigenState.NewEigenStateManager(l)
+		esm := stateManager.NewEigenStateManager(l)
 		avsOperatorState, err := NewAvsOperators(esm, grm, cfg.Network, cfg.Environment, l, cfg)
 		assert.Nil(t, err)
 		assert.NotNil(t, avsOperatorState)
 	})
 	t.Run("Should register AvsOperatorState", func(t *testing.T) {
-		esm := eigenState.NewEigenStateManager(l)
+		esm := stateManager.NewEigenStateManager(l)
 		blockNumber := uint64(200)
 		log := storage.TransactionLog{
 			TransactionHash:  "some hash",
@@ -75,7 +75,7 @@ func Test_AvsOperatorState(t *testing.T) {
 		teardown(avsOperatorState)
 	})
 	t.Run("Should register AvsOperatorState and generate the table for the block", func(t *testing.T) {
-		esm := eigenState.NewEigenStateManager(l)
+		esm := stateManager.NewEigenStateManager(l)
 		blockNumber := uint64(200)
 
 		log := storage.TransactionLog{
@@ -123,7 +123,7 @@ func Test_AvsOperatorState(t *testing.T) {
 		teardown(avsOperatorState)
 	})
 	t.Run("Should correctly generate state across multiple blocks", func(t *testing.T) {
-		esm := eigenState.NewEigenStateManager(l)
+		esm := stateManager.NewEigenStateManager(l)
 		blocks := []uint64{
 			300,
 			301,
diff --git a/internal/eigenState/baseEigenState.go b/internal/eigenState/base/baseEigenState.go
similarity index 98%
rename from internal/eigenState/baseEigenState.go
rename to internal/eigenState/base/baseEigenState.go
index c584d3e5..da85b722 100644
--- a/internal/eigenState/baseEigenState.go
+++ b/internal/eigenState/base/baseEigenState.go
@@ -1,4 +1,4 @@
-package eigenState
+package base
 
 import (
 	"encoding/json"
diff --git a/internal/eigenState/eigenstate_test.go b/internal/eigenState/eigenstate_test.go
new file mode 100644
index 00000000..bc00174a
--- /dev/null
+++ b/internal/eigenState/eigenstate_test.go
@@ -0,0 +1,68 @@
+package eigenState
+
+import (
+	"fmt"
+	"github.com/Layr-Labs/sidecar/internal/config"
+	"github.com/Layr-Labs/sidecar/internal/eigenState/avsOperators"
+	"github.com/Layr-Labs/sidecar/internal/eigenState/operatorShares"
+	"github.com/Layr-Labs/sidecar/internal/eigenState/stateManager"
+	"github.com/Layr-Labs/sidecar/internal/logger"
+	"github.com/Layr-Labs/sidecar/internal/tests"
+	"github.com/stretchr/testify/assert"
+	"go.uber.org/zap"
+	"gorm.io/gorm"
+	"testing"
+)
+
+func setup() (
+	*config.Config,
+	*gorm.DB,
+	*zap.Logger,
+	error,
+) {
+	cfg := tests.GetConfig()
+	l, _ := logger.NewLogger(&logger.LoggerConfig{Debug: cfg.Debug})
+
+	_, grm, err := tests.GetDatabaseConnection(cfg)
+
+	return cfg, grm, l, err
+}
+
+func teardown(grm *gorm.DB) {
+	grm.Exec("truncate table avs_operator_changes cascade")
+	grm.Exec("truncate table registered_avs_operators cascade")
+}
+
+func Test_EigenStateManager(t *testing.T) {
+	cfg, grm, l, err := setup()
+
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	t.Run("Should create a new EigenStateManager", func(t *testing.T) {
+		esm := stateManager.NewEigenStateManager(l)
+		assert.NotNil(t, esm)
+	})
+	t.Run("Should create a state root with states from models", func(t *testing.T) {
+		esm := stateManager.NewEigenStateManager(l)
+		avsOperatorsModel, err := avsOperators.NewAvsOperators(esm, grm, cfg.Network, cfg.Environment, l, cfg)
+		assert.Nil(t, err)
+		assert.NotNil(t, avsOperatorsModel)
+
+		operatorSharesModel, err := operatorShares.NewOperatorSharesModel(esm, grm, cfg.Network, cfg.Environment, l, cfg)
+		assert.Nil(t, err)
+		assert.NotNil(t, operatorSharesModel)
+
+		indexes := esm.GetSortedModelIndexes()
+		assert.Equal(t, 2, len(indexes))
+		assert.Equal(t, 0, indexes[0])
+		assert.Equal(t, 1, indexes[1])
+
+		root, err := esm.GenerateStateRoot(200)
+		assert.Nil(t, err)
+		assert.True(t, len(root) > 0)
+		fmt.Printf("Root: %+v\n", root)
+	})
+	teardown(grm)
+}
diff --git a/internal/eigenState/operatorShares/operatorShares.go b/internal/eigenState/operatorShares/operatorShares.go
index 0d9c16eb..1a29713b 100644
--- a/internal/eigenState/operatorShares/operatorShares.go
+++ b/internal/eigenState/operatorShares/operatorShares.go
@@ -5,7 +5,9 @@ import (
 	"encoding/json"
 	"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/parser"
 	"github.com/Layr-Labs/sidecar/internal/storage"
 	"github.com/Layr-Labs/sidecar/internal/utils"
@@ -46,8 +48,8 @@ type OperatorShares struct {
 
 // Implements IEigenStateModel
 type OperatorSharesModel struct {
-	eigenState.BaseEigenState
-	StateTransitions eigenState.StateTransitions[OperatorShareChange]
+	base.BaseEigenState
+	StateTransitions types.StateTransitions[OperatorShareChange]
 	Db               *gorm.DB
 	Network          config.Network
 	Environment      config.Environment
@@ -56,7 +58,7 @@ type OperatorSharesModel struct {
 }
 
 func NewOperatorSharesModel(
-	esm *eigenState.EigenStateManager,
+	esm *stateManager.EigenStateManager,
 	grm *gorm.DB,
 	Network config.Network,
 	Environment config.Environment,
@@ -64,12 +66,14 @@ func NewOperatorSharesModel(
 	globalConfig *config.Config,
 ) (*OperatorSharesModel, error) {
 	model := &OperatorSharesModel{
-		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(model, 1)
@@ -80,8 +84,8 @@ func (osm *OperatorSharesModel) GetModelName() string {
 	return "OperatorSharesModel"
 }
 
-func (osm *OperatorSharesModel) GetStateTransitions() (eigenState.StateTransitions[OperatorShareChange], []uint64) {
-	stateChanges := make(eigenState.StateTransitions[OperatorShareChange])
+func (osm *OperatorSharesModel) GetStateTransitions() (types.StateTransitions[OperatorShareChange], []uint64) {
+	stateChanges := make(types.StateTransitions[OperatorShareChange])
 
 	stateChanges[0] = func(log *storage.TransactionLog) (*OperatorShareChange, error) {
 		arguments := make([]parser.Argument, 0)
@@ -283,7 +287,7 @@ func (osm *OperatorSharesModel) getDifferencesInStates(currentBlock uint64) ([]O
 	return diffs, nil
 }
 
-func (osm *OperatorSharesModel) GenerateStateRoot(blockNumber uint64) (eigenState.StateRoot, error) {
+func (osm *OperatorSharesModel) GenerateStateRoot(blockNumber uint64) (types.StateRoot, error) {
 	diffs, err := osm.getDifferencesInStates(blockNumber)
 	if err != nil {
 		return "", err
@@ -293,7 +297,7 @@ func (osm *OperatorSharesModel) GenerateStateRoot(blockNumber uint64) (eigenStat
 	if err != nil {
 		return "", err
 	}
-	return eigenState.StateRoot(utils.ConvertBytesToString(fullTree.Root())), nil
+	return types.StateRoot(utils.ConvertBytesToString(fullTree.Root())), nil
 }
 
 func (osm *OperatorSharesModel) merkelizeState(blockNumber uint64, diffs []OperatorShares) (*merkletree.MerkleTree, error) {
diff --git a/internal/eigenState/operatorShares/operatorShares_test.go b/internal/eigenState/operatorShares/operatorShares_test.go
index 0571fbc8..12ec97b4 100644
--- a/internal/eigenState/operatorShares/operatorShares_test.go
+++ b/internal/eigenState/operatorShares/operatorShares_test.go
@@ -4,7 +4,7 @@ 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/stateManager"
 	"github.com/Layr-Labs/sidecar/internal/logger"
 	"github.com/Layr-Labs/sidecar/internal/storage"
 	"github.com/Layr-Labs/sidecar/internal/tests"
@@ -43,13 +43,13 @@ func Test_OperatorSharesState(t *testing.T) {
 	}
 
 	t.Run("Should create a new OperatorSharesState", func(t *testing.T) {
-		esm := eigenState.NewEigenStateManager(l)
+		esm := stateManager.NewEigenStateManager(l)
 		model, err := NewOperatorSharesModel(esm, grm, cfg.Network, cfg.Environment, l, cfg)
 		assert.Nil(t, err)
 		assert.NotNil(t, model)
 	})
 	t.Run("Should register OperatorSharesState", func(t *testing.T) {
-		esm := eigenState.NewEigenStateManager(l)
+		esm := stateManager.NewEigenStateManager(l)
 		blockNumber := uint64(200)
 		log := storage.TransactionLog{
 			TransactionHash:  "some hash",
@@ -75,7 +75,7 @@ func Test_OperatorSharesState(t *testing.T) {
 		teardown(model)
 	})
 	t.Run("Should register AvsOperatorState and generate the table for the block", func(t *testing.T) {
-		esm := eigenState.NewEigenStateManager(l)
+		esm := stateManager.NewEigenStateManager(l)
 		blockNumber := uint64(200)
 		log := storage.TransactionLog{
 			TransactionHash:  "some hash",
diff --git a/internal/eigenState/stateManager.go b/internal/eigenState/stateManager/stateManager.go
similarity index 75%
rename from internal/eigenState/stateManager.go
rename to internal/eigenState/stateManager/stateManager.go
index 1c0860b1..61c978c9 100644
--- a/internal/eigenState/stateManager.go
+++ b/internal/eigenState/stateManager/stateManager.go
@@ -1,7 +1,8 @@
-package eigenState
+package stateManager
 
 import (
 	"fmt"
+	"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"
@@ -11,19 +12,19 @@ import (
 )
 
 type EigenStateManager struct {
-	StateModels map[int]IEigenStateModel
+	StateModels map[int]types.IEigenStateModel
 	logger      *zap.Logger
 }
 
 func NewEigenStateManager(logger *zap.Logger) *EigenStateManager {
 	return &EigenStateManager{
-		StateModels: make(map[int]IEigenStateModel),
+		StateModels: make(map[int]types.IEigenStateModel),
 		logger:      logger,
 	}
 }
 
 // Allows a model to register itself with the state manager
-func (e *EigenStateManager) RegisterState(model IEigenStateModel, index int) {
+func (e *EigenStateManager) RegisterState(model types.IEigenStateModel, index int) {
 	if m, ok := e.StateModels[index]; ok {
 		e.logger.Sugar().Fatalf("Registering model model at index %d which already exists and belongs to %s", index, m.GetModelName())
 	}
@@ -32,7 +33,7 @@ func (e *EigenStateManager) RegisterState(model IEigenStateModel, index int) {
 
 // Given a log, allow each state model to determine if/how to process it
 func (e *EigenStateManager) HandleLogStateChange(log *storage.TransactionLog) error {
-	for _, index := range e.getSortedModelIndexes() {
+	for _, index := range e.GetSortedModelIndexes() {
 		state := e.StateModels[index]
 		if state.IsInterestingLog(log) {
 			_, err := state.HandleStateChange(log)
@@ -46,7 +47,7 @@ func (e *EigenStateManager) HandleLogStateChange(log *storage.TransactionLog) er
 
 // With all transactions/logs processed for a block, commit the final state to the table
 func (e *EigenStateManager) CommitFinalState(blockNumber uint64) error {
-	for _, index := range e.getSortedModelIndexes() {
+	for _, index := range e.GetSortedModelIndexes() {
 		state := e.StateModels[index]
 		err := state.WriteFinalState(blockNumber)
 		if err != nil {
@@ -56,8 +57,8 @@ func (e *EigenStateManager) CommitFinalState(blockNumber uint64) error {
 	return nil
 }
 
-func (e *EigenStateManager) GenerateStateRoot(blockNumber uint64) (StateRoot, error) {
-	sortedIndexes := e.getSortedModelIndexes()
+func (e *EigenStateManager) GenerateStateRoot(blockNumber uint64) (types.StateRoot, error) {
+	sortedIndexes := e.GetSortedModelIndexes()
 	roots := [][]byte{
 		[]byte(fmt.Sprintf("%d", blockNumber)),
 	}
@@ -79,10 +80,10 @@ func (e *EigenStateManager) GenerateStateRoot(blockNumber uint64) (StateRoot, er
 		return "", err
 	}
 
-	return StateRoot(utils.ConvertBytesToString(tree.Root())), nil
+	return types.StateRoot(utils.ConvertBytesToString(tree.Root())), nil
 }
 
-func (e *EigenStateManager) encodeModelLeaf(model IEigenStateModel, blockNumber uint64) ([]byte, error) {
+func (e *EigenStateManager) encodeModelLeaf(model types.IEigenStateModel, blockNumber uint64) ([]byte, error) {
 	root, err := model.GenerateStateRoot(blockNumber)
 	if err != nil {
 		return nil, err
@@ -90,7 +91,7 @@ func (e *EigenStateManager) encodeModelLeaf(model IEigenStateModel, blockNumber
 	return append([]byte(model.GetModelName()), []byte(root)[:]...), nil
 }
 
-func (e *EigenStateManager) getSortedModelIndexes() []int {
+func (e *EigenStateManager) GetSortedModelIndexes() []int {
 	indexes := make([]int, 0, len(e.StateModels))
 	for i := range e.StateModels {
 		indexes = append(indexes, i)
diff --git a/internal/eigenState/types.go b/internal/eigenState/types/types.go
similarity index 98%
rename from internal/eigenState/types.go
rename to internal/eigenState/types/types.go
index b0a984e9..82a657ed 100644
--- a/internal/eigenState/types.go
+++ b/internal/eigenState/types/types.go
@@ -1,4 +1,4 @@
-package eigenState
+package types
 
 import (
 	"github.com/Layr-Labs/sidecar/internal/storage"