Converting LeafVersion into an enum

Author: dr-orlovsky

src/util/psbt/serialize.rs

@@ -235,7 +235,7 @@ impl Serialize for (Script, LeafVersion) {
     fn serialize(&self) -> Vec<u8> {
         let mut buf = Vec::with_capacity(self.0.len() + 1);
         buf.extend(self.0.as_bytes());
-        buf.push(self.1.as_u8());
+        buf.push(self.1.into_consensus());
         buf
     }
 }
@@ -247,7 +247,7 @@ impl Deserialize for (Script, LeafVersion) {
         }
         // The last byte is LeafVersion.
         let script = Script::deserialize(&bytes[..bytes.len() - 1])?;
-        let leaf_ver = LeafVersion::from_u8(bytes[bytes.len() - 1])
+        let leaf_ver = LeafVersion::from_consensus(bytes[bytes.len() - 1])
             .map_err(|_| encode::Error::ParseFailed("invalid leaf version"))?;
         Ok((script, leaf_ver))
     }
@@ -283,7 +283,7 @@ impl Serialize for TapTree {
                     // TaprootMerkleBranch can only have len atmost 128(TAPROOT_CONTROL_MAX_NODE_COUNT).
                     // safe to cast from usize to u8
                     buf.push(leaf_info.merkle_branch.as_inner().len() as u8);
-                    buf.push(leaf_info.ver.as_u8());
+                    buf.push(leaf_info.ver.into_consensus());
                     leaf_info.script.consensus_encode(&mut buf).expect("Vecs dont err");
                 }
                 buf
@@ -305,7 +305,7 @@ impl Deserialize for TapTree {
                 bytes_iter.nth(consumed - 1);
             }
 
-            let leaf_version = LeafVersion::from_u8(*version)
+            let leaf_version = LeafVersion::from_consensus(*version)
                 .map_err(|_| encode::Error::ParseFailed("Leaf Version Error"))?;
             builder = builder.add_leaf_with_ver(usize::from(*depth), script, leaf_version)
                 .map_err(|_| encode::Error::ParseFailed("Tree not in DFS order"))?;

src/util/sighash.rs

@@ -229,13 +229,13 @@ impl<'s> ScriptPath<'s> {
     }
     /// Create a new ScriptPath structure using default leaf version value
     pub fn with_defaults(script: &'s Script) -> Self {
-        Self::new(script, LeafVersion::default())
+        Self::new(script, LeafVersion::TapScript)
     }
     /// Compute the leaf hash
     pub fn leaf_hash(&self) -> TapLeafHash {
         let mut enc = TapLeafHash::engine();
 
-        self.leaf_version.as_u8().consensus_encode(&mut enc).expect("Writing to hash enging should never fail");
+        self.leaf_version.into_consensus().consensus_encode(&mut enc).expect("Writing to hash enging should never fail");
         self.script.consensus_encode(&mut enc).expect("Writing to hash enging should never fail");
 
         TapLeafHash::from_engine(enc)

src/util/taproot.rs

@@ -15,6 +15,7 @@
 //!
 //! This module provides support for taproot tagged hashes.
 //!
+
 use prelude::*;
 use io;
 use secp256k1::{self, Secp256k1};
@@ -120,7 +121,7 @@ impl TapLeafHash {
     /// function to compute leaf hash from components
     pub fn from_script(script: &Script, ver: LeafVersion) -> TapLeafHash {
         let mut eng = TapLeafHash::engine();
-        ver.as_u8()
+        ver.into_consensus()
             .consensus_encode(&mut eng)
             .expect("engines don't error");
         script
@@ -142,6 +143,8 @@ pub const TAPROOT_LEAF_MASK: u8 = 0xfe;
 /// Tapscript leaf version
 // https://github.com/bitcoin/bitcoin/blob/e826b22da252e0599c61d21c98ff89f366b3120f/src/script/interpreter.h#L226
 pub const TAPROOT_LEAF_TAPSCRIPT: u8 = 0xc0;
+/// Taproot annex prefix
+pub const TAPROOT_ANNEX_PREFIX: u8 = 0x50;
 /// Tapscript control base size
 // https://github.com/bitcoin/bitcoin/blob/e826b22da252e0599c61d21c98ff89f366b3120f/src/script/interpreter.h#L227
 pub const TAPROOT_CONTROL_BASE_SIZE: usize = 33;
@@ -152,6 +155,7 @@ pub const TAPROOT_CONTROL_MAX_SIZE: usize =
 
 // type alias for versioned tap script corresponding merkle proof
 type ScriptMerkleProofMap = BTreeMap<(Script, LeafVersion), BTreeSet<TaprootMerkleBranch>>;
+
 /// Data structure for representing Taproot spending information.
 /// Taproot output corresponds to a combination of a
 /// single public key condition (known the internal key), and zero or more
@@ -216,7 +220,7 @@ impl TaprootSpendInfo {
     {
         let mut node_weights = BinaryHeap::<(Reverse<u64>, NodeInfo)>::new();
         for (p, leaf) in script_weights {
-            node_weights.push((Reverse(p as u64), NodeInfo::new_leaf_with_ver(leaf, LeafVersion::default())));
+            node_weights.push((Reverse(p as u64), NodeInfo::new_leaf_with_ver(leaf, LeafVersion::TapScript)));
         }
         if node_weights.is_empty() {
             return Err(TaprootBuilderError::IncompleteTree);
@@ -409,7 +413,7 @@ impl TaprootBuilder {
     /// See [`TaprootBuilder::add_leaf_with_ver`] for adding a leaf with specific version
     /// See [Uncyclopedia](https://en.wikipedia.org/wiki/Depth-first_search) for more details
     pub fn add_leaf(self, depth: usize, script: Script) -> Result<Self, TaprootBuilderError> {
-        self.add_leaf_with_ver(depth, script, LeafVersion::default())
+        self.add_leaf_with_ver(depth, script, LeafVersion::TapScript)
     }
 
     /// Add a hidden/omitted node at a depth `depth` to the builder.
@@ -680,7 +684,7 @@ impl ControlBlock {
             return Err(TaprootError::InvalidControlBlockSize(sl.len()));
         }
         let output_key_parity = secp256k1::Parity::from((sl[0] & 1) as i32);
-        let leaf_version = LeafVersion::from_u8(sl[0] & TAPROOT_LEAF_MASK)?;
+        let leaf_version = LeafVersion::from_consensus(sl[0] & TAPROOT_LEAF_MASK)?;
         let internal_key = UntweakedPublicKey::from_slice(&sl[1..TAPROOT_CONTROL_BASE_SIZE])
             .map_err(TaprootError::InvalidInternalKey)?;
         let merkle_branch = TaprootMerkleBranch::from_slice(&sl[TAPROOT_CONTROL_BASE_SIZE..])?;
@@ -700,7 +704,7 @@ impl ControlBlock {
 
     /// Serialize to a writer. Returns the number of bytes written
     pub fn encode<Write: io::Write>(&self, mut writer: Write) -> io::Result<usize> {
-        let first_byte: u8 = i32::from(self.output_key_parity) as u8 | self.leaf_version.as_u8();
+        let first_byte: u8 = i32::from(self.output_key_parity) as u8 | self.leaf_version.into_consensus();
         let mut bytes_written = 0;
         bytes_written += writer.write(&[first_byte])?;
         bytes_written += writer.write(&self.internal_key.serialize())?;
@@ -759,16 +763,16 @@ impl ControlBlock {
 /// The leaf version for tapleafs
 #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
-pub struct LeafVersion(u8);
+pub enum LeafVersion {
+    /// BIP-342 tapscript
+    TapScript,
 
-impl Default for LeafVersion {
-    fn default() -> Self {
-        LeafVersion(TAPROOT_LEAF_TAPSCRIPT)
-    }
+    /// Future leaf version
+    Future(u8)
 }
 
 impl LeafVersion {
-    /// Obtain LeafVersion from u8, will error when last bit of ver is even or
+    /// Obtain LeafVersion from u8, will error when last bit of ver is odd or
     /// when ver is 0x50 (ANNEX_TAG)
     // Text from BIP341:
     // In order to support some forms of static analysis that rely on
@@ -779,23 +783,32 @@ impl LeafVersion {
     // or an opcode that is not valid as the first opcode).
     // The values that comply to this rule are the 32 even values between
     // 0xc0 and 0xfe and also 0x66, 0x7e, 0x80, 0x84, 0x96, 0x98, 0xba, 0xbc, 0xbe
-    pub fn from_u8(ver: u8) -> Result<Self, TaprootError> {
-        if ver & TAPROOT_LEAF_MASK == ver && ver != 0x50 {
-            Ok(LeafVersion(ver))
-        } else {
-            Err(TaprootError::InvalidTaprootLeafVersion(ver))
+    pub fn from_consensus(version: u8) -> Result<Self, TaprootError> {
+        match version {
+            TAPROOT_LEAF_TAPSCRIPT => Ok(LeafVersion::TapScript),
+            TAPROOT_ANNEX_PREFIX => Err(TaprootError::InvalidTaprootLeafVersion(TAPROOT_ANNEX_PREFIX)),
+            odd if odd & TAPROOT_LEAF_MASK != odd => Err(TaprootError::InvalidTaprootLeafVersion(odd)),
+            future => Ok(LeafVersion::Future(future)),
         }
     }
 
     /// Get the inner version from LeafVersion
-    pub fn as_u8(&self) -> u8 {
-        self.0
+    pub fn into_consensus(self) -> u8 {
+        match self {
+            LeafVersion::TapScript => TAPROOT_LEAF_TAPSCRIPT,
+            LeafVersion::Future(version) => version,
+        }
     }
 }
 
-impl Into<u8> for LeafVersion {
-    fn into(self) -> u8 {
-        self.0
+impl fmt::Display for LeafVersion {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match (self, f.alternate()) {
+            (LeafVersion::TapScript, false) => f.write_str("tapscript"),
+            (LeafVersion::TapScript, true) => fmt::Display::fmt(&TAPROOT_LEAF_TAPSCRIPT, f),
+            (LeafVersion::Future(version), false) => write!(f, "future_script_{:#02x}", version),
+            (LeafVersion::Future(version), true) => fmt::Display::fmt(version, f),
+        }
     }
 }
 
@@ -1063,7 +1076,7 @@ mod test {
                 length,
                 tree_info
                     .script_map
-                    .get(&(Script::from_hex(script).unwrap(), LeafVersion::default()))
+                    .get(&(Script::from_hex(script).unwrap(), LeafVersion::TapScript))
                     .expect("Present Key")
                     .iter()
                     .next()
@@ -1078,7 +1091,7 @@ mod test {
 
         // Try to create and verify a control block from each path
         for (_weights, script) in script_weights {
-            let ver_script = (script, LeafVersion::default());
+            let ver_script = (script, LeafVersion::TapScript);
             let ctrl_block = tree_info.control_block(&ver_script).unwrap();
             assert!(ctrl_block.verify_taproot_commitment(&secp, &output_key, &ver_script.0))
         }
@@ -1114,7 +1127,7 @@ mod test {
         let output_key = tree_info.output_key();
 
         for script in vec![a, b, c, d, e] {
-            let ver_script = (script, LeafVersion::default());
+            let ver_script = (script, LeafVersion::TapScript);
             let ctrl_block = tree_info.control_block(&ver_script).unwrap();
             assert!(ctrl_block.verify_taproot_commitment(&secp, &output_key, &ver_script.0))
         }
@@ -1137,7 +1150,7 @@ mod test {
                 }
             } else {
                 let script = Script::from_str(v["script"].as_str().unwrap()).unwrap();
-                let ver = LeafVersion::from_u8(v["leafVersion"].as_u64().unwrap() as u8).unwrap();
+                let ver = LeafVersion::from_consensus(v["leafVersion"].as_u64().unwrap() as u8).unwrap();
                 leaves.push((script.clone(), ver));
                 builder = builder.add_leaf_with_ver(depth, script, ver).unwrap();
             }