save the dag state using the pruning point as the key and if it is 0,…

… use genesis id

chain/src/chain.rs

@@ -178,7 +178,7 @@ impl BlockChain {
     }
 
     fn init_dag(mut dag: BlockDAG, genesis_header: BlockHeader) -> Result<BlockDAG> {
-        match dag.get_dag_state(genesis_header.pruning_point()) {
+        match dag.get_dag_state(genesis_header.id()) {
             anyhow::Result::Ok(_dag_state) => (),
             Err(e) => match e.downcast::<StoreError>()? {
                 StoreError::KeyNotFound(_) => {
@@ -987,7 +987,12 @@ impl BlockChain {
     }
 
     pub fn get_dag_state(&self) -> Result<DagState> {
-        self.dag.get_dag_state(self.status().head().pruning_point())
+        let current_pruning_point = self.status().head().pruning_point();
+        if current_pruning_point == HashValue::zero() {
+            self.dag.get_dag_state(self.genesis_hash)
+        } else {
+            self.dag.get_dag_state(current_pruning_point)
+        }
     }
 }
 
@@ -1573,8 +1578,13 @@ impl BlockChain {
 
         if parent_header.pruning_point() == block.header().pruning_point() {
             info!("pruning point not changed, save dag state without prune. tips are {:?}, pruning point is {:?}", tips, block.header().pruning_point());
-            self.dag
-                .save_dag_state(block.header().pruning_point(), DagState { tips })?;
+            if block.header().pruning_point() == HashValue::zero() {
+                self.dag
+                    .save_dag_state(self.genesis_hash, DagState { tips })?;
+            } else {
+                self.dag
+                    .save_dag_state(block.header().pruning_point(), DagState { tips })?;
+            }
         } else {
             let new_tips = dag.pruning_point_manager().prune(
                 &DagState { tips: tips.clone() },

flexidag/src/blockdag.rs

@@ -113,10 +113,9 @@ impl BlockDAG {
             .write()
             .insert(origin, BlockHashes::new(vec![]))?;
 
-        let pruning_point = genesis.pruning_point();
         self.commit(genesis, origin)?;
         self.save_dag_state(
-            pruning_point,
+            genesis_id,
             DagState {
                 tips: vec![genesis_id],
             },
@@ -545,25 +544,30 @@ impl BlockDAG {
         self.ghost_dag_manager()
             .verify_and_ghostdata(blue_blocks, header)
     }
-    pub fn check_upgrade(&self, main: &BlockHeader) -> anyhow::Result<()> {
+    pub fn check_upgrade(&self, main: &BlockHeader, genesis_id: HashValue) -> anyhow::Result<()> {
         // set the state with key 0
-        if main.version() == 0 {
+        if main.version() == 0 || main.version() == 1 {
             let result_dag_state = self
                 .storage
                 .state_store
                 .read()
-                .get_state_by_hash(main.pruning_point());
+                .get_state_by_hash(genesis_id);
             match result_dag_state {
                 anyhow::Result::Ok(_dag_state) => (),
                 Err(_) => {
-                    let result_dag_state =
-                        self.storage.state_store.read().get_state_by_hash(0.into());
+                    let result_dag_state = self
+                        .storage
+                        .state_store
+                        .read()
+                        .get_state_by_hash(HashValue::zero());
+
                     match result_dag_state {
-                        anyhow::Result::Ok(dag_state) => self
-                            .storage
-                            .state_store
-                            .write()
-                            .insert(main.pruning_point(), dag_state)?,
+                        anyhow::Result::Ok(dag_state) => {
+                            self.storage
+                                .state_store
+                                .write()
+                                .insert(genesis_id, dag_state)?;
+                        }
                         Err(_) => {
                             let dag_state = self
                                 .storage
@@ -573,26 +577,15 @@ impl BlockDAG {
                             self.storage
                                 .state_store
                                 .write()
-                                .insert(0.into(), dag_state.clone())?;
+                                .insert(HashValue::zero(), dag_state.clone())?;
                             self.storage
                                 .state_store
                                 .write()
-                                .insert(HashValue::zero(), dag_state)?;
+                                .insert(genesis_id, dag_state)?;
                         }
                     }
                 }
             }
-            return Ok(());
-        } else if main.version() == 1 {
-            let dag_state = self
-                .storage
-                .state_store
-                .read()
-                .get_state_by_hash(0.into())?;
-            self.storage
-                .state_store
-                .write()
-                .insert(HashValue::zero(), dag_state)?;
         }
 
         anyhow::Ok(())

flexidag/src/prune/pruning_point_manager.rs

@@ -80,14 +80,8 @@ impl<T: ReachabilityStoreReader + Clone> PruningPointManagerT<T> {
         if min_required_blue_score_for_next_pruning_point + pruning_depth
             <= next_ghostdata.blue_score
         {
-            let ancestor = if previous_pruning_point == HashValue::zero() {
-                HashValue::new(ORIGIN)
-            } else {
-                previous_pruning_point
-            };
-
             for child in self.reachability_service().forward_chain_iterator(
-                ancestor,
+                previous_pruning_point,
                 next_ghostdata.selected_parent,
                 true,
             ) {

flexidag/tests/tests.rs

@@ -748,6 +748,7 @@ fn add_and_print(
         .with_parent_hash(parent)
         .with_parents_hash(parents)
         .with_number(number)
+        .with_pruning_point(Hash::zero())
         .build();
     let start = Instant::now();
     dag.commit(header.to_owned(), origin)?;
@@ -869,7 +870,6 @@ fn test_big_data_commit() -> anyhow::Result<()> {
     anyhow::Result::Ok(())
 }
 
-#[ignore = "pruning will be tested in next release"]
 #[test]
 fn test_prune() -> anyhow::Result<()> {
     // initialzie the dag firstly
@@ -968,26 +968,40 @@ fn test_prune() -> anyhow::Result<()> {
 
     // prunning process begins
     dag.save_dag_state(
-        Hash::zero(),
+        genesis.id(),
         DagState {
             tips: vec![block_red_3.id(), block_main_5.id()],
         },
     )?;
 
+    let (previous_ghostdata, previous_pruning_point) =
+        if block_main_5.pruning_point() == Hash::zero() {
+            (
+                dag.ghostdata_by_hash(genesis.id())?.ok_or_else(|| {
+                    format_err!("failed to get the ghostdata by genesis: {:?}", genesis.id())
+                })?,
+                genesis.id(),
+            )
+        } else {
+            (
+                dag.ghostdata_by_hash(block_main_5.pruning_point())?
+                    .ok_or_else(|| {
+                        format_err!(
+                            "failed to get the ghostdata by pruning point: {:?}",
+                            block_main_5.pruning_point()
+                        )
+                    })?,
+                block_main_5.pruning_point(),
+            )
+        };
+
     let MineNewDagBlockInfo {
         tips,
         blue_blocks: _,
         pruning_point,
     } = dag.calc_mergeset_and_tips(
-        block_main_5.pruning_point(),
-        dag.ghostdata_by_hash(block_main_5.pruning_point())?
-            .ok_or_else(|| {
-                format_err!(
-                    "failed to get the ghostdata by {:?}",
-                    block_main_5.pruning_point()
-                )
-            })?
-            .as_ref(),
+        previous_pruning_point,
+        previous_ghostdata.as_ref(),
         pruning_depth,
         pruning_finality,
     )?;

node/src/node.rs

@@ -336,7 +336,7 @@ impl NodeService {
             upgrade_time.as_secs()
         );
 
-        dag.check_upgrade(chain_info.status().head())?;
+        dag.check_upgrade(chain_info.status().head(), genesis.block().id())?;
 
         registry.put_shared(genesis).await?;