From fb1b5c6a2f39776e3dff503038b638781ffd6b96 Mon Sep 17 00:00:00 2001
From: jackzhhuang <jackzhhuang@gmail.com>
Date: Thu, 10 Oct 2024 10:23:38 +0800
Subject: [PATCH] save the dag state directly when resetting

---
 flexidag/src/blockdag.rs                      | 5 +++++
 sync/src/block_connector/write_block_chain.rs | 6 +++---
 2 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/flexidag/src/blockdag.rs b/flexidag/src/blockdag.rs
index 8635fce5d5..40493017bf 100644
--- a/flexidag/src/blockdag.rs
+++ b/flexidag/src/blockdag.rs
@@ -449,6 +449,11 @@ impl BlockDAG {
         Ok(self.storage.state_store.read().get_state_by_hash(hash)?)
     }
 
+    pub fn save_dag_state_directly(&self, hash: Hash, state: DagState) -> anyhow::Result<()> {
+        self.storage.state_store.write().insert(hash, state)?;
+        anyhow::Ok(())
+    }
+
     pub fn save_dag_state(&self, hash: Hash, state: DagState) -> anyhow::Result<()> {
         let writer = self.storage.state_store.write();
         match writer.get_state_by_hash(hash) {
diff --git a/sync/src/block_connector/write_block_chain.rs b/sync/src/block_connector/write_block_chain.rs
index 363de6e0b7..eaf90ec22b 100644
--- a/sync/src/block_connector/write_block_chain.rs
+++ b/sync/src/block_connector/write_block_chain.rs
@@ -414,20 +414,20 @@ where
                 .get_storage()
                 .get_genesis()?
                 .ok_or_else(|| format_err!("Cannot get the genesis in storage!"))?;
-            self.main.dag().save_dag_state(
+            self.main.dag().save_dag_state_directly(
                 genesis,
                 DagState {
                     tips: vec![new_head_block.header().id()],
                 },
             )?;
-            self.main.dag().save_dag_state(
+            self.main.dag().save_dag_state_directly(
                 HashValue::zero(),
                 DagState {
                     tips: vec![new_head_block.header().id()],
                 },
             )?;
         } else {
-            self.main.dag().save_dag_state(
+            self.main.dag().save_dag_state_directly(
                 new_head_block.header().pruning_point(),
                 DagState {
                     tips: vec![new_head_block.header().id()],