Simplify chain switching

chain/service/src/chain_service.rs

@@ -91,17 +91,8 @@ impl ActorService for ChainReaderService {
 impl EventHandler<Self, NewDagBlock> for ChainReaderService {
     fn handle_event(&mut self, event: NewDagBlock, _ctx: &mut ServiceContext<Self>) {
         info!("NewDagBlock in chain reader service");
-        let mut main = self
-            .inner
-            .get_main()
-            .fork(self.inner.main_head_header().id())
-            .unwrap_or_else(|e| {
-                panic!(
-                    "fork error when handle NewDagBlock in chain reader service: {:?}",
-                    e
-                )
-            });
-        self.inner.main = main
+        self.inner
+            .get_main_mut()
             .select_dag_state(event.executed_block.as_ref().header())
             .unwrap_or_else(|e| {
                 panic!(
@@ -362,6 +353,10 @@ impl ChainReaderServiceInner {
         &self.main
     }
 
+    pub fn get_main_mut(&mut self) -> &mut BlockChain {
+        &mut self.main
+    }
+
     pub fn get_storages(&self) -> (Arc<dyn Store>, Arc<dyn Store2>) {
         (self.storage.clone(), self.storage2.clone())
     }

chain/src/chain.rs

@@ -1964,7 +1964,79 @@ impl BlockChain {
         }
     }
 
-    pub fn select_dag_state(&mut self, header: &BlockHeader) -> Result<Self> {
+    pub fn switch_to_block(&mut self, block_id: HashValue) -> Result<()> {
+        let (storage, storage2) = &self.storage;
+        ensure!(
+            self.has_dag_block(block_id)?,
+            "Block with id {} does not exist in current chain.",
+            block_id
+        );
+        let head_block = storage
+            .get_block_by_hash(block_id)?
+            .ok_or_else(|| format_err!("Cannot find block by hash {:?}", block_id))?;
+
+        let block_info = storage
+            .get_block_info(head_block.id())?
+            .ok_or_else(|| format_err!("Can not find block info by hash {:?}", head_block.id()))?;
+
+        let txn_accumulator_info = block_info.get_txn_accumulator_info();
+        let block_accumulator_info = block_info.get_block_accumulator_info();
+        let vm_state_accumulator_info = block_info.get_vm_state_accumulator_info();
+
+        self.txn_accumulator = info_2_accumulator(
+            txn_accumulator_info.clone(),
+            AccumulatorStoreType::Transaction,
+            storage.as_ref(),
+        );
+        self.block_accumulator = info_2_accumulator(
+            block_accumulator_info.clone(),
+            AccumulatorStoreType::Block,
+            storage.as_ref(),
+        );
+        self.vm_state_accumulator = info_2_accumulator(
+            vm_state_accumulator_info.clone(),
+            AccumulatorStoreType::VMState,
+            storage.as_ref(),
+        );
+
+        let (state_root1, state_root2) = {
+            assert!(
+                self.vm_state_accumulator.num_leaves() > 1,
+                "vm_state_accumulator must have at least 2 leaves, but has {}",
+                self.vm_state_accumulator.num_leaves()
+            );
+
+            let leaf1_idx = self.vm_state_accumulator.num_leaves() - 2;
+            let leaf2_idx = self.vm_state_accumulator.num_leaves() - 1;
+
+            let state_root1 = self
+                .vm_state_accumulator
+                .get_leaf(leaf1_idx)?
+                .ok_or_else(|| format_err!("Can not find acc leaf at index {}", leaf1_idx))?;
+
+            let state_root2 = self
+                .vm_state_accumulator
+                .get_leaf(leaf2_idx)?
+                .ok_or_else(|| format_err!("Can not find acc leaf at index {}", leaf2_idx))?;
+
+            (state_root1, state_root2)
+        };
+
+        let chain_state = ChainStateDB::new(storage.clone().into_super_arc(), Some(state_root1));
+        let chain_state2 = ChainStateDB2::new(storage2.clone().into_super_arc(), Some(state_root2));
+        self.epoch = get_epoch_from_statedb(&chain_state2)?;
+
+        self.status = ChainStatusWithBlock {
+            status: ChainStatus::new(head_block.header.clone(), block_info),
+            head: head_block,
+            multi_state: MultiState::new(state_root1, state_root2),
+        };
+        self.statedb = (chain_state, chain_state2);
+
+        Ok(())
+    }
+
+    pub fn select_dag_state(&mut self, header: &BlockHeader) -> Result<()> {
         let new_pruning_point = if header.pruning_point() == HashValue::zero() {
             self.genesis_hash
         } else {
@@ -1976,15 +2048,13 @@ impl BlockChain {
             self.status().head().pruning_point()
         };
 
-        let chain = if current_pruning_point == new_pruning_point
+        let block_id = if current_pruning_point == new_pruning_point
             || current_pruning_point == HashValue::zero()
         {
             let state = self.dag().get_dag_state(new_pruning_point)?;
-            let block_id = self
-                .dag()
+            self.dag()
                 .ghost_dag_manager()
-                .find_selected_parent(state.tips.into_iter())?;
-            self.fork(block_id)?
+                .find_selected_parent(state.tips.into_iter())?
         } else {
             // Handle pruning point change: select best header from both states
             let new_state = self.dag().get_dag_state(new_pruning_point)?;
@@ -1999,15 +2069,12 @@ impl BlockChain {
                 .ghost_dag_manager()
                 .find_selected_parent(current_state.tips.into_iter())?;
 
-            let selected_header = self
-                .dag()
+            self.dag()
                 .ghost_dag_manager()
-                .find_selected_parent([new_header, current_header].into_iter())?;
-
-            self.fork(selected_header)?
+                .find_selected_parent([new_header, current_header].into_iter())?
         };
 
-        Ok(chain)
+        self.switch_to_block(block_id)
     }
 }
 

chain/tests/test_select_dag_state.rs

@@ -20,12 +20,12 @@ fn test_select_dag_state_same_pruning_point() -> Result<()> {
 
     // Select DAG state with same pruning point (should use simple path)
     let mut test_chain = mock_chain.fork_new_branch(Some(block_1.id()))?;
-    let new_chain = test_chain.select_dag_state(&block_2)?;
+    test_chain.select_dag_state(&block_2)?;
 
     // Verify: selected chain should have block_2 as head (same pruning point, simple case)
-    assert_eq!(new_chain.status().head().id(), block_2.id());
+    assert_eq!(test_chain.status().head().id(), block_2.id());
     assert_eq!(
-        new_chain.status().head().pruning_point(),
+        test_chain.status().head().pruning_point(),
         block_2.pruning_point()
     );
 
@@ -72,11 +72,11 @@ fn test_select_dag_state_different_pruning_points() -> Result<()> {
     // Now test select_dag_state with different pruning points
     // This should trigger the complex comparison logic in lines 1915-1934
     let mut test_chain = mock_chain.fork_new_branch(Some(blue_5.id()))?;
-    let selected_chain = test_chain.select_dag_state(blue_6.header())?;
+    test_chain.select_dag_state(blue_6.header())?;
 
     // Verify: the selected chain should handle the pruning point change correctly
     // The exact result depends on GHOSTDAG comparison, but it should not panic/fail
-    assert!(selected_chain.status().head().id() != HashValue::zero());
+    assert!(test_chain.status().head().id() != HashValue::zero());
 
     Ok(())
 }
@@ -99,10 +99,10 @@ fn test_select_dag_state_zero_pruning_point() -> Result<()> {
 
     // This should take the simple path (condition on line 1905-1907)
     let mut test_chain = mock_chain.fork_new_branch(Some(genesis_header.id()))?;
-    let selected_chain = test_chain.select_dag_state(&block_1)?;
+    test_chain.select_dag_state(&block_1)?;
 
     // Verify: should work without issues
-    assert_eq!(selected_chain.status().head().id(), block_1.id());
+    assert_eq!(test_chain.status().head().id(), block_1.id());
 
     Ok(())
 }
@@ -156,9 +156,9 @@ fn test_select_dag_state_regression_pruning_change() -> Result<()> {
         "select_dag_state should handle pruning point changes"
     );
 
-    let selected_chain = result?;
+    result?;
     // Should successfully create a chain, exact head depends on GHOSTDAG logic
-    assert!(selected_chain.status().head().id() != HashValue::zero());
+    assert!(test_chain.status().head().id() != HashValue::zero());
 
     Ok(())
 }

state/service/src/service.rs

@@ -17,7 +17,7 @@ use starcoin_state_api::{
 use starcoin_statedb::ChainStateDB;
 use starcoin_storage::{BlockStore, Storage, Store};
 use starcoin_types::state_set::AccountStateSet;
-use starcoin_types::system_events::NewHeadBlock;
+use starcoin_types::system_events::{NewDagBlock, NewHeadBlock};
 use starcoin_types::{
     access_path::AccessPath, account_address::AccountAddress, account_state::AccountState,
     state_set::ChainStateSet,
@@ -64,12 +64,14 @@ impl ServiceFactory<Self> for ChainStateService {
 impl ActorService for ChainStateService {
     fn started(&mut self, ctx: &mut ServiceContext<Self>) -> Result<()> {
         ctx.subscribe::<NewHeadBlock>();
+        ctx.subscribe::<NewDagBlock>();
         self.service.adjust_time();
         Ok(())
     }
 
     fn stopped(&mut self, ctx: &mut ServiceContext<Self>) -> Result<()> {
         ctx.unsubscribe::<NewHeadBlock>();
+        ctx.unsubscribe::<NewDagBlock>();
         Ok(())
     }
 }
@@ -139,6 +141,21 @@ impl EventHandler<Self, NewHeadBlock> for ChainStateService {
     }
 }
 
+impl EventHandler<Self, NewDagBlock> for ChainStateService {
+    fn handle_event(&mut self, msg: NewDagBlock, _ctx: &mut ServiceContext<ChainStateService>) {
+        let NewDagBlock {
+            executed_block: block,
+        } = msg;
+
+        let state_root = block.multi_state().state_root1();
+        debug!(
+            "ChainStateActor change StateRoot to (dag): {:?}",
+            state_root
+        );
+        self.service.change_root(state_root);
+    }
+}
+
 pub struct Inner {
     state_db: ChainStateDB,
     //for adjust local time by on chain time.

sync/src/block_connector/block_connector_service.rs

@@ -240,20 +240,30 @@ where
     TransactionPoolServiceT: TxPoolSyncService + 'static,
 {
     fn handle_event(&mut self, msg: NewDagBlock, _ctx: &mut ServiceContext<Self>) {
-        let block_header = match self
-            .chain_service
-            .switch_header(msg.executed_block.header())
-        {
+        let executed_block = msg.executed_block;
+        let block_header = match self.chain_service.switch_header(executed_block.header()) {
             std::result::Result::Ok(block_header) => block_header,
             Err(e) => {
                 error!(
                     "failed to switch header when processing NewDagBlock, error: {:?}, id: {:?}",
                     e,
-                    msg.executed_block.header().id()
+                    executed_block.header().id()
                 );
                 return;
             }
         };
+        if block_header.id() == executed_block.header().id() {
+            if let Err(e) = self
+                .chain_service
+                .apply_new_head_from_dag((*executed_block).clone())
+            {
+                error!(
+                    "failed to apply new head when processing NewDagBlock, error: {:?}, id: {:?}",
+                    e,
+                    executed_block.header().id()
+                );
+            }
+        }
 
         let _consume = self
             .pruning_point_channel

sync/src/block_connector/write_block_chain.rs

@@ -111,6 +111,11 @@ impl<TransactionPoolServiceT> WriteBlockChainService<TransactionPoolServiceT>
 where
     TransactionPoolServiceT: TxPoolSyncService + 'static,
 {
+    pub fn apply_new_head_from_dag(&mut self, executed_block: ExecutedBlock) -> Result<()> {
+        let enacted_blocks = vec![executed_block.block().clone()];
+        self.do_new_head(executed_block, 1, enacted_blocks, 0, vec![])
+    }
+
     pub fn new(
         config: Arc<NodeConfig>,
         startup_info: StartupInfo,
@@ -175,8 +180,7 @@ where
     }
 
     pub fn switch_header(&mut self, header: &BlockHeader) -> Result<BlockHeader> {
-        let new_branch = self.main.select_dag_state(header)?; // 1
-        self.select_head(new_branch)?;
+        self.main.select_dag_state(header)?;
         self.update_startup_info(&self.main.current_header())?;
         Ok(self.main.current_header())
     }

vm2/service/state/src/service.rs

@@ -9,7 +9,7 @@ use starcoin_service_registry::{
 };
 use starcoin_storage::Storage2;
 use starcoin_storage::{BlockStore, Storage, Store};
-use starcoin_types::system_events::NewHeadBlock;
+use starcoin_types::system_events::{NewDagBlock, NewHeadBlock};
 use starcoin_vm2_crypto::HashValue;
 use starcoin_vm2_state_api::{
     message::{StateRequest, StateResponse},
@@ -64,12 +64,14 @@ impl ServiceFactory<Self> for ChainStateService {
 impl ActorService for ChainStateService {
     fn started(&mut self, ctx: &mut ServiceContext<Self>) -> Result<()> {
         ctx.subscribe::<NewHeadBlock>();
+        ctx.subscribe::<NewDagBlock>();
         self.service.adjust_time();
         Ok(())
     }
 
     fn stopped(&mut self, ctx: &mut ServiceContext<Self>) -> Result<()> {
         ctx.unsubscribe::<NewHeadBlock>();
+        ctx.unsubscribe::<NewDagBlock>();
         Ok(())
     }
 }
@@ -146,6 +148,17 @@ impl EventHandler<Self, NewHeadBlock> for ChainStateService {
     }
 }
 
+impl EventHandler<Self, NewDagBlock> for ChainStateService {
+    fn handle_event(&mut self, msg: NewDagBlock, _ctx: &mut ServiceContext<Self>) {
+        let state_root = msg.executed_block.multi_state();
+        debug!(
+            "VM2 ChainStateActor change StateRoot to (dag): {:?}",
+            state_root
+        );
+        self.service.change_root(state_root.state_root2());
+    }
+}
+
 pub struct Inner {
     state_db: ChainStateDB,
     //for adjust local time by on chain time.