elements_miniscript/miniscript/
iter.rs

1// Written in 2022 by Dr Maxim Orlovsky <orlovsky@pandoracore.com>
2// SPDX-License-Identifier: CC0-1.0
3
4//! Miniscript Iterators
5//!
6//! Iterators for Miniscript with special functions for iterating
7//! over Public Keys, Public Key Hashes or both.
8use std::ops::Deref;
9use std::sync::Arc;
10
11use super::decode::Terminal;
12use super::{Miniscript, MiniscriptKey, ScriptContext};
13use crate::Extension;
14
15/// Iterator-related extensions for [Miniscript]
16impl<Pk: MiniscriptKey, Ctx: ScriptContext, Ext: Extension> Miniscript<Pk, Ctx, Ext> {
17    /// Creates a new [Iter] iterator that will iterate over all [Miniscript] items within
18    /// AST by traversing its branches. For the specific algorithm please see
19    /// [Iter::next] function.
20    pub fn iter(&self) -> Iter<'_, Pk, Ctx, Ext> {
21        Iter::new(self)
22    }
23
24    /// Creates a new [PkIter] iterator that will iterate over all plain public keys (and not
25    /// key hash values) present in [Miniscript] items within AST by traversing all its branches.
26    /// For the specific algorithm please see [PkIter::next] function.
27    pub fn iter_pk(&self) -> PkIter<'_, Pk, Ctx, Ext> {
28        PkIter::new(self)
29    }
30
31    /// Enumerates all child nodes of the current AST node (`self`) and returns a `Vec` referencing
32    /// them.
33    pub fn branches(&self) -> Vec<&Miniscript<Pk, Ctx, Ext>> {
34        match self.node {
35            Terminal::PkK(_) | Terminal::PkH(_) | Terminal::RawPkH(_) | Terminal::Multi(_, _) => {
36                vec![]
37            }
38
39            Terminal::Alt(ref node)
40            | Terminal::Swap(ref node)
41            | Terminal::Check(ref node)
42            | Terminal::DupIf(ref node)
43            | Terminal::Verify(ref node)
44            | Terminal::NonZero(ref node)
45            | Terminal::ZeroNotEqual(ref node) => vec![node],
46
47            Terminal::AndV(ref node1, ref node2)
48            | Terminal::AndB(ref node1, ref node2)
49            | Terminal::OrB(ref node1, ref node2)
50            | Terminal::OrD(ref node1, ref node2)
51            | Terminal::OrC(ref node1, ref node2)
52            | Terminal::OrI(ref node1, ref node2) => vec![node1, node2],
53
54            Terminal::AndOr(ref node1, ref node2, ref node3) => vec![node1, node2, node3],
55
56            Terminal::Thresh(_, ref node_vec) => node_vec.iter().map(Arc::deref).collect(),
57
58            _ => vec![],
59        }
60    }
61
62    /// Returns child node with given index, if any
63    pub fn get_nth_child(&self, n: usize) -> Option<&Miniscript<Pk, Ctx, Ext>> {
64        match (n, &self.node) {
65            (0, Terminal::Alt(node))
66            | (0, Terminal::Swap(node))
67            | (0, Terminal::Check(node))
68            | (0, Terminal::DupIf(node))
69            | (0, Terminal::Verify(node))
70            | (0, Terminal::NonZero(node))
71            | (0, Terminal::ZeroNotEqual(node))
72            | (0, Terminal::AndV(node, _))
73            | (0, Terminal::AndB(node, _))
74            | (0, Terminal::OrB(node, _))
75            | (0, Terminal::OrD(node, _))
76            | (0, Terminal::OrC(node, _))
77            | (0, Terminal::OrI(node, _))
78            | (1, Terminal::AndV(_, node))
79            | (1, Terminal::AndB(_, node))
80            | (1, Terminal::OrB(_, node))
81            | (1, Terminal::OrD(_, node))
82            | (1, Terminal::OrC(_, node))
83            | (1, Terminal::OrI(_, node))
84            | (0, Terminal::AndOr(node, _, _))
85            | (1, Terminal::AndOr(_, node, _))
86            | (2, Terminal::AndOr(_, _, node)) => Some(node),
87
88            (n, Terminal::Thresh(_, node_vec)) => node_vec.get(n).map(|x| &**x),
89
90            _ => None,
91        }
92    }
93
94    /// Returns `Option::Some` with cloned n'th public key from the current miniscript item,
95    /// if any. Otherwise returns `Option::None`.
96    ///
97    /// NB: The function analyzes only single miniscript item and not any of its descendants in AST.
98    pub fn get_nth_pk(&self, n: usize) -> Option<Pk> {
99        match (&self.node, n) {
100            (&Terminal::PkK(ref key), 0) | (&Terminal::PkH(ref key), 0) => Some(key.clone()),
101            (&Terminal::Multi(_, ref keys), _) | (&Terminal::MultiA(_, ref keys), _) => {
102                keys.get(n).cloned()
103            }
104            _ => None,
105        }
106    }
107}
108
109/// Iterator for traversing all [Miniscript] miniscript AST references starting from some specific
110/// node which constructs the iterator via [Miniscript::iter] method.
111pub struct Iter<'a, Pk: MiniscriptKey, Ctx: ScriptContext, Ext: Extension> {
112    next: Option<&'a Miniscript<Pk, Ctx, Ext>>,
113    // Here we store vec of path elements, where each element is a tuple, consisting of:
114    // 1. Miniscript node on the path
115    // 2. Index of the current branch
116    path: Vec<(&'a Miniscript<Pk, Ctx, Ext>, usize)>,
117}
118
119impl<'a, Pk: MiniscriptKey, Ctx: ScriptContext, Ext: Extension> Iter<'a, Pk, Ctx, Ext> {
120    fn new(miniscript: &'a Miniscript<Pk, Ctx, Ext>) -> Self {
121        Iter {
122            next: Some(miniscript),
123            path: vec![],
124        }
125    }
126}
127
128impl<'a, Pk: MiniscriptKey, Ctx: 'a + ScriptContext, Ext: 'a + Extension> Iterator
129    for Iter<'a, Pk, Ctx, Ext>
130{
131    type Item = &'a Miniscript<Pk, Ctx, Ext>;
132
133    /// First, the function returns `self`, then the first child of the self (if any),
134    /// then proceeds to the child of the child — down to a leaf of the tree in its first branch.
135    /// When the leaf is reached, it goes in the reverse direction on the same branch until it
136    /// founds a first branching node that had more than a single branch and returns it, traversing
137    /// it with the same algorithm again.
138    ///
139    /// For example, for the given AST
140    /// ```text
141    /// A --+--> B -----> C --+--> D -----> E
142    ///     |                 |
143    ///     |                 +--> F
144    ///     |                 |
145    ///     |                 +--> G --+--> H
146    ///     |                          |
147    ///     |                          +--> I -----> J
148    ///     +--> K
149    /// ```
150    /// `Iter::next()` will iterate over the nodes in the following order:
151    /// `A > B > C > D > E > F > G > H > I > J > K`
152    ///
153    /// To enumerate the branches iterator uses [Miniscript::branches] function.
154    fn next(&mut self) -> Option<Self::Item> {
155        let mut curr = self.next;
156        if curr.is_none() {
157            while let Some((node, child)) = self.path.pop() {
158                curr = node.get_nth_child(child);
159                if curr.is_some() {
160                    self.path.push((node, child + 1));
161                    break;
162                }
163            }
164        }
165        if let Some(node) = curr {
166            self.next = node.get_nth_child(0);
167            self.path.push((node, 1));
168        }
169        curr
170    }
171}
172
173/// Iterator for traversing all [MiniscriptKey]'s in AST starting from some specific node which
174/// constructs the iterator via [Miniscript::iter_pk] method.
175pub struct PkIter<'a, Pk: MiniscriptKey, Ctx: ScriptContext, Ext: Extension> {
176    node_iter: Iter<'a, Pk, Ctx, Ext>,
177    curr_node: Option<&'a Miniscript<Pk, Ctx, Ext>>,
178    key_index: usize,
179}
180
181impl<'a, Pk: MiniscriptKey, Ctx: ScriptContext, Ext: Extension> PkIter<'a, Pk, Ctx, Ext> {
182    fn new(miniscript: &'a Miniscript<Pk, Ctx, Ext>) -> Self {
183        let mut iter = Iter::new(miniscript);
184        PkIter {
185            curr_node: iter.next(),
186            node_iter: iter,
187            key_index: 0,
188        }
189    }
190}
191
192impl<'a, Pk: MiniscriptKey, Ctx: ScriptContext, Ext: Extension> Iterator
193    for PkIter<'a, Pk, Ctx, Ext>
194{
195    type Item = Pk;
196
197    fn next(&mut self) -> Option<Self::Item> {
198        loop {
199            match self.curr_node {
200                None => break None,
201                Some(node) => match node.get_nth_pk(self.key_index) {
202                    None => {
203                        self.curr_node = self.node_iter.next();
204                        self.key_index = 0;
205                        continue;
206                    }
207                    Some(pk) => {
208                        self.key_index += 1;
209                        break Some(pk);
210                    }
211                },
212            }
213        }
214    }
215}
216
217// Module is public since it export testcase generation which may be used in
218// dependent libraries for their own tasts based on Miniscript AST
219#[cfg(test)]
220pub mod test {
221    use bitcoin;
222    use elements::hashes::{hash160, ripemd160, sha256, sha256d, Hash};
223    use elements::secp256k1_zkp;
224
225    use super::Miniscript;
226    use crate::miniscript::context::Segwitv0;
227    use crate::NoExt;
228
229    pub type TestData = (
230        Miniscript<bitcoin::PublicKey, Segwitv0, NoExt>,
231        Vec<bitcoin::PublicKey>,
232        Vec<hash160::Hash>,
233        bool, // Indicates that the top-level contains public key or hashes
234    );
235
236    pub fn gen_secp_pubkeys(n: usize) -> Vec<secp256k1_zkp::PublicKey> {
237        let mut ret = Vec::with_capacity(n);
238        let secp = secp256k1_zkp::Secp256k1::new();
239        let mut sk = [0; 32];
240
241        for i in 1..n + 1 {
242            sk[0] = i as u8;
243            sk[1] = (i >> 8) as u8;
244            sk[2] = (i >> 16) as u8;
245
246            ret.push(secp256k1_zkp::PublicKey::from_secret_key(
247                &secp,
248                &secp256k1_zkp::SecretKey::from_slice(&sk[..]).unwrap(),
249            ));
250        }
251        ret
252    }
253
254    pub fn gen_bitcoin_pubkeys(n: usize, compressed: bool) -> Vec<bitcoin::PublicKey> {
255        gen_secp_pubkeys(n)
256            .into_iter()
257            .map(|inner| bitcoin::PublicKey { inner, compressed })
258            .collect()
259    }
260
261    pub fn gen_testcases() -> Vec<TestData> {
262        let k = gen_bitcoin_pubkeys(10, true);
263        let _h: Vec<hash160::Hash> = k
264            .iter()
265            .map(|pk| hash160::Hash::hash(&pk.to_bytes()))
266            .collect();
267
268        let preimage = vec![0xab; 32];
269        let sha256_hash = sha256::Hash::hash(&preimage);
270        let sha256d_hash_rev = sha256d::Hash::hash(&preimage);
271        let mut sha256d_hash_bytes = sha256d_hash_rev.to_byte_array();
272        sha256d_hash_bytes.reverse();
273        let sha256d_hash = sha256d::Hash::from_byte_array(sha256d_hash_bytes);
274        let hash160_hash = hash160::Hash::hash(&preimage);
275        let ripemd160_hash = ripemd160::Hash::hash(&preimage);
276
277        vec![
278            (ms_str!("after({})", 1000), vec![], vec![], false),
279            (ms_str!("older({})", 1000), vec![], vec![], false),
280            (ms_str!("sha256({})", sha256_hash), vec![], vec![], false),
281            (ms_str!("hash256({})", sha256d_hash), vec![], vec![], false),
282            (ms_str!("hash160({})", hash160_hash), vec![], vec![], false),
283            (
284                ms_str!("ripemd160({})", ripemd160_hash),
285                vec![],
286                vec![],
287                false,
288            ),
289            (ms_str!("c:pk_k({})", k[0]), vec![k[0]], vec![], true),
290            (ms_str!("c:pk_h({})", k[0]), vec![k[0]], vec![], true),
291            (
292                ms_str!("and_v(vc:pk_k({}),c:pk_h({}))", k[0], k[1]),
293                vec![k[0], k[1]],
294                vec![],
295                false,
296            ),
297            (
298                ms_str!("and_b(c:pk_k({}),sjtv:sha256({}))", k[0], sha256_hash),
299                vec![k[0]],
300                vec![],
301                false,
302            ),
303            (
304                ms_str!(
305                    "andor(c:pk_k({}),jtv:sha256({}),c:pk_h({}))",
306                    k[1],
307                    sha256_hash,
308                    k[2]
309                ),
310                vec![k[1], k[2]],
311                vec![],
312                false,
313            ),
314            (
315                ms_str!("multi(3,{},{},{},{},{})", k[9], k[8], k[7], k[0], k[1]),
316                vec![k[9], k[8], k[7], k[0], k[1]],
317                vec![],
318                true,
319            ),
320            (
321                ms_str!(
322                    "thresh(3,c:pk_k({}),sc:pk_k({}),sc:pk_k({}),sc:pk_k({}),sc:pk_k({}))",
323                    k[2],
324                    k[3],
325                    k[4],
326                    k[5],
327                    k[6]
328                ),
329                vec![k[2], k[3], k[4], k[5], k[6]],
330                vec![],
331                false,
332            ),
333            (
334                ms_str!(
335                    "or_d(multi(2,{},{}),and_v(v:multi(2,{},{}),older(10000)))",
336                    k[6],
337                    k[7],
338                    k[8],
339                    k[9]
340                ),
341                vec![k[6], k[7], k[8], k[9]],
342                vec![],
343                false,
344            ),
345            (
346                ms_str!(
347                    "or_d(multi(3,{},{},{},{},{}),\
348                      and_v(v:thresh(2,c:pk_h({}),\
349                      ac:pk_h({}),ac:pk_h({})),older(10000)))",
350                    k[0],
351                    k[2],
352                    k[4],
353                    k[6],
354                    k[9],
355                    k[1],
356                    k[3],
357                    k[5]
358                ),
359                vec![k[0], k[2], k[4], k[6], k[9], k[1], k[3], k[5]],
360                vec![],
361                false,
362            ),
363        ]
364    }
365
366    #[test]
367    fn find_keys() {
368        gen_testcases().into_iter().for_each(|(ms, k, _, _)| {
369            assert_eq!(ms.iter_pk().collect::<Vec<bitcoin::PublicKey>>(), k);
370        })
371    }
372}
elements_miniscript/miniscript/iter.rs

elements_miniscript/miniscript/
iter.rs
