miniscript: eliminate recursion in for_each_key

This is a breaking change because we no longer implement `ForEachKey` on `Terminal`. But IMHO we should never have implemented that trait (or any trait, really..) directly on `Terminal`. The fact that we did was just an implementation detail of how we used to do a lot of iteration.
apoelstra · Jul 14, 2023 · ea70d8a · ea70d8a
1 parent 935aa83
commit ea70d8a
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 50 deletions.
diff --git a/src/miniscript/astelem.rs b/src/miniscript/astelem.rs
@@ -21,7 +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,
+    errstr, expression, script_num_size, AbsLockTime, Error, Miniscript, MiniscriptKey,
     Terminal, ToPublicKey, TranslateErr, TranslatePk, Translator,
 };
 
@@ -64,44 +64,6 @@ where
 }
 
 impl<Pk: MiniscriptKey, Ctx: ScriptContext> Terminal<Pk, Ctx> {
-    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<Q, CtxQ, T, E>(
         &self,
         t: &mut T,
@@ -250,12 +212,6 @@ impl<Pk: MiniscriptKey, Ctx: ScriptContext> Terminal<Pk, Ctx> {
     }
 }
 
-impl<Pk: MiniscriptKey, Ctx: ScriptContext> ForEachKey<Pk> for Terminal<Pk, Ctx> {
-    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, Ctx: ScriptContext> fmt::Debug for Terminal<Pk, Ctx> {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         f.write_str("[")?;

diff --git a/src/miniscript/mod.rs b/src/miniscript/mod.rs
@@ -23,6 +23,7 @@ use bitcoin::taproot::{LeafVersion, TapLeafHash};
 use self::analyzable::ExtParams;
 pub use self::context::{BareCtx, Legacy, Segwitv0, Tap};
 use crate::prelude::*;
+use crate::iter::TreeLike;
 use crate::TranslateErr;
 
 pub mod analyzable;
@@ -296,7 +297,15 @@ impl<Pk: MiniscriptKey, Ctx: ScriptContext> Miniscript<Pk, Ctx> {
 
 impl<Pk: MiniscriptKey, Ctx: ScriptContext> ForEachKey<Pk> for Miniscript<Pk, Ctx> {
     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
     }
 }
 
@@ -319,10 +328,6 @@ where
 }
 
 impl<Pk: MiniscriptKey, Ctx: ScriptContext> Miniscript<Pk, Ctx> {
-    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<Q, CtxQ, T, FuncError>(
         &self,
         t: &mut T,