Skip to content

Commit

Permalink
Merge pull request #40 from BP-WG/taptree-leafscripts
Browse files Browse the repository at this point in the history 
Use TapTree::leaf_scripts iterator
  • Loading branch information
@dr-orlovsky
dr-orlovsky committed Jun 30, 2022
2 parents a72b7bf + d53639f commit 47aa3b5
Show file tree
Hide file tree
Showing 3 changed files with 109 additions and 75 deletions.
2 changes: 1 addition & 1 deletion Cargo.lock
View file

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

4 changes: 2 additions & 2 deletions scripts/Cargo.toml
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
[package]
name = "bitcoin_scripts"
version = "0.8.1"
version = "0.8.2"
license = "Apache-2.0"
authors = ["Dr. Maxim Orlovsky <[email protected]>"]
description = "Bitcoin extended script types"
Expand All @@ -15,7 +15,7 @@ exclude = []

[dependencies]
amplify = "3.12.1"
bitcoin = "0.28.1"
bitcoin = "=0.28.1" # We pin this version due to an upstream bug. See setails: <https://github.com/BP-WG/descriptor-wallet/pull/40>
secp256k1 = { version = "0.22.1", features = ["global-context"] } # Remove after bitcoin #922
miniscript = { version = "7.0.0", optional = true }
strict_encoding = "0.8.0"
Expand Down
178 changes: 106 additions & 72 deletions scripts/src/taproot.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,6 @@ use bitcoin::hashes::Hash;
use bitcoin::psbt::TapTree;
use bitcoin::util::taproot::{LeafVersion, TapBranchHash, TapLeafHash, TaprootBuilder};
use bitcoin::Script;
use secp256k1::{KeyPair, SECP256K1};
use strict_encoding::{StrictDecode, StrictEncode};

use crate::types::IntoNodeHash;
Expand Down Expand Up @@ -981,15 +980,16 @@ impl TaprootScriptTree {

/// Returns iterator over known scripts stored in the tree.
///
/// NB: the iterator ignores scripts behind hidden nodes.
/// NB: the iterator ignores scripts behind hidden nodes. It iterates the
/// scripts in DFS (and not consensus) order.
#[inline]
pub fn scripts(&self) -> TreeScriptIter { TreeScriptIter::from(self) }

/// Returns iterator over all known nodes of the tree.
/// Returns iterator over all known nodes of the tree in DFS order.
#[inline]
pub fn nodes(&self) -> TreeNodeIter { TreeNodeIter::from(self) }

/// Returns mutable iterator over all known nodes of the tree.
/// Returns mutable iterator over all known nodes of the tree in DFS order.
#[inline]
pub(self) fn nodes_mut(&mut self) -> TreeNodeIterMut { TreeNodeIterMut::from(self) }

Expand Down Expand Up @@ -1233,65 +1233,63 @@ impl TaprootScriptTree {

impl From<TapTree> for TaprootScriptTree {
fn from(tree: TapTree) -> Self {
// TODO: Do via iterator once #922 will be merged
let dumb_key = KeyPair::from_secret_key(SECP256K1, secp256k1::ONE_KEY).public_key();
let spent_info = tree
.into_builder()
.finalize(SECP256K1, dumb_key)
.expect("non-final taptree");

let mut root: Option<PartialTreeNode> = None;
for ((script, leaf_version), map) in spent_info.as_script_map() {
for merkle_branch in map {
let merkle_branch = merkle_branch.as_inner();
let leaf_depth = merkle_branch.len() as u8;

let mut curr_hash =
TapLeafHash::from_script(script, *leaf_version).into_node_hash();
let merkle_branch = merkle_branch
.iter()
.map(|step| {
curr_hash =
TapBranchHash::from_node_hashes(*step, curr_hash).into_node_hash();
curr_hash
})
.collect::<Vec<_>>();
let mut hash_iter = merkle_branch.iter().rev();

match (root.is_some(), hash_iter.next()) {
(false, None) => {
root = Some(PartialTreeNode::with_leaf(*leaf_version, script.clone(), 0))
}
(false, Some(hash)) => {
root = Some(PartialTreeNode::with_branch(
TapBranchHash::from_inner(hash.into_inner()),
0,
))
}
(true, None) => unreachable!("broken TapTree structure"),
(true, Some(_)) => {}
// TODO: This is a bugfix, which should be reversed once <https://github.com/rust-bitcoin/rust-bitcoin/issues/1069> is fixed upstream
let mut script_leaves = tree.script_leaves().collect::<Vec<_>>();
script_leaves.reverse();
for leaf in script_leaves {
let merkle_branch = leaf.merkle_branch().as_inner();
let leaf_depth = merkle_branch.len() as u8;

let mut curr_hash =
TapLeafHash::from_script(leaf.script(), leaf.leaf_version()).into_node_hash();
let merkle_branch = merkle_branch
.iter()
.map(|step| {
curr_hash = TapBranchHash::from_node_hashes(*step, curr_hash).into_node_hash();
curr_hash
})
.collect::<Vec<_>>();
let mut hash_iter = merkle_branch.iter().rev();

match (root.is_some(), hash_iter.next()) {
(false, None) => {
root = Some(PartialTreeNode::with_leaf(
leaf.leaf_version(),
leaf.script().clone(),
0,
))
}
let mut node = root.as_mut().expect("unreachable");
for (depth, hash) in hash_iter.enumerate() {
match node {
PartialTreeNode::Leaf(..) => unreachable!("broken TapTree structure"),
PartialTreeNode::Branch(branch, _) => {
let child = PartialTreeNode::with_branch(
TapBranchHash::from_inner(hash.into_inner()),
depth as u8 + 1,
);
node = branch.push_child(child).expect("broken TapTree structure");
}
}
(false, Some(hash)) => {
root = Some(PartialTreeNode::with_branch(
TapBranchHash::from_inner(hash.into_inner()),
0,
))
}
let leaf = PartialTreeNode::with_leaf(*leaf_version, script.clone(), leaf_depth);
(true, None) => unreachable!("broken TapTree structure"),
(true, Some(_)) => {}
}
let mut node = root.as_mut().expect("unreachable");
for (depth, hash) in hash_iter.enumerate() {
match node {
PartialTreeNode::Leaf(..) => { /* nothing to do here */ }
PartialTreeNode::Leaf(..) => unreachable!("broken TapTree structure"),
PartialTreeNode::Branch(branch, _) => {
branch.push_child(leaf);
let child = PartialTreeNode::with_branch(
TapBranchHash::from_inner(hash.into_inner()),
depth as u8 + 1,
);
node = branch.push_child(child).expect("broken TapTree structure");
}
}
}
let leaf =
PartialTreeNode::with_leaf(leaf.leaf_version(), leaf.script().clone(), leaf_depth);
match node {
PartialTreeNode::Leaf(..) => { /* nothing to do here */ }
PartialTreeNode::Branch(branch, _) => {
branch.push_child(leaf);
}
}
}

let root = root
Expand Down Expand Up @@ -1426,6 +1424,8 @@ enum BranchDirection {

/// Iterator over leaf scripts stored in the leaf nodes of the taproot script
/// tree.
///
/// NB: The scripts are iterated in the DFS order (not consensus).
pub struct TreeScriptIter<'tree> {
// Here we store vec of path elements, where each element is a tuple, consisting of:
// 1. Tree node on the path
Expand Down Expand Up @@ -1507,11 +1507,10 @@ impl From<TaprootScriptTree> for TapTree {
mod test {
use std::collections::BTreeSet;

use amplify::Wrapper;
use bitcoin::blockdata::opcodes::all;
use bitcoin::hashes::hex::FromHex;
use bitcoin::schnorr::UntweakedPublicKey;
use bitcoin::util::taproot::TaprootBuilder;
use secp256k1::ONE_KEY;

use super::*;

Expand All @@ -1534,21 +1533,13 @@ mod test {
let taptree = compose_tree(opcode, depth_map);
let script_tree = TaprootScriptTree::from(taptree.clone());

let dumb = KeyPair::from_secret_key(SECP256K1, ONE_KEY);

let map = taptree
.clone()
.into_builder()
.finalize(SECP256K1, UntweakedPublicKey::from_keypair(&dumb))
.unwrap();
let scripts = map
.as_script_map()
.iter()
.map(|((script, version), path)| {
let scripts = taptree
.script_leaves()
.map(|leaf| {
(
path.iter().next().unwrap().as_inner().len() as u8,
*version,
script,
leaf.merkle_branch().as_inner().len() as u8,
leaf.leaf_version(),
leaf.script(),
)
})
.collect::<BTreeSet<_>>();
Expand Down Expand Up @@ -1748,7 +1739,7 @@ mod test {
}

#[test]
fn instll_path_proof() {
fn instill_path_proof() {
let path = DfsPath::from_str("00101").unwrap();

let taptree = compose_tree(0x51, [3, 5, 5, 4, 3, 3, 2, 3, 4, 5, 6, 8, 8, 7]);
Expand Down Expand Up @@ -1804,4 +1795,47 @@ mod test {
PartnerNode::Script(s!("Script(OP_PUSHNUM_3)")),
]);
}

#[test]
fn tapscripttree_roudtrip() {
let taptree = compose_tree(0x51, [3, 5, 5, 4, 3, 3, 2, 3, 4, 5, 6, 8, 8, 7]);
let script_tree = TaprootScriptTree::from(taptree.clone());
let taptree_roundtrip = TapTree::from(script_tree);
assert_eq!(taptree, taptree_roundtrip);
}

#[test]
fn tapscripttree_taptree_eq() {
let taptree = compose_tree(0x51, [3, 5, 5, 4, 3, 3, 2, 3, 4, 5, 6, 8, 8, 7]);
let script_tree = TaprootScriptTree::from(taptree.clone());
assert!(script_tree.check().is_ok());

// TODO: This is a bugfix, which should be reversed once <https://github.com/rust-bitcoin/rust-bitcoin/issues/1069> is fixed upstream
let mut script_leaves = taptree.script_leaves().collect::<Vec<_>>();
script_leaves.reverse();

for (leaf, (_, leaf_script)) in script_leaves.iter().zip(script_tree.scripts()) {
assert_eq!(leaf.script(), leaf_script.script.as_inner());
}
}

#[test]
fn tapscripttree_dfs() {
let depth_map = [3, 5, 5, 4, 3, 3, 2, 3, 4, 5, 6, 8, 8, 7];
let mut val = 0x51;

let taptree = compose_tree(val, depth_map);
let script_tree = TaprootScriptTree::from(taptree.clone());
assert!(script_tree.check().is_ok());

for (depth, leaf_script) in script_tree.scripts() {
let script = Script::from_hex(&format!("{:02x}", val)).unwrap();

assert_eq!(depth, depth_map[(val - 0x51) as usize]);
assert_eq!(script, leaf_script.script.to_inner());

let (new_val, _) = val.overflowing_add(1);
val = new_val;
}
}
}

0 comments on commit 47aa3b5

Please sign in to comment.