Add helpers for setting work on fake chains

zebra-chain/src/work/difficulty.rs

@@ -239,20 +239,60 @@ impl CompactDifficulty {
     /// [Zcash Specification]: https://zips.z.cash/protocol/protocol.pdf#workdef
     pub fn to_work(&self) -> Option<Work> {
         let expanded = self.to_expanded()?;
+        Work::try_from(expanded).ok()
+    }
+}
 
+impl TryFrom<ExpandedDifficulty> for Work {
+    type Error = ();
+
+    fn try_from(expanded: ExpandedDifficulty) -> Result<Self, Self::Error> {
         // We need to compute `2^256 / (expanded + 1)`, but we can't represent
         // 2^256, as it's too large for a u256. However, as 2^256 is at least as
         // large as `expanded + 1`, it is equal to
         // `((2^256 - expanded - 1) / (expanded + 1)) + 1`, or
         let result = (!expanded.0 / (expanded.0 + 1)) + 1;
         if result <= u128::MAX.into() {
-            Some(Work(result.as_u128()))
+            Ok(Work(result.as_u128()))
         } else {
-            None
+            Err(())
         }
     }
 }
 
+impl TryFrom<u128> for Work {
+    type Error = ();
+
+    fn try_from(value: u128) -> Result<Self, Self::Error> {
+        if value == 0 {
+            Err(())
+        } else {
+            Ok(Work(value))
+        }
+    }
+}
+
+impl From<ExpandedDifficulty> for CompactDifficulty {
+    fn from(value: ExpandedDifficulty) -> Self {
+        value.to_compact()
+    }
+}
+
+impl From<Work> for CompactDifficulty {
+    fn from(value: Work) -> Self {
+        let expanded = ExpandedDifficulty::from(value);
+        Self::from(expanded)
+    }
+}
+
+impl From<Work> for ExpandedDifficulty {
+    fn from(value: Work) -> Self {
+        let work: U256 = value.0.into();
+        let expanded = (!work + 1) / work;
+        ExpandedDifficulty(expanded)
+    }
+}
+
 impl ExpandedDifficulty {
     /// Returns the difficulty of the hash.
     ///

zebra-state/src/service/memory_state/non_finalized_state.rs

@@ -235,13 +235,10 @@ mod tests {
     fn best_chain_wins() -> Result<()> {
         zebra_test::init();
 
-        let block2: Arc<Block> =
-            zebra_test::vectors::BLOCK_MAINNET_419201_BYTES.zcash_deserialize_into()?;
         let block1: Arc<Block> =
             zebra_test::vectors::BLOCK_MAINNET_419200_BYTES.zcash_deserialize_into()?;
-        // Create a random block which will have a much worse difficulty threshold
-        // than an intentionally mined block from the mainnet
-        let child = block1.make_fake_child();
+        let block2 = block1.make_fake_child().set_work(10);
+        let child = block1.make_fake_child().set_work(1);
 
         let expected_hash = block2.hash();
 
@@ -259,13 +256,10 @@ mod tests {
     fn finalize_pops_from_best_chain() -> Result<()> {
         zebra_test::init();
 
-        let block2: Arc<Block> =
-            zebra_test::vectors::BLOCK_MAINNET_419201_BYTES.zcash_deserialize_into()?;
         let block1: Arc<Block> =
             zebra_test::vectors::BLOCK_MAINNET_419200_BYTES.zcash_deserialize_into()?;
-        // Create a random block which will have a much worse difficulty threshold
-        // than an intentionally mined block from the mainnet
-        let child = block1.make_fake_child();
+        let block2 = block1.make_fake_child().set_work(10);
+        let child = block1.make_fake_child().set_work(1);
 
         let mut state = NonFinalizedState::default();
         state.commit_new_chain(block1.clone());
@@ -288,14 +282,11 @@ mod tests {
     fn commit_block_extending_best_chain_doesnt_drop_worst_chains() -> Result<()> {
         zebra_test::init();
 
-        let block2: Arc<Block> =
-            zebra_test::vectors::BLOCK_MAINNET_419201_BYTES.zcash_deserialize_into()?;
         let block1: Arc<Block> =
             zebra_test::vectors::BLOCK_MAINNET_419200_BYTES.zcash_deserialize_into()?;
-        // Create a random block which will have a much worse difficulty threshold
-        // than an intentionally mined block from the mainnet
-        let child1 = block1.make_fake_child();
-        let child2 = block2.make_fake_child();
+        let block2 = block1.make_fake_child().set_work(10);
+        let child1 = block1.make_fake_child().set_work(1);
+        let child2 = block2.make_fake_child().set_work(1);
 
         let mut state = NonFinalizedState::default();
         assert_eq!(0, state.chain_set.len());
@@ -310,4 +301,79 @@ mod tests {
 
         Ok(())
     }
+
+    #[test]
+    fn shorter_chain_can_be_best_chain() -> Result<()> {
+        zebra_test::init();
+
+        let block1: Arc<Block> =
+            zebra_test::vectors::BLOCK_MAINNET_419200_BYTES.zcash_deserialize_into()?;
+
+        let long_chain_block1 = block1.make_fake_child().set_work(1);
+        let long_chain_block2 = long_chain_block1.make_fake_child().set_work(1);
+
+        let short_chain_block = block1.make_fake_child().set_work(3);
+
+        let mut state = NonFinalizedState::default();
+        state.commit_new_chain(block1);
+        state.commit_block(long_chain_block1);
+        state.commit_block(long_chain_block2);
+        state.commit_block(short_chain_block);
+        assert_eq!(2, state.chain_set.len());
+
+        assert_eq!(2, state.best_chain_len());
+
+        Ok(())
+    }
+
+    #[test]
+    fn longer_chain_with_more_work_wins() -> Result<()> {
+        zebra_test::init();
+
+        let block1: Arc<Block> =
+            zebra_test::vectors::BLOCK_MAINNET_419200_BYTES.zcash_deserialize_into()?;
+
+        let long_chain_block1 = block1.make_fake_child().set_work(1);
+        let long_chain_block2 = long_chain_block1.make_fake_child().set_work(1);
+        let long_chain_block3 = long_chain_block2.make_fake_child().set_work(1);
+        let long_chain_block4 = long_chain_block3.make_fake_child().set_work(1);
+
+        let short_chain_block = block1.make_fake_child().set_work(3);
+
+        let mut state = NonFinalizedState::default();
+        state.commit_new_chain(block1);
+        state.commit_block(long_chain_block1);
+        state.commit_block(long_chain_block2);
+        state.commit_block(long_chain_block3);
+        state.commit_block(long_chain_block4);
+        state.commit_block(short_chain_block);
+        assert_eq!(2, state.chain_set.len());
+
+        assert_eq!(5, state.best_chain_len());
+
+        Ok(())
+    }
+
+    #[test]
+    fn equal_length_goes_to_more_work() -> Result<()> {
+        zebra_test::init();
+
+        let block1: Arc<Block> =
+            zebra_test::vectors::BLOCK_MAINNET_419200_BYTES.zcash_deserialize_into()?;
+
+        let less_work_child = block1.make_fake_child().set_work(1);
+        let more_work_child = block1.make_fake_child().set_work(3);
+        let expected_hash = more_work_child.hash();
+
+        let mut state = NonFinalizedState::default();
+        state.commit_new_chain(block1);
+        state.commit_block(less_work_child);
+        state.commit_block(more_work_child);
+        assert_eq!(2, state.chain_set.len());
+
+        let tip_hash = state.tip().unwrap().1;
+        assert_eq!(expected_hash, tip_hash);
+
+        Ok(())
+    }
 }

zebra-state/src/tests.rs

@@ -1,9 +1,10 @@
-use std::{mem, sync::Arc};
+use std::{convert::TryFrom, mem, sync::Arc};
 
 use zebra_chain::{
     block::{self, Block},
     transaction::Transaction,
     transparent,
+    work::difficulty::Work,
 };
 
 use super::*;
@@ -11,9 +12,11 @@ use super::*;
 /// Helper trait for constructing "valid" looking chains of blocks
 pub trait FakeChainHelper {
     fn make_fake_child(&self) -> Arc<Block>;
+
+    fn set_work(self, work: u128) -> Arc<Block>;
 }
 
-impl FakeChainHelper for Block {
+impl FakeChainHelper for Arc<Block> {
     fn make_fake_child(&self) -> Arc<Block> {
         let parent_hash = self.hash();
         let mut child = Block::clone(self);
@@ -37,6 +40,13 @@ impl FakeChainHelper for Block {
 
         Arc::new(child)
     }
+
+    fn set_work(mut self, work: u128) -> Arc<Block> {
+        let work = Work::try_from(work).expect("tests should only pass in valid work values");
+        let block = Arc::make_mut(&mut self);
+        block.header.difficulty_threshold = work.into();
+        self
+    }
 }
 
 /// Block heights, and the expected minimum block locator height