Implement std traits for `TaprootMerkleBranch`

The type is naturally a collection of hashes so make it behave that way
by implementing `Deref`, `AsRef`, `Borrow` and their mutable versions as
well as `IntoIterator` for its reference. `IntoIterator` for itself is
not yet implemented because it's a bit more complicated.
This commit is contained in:
Martin Habovstiak 2023-12-07 23:45:15 +01:00
parent 93b415589d
commit 9d23c1d0a8
1 changed files with 61 additions and 5 deletions

View File

@ -1118,20 +1118,20 @@ impl TaprootMerkleBranch {
/// ///
/// The number of bytes written to the writer. /// The number of bytes written to the writer.
pub fn encode<Write: io::Write + ?Sized>(&self, writer: &mut Write) -> io::Result<usize> { pub fn encode<Write: io::Write + ?Sized>(&self, writer: &mut Write) -> io::Result<usize> {
for hash in self.0.iter() { for hash in self {
writer.write_all(hash.as_ref())?; writer.write_all(hash.as_ref())?;
} }
Ok(self.0.len() * TapNodeHash::LEN) Ok(self.len() * TapNodeHash::LEN)
} }
/// Serializes `self` as bytes. /// Serializes `self` as bytes.
pub fn serialize(&self) -> Vec<u8> { pub fn serialize(&self) -> Vec<u8> {
self.0.iter().flat_map(|e| e.as_byte_array()).copied().collect::<Vec<u8>>() self.iter().flat_map(|e| e.as_byte_array()).copied().collect::<Vec<u8>>()
} }
/// Appends elements to proof. /// Appends elements to proof.
fn push(&mut self, h: TapNodeHash) -> Result<(), TaprootBuilderError> { fn push(&mut self, h: TapNodeHash) -> Result<(), TaprootBuilderError> {
if self.0.len() >= TAPROOT_CONTROL_MAX_NODE_COUNT { if self.len() >= TAPROOT_CONTROL_MAX_NODE_COUNT {
Err(TaprootBuilderError::InvalidMerkleTreeDepth(self.0.len())) Err(TaprootBuilderError::InvalidMerkleTreeDepth(self.0.len()))
} else { } else {
self.0.push(h); self.0.push(h);
@ -1194,6 +1194,62 @@ impl From<TaprootMerkleBranch> for Vec<TapNodeHash> {
fn from(branch: TaprootMerkleBranch) -> Self { branch.0 } fn from(branch: TaprootMerkleBranch) -> Self { branch.0 }
} }
impl<'a> IntoIterator for &'a TaprootMerkleBranch {
type IntoIter = core::slice::Iter<'a, TapNodeHash>;
type Item = &'a TapNodeHash;
fn into_iter(self) -> Self::IntoIter {
self.0.iter()
}
}
impl<'a> IntoIterator for &'a mut TaprootMerkleBranch {
type IntoIter = core::slice::IterMut<'a, TapNodeHash>;
type Item = &'a mut TapNodeHash;
fn into_iter(self) -> Self::IntoIter {
self.0.iter_mut()
}
}
impl core::ops::Deref for TaprootMerkleBranch {
type Target = [TapNodeHash];
fn deref(&self) -> &Self::Target {
&self.0
}
}
impl core::ops::DerefMut for TaprootMerkleBranch {
fn deref_mut(&mut self) -> &mut Self::Target {
&mut self.0
}
}
impl AsRef<[TapNodeHash]> for TaprootMerkleBranch {
fn as_ref(&self) -> &[TapNodeHash] {
&self.0
}
}
impl AsMut<[TapNodeHash]> for TaprootMerkleBranch {
fn as_mut(&mut self) -> &mut [TapNodeHash] {
&mut self.0
}
}
impl Borrow<[TapNodeHash]> for TaprootMerkleBranch {
fn borrow(&self) -> &[TapNodeHash] {
&self.0
}
}
impl BorrowMut<[TapNodeHash]> for TaprootMerkleBranch {
fn borrow_mut(&mut self) -> &mut [TapNodeHash] {
&mut self.0
}
}
/// Control block data structure used in Tapscript satisfaction. /// Control block data structure used in Tapscript satisfaction.
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
@ -1283,7 +1339,7 @@ impl ControlBlock {
// Initially the curr_hash is the leaf hash // Initially the curr_hash is the leaf hash
let mut curr_hash = TapNodeHash::from_script(script, self.leaf_version); let mut curr_hash = TapNodeHash::from_script(script, self.leaf_version);
// Verify the proof // Verify the proof
for elem in self.merkle_branch.as_slice() { for elem in &self.merkle_branch {
// Recalculate the curr hash as parent hash // Recalculate the curr hash as parent hash
curr_hash = TapNodeHash::from_node_hashes(curr_hash, *elem); curr_hash = TapNodeHash::from_node_hashes(curr_hash, *elem);
} }