diff --git a/src/interpreter/inner.rs b/src/interpreter/inner.rs index 538cacb7e..59de53324 100644 --- a/src/interpreter/inner.rs +++ b/src/interpreter/inner.rs @@ -380,7 +380,7 @@ impl ToNoChecks for Miniscript { translate_hash_clone!(bitcoin::PublicKey, BitcoinKey, ()); } - self.real_translate_pk(&mut TranslateFullPk) + self.translate_pk_ctx(&mut TranslateFullPk) .expect("Translation should succeed") } } @@ -397,7 +397,7 @@ impl ToNoChecks for Miniscript +// SPDX-License-Identifier: CC0-1.0 + +//! Abstract Tree Iteration +//! +//! This module provides functionality to treat Miniscript objects abstractly +//! as trees, iterating over them in various orders. The iterators in this +//! module can be used to avoid explicitly recursive algorithms. +//! + +mod tree; + +pub use tree::{ + PostOrderIter, PostOrderIterItem, PreOrderIter, PreOrderIterItem, Tree, TreeLike, + VerbosePreOrderIter, +}; + +use crate::sync::Arc; +use crate::{Miniscript, MiniscriptKey, ScriptContext, Terminal}; + +impl<'a, Pk: MiniscriptKey, Ctx: ScriptContext> TreeLike for &'a Miniscript { + fn as_node(&self) -> Tree { + match self.node { + Terminal::PkK(..) + | Terminal::PkH(..) + | Terminal::RawPkH(..) + | Terminal::After(..) + | Terminal::Older(..) + | Terminal::Sha256(..) + | Terminal::Hash256(..) + | Terminal::Ripemd160(..) + | Terminal::Hash160(..) + | Terminal::True + | Terminal::False + | Terminal::Multi(..) + | Terminal::MultiA(..) => Tree::Nullary, + Terminal::Alt(ref sub) + | Terminal::Swap(ref sub) + | Terminal::Check(ref sub) + | Terminal::DupIf(ref sub) + | Terminal::Verify(ref sub) + | Terminal::NonZero(ref sub) + | Terminal::ZeroNotEqual(ref sub) => Tree::Unary(sub), + Terminal::AndV(ref left, ref right) + | Terminal::AndB(ref left, ref right) + | Terminal::OrB(ref left, ref right) + | Terminal::OrD(ref left, ref right) + | Terminal::OrC(ref left, ref right) + | Terminal::OrI(ref left, ref right) => Tree::Binary(left, right), + Terminal::AndOr(ref a, ref b, ref c) => Tree::Nary(Arc::from([a.as_ref(), b, c])), + Terminal::Thresh(_, ref subs) => Tree::Nary(subs.iter().map(Arc::as_ref).collect()), + } + } +} + +impl TreeLike for Arc> { + fn as_node(&self) -> Tree { + match self.node { + Terminal::PkK(..) + | Terminal::PkH(..) + | Terminal::RawPkH(..) + | Terminal::After(..) + | Terminal::Older(..) + | Terminal::Sha256(..) + | Terminal::Hash256(..) + | Terminal::Ripemd160(..) + | Terminal::Hash160(..) + | Terminal::True + | Terminal::False + | Terminal::Multi(..) + | Terminal::MultiA(..) => Tree::Nullary, + Terminal::Alt(ref sub) + | Terminal::Swap(ref sub) + | Terminal::Check(ref sub) + | Terminal::DupIf(ref sub) + | Terminal::Verify(ref sub) + | Terminal::NonZero(ref sub) + | Terminal::ZeroNotEqual(ref sub) => Tree::Unary(Arc::clone(sub)), + Terminal::AndV(ref left, ref right) + | Terminal::AndB(ref left, ref right) + | Terminal::OrB(ref left, ref right) + | Terminal::OrD(ref left, ref right) + | Terminal::OrC(ref left, ref right) + | Terminal::OrI(ref left, ref right) => { + Tree::Binary(Arc::clone(left), Arc::clone(right)) + } + Terminal::AndOr(ref a, ref b, ref c) => { + Tree::Nary(Arc::from([Arc::clone(a), Arc::clone(b), Arc::clone(c)])) + } + Terminal::Thresh(_, ref subs) => Tree::Nary(subs.iter().map(Arc::clone).collect()), + } + } +} diff --git a/src/iter/tree.rs b/src/iter/tree.rs new file mode 100644 index 000000000..4ff0e1ee9 --- /dev/null +++ b/src/iter/tree.rs @@ -0,0 +1,324 @@ +// Written in 2023 by Andrew Poelstra +// SPDX-License-Identifier: CC0-1.0 + +//! Abstract Trees +//! +//! This module provides the [`TreeLike`] trait which represents a node in a +//! tree, and several iterators over trees whose nodes implement this trait. +//! + +use crate::prelude::*; +use crate::sync::Arc; + +/// Abstract node of a tree. +/// +/// Tracks the arity (out-degree) of a node, which is the only thing that +/// is needed for iteration purposes. +pub enum Tree { + /// Combinator with no children. + Nullary, + /// Combinator with one child. + Unary(T), + /// Combinator with two children. + Binary(T, T), + /// Combinator with more than two children. + Nary(Arc<[T]>), +} + +/// A trait for any structure which has the shape of a Miniscript tree. +/// +/// As a general rule, this should be implemented on references to nodes, +/// rather than nodes themselves, because it provides algorithms that +/// assume copying is cheap. +/// +/// To implement this trait, you only need to implement the [`TreeLike::as_node`] +/// method, which will usually be very mechanical. Everything else is provided. +/// However, to avoid allocations, it may make sense to also implement +/// [`TreeLike::n_children`] and [`TreeLike::nth_child`] because the default +/// implementations will allocate vectors for n-ary nodes. +pub trait TreeLike: Clone + Sized { + /// Interpret the node as an abstract node. + fn as_node(&self) -> Tree; + + /// Accessor for the number of children this node has. + fn n_children(&self) -> usize { + match self.as_node() { + Tree::Nullary => 0, + Tree::Unary(..) => 1, + Tree::Binary(..) => 2, + Tree::Nary(children) => children.len(), + } + } + + /// Accessor for the nth child of the node, if a child with that index exists. + fn nth_child(&self, n: usize) -> Option { + match (n, self.as_node()) { + (_, Tree::Nullary) => None, + (0, Tree::Unary(sub)) => Some(sub), + (_, Tree::Unary(..)) => None, + (0, Tree::Binary(sub, _)) => Some(sub), + (1, Tree::Binary(_, sub)) => Some(sub), + (_, Tree::Binary(..)) => None, + (n, Tree::Nary(children)) => children.get(n).cloned(), + } + } + + /// Obtains an iterator of all the nodes rooted at the node, in pre-order. + fn pre_order_iter(self) -> PreOrderIter { + PreOrderIter { stack: vec![self] } + } + + /// Obtains a verbose iterator of all the nodes rooted at the DAG, in pre-order. + /// + /// See the documentation of [`VerbosePreOrderIter`] for more information about what + /// this does. Essentially, if you find yourself using [`Self::pre_order_iter`] and + /// then adding a stack to manually track which items and their children have been + /// yielded, you may be better off using this iterator instead. + fn verbose_pre_order_iter(self) -> VerbosePreOrderIter { + VerbosePreOrderIter { + stack: vec![PreOrderIterItem::initial(self, None)], + index: 0, + } + } + + /// Obtains an iterator of all the nodes rooted at the DAG, in post order. + /// + /// Each node is only yielded once, at the leftmost position that it + /// appears in the DAG. + fn post_order_iter(self) -> PostOrderIter { + PostOrderIter { + index: 0, + stack: vec![IterStackItem::unprocessed(self, None)], + } + } +} + +/// Element stored internally on the stack of a [`PostOrderIter`]. +/// +/// This is **not** the type that is yielded by the [`PostOrderIter`]; +/// in fact, this type is not even exported. +#[derive(Clone, Debug)] +struct IterStackItem { + /// The element on the stack + elem: T, + /// Whether we have dealt with this item (and pushed its children, + /// if any) yet. + processed: bool, + /// If the item has been processed, the index of its children. + child_indices: Vec, + /// Whether the element is a left- or right-child of its parent. + parent_stack_idx: Option, +} + +impl IterStackItem { + /// Constructor for a new stack item with a given element and relationship + /// to its parent. + fn unprocessed(elem: T, parent_stack_idx: Option) -> Self { + IterStackItem { + processed: false, + child_indices: Vec::with_capacity(elem.n_children()), + parent_stack_idx, + elem, + } + } +} + +/// Iterates over a DAG in _post order_. +/// +/// That means nodes are yielded in the order (left child, right child, parent). +#[derive(Clone, Debug)] +pub struct PostOrderIter { + /// The index of the next item to be yielded + index: usize, + /// A stack of elements to be yielded; each element is a node, then its left + /// and right children (if they exist and if they have been yielded already) + stack: Vec>, +} + +/// A set of data yielded by a `PostOrderIter`. +pub struct PostOrderIterItem { + /// The actual node data + pub node: T, + /// The index of this node (equivalent to if you'd called `.enumerate()` on + /// the iterator) + pub index: usize, + /// The indices of this node's children. + pub child_indices: Vec, +} + +impl Iterator for PostOrderIter { + type Item = PostOrderIterItem; + + fn next(&mut self) -> Option { + let mut current = self.stack.pop()?; + + if !current.processed { + current.processed = true; + + // When we first encounter an item, it is completely unknown; it is + // nominally the next item to be yielded, but it might have children, + // and if so, they come first + let current_stack_idx = self.stack.len(); + let n_children = current.elem.n_children(); + self.stack.push(current); + for idx in (0..n_children).rev() { + self.stack.push(IterStackItem::unprocessed( + self.stack[current_stack_idx].elem.nth_child(idx).unwrap(), + Some(current_stack_idx), + )); + } + self.next() + } else { + // The second time we encounter an item, we have dealt with its children, + // updated the child indices for this item, and are now ready to yield it + // rather than putting it back in the stack. + // + // Before yielding though, we must the item's parent's child indices with + // this item's index. + if let Some(idx) = current.parent_stack_idx { + self.stack[idx].child_indices.push(self.index); + } + + self.index += 1; + Some(PostOrderIterItem { + node: current.elem, + index: self.index - 1, + child_indices: current.child_indices, + }) + } + } +} + +/// Iterates over a [`TreeLike`] in _pre order_. +/// +/// Unlike the post-order iterator, this one does not keep track of indices +/// (this would be impractical since when we yield a node we have not yet +/// yielded its children, so we cannot know their indices). If you do need +/// the indices for some reason, the best strategy may be to run the +/// post-order iterator, collect into a vector, then iterate through that +/// backward. +#[derive(Clone, Debug)] +pub struct PreOrderIter { + /// A stack of elements to be yielded. As items are yielded, their right + /// children are put onto the stack followed by their left, so that the + /// appropriate one will be yielded on the next iteration. + stack: Vec, +} + +impl Iterator for PreOrderIter { + type Item = T; + + fn next(&mut self) -> Option { + // This algorithm is _significantly_ simpler than the post-order one, + // mainly because we don't care about child indices. + let top = self.stack.pop()?; + match top.as_node() { + Tree::Nullary => {} + Tree::Unary(next) => self.stack.push(next), + Tree::Binary(left, right) => { + self.stack.push(right); + self.stack.push(left); + } + Tree::Nary(children) => { + self.stack.extend(children.into_iter().rev().cloned()); + } + } + Some(top) + } +} + +/// Iterates over a [`TreeLike`] in "verbose pre order", yielding extra state changes. +/// +/// This yields nodes followed by their children, followed by the node *again* +/// after each child. This means that each node will be yielded a total of +/// (n+1) times, where n is its number of children. +/// +/// The different times that a node is yielded can be distinguished by looking +/// at the [`PreOrderIterItem::n_children_yielded`] (which, in particular, +/// will be 0 on the first yield) and [`PreOrderIterItem::is_complete`] (which +/// will be true on the last yield) fields of the yielded item. +#[derive(Clone, Debug)] +pub struct VerbosePreOrderIter { + /// A stack of elements to be yielded. As items are yielded, their right + /// children are put onto the stack followed by their left, so that the + /// appropriate one will be yielded on the next iteration. + stack: Vec>, + /// The index of the next item to be yielded. + /// + /// Note that unlike the [`PostOrderIter`], this value is not monotonic + /// and not equivalent to just using `enumerate` on the iterator, because + /// elements may be yielded multiple times. + index: usize, +} + +impl Iterator for VerbosePreOrderIter { + type Item = PreOrderIterItem; + + fn next(&mut self) -> Option { + // This algorithm is still simpler than the post-order one, because while + // we care about node indices, we don't care about their childrens' indices. + let mut top = self.stack.pop()?; + + // If this is the first time we're be yielding this element, set its index. + if top.n_children_yielded == 0 { + top.index = self.index; + self.index += 1; + } + // Push the next child. + let n_children = top.node.n_children(); + if top.n_children_yielded < n_children { + self.stack.push(top.clone().increment(n_children)); + let child = top.node.nth_child(top.n_children_yielded).unwrap(); + self.stack + .push(PreOrderIterItem::initial(child, Some(top.node.clone()))); + } + + // Then yield the element. + Some(top) + } +} + +/// A set of data yielded by a [`VerbosePreOrderIter`]. +#[derive(Clone, Debug)] +pub struct PreOrderIterItem { + /// The actual element being yielded. + pub node: T, + /// The parent of this node. `None` for the initial node, but will be + /// populated for all other nodes. + pub parent: Option, + /// The index when the element was first yielded. + pub index: usize, + /// How many of this item's children have been yielded. + /// + /// This can also be interpreted as a count of how many times this + /// item has been yielded before. + pub n_children_yielded: usize, + /// Whether this item is done (will not be yielded again). + pub is_complete: bool, +} + +impl PreOrderIterItem { + /// Creates a `PreOrderIterItem` which yields a given element for the first time. + /// + /// Marks the index as 0. The index must be manually set before yielding. + fn initial(node: T, parent: Option) -> Self { + PreOrderIterItem { + is_complete: node.n_children() == 0, + node, + parent, + index: 0, + n_children_yielded: 0, + } + } + + /// Creates a `PreOrderIterItem` which yields a given element again. + fn increment(self, n_children: usize) -> Self { + PreOrderIterItem { + node: self.node, + index: self.index, + parent: self.parent, + n_children_yielded: self.n_children_yielded + 1, + is_complete: self.n_children_yielded + 1 == n_children, + } + } +} diff --git a/src/lib.rs b/src/lib.rs index 69a551018..084f4ac9a 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -122,6 +122,7 @@ pub use pub_macros::*; pub mod descriptor; pub mod expression; pub mod interpreter; +pub mod iter; pub mod miniscript; pub mod policy; pub mod psbt; diff --git a/src/miniscript/astelem.rs b/src/miniscript/astelem.rs index 62dddf633..5c85b53d5 100644 --- a/src/miniscript/astelem.rs +++ b/src/miniscript/astelem.rs @@ -21,8 +21,7 @@ use crate::miniscript::ScriptContext; use crate::prelude::*; use crate::util::MsKeyBuilder; use crate::{ - errstr, expression, script_num_size, AbsLockTime, Error, ForEachKey, Miniscript, MiniscriptKey, - Terminal, ToPublicKey, TranslateErr, TranslatePk, Translator, + errstr, expression, AbsLockTime, Error, Miniscript, MiniscriptKey, Terminal, ToPublicKey, }; impl Terminal { @@ -46,216 +45,6 @@ impl Terminal { } } -impl TranslatePk for Terminal -where - Pk: MiniscriptKey, - Q: MiniscriptKey, - Ctx: ScriptContext, -{ - type Output = Terminal; - - /// Converts an AST element with one public key type to one of another public key type. - fn translate_pk(&self, translate: &mut T) -> Result> - where - T: Translator, - { - self.real_translate_pk(translate) - } -} - -impl Terminal { - pub(super) fn real_for_each_key<'a, F: FnMut(&'a Pk) -> bool>(&'a self, pred: &mut F) -> bool { - match *self { - Terminal::PkK(ref p) => pred(p), - Terminal::PkH(ref p) => pred(p), - Terminal::RawPkH(..) - | Terminal::After(..) - | Terminal::Older(..) - | Terminal::Sha256(..) - | Terminal::Hash256(..) - | Terminal::Ripemd160(..) - | Terminal::Hash160(..) - | Terminal::True - | Terminal::False => true, - Terminal::Alt(ref sub) - | Terminal::Swap(ref sub) - | Terminal::Check(ref sub) - | Terminal::DupIf(ref sub) - | Terminal::Verify(ref sub) - | Terminal::NonZero(ref sub) - | Terminal::ZeroNotEqual(ref sub) => sub.real_for_each_key(pred), - Terminal::AndV(ref left, ref right) - | Terminal::AndB(ref left, ref right) - | Terminal::OrB(ref left, ref right) - | Terminal::OrD(ref left, ref right) - | Terminal::OrC(ref left, ref right) - | Terminal::OrI(ref left, ref right) => { - left.real_for_each_key(&mut *pred) && right.real_for_each_key(pred) - } - Terminal::AndOr(ref a, ref b, ref c) => { - a.real_for_each_key(&mut *pred) - && b.real_for_each_key(&mut *pred) - && c.real_for_each_key(pred) - } - Terminal::Thresh(_, ref subs) => subs.iter().all(|sub| sub.real_for_each_key(pred)), - Terminal::Multi(_, ref keys) | Terminal::MultiA(_, ref keys) => keys.iter().all(pred), - } - } - - pub(super) fn real_translate_pk( - &self, - t: &mut T, - ) -> Result, TranslateErr> - where - Q: MiniscriptKey, - CtxQ: ScriptContext, - T: Translator, - { - let frag: Terminal = match *self { - Terminal::PkK(ref p) => Terminal::PkK(t.pk(p)?), - Terminal::PkH(ref p) => Terminal::PkH(t.pk(p)?), - Terminal::RawPkH(ref p) => Terminal::RawPkH(*p), - Terminal::After(n) => Terminal::After(n), - Terminal::Older(n) => Terminal::Older(n), - Terminal::Sha256(ref x) => Terminal::Sha256(t.sha256(x)?), - Terminal::Hash256(ref x) => Terminal::Hash256(t.hash256(x)?), - Terminal::Ripemd160(ref x) => Terminal::Ripemd160(t.ripemd160(x)?), - Terminal::Hash160(ref x) => Terminal::Hash160(t.hash160(x)?), - Terminal::True => Terminal::True, - Terminal::False => Terminal::False, - Terminal::Alt(ref sub) => Terminal::Alt(Arc::new(sub.real_translate_pk(t)?)), - Terminal::Swap(ref sub) => Terminal::Swap(Arc::new(sub.real_translate_pk(t)?)), - Terminal::Check(ref sub) => Terminal::Check(Arc::new(sub.real_translate_pk(t)?)), - Terminal::DupIf(ref sub) => Terminal::DupIf(Arc::new(sub.real_translate_pk(t)?)), - Terminal::Verify(ref sub) => Terminal::Verify(Arc::new(sub.real_translate_pk(t)?)), - Terminal::NonZero(ref sub) => Terminal::NonZero(Arc::new(sub.real_translate_pk(t)?)), - Terminal::ZeroNotEqual(ref sub) => { - Terminal::ZeroNotEqual(Arc::new(sub.real_translate_pk(t)?)) - } - Terminal::AndV(ref left, ref right) => Terminal::AndV( - Arc::new(left.real_translate_pk(t)?), - Arc::new(right.real_translate_pk(t)?), - ), - Terminal::AndB(ref left, ref right) => Terminal::AndB( - Arc::new(left.real_translate_pk(t)?), - Arc::new(right.real_translate_pk(t)?), - ), - Terminal::AndOr(ref a, ref b, ref c) => Terminal::AndOr( - Arc::new(a.real_translate_pk(t)?), - Arc::new(b.real_translate_pk(t)?), - Arc::new(c.real_translate_pk(t)?), - ), - Terminal::OrB(ref left, ref right) => Terminal::OrB( - Arc::new(left.real_translate_pk(t)?), - Arc::new(right.real_translate_pk(t)?), - ), - Terminal::OrD(ref left, ref right) => Terminal::OrD( - Arc::new(left.real_translate_pk(t)?), - Arc::new(right.real_translate_pk(t)?), - ), - Terminal::OrC(ref left, ref right) => Terminal::OrC( - Arc::new(left.real_translate_pk(t)?), - Arc::new(right.real_translate_pk(t)?), - ), - Terminal::OrI(ref left, ref right) => Terminal::OrI( - Arc::new(left.real_translate_pk(t)?), - Arc::new(right.real_translate_pk(t)?), - ), - Terminal::Thresh(k, ref subs) => { - let subs: Result>>, _> = subs - .iter() - .map(|s| s.real_translate_pk(t).map(Arc::new)) - .collect(); - Terminal::Thresh(k, subs?) - } - Terminal::Multi(k, ref keys) => { - let keys: Result, _> = keys.iter().map(|k| t.pk(k)).collect(); - Terminal::Multi(k, keys?) - } - Terminal::MultiA(k, ref keys) => { - let keys: Result, _> = keys.iter().map(|k| t.pk(k)).collect(); - Terminal::MultiA(k, keys?) - } - }; - Ok(frag) - } - - /// Substitutes raw public keys hashes with the public keys as provided by map. - pub fn substitute_raw_pkh(&self, pk_map: &BTreeMap) -> Terminal { - match self { - Terminal::RawPkH(ref p) => match pk_map.get(p) { - Some(pk) => Terminal::PkH(pk.clone()).into(), - None => Terminal::RawPkH(*p).into(), - }, - Terminal::PkK(..) - | Terminal::PkH(..) - | Terminal::Multi(..) - | Terminal::MultiA(..) - | Terminal::After(..) - | Terminal::Older(..) - | Terminal::Sha256(..) - | Terminal::Hash256(..) - | Terminal::Ripemd160(..) - | Terminal::Hash160(..) - | Terminal::True - | Terminal::False => self.clone().into(), - Terminal::Alt(ref sub) => Terminal::Alt(Arc::new(sub.substitute_raw_pkh(pk_map))), - Terminal::Swap(ref sub) => Terminal::Swap(Arc::new(sub.substitute_raw_pkh(pk_map))), - Terminal::Check(ref sub) => Terminal::Check(Arc::new(sub.substitute_raw_pkh(pk_map))), - Terminal::DupIf(ref sub) => Terminal::DupIf(Arc::new(sub.substitute_raw_pkh(pk_map))), - Terminal::Verify(ref sub) => Terminal::Verify(Arc::new(sub.substitute_raw_pkh(pk_map))), - Terminal::NonZero(ref sub) => { - Terminal::NonZero(Arc::new(sub.substitute_raw_pkh(pk_map))) - } - Terminal::ZeroNotEqual(ref sub) => { - Terminal::ZeroNotEqual(Arc::new(sub.substitute_raw_pkh(pk_map))) - } - Terminal::AndV(ref left, ref right) => Terminal::AndV( - Arc::new(left.substitute_raw_pkh(pk_map)), - Arc::new(right.substitute_raw_pkh(pk_map)), - ), - Terminal::AndB(ref left, ref right) => Terminal::AndB( - Arc::new(left.substitute_raw_pkh(pk_map)), - Arc::new(right.substitute_raw_pkh(pk_map)), - ), - Terminal::AndOr(ref a, ref b, ref c) => Terminal::AndOr( - Arc::new(a.substitute_raw_pkh(pk_map)), - Arc::new(b.substitute_raw_pkh(pk_map)), - Arc::new(c.substitute_raw_pkh(pk_map)), - ), - Terminal::OrB(ref left, ref right) => Terminal::OrB( - Arc::new(left.substitute_raw_pkh(pk_map)), - Arc::new(right.substitute_raw_pkh(pk_map)), - ), - Terminal::OrD(ref left, ref right) => Terminal::OrD( - Arc::new(left.substitute_raw_pkh(pk_map)), - Arc::new(right.substitute_raw_pkh(pk_map)), - ), - Terminal::OrC(ref left, ref right) => Terminal::OrC( - Arc::new(left.substitute_raw_pkh(pk_map)), - Arc::new(right.substitute_raw_pkh(pk_map)), - ), - Terminal::OrI(ref left, ref right) => Terminal::OrI( - Arc::new(left.substitute_raw_pkh(pk_map)), - Arc::new(right.substitute_raw_pkh(pk_map)), - ), - Terminal::Thresh(k, ref subs) => { - let subs: Vec>> = subs - .iter() - .map(|s| Arc::new(s.substitute_raw_pkh(pk_map))) - .collect(); - Terminal::Thresh(*k, subs) - } - } - } -} - -impl ForEachKey for Terminal { - fn for_each_key<'a, F: FnMut(&'a Pk) -> bool>(&'a self, mut pred: F) -> bool { - self.real_for_each_key(&mut pred) - } -} - impl fmt::Debug for Terminal { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { f.write_str("[")?; @@ -804,64 +593,4 @@ impl Terminal { } } } - - /// Size, in bytes of the script-pubkey. If this Miniscript is used outside - /// of segwit (e.g. in a bare or P2SH descriptor), this quantity should be - /// multiplied by 4 to compute the weight. - /// - /// In general, it is not recommended to use this function directly, but - /// to instead call the corresponding function on a `Descriptor`, which - /// will handle the segwit/non-segwit technicalities for you. - pub fn script_size(&self) -> usize { - match *self { - Terminal::PkK(ref pk) => Ctx::pk_len(pk), - Terminal::PkH(..) | Terminal::RawPkH(..) => 24, - Terminal::After(n) => script_num_size(n.to_consensus_u32() as usize) + 1, - Terminal::Older(n) => script_num_size(n.to_consensus_u32() as usize) + 1, - Terminal::Sha256(..) => 33 + 6, - Terminal::Hash256(..) => 33 + 6, - Terminal::Ripemd160(..) => 21 + 6, - Terminal::Hash160(..) => 21 + 6, - Terminal::True => 1, - Terminal::False => 1, - Terminal::Alt(ref sub) => sub.node.script_size() + 2, - Terminal::Swap(ref sub) => sub.node.script_size() + 1, - Terminal::Check(ref sub) => sub.node.script_size() + 1, - Terminal::DupIf(ref sub) => sub.node.script_size() + 3, - Terminal::Verify(ref sub) => { - sub.node.script_size() + usize::from(!sub.ext.has_free_verify) - } - Terminal::NonZero(ref sub) => sub.node.script_size() + 4, - Terminal::ZeroNotEqual(ref sub) => sub.node.script_size() + 1, - Terminal::AndV(ref l, ref r) => l.node.script_size() + r.node.script_size(), - Terminal::AndB(ref l, ref r) => l.node.script_size() + r.node.script_size() + 1, - Terminal::AndOr(ref a, ref b, ref c) => { - a.node.script_size() + b.node.script_size() + c.node.script_size() + 3 - } - Terminal::OrB(ref l, ref r) => l.node.script_size() + r.node.script_size() + 1, - Terminal::OrD(ref l, ref r) => l.node.script_size() + r.node.script_size() + 3, - Terminal::OrC(ref l, ref r) => l.node.script_size() + r.node.script_size() + 2, - Terminal::OrI(ref l, ref r) => l.node.script_size() + r.node.script_size() + 3, - Terminal::Thresh(k, ref subs) => { - assert!(!subs.is_empty(), "threshold must be nonempty"); - script_num_size(k) // k - + 1 // EQUAL - + subs.iter().map(|s| s.node.script_size()).sum::() - + subs.len() // ADD - - 1 // no ADD on first element - } - Terminal::Multi(k, ref pks) => { - script_num_size(k) - + 1 - + script_num_size(pks.len()) - + pks.iter().map(|pk| Ctx::pk_len(pk)).sum::() - } - Terminal::MultiA(k, ref pks) => { - script_num_size(k) - + 1 // NUMEQUAL - + pks.iter().map(|pk| Ctx::pk_len(pk)).sum::() // n keys - + pks.len() // n times CHECKSIGADD - } - } - } } diff --git a/src/miniscript/mod.rs b/src/miniscript/mod.rs index c31199d6b..db0b2baa6 100644 --- a/src/miniscript/mod.rs +++ b/src/miniscript/mod.rs @@ -22,8 +22,9 @@ use bitcoin::taproot::{LeafVersion, TapLeafHash}; use self::analyzable::ExtParams; pub use self::context::{BareCtx, Legacy, Segwitv0, Tap}; +use crate::iter::TreeLike; use crate::prelude::*; -use crate::TranslateErr; +use crate::{script_num_size, TranslateErr}; pub mod analyzable; pub mod astelem; @@ -258,7 +259,55 @@ where /// to instead call the corresponding function on a `Descriptor`, which /// will handle the segwit/non-segwit technicalities for you. pub fn script_size(&self) -> usize { - self.node.script_size() + let mut len = 0; + for ms in self.pre_order_iter() { + len += match ms.node { + Terminal::PkK(ref pk) => Ctx::pk_len(pk), + Terminal::PkH(..) | Terminal::RawPkH(..) => 24, + Terminal::After(n) => script_num_size(n.to_consensus_u32() as usize) + 1, + Terminal::Older(n) => script_num_size(n.to_consensus_u32() as usize) + 1, + Terminal::Sha256(..) => 33 + 6, + Terminal::Hash256(..) => 33 + 6, + Terminal::Ripemd160(..) => 21 + 6, + Terminal::Hash160(..) => 21 + 6, + Terminal::True => 1, + Terminal::False => 1, + Terminal::Alt(..) => 2, + Terminal::Swap(..) => 1, + Terminal::Check(..) => 1, + Terminal::DupIf(..) => 3, + Terminal::Verify(ref sub) => usize::from(!sub.ext.has_free_verify), + Terminal::NonZero(..) => 4, + Terminal::ZeroNotEqual(..) => 1, + Terminal::AndV(..) => 0, + Terminal::AndB(..) => 1, + Terminal::AndOr(..) => 3, + Terminal::OrB(..) => 1, + Terminal::OrD(..) => 3, + Terminal::OrC(..) => 2, + Terminal::OrI(..) => 3, + Terminal::Thresh(k, ref subs) => { + assert!(!subs.is_empty(), "threshold must be nonempty"); + script_num_size(k) // k + + 1 // EQUAL + + subs.len() // ADD + - 1 // no ADD on first element + } + Terminal::Multi(k, ref pks) => { + script_num_size(k) + + 1 + + script_num_size(pks.len()) + + pks.iter().map(|pk| Ctx::pk_len(pk)).sum::() + } + Terminal::MultiA(k, ref pks) => { + script_num_size(k) + + 1 // NUMEQUAL + + pks.iter().map(|pk| Ctx::pk_len(pk)).sum::() // n keys + + pks.len() // n times CHECKSIGADD + } + } + } + len } } @@ -296,7 +345,27 @@ impl Miniscript { impl ForEachKey for Miniscript { fn for_each_key<'a, F: FnMut(&'a Pk) -> bool>(&'a self, mut pred: F) -> bool { - self.real_for_each_key(&mut pred) + for ms in self.pre_order_iter() { + match ms.node { + Terminal::PkK(ref p) => { + if !pred(p) { + return false; + } + } + Terminal::PkH(ref p) => { + if !pred(p) { + return false; + } + } + Terminal::Multi(_, ref keys) | Terminal::MultiA(_, ref keys) => { + if !keys.iter().all(&mut pred) { + return false; + } + } + _ => {} + } + } + true } } @@ -310,20 +379,16 @@ where /// Translates a struct from one generic to another where the translation /// for Pk is provided by [`Translator`] - fn translate_pk(&self, translate: &mut T) -> Result> + fn translate_pk(&self, t: &mut T) -> Result> where T: Translator, { - self.real_translate_pk(translate) + self.translate_pk_ctx(t) } } impl Miniscript { - fn real_for_each_key<'a, F: FnMut(&'a Pk) -> bool>(&'a self, pred: &mut F) -> bool { - self.node.real_for_each_key(pred) - } - - pub(super) fn real_translate_pk( + pub(super) fn translate_pk_ctx( &self, t: &mut T, ) -> Result, TranslateErr> @@ -332,13 +397,74 @@ impl Miniscript { CtxQ: ScriptContext, T: Translator, { - let inner = self.node.real_translate_pk(t)?; - Miniscript::from_ast(inner).map_err(TranslateErr::OuterError) + let mut translated = vec![]; + for data in Arc::new(self.clone()).post_order_iter() { + // convenience method to reduce typing + let child_n = |n| Arc::clone(&translated[data.child_indices[n]]); + + let new_term = match data.node.node { + Terminal::PkK(ref p) => Terminal::PkK(t.pk(p)?), + Terminal::PkH(ref p) => Terminal::PkH(t.pk(p)?), + Terminal::RawPkH(ref p) => Terminal::RawPkH(*p), + Terminal::After(n) => Terminal::After(n), + Terminal::Older(n) => Terminal::Older(n), + Terminal::Sha256(ref x) => Terminal::Sha256(t.sha256(x)?), + Terminal::Hash256(ref x) => Terminal::Hash256(t.hash256(x)?), + Terminal::Ripemd160(ref x) => Terminal::Ripemd160(t.ripemd160(x)?), + Terminal::Hash160(ref x) => Terminal::Hash160(t.hash160(x)?), + Terminal::True => Terminal::True, + Terminal::False => Terminal::False, + Terminal::Alt(..) => Terminal::Alt(child_n(0)), + Terminal::Swap(..) => Terminal::Swap(child_n(0)), + Terminal::Check(..) => Terminal::Check(child_n(0)), + Terminal::DupIf(..) => Terminal::DupIf(child_n(0)), + Terminal::Verify(..) => Terminal::Verify(child_n(0)), + Terminal::NonZero(..) => Terminal::NonZero(child_n(0)), + Terminal::ZeroNotEqual(..) => Terminal::ZeroNotEqual(child_n(0)), + Terminal::AndV(..) => Terminal::AndV(child_n(0), child_n(1)), + Terminal::AndB(..) => Terminal::AndB(child_n(0), child_n(1)), + Terminal::AndOr(..) => Terminal::AndOr(child_n(0), child_n(1), child_n(2)), + Terminal::OrB(..) => Terminal::OrB(child_n(0), child_n(1)), + Terminal::OrD(..) => Terminal::OrD(child_n(0), child_n(1)), + Terminal::OrC(..) => Terminal::OrC(child_n(0), child_n(1)), + Terminal::OrI(..) => Terminal::OrI(child_n(0), child_n(1)), + Terminal::Thresh(k, ref subs) => { + Terminal::Thresh(k, (0..subs.len()).map(child_n).collect()) + } + Terminal::Multi(k, ref keys) => { + let keys: Result, _> = keys.iter().map(|k| t.pk(k)).collect(); + Terminal::Multi(k, keys?) + } + Terminal::MultiA(k, ref keys) => { + let keys: Result, _> = keys.iter().map(|k| t.pk(k)).collect(); + Terminal::MultiA(k, keys?) + } + }; + let new_ms = Miniscript::from_ast(new_term).map_err(TranslateErr::OuterError)?; + translated.push(Arc::new(new_ms)); + } + + Ok(Arc::try_unwrap(translated.pop().unwrap()).unwrap()) } /// Substitutes raw public keys hashes with the public keys as provided by map. pub fn substitute_raw_pkh(&self, pk_map: &BTreeMap) -> Miniscript { - Miniscript::from_ast(self.node.substitute_raw_pkh(pk_map)).expect("type check failed") + let mut translated = vec![]; + for data in Arc::new(self.clone()).post_order_iter() { + let new_term = if let Terminal::RawPkH(ref p) = data.node.node { + match pk_map.get(p) { + Some(pk) => Terminal::PkH(pk.clone()).into(), + None => Terminal::RawPkH(*p).into(), + } + } else { + data.node.node.clone() + }; + + let new_ms = Miniscript::from_ast(new_term).expect("typeck"); + translated.push(Arc::new(new_ms)); + } + + Arc::try_unwrap(translated.pop().unwrap()).unwrap() } }