From 3fdc1b181e3cc2bd148489cbd30c073abacbc532 Mon Sep 17 00:00:00 2001
From: yumiel <yumiel.ko@kaia.io>
Date: Thu, 19 Dec 2024 08:56:15 +0800
Subject: [PATCH] add valsetModule engine at SetupKaiaxModules

---
 blockchain/blockchain.go |  5 -----
 node/cn/backend.go       | 16 +++++++++++++---
 2 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/blockchain/blockchain.go b/blockchain/blockchain.go
index 8975a8dbb..e18e8ae96 100644
--- a/blockchain/blockchain.go
+++ b/blockchain/blockchain.go
@@ -2151,11 +2151,6 @@ func (bc *BlockChain) insertChain(chain types.Blocks) (int, []interface{}, []*ty
 			}
 		}
 
-		// update governance CurrentSet if it is at an epoch block
-		if bc.engine.CreateSnapshot(bc, block.NumberU64(), block.Hash(), nil) != nil {
-			return i, events, coalescedLogs, err
-		}
-
 		// update governance parameters
 		if istanbul, ok := bc.engine.(consensus.Istanbul); ok {
 			if err = istanbul.UpdateParam(block.NumberU64()); err != nil {
diff --git a/node/cn/backend.go b/node/cn/backend.go
index df11aa5f4..371011b63 100644
--- a/node/cn/backend.go
+++ b/node/cn/backend.go
@@ -52,6 +52,7 @@ import (
 	"github.com/kaiachain/kaia/kaiax/staking"
 	staking_impl "github.com/kaiachain/kaia/kaiax/staking/impl"
 	supply_impl "github.com/kaiachain/kaia/kaiax/supply/impl"
+	valset_impl "github.com/kaiachain/kaia/kaiax/valset/impl"
 	"github.com/kaiachain/kaia/networks/p2p"
 	"github.com/kaiachain/kaia/networks/rpc"
 	"github.com/kaiachain/kaia/node"
@@ -538,6 +539,7 @@ func (s *CN) SetupKaiaxModules() error {
 		mReward  = reward_impl.NewRewardModule()
 		mSupply  = supply_impl.NewSupplyModule()
 		mGov     = gov_impl.NewGovModule()
+		mValset  = valset_impl.NewValsetModule()
 	)
 
 	// Initialize modules
@@ -565,6 +567,13 @@ func (s *CN) SetupKaiaxModules() error {
 			Chain:       s.blockchain,
 			NodeAddress: s.nodeAddress,
 		}),
+		mValset.Init(&valset_impl.InitOpts{
+			ChainKv:     s.chainDB.GetMiscDB(),
+			Chain:       s.blockchain,
+			Governance:  mGov,
+			StakingInfo: mStaking,
+			NodeAddress: s.nodeAddress,
+		}),
 	)
 	if err != nil {
 		return err
@@ -572,12 +581,13 @@ func (s *CN) SetupKaiaxModules() error {
 
 	// Register modules to respective components
 	// TODO-kaiax: Organize below lines.
-	s.RegisterBaseModules(mStaking, mReward, mSupply, mGov)
+	s.RegisterBaseModules(mStaking, mReward, mSupply, mGov, mValset)
 	s.RegisterJsonRpcModules(mStaking, mReward, mSupply, mGov)
 	s.miner.RegisterExecutionModule(mStaking, mSupply, mGov)
-	s.blockchain.RegisterExecutionModule(mSupply, mGov)
-	s.blockchain.RegisterRewindableModule(mStaking, mSupply, mGov)
+	s.blockchain.RegisterExecutionModule(mSupply, mGov, mValset)
+	s.blockchain.RegisterRewindableModule(mStaking, mSupply, mGov, mValset)
 	if engine, ok := s.engine.(consensus.Istanbul); ok {
+		engine.RegisterValsetModule(mValset)
 		engine.RegisterStakingModule(mStaking)
 		engine.RegisterConsensusModule(mReward, mGov)
 	}