types: use GhostCell in union_bound rather than mutexes

This eliminates all the locking code from union_bound.rs, turning any
lifetime mistake into compiler errors rather than deadlocks :). It also
lets us simplify a fair bit of code, removing all the drops, all the
lock().unwrap()s, and one shallow_clone.

There are some borrow()s and borrow_mut()s, but if these compile then
they are no-ops, so the reader can basically ignore them.

The previous benchmark was: 70,972.00 ns/iter (+/- 662.39)
            The new one is: 69,947.54 ns/iter (+/- 286.68)

which is a ~1.5% improvement, if these numbers are to be believed.

Author: apoelstra

src/types/context.rs

@@ -22,7 +22,9 @@ use ghost_cell::GhostToken;
 
 use crate::dag::{Dag, DagLike};
 
-use super::{Bound, CompleteBound, Error, Final, Incomplete, Type, TypeInner, WithGhostToken};
+use super::{
+    Bound, CompleteBound, Error, Final, Incomplete, Type, TypeInner, UbElement, WithGhostToken,
+};
 
 // Copied from ghost_cell source. See
 //     https://arhan.sh/blog/the-generativity-pattern-in-rust/
@@ -78,7 +80,7 @@ impl<'brand> Context<'brand> {
     pub fn new(token: GhostToken<'brand>) -> Self {
         Context {
             inner: Arc::new(Mutex::new(WithGhostToken {
-                _token: token,
+                token,
                 inner: ContextInner { slab: vec![] },
             })),
         }
@@ -151,6 +153,15 @@ impl<'brand> Context<'brand> {
         lock.inner.slab[bound.index].shallow_clone()
     }
 
+    /// Accesses a bound through a union-bound element.
+    pub(super) fn get_root_ref(
+        &self,
+        bound: &UbElement<'brand, BoundRef<'brand>>,
+    ) -> BoundRef<'brand> {
+        let mut lock = self.lock();
+        bound.root(&mut lock.token)
+    }
+
     /// Reassigns a bound to a different bound.
     ///
     /// # Panics
@@ -177,10 +188,10 @@ impl<'brand> Context<'brand> {
         prod_r: &Type<'brand>,
         hint: &'static str,
     ) -> Result<(), Error> {
-        let existing_root = existing.inner.bound.root();
+        let mut lock = self.lock();
+        let existing_root = existing.inner.bound.root(&mut lock.token);
         let new_bound = Bound::Product(prod_l.inner.shallow_clone(), prod_r.inner.shallow_clone());
 
-        let mut lock = self.lock();
         lock.bind(existing_root, new_bound).map_err(|e| {
             let new_bound = lock.alloc_bound(e.new);
             drop(lock);
@@ -260,7 +271,9 @@ impl<'brand> DagLike for (&'_ Context<'brand>, BoundRef<'brand>) {
         match self.0.get(&self.1) {
             Bound::Free(..) | Bound::Complete(..) => Dag::Nullary,
             Bound::Sum(ref ty1, ref ty2) | Bound::Product(ref ty1, ref ty2) => {
-                Dag::Binary((self.0, ty1.bound.root()), (self.0, ty2.bound.root()))
+                let root1 = self.0.get_root_ref(&ty1.bound);
+                let root2 = self.0.get_root_ref(&ty2.bound);
+                Dag::Binary((self.0, root1), (self.0, root2))
             }
         }
     }
@@ -295,26 +308,29 @@ impl<'brand> ContextInner<'brand> {
         );
         self.slab[bound.index] = new;
     }
+}
 
+impl<'brand> WithGhostToken<'brand, ContextInner<'brand>> {
     /// It is a common situation that we are pairing two types, and in the
     /// case that they are both complete, we want to pair the complete types.
     ///
     /// This method deals with all the annoying/complicated member variable
     /// paths to get the actual complete data out.
     fn complete_pair_data(
-        &self,
+        &mut self,
         inn1: &TypeInner<'brand>,
         inn2: &TypeInner<'brand>,
     ) -> Option<(Arc<Final>, Arc<Final>)> {
-        let bound1 = &self.slab[inn1.bound.root().index];
-        let bound2 = &self.slab[inn2.bound.root().index];
+        let idx1 = inn1.bound.root(&mut self.token).index;
+        let idx2 = inn2.bound.root(&mut self.token).index;
+        let bound1 = &self.slab[idx1];
+        let bound2 = &self.slab[idx2];
         if let (Bound::Complete(ref data1), Bound::Complete(ref data2)) = (bound1, bound2) {
             Some((Arc::clone(data1), Arc::clone(data2)))
         } else {
             None
         }
     }
-
     /// Unify the type with another one.
     ///
     /// Fails if the bounds on the two types are incompatible
@@ -323,9 +339,11 @@ impl<'brand> ContextInner<'brand> {
         existing: &TypeInner<'brand>,
         other: &TypeInner<'brand>,
     ) -> Result<(), BindError<'brand>> {
-        existing.bound.unify(&other.bound, |x_bound, y_bound| {
-            self.bind(x_bound, self.slab[y_bound.index].shallow_clone())
-        })
+        existing
+            .bound
+            .unify(self, &other.bound, |self_, x_bound, y_bound| {
+                self_.bind(x_bound, self_.slab[y_bound.index].shallow_clone())
+            })
     }
 
     fn bind(
@@ -368,8 +386,8 @@ impl<'brand> ContextInner<'brand> {
                         Bound::Product(ref ty1, ref ty2),
                     )
                     | (CompleteBound::Sum(ref comp1, ref comp2), Bound::Sum(ref ty1, ref ty2)) => {
-                        let bound1 = ty1.bound.root();
-                        let bound2 = ty2.bound.root();
+                        let bound1 = ty1.bound.root(&mut self.token);
+                        let bound2 = ty2.bound.root(&mut self.token);
                         self.bind(bound1, Bound::Complete(Arc::clone(comp1)))?;
                         self.bind(bound2, Bound::Complete(Arc::clone(comp2)))
                     }

src/types/mod.rs

@@ -192,7 +192,7 @@ pub struct Type<'brand> {
 struct TypeInner<'brand> {
     /// A set of constraints, which maintained by the union-bound algorithm and
     /// is progressively tightened as type inference proceeds.
-    bound: UbElement<BoundRef<'brand>>,
+    bound: UbElement<'brand, BoundRef<'brand>>,
 }
 
 impl TypeInner<'_> {
@@ -258,7 +258,9 @@ impl<'brand> Type<'brand> {
 
     /// Accessor for the data of this type, if it is complete
     pub fn final_data(&self) -> Option<Arc<Final>> {
-        if let Bound::Complete(ref data) = self.ctx.get(&self.inner.bound.root()) {
+        let root = self.ctx.get_root_ref(&self.inner.bound);
+        let bound = self.ctx.get(&root);
+        if let Bound::Complete(ref data) = bound {
             Some(Arc::clone(data))
         } else {
             None
@@ -276,27 +278,27 @@ impl<'brand> Type<'brand> {
 
     /// Converts a type as-is to an `Incomplete` type for use in an error.
     pub fn to_incomplete(&self) -> Arc<Incomplete> {
-        let root = self.inner.bound.root();
+        let root = self.ctx.get_root_ref(&self.inner.bound);
         Incomplete::from_bound_ref(&self.ctx, root)
     }
 
     /// Attempts to finalize the type. Returns its TMR on success.
     pub fn finalize(&self) -> Result<Arc<Final>, Error> {
-        let root = self.inner.bound.root();
+        let root = self.ctx.get_root_ref(&self.inner.bound);
         let bound = self.ctx.get(&root);
         if let Bound::Complete(ref data) = bound {
             return Ok(Arc::clone(data));
         }
 
         // First, do occurs-check to ensure that we have no infinitely sized types.
-        if let Some(infinite_bound) = Incomplete::occurs_check(&self.ctx, root) {
+        if let Some(infinite_bound) = Incomplete::occurs_check(&self.ctx, root.shallow_clone()) {
             return Err(Error::OccursCheck { infinite_bound });
         }
 
         // Now that we know our types have finite size, we can safely use a
         // post-order iterator to finalize them.
         let mut finalized = vec![];
-        for data in (&self.ctx, self.inner.bound.root()).post_order_iter::<NoSharing>() {
+        for data in (&self.ctx, root).post_order_iter::<NoSharing>() {
             let bound_get = data.node.0.get(&data.node.1);
             let final_data = match bound_get {
                 Bound::Free(_) => Final::unit(),
@@ -339,9 +341,8 @@ const MAX_DISPLAY_LENGTH: usize = 10000;
 
 impl fmt::Debug for Type<'_> {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        for data in (&self.ctx, self.inner.bound.root())
-            .verbose_pre_order_iter::<NoSharing>(Some(MAX_DISPLAY_DEPTH))
-        {
+        let root = self.ctx.get_root_ref(&self.inner.bound);
+        for data in (&self.ctx, root).verbose_pre_order_iter::<NoSharing>(Some(MAX_DISPLAY_DEPTH)) {
             if data.index > MAX_DISPLAY_LENGTH {
                 write!(f, "... [truncated type after {} nodes]", MAX_DISPLAY_LENGTH)?;
                 return Ok(());
@@ -376,9 +377,8 @@ impl fmt::Debug for Type<'_> {
 
 impl fmt::Display for Type<'_> {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        for data in (&self.ctx, self.inner.bound.root())
-            .verbose_pre_order_iter::<NoSharing>(Some(MAX_DISPLAY_DEPTH))
-        {
+        let root = self.ctx.get_root_ref(&self.inner.bound);
+        for data in (&self.ctx, root).verbose_pre_order_iter::<NoSharing>(Some(MAX_DISPLAY_DEPTH)) {
             if data.index > MAX_DISPLAY_LENGTH {
                 write!(f, "... [truncated type after {} nodes]", MAX_DISPLAY_LENGTH)?;
                 return Ok(());