TreeStore module implementation patch

Author: h5law

persistence/module.go

@@ -2,6 +2,7 @@ package persistence
 
 import (
 	"context"
+	"errors"
 	"fmt"
 
 	"github.com/jackc/pgx/v5"
@@ -107,14 +108,18 @@ func (*persistenceModule) Create(bus modules.Bus, options ...modules.ModuleOptio
 	if err != nil {
 		return nil, err
 	}
+	treeStoreModule, ok := treeModule.(modules.TreeStoreModule)
+	if !ok {
+		return nil, errors.New("error casting TreeStoreModule")
+	}
 
 	m.config = persistenceCfg
 	m.genesisState = genesisState
 	m.networkId = runtimeMgr.GetConfig().NetworkId
 
 	m.blockStore = blockStore
 	m.txIndexer = txIndexer
-	m.stateTrees = treeModule
+	m.stateTrees = treeStoreModule
 
 	// TECHDEBT: reconsider if this is the best place to call `populateGenesisState`. Note that
 	// 		     this forces the genesis state to be reloaded on every node startup until state
@@ -235,10 +240,6 @@ func (m *persistenceModule) GetTxIndexer() indexer.TxIndexer {
 	return m.txIndexer
 }
 
-func (m *persistenceModule) GetTreeStore() modules.TreeStoreModule {
-	return m.stateTrees
-}
-
 func (m *persistenceModule) GetNetworkID() string {
 	return m.networkId
 }

persistence/trees/module.go

@@ -8,14 +8,16 @@ import (
 	"github.com/pokt-network/smt"
 )
 
-func (*treeStore) Create(bus modules.Bus, options ...modules.TreeStoreOption) (modules.TreeStoreModule, error) {
+var _ modules.Module = &treeStore{}
+
+func (*treeStore) Create(bus modules.Bus, options ...modules.ModuleOption) (modules.Module, error) {
 	m := &treeStore{}
 
 	for _, option := range options {
 		option(m)
 	}
 
-	m.SetBus(bus)
+	bus.RegisterModule(m)
 
 	if err := m.setupTrees(); err != nil {
 		return nil, err
@@ -24,14 +26,20 @@ func (*treeStore) Create(bus modules.Bus, options ...modules.TreeStoreOption) (m
 	return m, nil
 }
 
-func Create(bus modules.Bus, options ...modules.TreeStoreOption) (modules.TreeStoreModule, error) {
+func Create(bus modules.Bus, options ...modules.ModuleOption) (modules.Module, error) {
 	return new(treeStore).Create(bus, options...)
 }
 
+func (t *treeStore) GetModuleName() string  { return modules.TreeStoreModuleName }
+func (t *treeStore) Start() error           { return nil }
+func (t *treeStore) Stop() error            { return nil }
+func (t *treeStore) GetBus() modules.Bus    { return t.bus }
+func (t *treeStore) SetBus(bus modules.Bus) { t.bus = bus }
+
 // WithTreeStoreDirectory assigns the path where the tree store
 // saves its data.
-func WithTreeStoreDirectory(path string) modules.TreeStoreOption {
-	return func(m modules.TreeStoreModule) {
+func WithTreeStoreDirectory(path string) modules.ModuleOption {
+	return func(m modules.InitializableModule) {
 		mod, ok := m.(*treeStore)
 		if ok {
 			mod.treeStoreDir = path

persistence/trees/trees.go

@@ -81,6 +81,8 @@ var _ modules.TreeStoreModule = &treeStore{}
 type treeStore struct {
 	base_modules.IntegratableModule
 
+	bus modules.Bus
+
 	treeStoreDir string
 	rootTree     *stateTree
 	merkleTrees  map[string]*stateTree

runtime/bus.go

@@ -71,6 +71,10 @@ func (m *bus) GetPersistenceModule() modules.PersistenceModule {
 	return getModuleFromRegistry[modules.PersistenceModule](m, modules.PersistenceModuleName)
 }
 
+func (m *bus) GetTreeStoreModule() modules.TreeStoreModule {
+	return getModuleFromRegistry[modules.TreeStoreModule](m, modules.TreeStoreModuleName)
+}
+
 func (m *bus) GetP2PModule() modules.P2PModule {
 	return getModuleFromRegistry[modules.P2PModule](m, modules.P2PModuleName)
 }

shared/modules/bus_module.go

@@ -25,6 +25,7 @@ type Bus interface {
 
 	// Pocket modules
 	GetPersistenceModule() PersistenceModule
+	GetTreeStoreModule() TreeStoreModule
 	GetP2PModule() P2PModule
 	GetUtilityModule() UtilityModule
 	GetConsensusModule() ConsensusModule

shared/modules/persistence_module.go

@@ -30,9 +30,6 @@ type PersistenceModule interface {
 	GetTxIndexer() indexer.TxIndexer
 	TransactionExists(transactionHash string) (bool, error)
 
-	// TreeStore operations
-	GetTreeStore() TreeStoreModule
-
 	// Debugging / development only
 	HandleDebugMessage(*messaging.DebugMessage) error
 }

shared/modules/treestore_module.go

@@ -16,7 +16,7 @@ type TreeStoreFactory = FactoryWithOptions[TreeStoreModule, TreeStoreOption]
 // TreeStoreModules defines the interface for atomic updates and rollbacks to the internal
 // merkle trees that compose the state hash of pocket.
 type TreeStoreModule interface {
-	IntegratableModule
+	Module
 
 	// Update returns the new state hash for a given height.
 	// * Height is passed through to the Update function and is used to query the TxIndexer for transactions