Separate Jet env from Jet trait

Author: ivanlele

jets-bench/benches/elements/main.rs

@@ -5,13 +5,13 @@ use elements::confidential;
 use rand::rngs::ThreadRng;
 use simplicity::elements;
 use simplicity::jet::elements::ElementsEnv;
-use simplicity::jet::{Elements, Jet};
+use simplicity::jet::{Elements, ElementsTxEnv, Jet, JetEnvironment};
 use simplicity::types;
 use simplicity::types::Final;
 use simplicity::Value;
 use simplicity_bench::input::{
-    self, EqProduct, GenericProduct, InputSample, PrefixBit, Sha256Ctx, UniformBits,
-    DivMod12864Input,
+    self, DivMod12864Input, EqProduct, GenericProduct, InputSample, PrefixBit, Sha256Ctx,
+    UniformBits,
 };
 use simplicity_bench::{
     genesis_pegin, BenchSample, EnvSampling, InputSampling, JetBuffer, JetParams, SimplicityCtx8,
@@ -625,7 +625,7 @@ fn bench(c: &mut Criterion) {
                             let (src, dst) = buffer.write(&src_ty, params, &mut rng);
                             (dst, src, &env, buffer)
                         },
-                        |(mut dst, src, env, _buffer)| jet.c_jet_ptr()(&mut dst, src, env.c_tx_env()),
+                        |(mut dst, src, env, _buffer)| ElementsTxEnv::c_jet_ptr(&jet)(&mut dst, src, env.c_tx_env()),
                         BatchSize::SmallInput,
                     )
                 });
@@ -743,7 +743,7 @@ fn bench(c: &mut Criterion) {
                         let (src, dst) = buffer.write(&src_ty, params, &mut rng);
                         (dst, src, buffer)
                     },
-                    |(mut dst, src, _buffer)| jet.c_jet_ptr()(&mut dst, src, env.c_tx_env()),
+                    |(mut dst, src, _buffer)| ElementsTxEnv::c_jet_ptr(&jet)(&mut dst, src, env.c_tx_env()),
                     BatchSize::SmallInput,
                 )
             });
@@ -806,7 +806,7 @@ fn bench(c: &mut Criterion) {
                         let (src, dst) = buffer.write(&src_ty, params, &mut rng);
                         (dst, src, buffer)
                     },
-                    |(mut dst, src, _buffer)| jet.c_jet_ptr()(&mut dst, src, env.c_tx_env()),
+                    |(mut dst, src, _buffer)| ElementsTxEnv::c_jet_ptr(&jet)(&mut dst, src, env.c_tx_env()),
                     BatchSize::SmallInput,
                 )
             });
@@ -903,7 +903,7 @@ fn bench(c: &mut Criterion) {
                         let (src, dst) = buffer.write(&src_ty, params, &mut rng);
                         (dst, src, buffer)
                     },
-                    |(mut dst, src, _buffer)| jet.c_jet_ptr()(&mut dst, src, env.c_tx_env()),
+                    |(mut dst, src, _buffer)| ElementsTxEnv::c_jet_ptr(&jet)(&mut dst, src, env.c_tx_env()),
                     BatchSize::SmallInput,
                 )
             });

jets-bench/src/input.rs

@@ -8,7 +8,7 @@ use simplicity::ffi::c_jets::frame_ffi::c_writeBit;
 use simplicity::ffi::ffi::UWORD;
 use simplicity::ffi::CFrameItem;
 use simplicity::hashes::Hash;
-use simplicity::jet::Elements;
+use simplicity::jet::{Elements, ElementsTxEnv, JetEnvironment};
 use simplicity::types::{self, CompleteBound};
 use simplicity::Value;
 
@@ -296,7 +296,6 @@ impl FlatValue {
     fn call_jet(&self, jet: Elements, dest_bits: usize) -> Self {
         use core::{mem, ptr};
         use simplicity::ffi::c_jets::uword_width;
-        use simplicity::jet::Jet as _;
 
         assert!(dest_bits <= 8 * MAX_VALUE_BYTES);
         let mut ret = Self::zero_n_bits(dest_bits);
@@ -354,10 +353,11 @@ impl FlatValue {
 
             // We can assert this because in our sampling code jets should never
             // fail. In the benchmarking code they might.
-            assert!(jet.c_jet_ptr()(
+
+            assert!(ElementsTxEnv::c_jet_ptr(&jet)(
                 &mut dst_write_frame,
                 src_read_frame,
-                Elements::c_jet_env(&env)
+                env.c_tx_env()
             ));
             // The write frame winds up as an array of usizes with all bytes in
             // reverse order. (The bytes of the usizes are in reverse order due
@@ -1366,8 +1366,12 @@ impl InputSample for DivMod12864Input {
             for (bit1, bit2) in sample_1.bit_iter().zip(sample_2.bit_iter()) {
                 match (bit1, bit2) {
                     (false, false) | (true, true) => {} // both equal
-                    (true, false) => return FlatValue::product(&[sample_2, UniformBits.sample(0, 64), sample_1]),
-                    (false, true) => return FlatValue::product(&[sample_1, UniformBits.sample(0, 64), sample_2]),
+                    (true, false) => {
+                        return FlatValue::product(&[sample_2, UniformBits.sample(0, 64), sample_1])
+                    }
+                    (false, true) => {
+                        return FlatValue::product(&[sample_1, UniformBits.sample(0, 64), sample_2])
+                    }
                 }
             }
             unreachable!("if we get here, two uniform 63-bit samples were exactly equal")
@@ -1428,7 +1432,7 @@ mod tests {
 
         let (src, mut dst) = buffer.write(&src_ty, &params, &mut rand::thread_rng());
 
-        jet.c_jet_ptr()(&mut dst, src, env.c_tx_env());
+        ElementsTxEnv::c_jet_ptr(&jet)(&mut dst, src, env.c_tx_env());
     }
 
     #[test]

src/bit_machine/mod.rs

@@ -15,6 +15,7 @@ use std::fmt;
 use std::sync::Arc;
 
 use crate::analysis;
+use crate::jet::JetEnvironment;
 use crate::jet::{Jet, JetFailed};
 use crate::node::{self, RedeemNode};
 use crate::types::Final;
@@ -60,9 +61,9 @@ impl BitMachine {
     }
 
     #[cfg(test)]
-    pub fn test_exec<J: Jet>(
-        program: Arc<crate::node::ConstructNode<J>>,
-        env: &J::Environment,
+    pub fn test_exec<JE: JetEnvironment>(
+        program: Arc<crate::node::ConstructNode<JE::Jet>>,
+        env: &JE,
     ) -> Result<Value, ExecutionError> {
         use crate::node::SimpleFinalizer;
 
@@ -220,10 +221,10 @@ impl BitMachine {
     ///  ## Precondition
     ///
     /// The Bit Machine is constructed via [`Self::for_program()`] to ensure enough space.
-    pub fn exec<J: Jet>(
+    pub fn exec<JE: JetEnvironment>(
         &mut self,
-        program: &RedeemNode<J>,
-        env: &J::Environment,
+        program: &RedeemNode<JE::Jet>,
+        env: &JE,
     ) -> Result<Value, ExecutionError> {
         self.exec_with_tracker(program, env, &mut NoTracker)
     }
@@ -236,22 +237,22 @@ impl BitMachine {
     ///  ## Precondition
     ///
     /// The Bit Machine is constructed via [`Self::for_program()`] to ensure enough space.
-    pub fn exec_with_tracker<J: Jet, T: ExecTracker<J>>(
+    pub fn exec_with_tracker<JE: JetEnvironment, T: ExecTracker<JE::Jet>>(
         &mut self,
-        program: &RedeemNode<J>,
-        env: &J::Environment,
+        program: &RedeemNode<JE::Jet>,
+        env: &JE,
         tracker: &mut T,
     ) -> Result<Value, ExecutionError> {
-        enum CallStack<'a, J: Jet> {
-            Goto(&'a RedeemNode<J>),
+        enum CallStack<'a, JE: JetEnvironment> {
+            Goto(&'a RedeemNode<JE::Jet>),
             MoveWriteFrameToRead,
             DropReadFrame,
             CopyFwd(usize),
             Back(usize),
         }
 
         // Not used, but useful for debugging, so keep it around
-        impl<J: Jet> fmt::Debug for CallStack<'_, J> {
+        impl<JE: JetEnvironment> fmt::Debug for CallStack<'_, JE> {
             fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
                 match self {
                     CallStack::Goto(ins) => write!(f, "goto {}", ins.inner()),
@@ -268,7 +269,7 @@ impl BitMachine {
         }
 
         let mut ip = program;
-        let mut call_stack = vec![];
+        let mut call_stack: Vec<CallStack<'_, JE>> = vec![];
 
         let output_width = ip.arrow().target.bit_width();
         if output_width > 0 {
@@ -435,7 +436,7 @@ impl BitMachine {
         }
     }
 
-    fn exec_jet<J: Jet>(&mut self, jet: J, env: &J::Environment) -> Result<(), JetFailed> {
+    fn exec_jet<JE: JetEnvironment>(&mut self, jet: JE::Jet, env: &JE) -> Result<(), JetFailed> {
         use crate::ffi::c_jets::frame_ffi::{c_readBit, c_writeBit, CFrameItem};
         use crate::ffi::c_jets::uword_width;
         use crate::ffi::ffi::UWORD;
@@ -524,8 +525,8 @@ impl BitMachine {
         let (input_read_frame, _input_buffer) = unsafe { get_input_frame(self, input_width) };
         let (mut output_write_frame, output_buffer) = unsafe { get_output_frame(output_width) };
 
-        let jet_fn = jet.c_jet_ptr();
-        let c_env = J::c_jet_env(env);
+        let jet_fn = JE::c_jet_ptr(&jet);
+        let c_env = env.c_jet_env();
         let success = jet_fn(&mut output_write_frame, input_read_frame, c_env);
 
         if !success {

src/human_encoding/mod.rs

@@ -224,19 +224,19 @@ impl<J: Jet> Forest<J> {
 #[cfg(test)]
 mod tests {
     use crate::human_encoding::Forest;
-    use crate::jet::{Core, Jet};
+    use crate::jet::{CoreEnv, JetEnvironment};
     use crate::types;
     use crate::{BitMachine, Value};
     use std::collections::HashMap;
     use std::sync::Arc;
 
-    fn assert_finalize_ok<J: Jet>(
+    fn assert_finalize_ok<JE: JetEnvironment>(
         s: &str,
         witness: &HashMap<Arc<str>, Value>,
-        env: &J::Environment,
+        env: &JE,
     ) {
         types::Context::with_context(|ctx| {
-            let program = Forest::<J>::parse(s)
+            let program = Forest::<JE::Jet>::parse(s)
                 .expect("Failed to parse human encoding")
                 .to_witness_node(&ctx, witness)
                 .expect("Forest is missing expected root")
@@ -247,14 +247,14 @@ mod tests {
         });
     }
 
-    fn assert_finalize_err<J: Jet>(
+    fn assert_finalize_err<JE: JetEnvironment>(
         s: &str,
         witness: &HashMap<Arc<str>, Value>,
-        env: &J::Environment,
+        env: &JE,
         err_msg: &'static str,
     ) {
         types::Context::with_context(|ctx| {
-            let program = match Forest::<J>::parse(s)
+            let program = match Forest::<JE::Jet>::parse(s)
                 .expect("Failed to parse human encoding")
                 .to_witness_node(&ctx, witness)
                 .expect("Forest is missing expected root")
@@ -290,19 +290,19 @@ mod tests {
             (Arc::from("a"), Value::u8(0x00)),
             (Arc::from("b"), Value::u8(0x01)),
         ]);
-        assert_finalize_ok::<Core>(s, &a_less_than_b, &());
+        assert_finalize_ok::<CoreEnv>(s, &a_less_than_b, &());
 
         let b_greater_equal_a = HashMap::from([
             (Arc::from("a"), Value::u8(0x01)),
             (Arc::from("b"), Value::u8(0x01)),
         ]);
-        assert_finalize_err::<Core>(s, &b_greater_equal_a, &(), "Jet failed during execution");
+        assert_finalize_err::<CoreEnv>(s, &b_greater_equal_a, &(), "Jet failed during execution");
     }
 
     #[test]
     fn executed_witness_without_value() {
         let witness = HashMap::from([(Arc::from("wit1"), Value::u32(1337))]);
-        assert_finalize_err::<Core>(
+        assert_finalize_err::<CoreEnv>(
             "
                 wit1 := witness : 1 -> 2^32
                 wit2 := witness : 1 -> 2^32
@@ -326,22 +326,22 @@ mod tests {
             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, &());
+        assert_finalize_ok::<CoreEnv>(s, &wit2_is_pruned, &());
 
         let wit2_is_missing = HashMap::from([(Arc::from("wit1"), Value::u1(1))]);
-        assert_finalize_err::<Core>(s, &wit2_is_missing, &(), "Jet failed during execution");
+        assert_finalize_err::<CoreEnv>(s, &wit2_is_missing, &(), "Jet failed during execution");
 
         let wit2_is_present = HashMap::from([
             (Arc::from("wit1"), Value::u1(1)),
             (Arc::from("wit2"), Value::u64(u64::MAX)),
         ]);
-        assert_finalize_ok::<Core>(s, &wit2_is_present, &());
+        assert_finalize_ok::<CoreEnv>(s, &wit2_is_present, &());
     }
 
     #[test]
     fn executed_hole_with_value() {
         let empty = HashMap::new();
-        assert_finalize_ok::<Core>(
+        assert_finalize_ok::<CoreEnv>(
             "
                 id1 := iden : 2^256 * 1 -> 2^256 * 1
                 main := comp (disconnect id1 ?hole) unit
@@ -355,7 +355,7 @@ mod tests {
     #[test]
     fn executed_hole_without_value() {
         let empty = HashMap::new();
-        assert_finalize_err::<Core>(
+        assert_finalize_err::<CoreEnv>(
             "
                 wit1 := witness
                 main := comp wit1 comp disconnect iden ?dis2 unit
@@ -374,9 +374,9 @@ mod tests {
             main := comp wit2 jet_verify : 1 -> 1
         ";
         let wit1_populated = HashMap::from([(Arc::from("wit1"), Value::u1(1))]);
-        assert_finalize_err::<Core>(s, &wit1_populated, &(), "Jet failed during execution");
+        assert_finalize_err::<CoreEnv>(s, &wit1_populated, &(), "Jet failed during execution");
 
         let wit2_populated = HashMap::from([(Arc::from("wit2"), Value::u1(1))]);
-        assert_finalize_ok::<Core>(s, &wit2_populated, &());
+        assert_finalize_ok::<CoreEnv>(s, &wit2_populated, &());
     }
 }

src/human_encoding/parse/mod.rs

@@ -586,18 +586,18 @@ mod tests {
 
     use crate::dag::MaxSharing;
     use crate::human_encoding::Forest;
-    use crate::jet::{Core, Jet};
+    use crate::jet::{Core, CoreEnv, Jet, JetEnvironment};
     use crate::node::Inner;
     use crate::value::Word;
     use crate::{BitMachine, Value};
 
-    fn assert_cmr_witness<J: Jet>(
+    fn assert_cmr_witness<JE: JetEnvironment>(
         s: &str,
         cmr: &str,
         witness: &HashMap<Arc<str>, Value>,
-        env: &J::Environment,
+        env: &JE,
     ) {
-        match parse::<J>(s) {
+        match parse::<JE::Jet>(s) {
             Ok(forest) => {
                 assert_eq!(forest.len(), 1);
                 let main = &forest["main"];
@@ -673,7 +673,7 @@ mod tests {
     #[test]
     fn simple_program() {
         let empty = HashMap::new();
-        assert_cmr_witness::<Core>(
+        assert_cmr_witness::<CoreEnv>(
             "main := unit",
             "c40a10263f7436b4160acbef1c36fba4be4d95df181a968afeab5eac247adff7",
             &empty,
@@ -684,7 +684,7 @@ mod tests {
             (Arc::from("wit1"), Value::u32(0x00010203)),
             (Arc::from("wit2"), Value::u32(0x00010203)),
         ]);
-        assert_cmr_witness::<Core>(
+        assert_cmr_witness::<CoreEnv>(
             "
                 wit1 := witness : 1 -> 2^32
                 wit2 := witness : 1 -> 2^32
@@ -716,11 +716,11 @@ mod tests {
     #[cfg(feature = "elements")]
     fn bip340_program() {
         use crate::jet::elements::ElementsEnv;
-        use crate::jet::Elements;
+        use crate::jet::ElementsTxEnv;
 
         let empty = HashMap::new();
         let dummy = ElementsEnv::dummy();
-        assert_cmr_witness::<Elements>(
+        assert_cmr_witness::<ElementsTxEnv>(
             "main := unit",
             "c40a10263f7436b4160acbef1c36fba4be4d95df181a968afeab5eac247adff7",
             &empty,
@@ -737,7 +737,7 @@ mod tests {
         ];
 
         let signature = HashMap::from([(Arc::from("wit1"), Value::u512(sig))]);
-        assert_cmr_witness::<Elements>(
+        assert_cmr_witness::<ElementsTxEnv>(
             "
                 -- Witnesses
                 wit1 := witness : 1 -> 2^512