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

miniscript/miniscript/
iter.rs
