From 9d9872711f8b3b846c38998e67caa96bc1783c14 Mon Sep 17 00:00:00 2001
From: Shing Him Ng <shinghim@protonmail.com>
Date: Mon, 23 Sep 2024 08:33:28 -0500
Subject: [PATCH] Refine Sequence and Version Arbitrary impls

---
 primitives/src/sequence.rs    | 24 ++++++++++++++++++++++--
 primitives/src/transaction.rs |  9 +++++++--
 2 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/primitives/src/sequence.rs b/primitives/src/sequence.rs
index d22b77b11..dfa872c8e 100644
--- a/primitives/src/sequence.rs
+++ b/primitives/src/sequence.rs
@@ -244,7 +244,27 @@ units::impl_parse_str_from_int_infallible!(Sequence, u32, from_consensus);
 #[cfg(feature = "arbitrary")]
 impl<'a> Arbitrary<'a> for Sequence {
     fn arbitrary(u: &mut Unstructured<'a>) -> arbitrary::Result<Self> {
-        let s = u32::arbitrary(u)?;
-        Ok(Sequence(s))
+        let mut choice_range = 4;
+        if cfg!(feature = "alloc") {
+            choice_range = 8;
+        }
+
+        // Equally weight the cases of meaningful sequence numbers
+        let choice = u.int_in_range(0..=choice_range)?;
+        match choice {
+            0 => Ok(Sequence::MAX),
+            1 => Ok(Sequence::ZERO),
+            2 => Ok(Sequence::MIN_NO_RBF),
+            3 => Ok(Sequence::ENABLE_RBF_NO_LOCKTIME),
+            #[cfg(feature = "alloc")]
+            4 => Ok(Sequence::from_consensus(relative::Height::MIN.to_consensus_u32())),
+            #[cfg(feature = "alloc")]
+            5 => Ok(Sequence::from_consensus(relative::Height::MAX.to_consensus_u32())),
+            #[cfg(feature = "alloc")]
+            6 => Ok(Sequence::from_consensus(relative::Time::MIN.to_consensus_u32())),
+            #[cfg(feature = "alloc")]
+            7 => Ok(Sequence::from_consensus(relative::Time::MAX.to_consensus_u32())),
+            _ => Ok(Sequence(u.arbitrary()?))
+        }
     }
 }
diff --git a/primitives/src/transaction.rs b/primitives/src/transaction.rs
index cd8acd0ed..98a741a00 100644
--- a/primitives/src/transaction.rs
+++ b/primitives/src/transaction.rs
@@ -76,8 +76,13 @@ impl fmt::Display for Version {
 #[cfg(feature = "arbitrary")]
 impl<'a> Arbitrary<'a> for Version {
     fn arbitrary(u: &mut Unstructured<'a>) -> arbitrary::Result<Self> {
-        let v = i32::arbitrary(u)?;
-        Ok(Version(v))
+        // Equally weight the case of normal version numbers
+        let choice = u.int_in_range(0..=2)?;
+        match choice {
+            0 => Ok(Version::ONE),
+            1 => Ok(Version::TWO),
+            _ => Ok(Version(u.arbitrary()?))
+        }
     }
 }