types: fix bug in "halving" variant of the union-bound algorithm

In the union-bound algorithm we track our unified type variables as
sets, which are represented by trees of pointers, where the roots of the
trees are used as set representatives for determining equality etc. When
unifying variables, we move one's tree into the tree of the other, and
then copy all its bounds to the root of the new unified tree.

Most operations within union-bound are given an arbitrary set element
and immediately locate its root before doing anything. We obtain better
asymptotics by using the "halving" variant of union-bound, wherein on
each root lookup, as we follow the path from the element to its root, we
update the links to skip every other one. Basically, if node x points to
its parent which points to a grandparent, we update x to point to the
grandparent directly.

*However* our existing code, rather than making x point to the
grandparent, would just copy the grandparent over the parent. I'm not
sure what I was thinking here -- I think I got tangled up trying to deal
with Rust's multiple-aliasing rules without deadlocking and lost track
of the actual algorithm, or maybe I misread Russell's C code which I was
mostly copying from -- but the result is both weird and incorrect. In
particular, if the grandparent is a root, then by copying its data into
parent, we duplicate the root, effectively un-unifying the variable at
some hard-to-predict point.

Once the un-unification happens, then future unifications will appear to
only get half-applied, the resulting type inference will be slightly
wrong, we will produce programs which fail type-checking or
sharing-checking after being encoded, even though they were successfully
constructed.

The fix, amusingly, actually reduces the amount of locking, which might
provide a slight speedup, and will certainly make it easier to convince
yourself that this code is deadlock-free. (Probably we should drop all
these mutexes and just use raw pointers and unsafe code, which would be
easier to follow; the mutexes serve no purpose other than to mollify the
Rust typechecker; we assure type-safety by locking the entire slab, and
locking individual set elements is neither necessary nor sufficient to
get a thread-safe type inference algorithm.)

Thanks to Russell O'Connor for identifying this bug, and some thanks to
Claude 3.7-sonnet for producing weird not-quite-sensible "fixes" that
would change observed behavior on our bad test cases, and pointed us
to roughly the correct location.

Author: apoelstra

src/types/union_bound.rs

@@ -93,15 +93,6 @@ impl<T> UbData<T> {
     }
 }
 
-impl<T: PointerLike> UbData<T> {
-    fn shallow_clone(&self) -> Self {
-        match self {
-            UbData::Root(x) => UbData::Root(x.shallow_clone()),
-            UbData::EqualTo(eq) => UbData::EqualTo(eq.shallow_clone()),
-        }
-    }
-}
-
 impl<T> UbElement<T> {
     /// Turns an existing piece of data into a singleton union-bound set.
     pub fn new(data: T) -> Self {
@@ -148,21 +139,19 @@ impl<T: PointerLike> UbElement<T> {
                 }
             };
 
-            let mut parent_inner_lock = parent.inner.lock().unwrap();
+            let parent_inner_lock = parent.inner.lock().unwrap();
             // If the parent has a parent, remove the intermediate link. (This is
             // the "halving" variant of union-bound.)
             match parent_inner_lock.data {
                 // If there is no grandparent, return the parent.
                 UbData::Root(..) => return parent.shallow_clone(),
                 UbData::EqualTo(ref grandparent) => {
                     let grandparent = grandparent.shallow_clone();
-                    // Okay to lock both grandparent and parent at once because the
-                    // algorithm guarantees there are no cycles.
-                    let gp_inner_lock = grandparent.inner.lock().unwrap();
-                    parent_inner_lock.data = gp_inner_lock.data.shallow_clone();
-                    drop(gp_inner_lock);
-                    drop(parent_inner_lock);
-                    x = grandparent.shallow_clone();
+                    // The previous `x_inner_lock` was dropped above, so this will not deadlock.
+                    let mut x_inner_lock = x.inner.lock().unwrap();
+                    x_inner_lock.data = UbData::EqualTo(grandparent.shallow_clone());
+                    drop(x_inner_lock);
+                    x = grandparent;
                 }
             }
         }