bit_machine: add PruneTracker trait and RedeemNode::prune_with_tracker

This eliminates the private `exec_prune` method which was weirdly-named
and no longer makes sense now that we're using a generic tracker rather
than a SetTracker.

Author: apoelstra

src/bit_machine/mod.rs

@@ -22,7 +22,7 @@ use crate::{Cmr, FailEntropy, Value};
 use frame::Frame;
 
 pub use self::limits::LimitError;
-pub use self::tracker::{ExecTracker, NoTracker, SetTracker};
+pub use self::tracker::{ExecTracker, NoTracker, PruneTracker, SetTracker};
 
 /// An iterator over the contents of a read or write frame which yields bits.
 pub type FrameIter<'a> = crate::BitIter<core::iter::Copied<core::slice::Iter<'a, u8>>>;
@@ -226,26 +226,6 @@ impl BitMachine {
         self.exec_with_tracker(program, env, &mut NoTracker)
     }
 
-    /// Execute the given `program` on the Bit Machine and track executed case branches.
-    ///
-    /// If the program runs successfully, then two sets of IHRs are returned:
-    ///
-    /// 1) The IHRs of case nodes whose _left_ branch was executed.
-    /// 2) The IHRs of case nodes whose _right_ branch was executed.
-    ///
-    /// ## Precondition
-    ///
-    /// The Bit Machine is constructed via [`Self::for_program()`] to ensure enough space.
-    pub(crate) fn exec_prune<J: Jet>(
-        &mut self,
-        program: &RedeemNode<J>,
-        env: &J::Environment,
-    ) -> Result<SetTracker, ExecutionError> {
-        let mut tracker = SetTracker::default();
-        self.exec_with_tracker(program, env, &mut tracker)?;
-        Ok(tracker)
-    }
-
     /// Execute the given `program` on the Bit Machine, using the given environment and tracker.
     ///
     ///  ## Precondition

src/bit_machine/tracker.rs

@@ -44,29 +44,21 @@ pub trait ExecTracker<J: Jet> {
     fn is_track_debug_enabled(&self) -> bool;
 }
 
+pub trait PruneTracker<J: Jet>: ExecTracker<J> {
+    /// Returns true if the left branch of the of the `Case` node with the IHR `ihr` was taken.
+    fn contains_left(&self, ihr: Ihr) -> bool;
+
+    /// Returns true if the right branch of the of the `Case` node with the IHR `ihr` was taken.
+    fn contains_right(&self, ihr: Ihr) -> bool;
+}
+
 /// Tracker of executed left and right branches for each case node.
 #[derive(Clone, Debug, Default)]
 pub struct SetTracker {
     left: HashSet<Ihr>,
     right: HashSet<Ihr>,
 }
 
-impl SetTracker {
-    /// Access the set of IHRs of case nodes whose left branch was executed.
-    pub fn left(&self) -> &HashSet<Ihr> {
-        &self.left
-    }
-
-    /// Access the set of IHRs of case nodes whose right branch was executed.
-    pub fn right(&self) -> &HashSet<Ihr> {
-        &self.right
-    }
-}
-
-/// Tracker that does not do anything (noop).
-#[derive(Copy, Clone, Debug)]
-pub struct NoTracker;
-
 impl<J: Jet> ExecTracker<J> for SetTracker {
     fn track_left(&mut self, ihr: Ihr) {
         self.left.insert(ihr);
@@ -85,6 +77,20 @@ impl<J: Jet> ExecTracker<J> for SetTracker {
     }
 }
 
+impl<J: Jet> PruneTracker<J> for SetTracker {
+    fn contains_left(&self, ihr: Ihr) -> bool {
+        self.left.contains(&ihr)
+    }
+
+    fn contains_right(&self, ihr: Ihr) -> bool {
+        self.right.contains(&ihr)
+    }
+}
+
+/// Tracker that does not do anything (noop).
+#[derive(Copy, Clone, Debug)]
+pub struct NoTracker;
+
 impl<J: Jet> ExecTracker<J> for NoTracker {
     fn track_left(&mut self, _: Ihr) {}
 

src/node/redeem.rs

@@ -1,7 +1,7 @@
 // SPDX-License-Identifier: CC0-1.0
 
 use crate::analysis::NodeBounds;
-use crate::bit_machine::{ExecutionError, SetTracker};
+use crate::bit_machine::{ExecutionError, PruneTracker, SetTracker};
 use crate::dag::{DagLike, InternalSharing, MaxSharing, PostOrderIterItem};
 use crate::jet::Jet;
 use crate::types::{self, arrow::FinalArrow};
@@ -290,13 +290,25 @@ impl<J: Jet> RedeemNode<J> {
     /// In this case, the witness data needs to be revised.
     /// The other pruning steps (2 & 3) never fail.
     pub fn prune(&self, env: &J::Environment) -> Result<Arc<RedeemNode<J>>, ExecutionError> {
-        struct Pruner<'brand, J> {
+        self.prune_with_tracker(env, &mut SetTracker::default())
+    }
+
+    pub fn prune_with_tracker<T: PruneTracker<J>>(
+        &self,
+        env: &J::Environment,
+        tracker: &mut T,
+    ) -> Result<Arc<RedeemNode<J>>, ExecutionError> {
+        struct Pruner<'brand, 't, J, T> {
             inference_context: types::Context<'brand>,
-            tracker: SetTracker,
+            tracker: &'t mut T,
             phantom: PhantomData<J>,
         }
 
-        impl<'brand, J: Jet> Converter<Redeem<J>, Construct<'brand, J>> for Pruner<'brand, J> {
+        impl<'brand, 't, J, T> Converter<Redeem<J>, Construct<'brand, J>> for Pruner<'brand, 't, J, T>
+        where
+            J: Jet,
+            T: PruneTracker<J>,
+        {
             type Error = std::convert::Infallible;
 
             fn convert_witness(
@@ -332,8 +344,8 @@ impl<J: Jet> RedeemNode<J> {
                 // but the Converter trait gives us access to the unpruned node (`data`).
                 // The Bit Machine tracked (un)used case branches based on the unpruned IHR.
                 match (
-                    self.tracker.left().contains(&data.node.ihr()),
-                    self.tracker.right().contains(&data.node.ihr()),
+                    self.tracker.contains_left(data.node.ihr()),
+                    self.tracker.contains_right(data.node.ihr()),
                 ) {
                     (true, true) => Ok(Hide::Neither),
                     (false, true) => Ok(Hide::Left),
@@ -418,7 +430,7 @@ impl<J: Jet> RedeemNode<J> {
         // 1) Run the Bit Machine and mark (un)used branches.
         // This is the only fallible step in the pruning process.
         let mut mac = BitMachine::for_program(self)?;
-        let tracker = mac.exec_prune(self, env)?;
+        mac.exec_with_tracker(self, env, tracker)?;
 
         // 2) Prune out unused case branches.
         // Because the types of the pruned program may change,