Remove Concrete translation recursion

Remove recursion in the two translation functions for
`policy::Concrete`. Requires implementing `TreeLike` for
`policy::Concrete` as we do for `Miniscript`.

Note this does not remove _all_ recursion in the `concrete` module (see
for example `num_tap_leaves`).

Author: tcharding

src/iter/mod.rs

@@ -15,7 +15,7 @@ pub use tree::{
 };
 
 use crate::sync::Arc;
-use crate::{Miniscript, MiniscriptKey, ScriptContext, Terminal};
+use crate::{policy, Miniscript, MiniscriptKey, ScriptContext, Terminal};
 
 impl<'a, Pk: MiniscriptKey, Ctx: ScriptContext> TreeLike for &'a Miniscript<Pk, Ctx> {
     fn as_node(&self) -> Tree<Self> {
@@ -68,3 +68,29 @@ impl<Pk: MiniscriptKey, Ctx: ScriptContext> TreeLike for Arc<Miniscript<Pk, Ctx>
         }
     }
 }
+
+impl<'a, Pk: MiniscriptKey> TreeLike for &'a policy::Concrete<Pk> {
+    fn as_node(&self) -> Tree<Self> {
+        use policy::Concrete::*;
+        match *self {
+            Unsatisfiable | Trivial | Key(_) | After(_) | Older(_) | Sha256(_) | Hash256(_)
+            | Ripemd160(_) | Hash160(_) => Tree::Nullary,
+            And(ref subs) => Tree::Nary(subs.iter().map(Arc::as_ref).collect()),
+            Or(ref v) => Tree::Nary(v.iter().map(|(_, p)| Arc::as_ref(p)).collect()),
+            Threshold(_, ref subs) => Tree::Nary(subs.iter().map(Arc::as_ref).collect()),
+        }
+    }
+}
+
+impl<'a, Pk: MiniscriptKey> TreeLike for Arc<policy::Concrete<Pk>> {
+    fn as_node(&self) -> Tree<Self> {
+        use policy::Concrete::*;
+        match self.as_ref() {
+            Unsatisfiable | Trivial | Key(_) | After(_) | Older(_) | Sha256(_) | Hash256(_)
+            | Ripemd160(_) | Hash160(_) => Tree::Nullary,
+            And(ref subs) => Tree::Nary(subs.iter().map(Arc::clone).collect()),
+            Or(ref v) => Tree::Nary(v.iter().map(|(_, p)| Arc::clone(p)).collect()),
+            Threshold(_, ref subs) => Tree::Nary(subs.iter().map(Arc::clone).collect()),
+        }
+    }
+}

src/policy/concrete.rs

@@ -8,6 +8,7 @@ use core::{fmt, str};
 use std::error;
 
 use bitcoin::{absolute, Sequence};
+use sync::Arc;
 #[cfg(feature = "compiler")]
 use {
     crate::descriptor::TapTree,
@@ -19,11 +20,11 @@ use {
     crate::Miniscript,
     crate::Tap,
     core::cmp::Reverse,
-    sync::Arc,
 };
 
 use super::ENTAILMENT_MAX_TERMINALS;
 use crate::expression::{self, FromTree};
+use crate::iter::TreeLike;
 use crate::miniscript::types::extra_props::TimelockInfo;
 use crate::prelude::*;
 #[cfg(all(doc, not(feature = "compiler")))]
@@ -58,12 +59,12 @@ pub enum Policy<Pk: MiniscriptKey> {
     /// A HASH160 whose preimage must be provided to satisfy the descriptor.
     Hash160(Pk::Hash160),
     /// A list of sub-policies, all of which must be satisfied.
-    And(Vec<Policy<Pk>>),
+    And(Vec<Arc<Policy<Pk>>>),
     /// A list of sub-policies, one of which must be satisfied, along with
     /// relative probabilities for each one.
-    Or(Vec<(usize, Policy<Pk>)>),
+    Or(Vec<(usize, Arc<Policy<Pk>>)>),
     /// A set of descriptors, satisfactions must be provided for `k` of them.
-    Threshold(usize, Vec<Policy<Pk>>),
+    Threshold(usize, Vec<Arc<Policy<Pk>>>),
 }
 
 impl<Pk> Policy<Pk>
@@ -656,30 +657,21 @@ impl<Pk: MiniscriptKey> PolicyArc<Pk> {
 
 impl<Pk: MiniscriptKey> ForEachKey<Pk> for Policy<Pk> {
     fn for_each_key<'a, F: FnMut(&'a Pk) -> bool>(&'a self, mut pred: F) -> bool {
-        self.real_for_each_key(&mut pred)
-    }
-}
-
-impl<Pk: MiniscriptKey> Policy<Pk> {
-    fn real_for_each_key<'a, F: FnMut(&'a Pk) -> bool>(&'a self, pred: &mut F) -> bool {
-        match *self {
-            Policy::Unsatisfiable | Policy::Trivial => true,
-            Policy::Key(ref pk) => pred(pk),
-            Policy::Sha256(..)
-            | Policy::Hash256(..)
-            | Policy::Ripemd160(..)
-            | Policy::Hash160(..)
-            | Policy::After(..)
-            | Policy::Older(..) => true,
-            Policy::Threshold(_, ref subs) | Policy::And(ref subs) => {
-                subs.iter().all(|sub| sub.real_for_each_key(&mut *pred))
+        for policy in self.pre_order_iter() {
+            match policy {
+                Policy::Key(ref pk) => {
+                    if !pred(pk) {
+                        return false;
+                    }
+                }
+                _ => {}
             }
-            Policy::Or(ref subs) => subs
-                .iter()
-                .all(|(_, sub)| sub.real_for_each_key(&mut *pred)),
         }
+        true
     }
+}
 
+impl<Pk: MiniscriptKey> Policy<Pk> {
     /// Converts a policy using one kind of public key to another type of public key.
     ///
     /// For example usage please see [`crate::policy::semantic::Policy::translate_pk`].
@@ -688,81 +680,70 @@ impl<Pk: MiniscriptKey> Policy<Pk> {
         T: Translator<Pk, Q, E>,
         Q: MiniscriptKey,
     {
-        self._translate_pk(t)
-    }
-
-    fn _translate_pk<Q, E, T>(&self, t: &mut T) -> Result<Policy<Q>, E>
-    where
-        T: Translator<Pk, Q, E>,
-        Q: MiniscriptKey,
-    {
-        match *self {
-            Policy::Unsatisfiable => Ok(Policy::Unsatisfiable),
-            Policy::Trivial => Ok(Policy::Trivial),
-            Policy::Key(ref pk) => t.pk(pk).map(Policy::Key),
-            Policy::Sha256(ref h) => t.sha256(h).map(Policy::Sha256),
-            Policy::Hash256(ref h) => t.hash256(h).map(Policy::Hash256),
-            Policy::Ripemd160(ref h) => t.ripemd160(h).map(Policy::Ripemd160),
-            Policy::Hash160(ref h) => t.hash160(h).map(Policy::Hash160),
-            Policy::Older(n) => Ok(Policy::Older(n)),
-            Policy::After(n) => Ok(Policy::After(n)),
-            Policy::Threshold(k, ref subs) => {
-                let new_subs: Result<Vec<Policy<Q>>, _> =
-                    subs.iter().map(|sub| sub._translate_pk(t)).collect();
-                new_subs.map(|ok| Policy::Threshold(k, ok))
-            }
-            Policy::And(ref subs) => Ok(Policy::And(
-                subs.iter()
-                    .map(|sub| sub._translate_pk(t))
-                    .collect::<Result<Vec<Policy<Q>>, E>>()?,
-            )),
-            Policy::Or(ref subs) => Ok(Policy::Or(
-                subs.iter()
-                    .map(|(prob, sub)| Ok((*prob, sub._translate_pk(t)?)))
-                    .collect::<Result<Vec<(usize, Policy<Q>)>, E>>()?,
-            )),
+        use Policy::*;
+
+        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_policy = match data.node.as_ref() {
+                Unsatisfiable => Unsatisfiable,
+                Trivial => Trivial,
+                Key(ref pk) => t.pk(pk).map(Key)?,
+                Sha256(ref h) => t.sha256(h).map(Sha256)?,
+                Hash256(ref h) => t.hash256(h).map(Hash256)?,
+                Ripemd160(ref h) => t.ripemd160(h).map(Ripemd160)?,
+                Hash160(ref h) => t.hash160(h).map(Hash160)?,
+                Older(n) => Older(*n),
+                After(n) => After(*n),
+                Threshold(k, ref subs) => Threshold(*k, (0..subs.len()).map(child_n).collect()),
+                And(ref subs) => And((0..subs.len()).map(child_n).collect()),
+                Or(ref subs) => Or((0..subs.len()).map(|i| (i, child_n(i))).collect()),
+            };
+            translated.push(Arc::new(new_policy));
         }
+
+        Ok(Arc::try_unwrap(translated.pop().unwrap()).unwrap())
     }
 
     /// Translates `Concrete::Key(key)` to `Concrete::Unsatisfiable` when extracting `TapKey`.
     pub fn translate_unsatisfiable_pk(self, key: &Pk) -> Policy<Pk> {
-        match self {
-            Policy::Key(ref k) if k.clone() == *key => Policy::Unsatisfiable,
-            Policy::And(subs) => Policy::And(
-                subs.into_iter()
-                    .map(|sub| sub.translate_unsatisfiable_pk(key))
-                    .collect::<Vec<_>>(),
-            ),
-            Policy::Or(subs) => Policy::Or(
-                subs.into_iter()
-                    .map(|(k, sub)| (k, sub.translate_unsatisfiable_pk(key)))
-                    .collect::<Vec<_>>(),
-            ),
-            Policy::Threshold(k, subs) => Policy::Threshold(
-                k,
-                subs.into_iter()
-                    .map(|sub| sub.translate_unsatisfiable_pk(key))
-                    .collect::<Vec<_>>(),
-            ),
-            x => x,
+        use Policy::*;
+
+        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_policy = match data.node.as_ref() {
+                Policy::Key(ref k) if k.clone() == *key => Some(Policy::Unsatisfiable),
+                Threshold(k, ref subs) => {
+                    Some(Threshold(*k, (0..subs.len()).map(child_n).collect()))
+                }
+                And(ref subs) => Some(And((0..subs.len()).map(child_n).collect())),
+                Or(ref subs) => Some(Or((0..subs.len()).map(|i| (i, child_n(i))).collect())),
+                _ => None,
+            };
+            match new_policy {
+                Some(new_policy) => translated.push(Arc::new(new_policy)),
+                None => translated.push(Arc::clone(&data.node)),
+            }
         }
+
+        Arc::try_unwrap(translated.pop().unwrap()).unwrap()
     }
 
     /// Gets all keys in the policy.
     pub fn keys(&self) -> Vec<&Pk> {
-        match *self {
-            Policy::Key(ref pk) => vec![pk],
-            Policy::Threshold(_k, ref subs) => {
-                subs.iter().flat_map(|sub| sub.keys()).collect::<Vec<_>>()
+        let mut keys = vec![];
+        for policy in self.pre_order_iter() {
+            match policy {
+                Policy::Key(ref pk) => keys.push(pk),
+                _ => {}
             }
-            Policy::And(ref subs) => subs.iter().flat_map(|sub| sub.keys()).collect::<Vec<_>>(),
-            Policy::Or(ref subs) => subs
-                .iter()
-                .flat_map(|(ref _k, ref sub)| sub.keys())
-                .collect::<Vec<_>>(),
-            // map all hashes and time
-            _ => vec![],
         }
+        keys
     }
 
     /// Gets the number of [TapLeaf](`TapTree::Leaf`)s considering exhaustive root-level [`Policy::Or`]
@@ -1144,7 +1125,7 @@ impl_block_str!(
                 for arg in &top.args {
                     subs.push(Policy::from_tree(arg)?);
                 }
-                Ok(Policy::And(subs))
+                Ok(Policy::And(subs.into_iter().map(|p| Arc::new(p)).collect()))
             }
             ("or", _) => {
                 if top.args.len() != 2 {
@@ -1154,7 +1135,7 @@ impl_block_str!(
                 for arg in &top.args {
                     subs.push(Policy::from_tree_prob(arg, true)?);
                 }
-                Ok(Policy::Or(subs))
+                Ok(Policy::Or(subs.into_iter().map(|(probability, policy)| (probability, Arc::new(policy))).collect()))
             }
             ("thresh", nsubs) => {
                 if top.args.is_empty() || !top.args[0].args.is_empty() {
@@ -1170,7 +1151,7 @@ impl_block_str!(
                 for arg in &top.args[1..] {
                     subs.push(Policy::from_tree(arg)?);
                 }
-                Ok(Policy::Threshold(thresh as usize, subs))
+                Ok(Policy::Threshold(thresh as usize, subs.into_iter().map(|p| Arc::new(p)).collect()))
             }
             _ => Err(errstr(top.name)),
         }

src/policy/mod.rs

@@ -198,7 +198,8 @@ impl<Pk: MiniscriptKey> Liftable<Pk> for Concrete<Pk> {
             Concrete::Ripemd160(ref h) => Semantic::Ripemd160(h.clone()),
             Concrete::Hash160(ref h) => Semantic::Hash160(h.clone()),
             Concrete::And(ref subs) => {
-                let semantic_subs: Result<_, Error> = subs.iter().map(Liftable::lift).collect();
+                let semantic_subs: Result<_, Error> =
+                    subs.iter().map(|p| Liftable::lift(p.as_ref())).collect();
                 Semantic::Threshold(2, semantic_subs?)
             }
             Concrete::Or(ref subs) => {
@@ -207,7 +208,8 @@ impl<Pk: MiniscriptKey> Liftable<Pk> for Concrete<Pk> {
                 Semantic::Threshold(1, semantic_subs?)
             }
             Concrete::Threshold(k, ref subs) => {
-                let semantic_subs: Result<_, Error> = subs.iter().map(Liftable::lift).collect();
+                let semantic_subs: Result<_, Error> =
+                    subs.iter().map(|p| Liftable::lift(p.as_ref())).collect();
                 Semantic::Threshold(k, semantic_subs?)
             }
         }