Merge BlockstreamResearch#361: Runtime jet type check

24ce63a add the Context jet type check (ivanlele)

Pull request description:

  This PR introduces a runtime jet type check in the Context to prevent unintended behavior.


ACKs for top commit:
  apoelstra:
    ACK 24ce63a; successfully ran local tests


Tree-SHA512: d162294d9acbda40c8f28f9da5ff3f0841ecaaa48acd9104a4b9d920df79bc0ed577ea0f472084f278ad5b10de06ea481c44b50bc384e1027cfaf2b6a442a508

Author: apoelstra

src/types/arrow.rs

@@ -356,6 +356,8 @@ impl<'brand> DisconnectConstructible<'brand, Option<&Arrow<'brand>>> for Arrow<'
 
 impl<'brand> JetConstructible<'brand> for Arrow<'brand> {
     fn jet(inference_context: &Context<'brand>, jet: &dyn Jet) -> Self {
+        inference_context.check_jet(jet);
+
         Arrow {
             source: jet.source_ty().to_type(inference_context),
             target: jet.target_ty().to_type(inference_context),

src/types/context.rs

@@ -14,13 +14,15 @@
 //! the other.
 //!
 
+use std::any::TypeId;
 use std::fmt;
 use std::marker::PhantomData;
 use std::sync::{Arc, Mutex, MutexGuard};
 
 use ghost_cell::GhostToken;
 
 use crate::dag::{Dag, DagLike};
+use crate::jet::Jet;
 
 use super::{
     Bound, CompleteBound, Error, Final, Incomplete, Type, TypeInner, UbElement, WithGhostToken,
@@ -48,6 +50,8 @@ pub struct Context<'brand> {
 
 struct ContextInner<'brand> {
     slab: Vec<Bound<'brand>>,
+    /// Concrete jet type registered in this context, if any.
+    jet_type: Option<TypeId>,
 }
 
 impl fmt::Debug for Context<'_> {
@@ -81,7 +85,10 @@ impl<'brand> Context<'brand> {
         Context {
             inner: Arc::new(Mutex::new(WithGhostToken {
                 token,
-                inner: ContextInner { slab: vec![] },
+                inner: ContextInner {
+                    slab: vec![],
+                    jet_type: None,
+                },
             })),
         }
     }
@@ -147,6 +154,23 @@ impl<'brand> Context<'brand> {
         }
     }
 
+    /// Asserts that all jets in this context have the same concrete type.
+    ///
+    /// Records the jet's type on first call, panics on subsequent calls with
+    /// a different concrete type.
+    pub fn check_jet(&self, jet: &dyn Jet) {
+        let new_id = jet.as_any().type_id();
+        let mut lock = self.lock();
+
+        if let Some(existing_id) = lock.inner.jet_type {
+            assert!(existing_id == new_id, "mixed jet types in context");
+
+            return;
+        }
+
+        lock.inner.jet_type = Some(new_id);
+    }
+
     /// Accesses a bound.
     pub(super) fn get(&self, bound: &BoundRef<'brand>) -> Bound<'brand> {
         let lock = self.lock();

src/types/mod.rs

@@ -415,6 +415,8 @@ impl fmt::Display for Type<'_> {
 mod tests {
     use super::*;
 
+    use crate::jet::Core;
+    use crate::node::JetConstructible;
     use crate::node::{ConstructNode, CoreConstructible};
 
     #[test]
@@ -446,4 +448,24 @@ mod tests {
             let _ = format!("{:?}", case.arrow().source);
         });
     }
+
+    #[test]
+    fn check_jet_same_type_ok() {
+        Context::with_context(|ctx| {
+            let _ = Arc::<ConstructNode>::jet(&ctx, &Core::Add32);
+            let _ = Arc::<ConstructNode>::jet(&ctx, &Core::Subtract32);
+        });
+    }
+
+    #[cfg(feature = "elements")]
+    #[test]
+    #[should_panic(expected = "mixed jet types in context")]
+    fn check_jet_different_types_panics() {
+        use crate::jet::Elements;
+
+        Context::with_context(|ctx| {
+            let _ = Arc::<ConstructNode>::jet(&ctx, &Core::Add32);
+            let _ = Arc::<ConstructNode>::jet(&ctx, &Elements::Add32);
+        });
+    }
 }