Implement pruning

src/bit_machine/mod.rs

@@ -8,14 +8,15 @@
 
 mod frame;
 
+use std::collections::HashSet;
 use std::error;
 use std::fmt;
 use std::sync::Arc;
 
-use crate::analysis;
 use crate::jet::{Jet, JetFailed};
 use crate::node::{self, RedeemNode};
 use crate::types::Final;
+use crate::{analysis, Imr};
 use crate::{Cmr, FailEntropy, Value};
 use frame::Frame;
 
@@ -204,13 +205,44 @@ impl BitMachine {
         Ok(())
     }
 
-    /// Execute the given program on the Bit Machine, using the given environment.
+    /// Execute the given `program` on the Bit Machine, using the given environment.
     ///
-    /// Make sure the Bit Machine has enough space by constructing it via [`Self::for_program()`].
-    pub fn exec<J: Jet + std::fmt::Debug>(
+    ///  ## Precondition
+    ///
+    /// The Bit Machine is constructed via [`Self::for_program()`] to ensure enough space.
+    pub fn exec<J: Jet>(
         &mut self,
         program: &RedeemNode<J>,
         env: &J::Environment,
+    ) -> Result<Value, ExecutionError> {
+        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 IMRs are returned:
+    ///
+    /// 1) The IMRs of case nodes whose _left_ branch was executed.
+    /// 2) The IMRs 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)
+    }
+
+    fn exec_with_tracker<J: Jet, T: CaseTracker>(
+        &mut self,
+        program: &RedeemNode<J>,
+        env: &J::Environment,
+        tracker: &mut T,
     ) -> Result<Value, ExecutionError> {
         enum CallStack<'a, J: Jet> {
             Goto(&'a RedeemNode<J>),
@@ -316,12 +348,14 @@ impl BitMachine {
                             self.fwd(1 + a.pad_right(b));
                             call_stack.push(CallStack::Back(1 + a.pad_right(b)));
                             call_stack.push(CallStack::Goto(right));
+                            tracker.track_right(ip.imr());
                         }
                         (node::Inner::Case(left, _), false)
                         | (node::Inner::AssertL(left, _), false) => {
                             self.fwd(1 + a.pad_left(b));
                             call_stack.push(CallStack::Back(1 + a.pad_left(b)));
                             call_stack.push(CallStack::Goto(left));
+                            tracker.track_left(ip.imr());
                         }
                         (node::Inner::AssertL(_, r_cmr), true) => {
                             return Err(ExecutionError::ReachedPrunedBranch(*r_cmr))
@@ -472,6 +506,60 @@ impl BitMachine {
     }
 }
 
+/// A type that keeps track of which case branches were executed
+/// during the execution of the Bit Machine.
+///
+/// The trait is implemented for [`SetTracker`], which does the actual tracking,
+/// and it is implemented for [`NoTracker`], which is a dummy tracker that is
+/// optimized out by the compiler.
+///
+/// The trait enables us to turn tracking on or off depending on a generic parameter.
+trait CaseTracker {
+    /// Track the execution of the left branch of the case node with the given `imr`.
+    fn track_left(&mut self, imr: Imr);
+
+    /// Track the execution of the right branch of the case node with the given `imr`.
+    fn track_right(&mut self, imr: Imr);
+}
+
+/// Tracker of executed left and right branches for each case node.
+#[derive(Clone, Debug, Default)]
+pub(crate) struct SetTracker {
+    left: HashSet<Imr>,
+    right: HashSet<Imr>,
+}
+
+impl SetTracker {
+    /// Access the set of IMRs of case nodes whose left branch was executed.
+    pub fn left(&self) -> &HashSet<Imr> {
+        &self.left
+    }
+
+    /// Access the set of IMRs of case nodes whose right branch was executed.
+    pub fn right(&self) -> &HashSet<Imr> {
+        &self.right
+    }
+}
+
+#[derive(Copy, Clone, Debug)]
+struct NoTracker;
+
+impl CaseTracker for SetTracker {
+    fn track_left(&mut self, imr: Imr) {
+        self.left.insert(imr);
+    }
+
+    fn track_right(&mut self, imr: Imr) {
+        self.right.insert(imr);
+    }
+}
+
+impl CaseTracker for NoTracker {
+    fn track_left(&mut self, _: Imr) {}
+
+    fn track_right(&mut self, _: Imr) {}
+}
+
 /// Errors related to simplicity Execution
 #[derive(Debug)]
 pub enum ExecutionError {

src/human_encoding/mod.rs

@@ -236,7 +236,7 @@ mod tests {
             .expect("Failed to parse human encoding")
             .to_witness_node(witness)
             .expect("Forest is missing expected root")
-            .finalize()
+            .finalize_pruned(env)
             .expect("Failed to finalize");
         let mut mac = BitMachine::for_program(&program);
         mac.exec(&program, env).expect("Failed to run program");
@@ -252,7 +252,7 @@ mod tests {
             .expect("Failed to parse human encoding")
             .to_witness_node(witness)
             .expect("Forest is missing expected root")
-            .finalize()
+            .finalize_pruned(env)
         {
             Ok(program) => program,
             Err(error) => {
@@ -268,7 +268,7 @@ mod tests {
     }
 
     #[test]
-    fn filled_witness() {
+    fn executed_witness_with_value() {
         let s = "
             a := witness
             b := witness
@@ -293,7 +293,7 @@ mod tests {
     }
 
     #[test]
-    fn unfilled_witness() {
+    fn executed_witness_without_value() {
         let witness = HashMap::from([(Arc::from("wit1"), Value::u32(1337))]);
         assert_finalize_err::<Core>(
             "
@@ -305,40 +305,34 @@ mod tests {
             ",
             &witness,
             &(),
-            "unable to satisfy program",
+            "Jet failed during execution",
         );
     }
 
     #[test]
-    fn unfilled_witness_pruned() {
+    fn pruned_witness_without_value() {
         let s = "
-            wit1 := witness
-            wit2 := witness
-            main := comp (pair wit1 unit) case unit wit2
+            wit1 := witness : 1 -> 2
+            wit2 := witness : 1 -> 2^64
+            input := pair wit1 unit : 1 -> 2 * 1
+            process := case (drop injr unit) (drop comp wit2 jet_all_64) : 2 * 1 -> 2
+            main := comp input comp process jet_verify : 1 -> 1
         ";
         let wit2_is_pruned = HashMap::from([(Arc::from("wit1"), Value::u1(0))]);
         assert_finalize_ok::<Core>(s, &wit2_is_pruned, &());
 
         let wit2_is_missing = HashMap::from([(Arc::from("wit1"), Value::u1(1))]);
-        // FIXME The finalization should fail
-        // This doesn't happen because we don't run the program,
-        // so we cannot always determine which nodes must be pruned
-        assert_finalize_err::<Core>(
-            s,
-            &wit2_is_missing,
-            &(),
-            "Execution reached a pruned branch: a0fc8debd6796917c86b77aded82e6c61649889ae8f2ed65b57b41aa9d90e375"
-        );
+        assert_finalize_err::<Core>(s, &wit2_is_missing, &(), "Jet failed during execution");
 
         let wit2_is_present = HashMap::from([
             (Arc::from("wit1"), Value::u1(1)),
-            (Arc::from("wit2"), Value::unit()),
+            (Arc::from("wit2"), Value::u64(u64::MAX)),
         ]);
         assert_finalize_ok::<Core>(s, &wit2_is_present, &());
     }
 
     #[test]
-    fn filled_hole() {
+    fn executed_hole_with_value() {
         let empty = HashMap::new();
         assert_finalize_ok::<Core>(
             "
@@ -352,7 +346,7 @@ mod tests {
     }
 
     #[test]
-    fn unfilled_hole() {
+    fn executed_hole_without_value() {
         let empty = HashMap::new();
         assert_finalize_err::<Core>(
             "
@@ -361,21 +355,21 @@ mod tests {
             ",
             &empty,
             &(),
-            "unable to satisfy program",
+            "disconnect node had one child (redeem time); must have two",
         );
     }
 
     #[test]
     fn witness_name_override() {
         let s = "
-            wit1 := witness
-            wit2 := wit1
-            main := comp wit2 iden
+            wit1 := witness : 1 -> 2
+            wit2 := wit1 : 1 -> 2
+            main := comp wit2 jet_verify : 1 -> 1
         ";
-        let wit1_populated = HashMap::from([(Arc::from("wit1"), Value::unit())]);
-        assert_finalize_err::<Core>(s, &wit1_populated, &(), "unable to satisfy program");
+        let wit1_populated = HashMap::from([(Arc::from("wit1"), Value::u1(1))]);
+        assert_finalize_err::<Core>(s, &wit1_populated, &(), "Jet failed during execution");
 
-        let wit2_populated = HashMap::from([(Arc::from("wit2"), Value::unit())]);
+        let wit2_populated = HashMap::from([(Arc::from("wit2"), Value::u1(1))]);
         assert_finalize_ok::<Core>(s, &wit2_populated, &());
     }
 }

src/human_encoding/parse/mod.rs

@@ -592,7 +592,7 @@ mod tests {
 
                 let program = main
                     .to_witness_node(witness, &forest)
-                    .finalize()
+                    .finalize_unpruned()
                     .expect("finalize");
 
                 let mut mac = BitMachine::for_program(&program);

src/lib.rs

@@ -77,17 +77,17 @@ pub fn leaf_version() -> elements::taproot::LeafVersion {
 #[derive(Debug)]
 pub enum Error {
     /// Decoder error
-    Decode(crate::decode::Error),
+    Decode(decode::Error),
     /// A disconnect node was populated at commitment time
     DisconnectCommitTime,
     /// A disconnect node was *not* populated at redeem time
     DisconnectRedeemTime,
     /// Type-checking error
-    Type(crate::types::Error),
+    Type(types::Error),
+    // Execution error
+    Execution(bit_machine::ExecutionError),
     /// Witness iterator ended early
     NoMoreWitnesses,
-    /// Finalization failed; did not have enough witness data to satisfy program.
-    IncompleteFinalization,
     /// Tried to parse a jet but the name wasn't recognized
     InvalidJetName(String),
     /// Policy error
@@ -106,7 +106,7 @@ impl fmt::Display for Error {
                 f.write_str("disconnect node had one child (redeem time); must have two")
             }
             Error::Type(ref e) => fmt::Display::fmt(e, f),
-            Error::IncompleteFinalization => f.write_str("unable to satisfy program"),
+            Error::Execution(ref e) => fmt::Display::fmt(e, f),
             Error::InvalidJetName(s) => write!(f, "unknown jet `{}`", s),
             Error::NoMoreWitnesses => f.write_str("no more witness data available"),
             #[cfg(feature = "elements")]
@@ -122,8 +122,8 @@ impl std::error::Error for Error {
             Error::DisconnectCommitTime => None,
             Error::DisconnectRedeemTime => None,
             Error::Type(ref e) => Some(e),
+            Error::Execution(ref e) => Some(e),
             Error::NoMoreWitnesses => None,
-            Error::IncompleteFinalization => None,
             Error::InvalidJetName(..) => None,
             #[cfg(feature = "elements")]
             Error::Policy(ref e) => Some(e),

src/node/convert.rs

@@ -181,7 +181,7 @@ impl<W: Iterator<Item = Value>, J: Jet> Converter<Commit<J>, Redeem<J>> for Simp
         _: Option<&Arc<RedeemNode<J>>>,
         _: &NoDisconnect,
     ) -> Result<Arc<RedeemNode<J>>, Self::Error> {
-        Err(crate::Error::IncompleteFinalization)
+        Err(crate::Error::DisconnectRedeemTime)
     }
 
     fn convert_data(

src/node/display.rs

@@ -148,7 +148,7 @@ mod tests {
             .unwrap()
             .to_witness_node(&empty_witness)
             .unwrap()
-            .finalize()
+            .finalize_unpruned()
             .unwrap()
     }