Improve time complexity of equality relations

src/librustc/middle/infer/bivariate.rs

@@ -77,8 +77,8 @@ impl<'a, 'tcx> TypeRelation<'a, 'tcx> for Bivariate<'a, 'tcx> {
         if a == b { return Ok(a); }
 
         let infcx = self.fields.infcx;
-        let a = infcx.type_variables.borrow().replace_if_possible(a);
-        let b = infcx.type_variables.borrow().replace_if_possible(b);
+        let a = infcx.type_variables.borrow_mut().replace_if_possible(a);
+        let b = infcx.type_variables.borrow_mut().replace_if_possible(b);
         match (&a.sty, &b.sty) {
             (&ty::TyInfer(TyVar(a_id)), &ty::TyInfer(TyVar(b_id))) => {
                 infcx.type_variables.borrow_mut().relate_vars(a_id, BiTo, b_id);

src/librustc/middle/infer/combine.rs

@@ -210,6 +210,12 @@ impl<'a, 'tcx> CombineFields<'a, 'tcx> {
                 None => break,
                 Some(e) => e,
             };
+            // Get the actual variable that b_vid has been inferred to
+            let (b_vid, b_ty) = {
+                let mut variables = self.infcx.type_variables.borrow_mut();
+                let b_vid = variables.root_var(b_vid);
+                (b_vid, variables.probe_root(b_vid))
+            };
 
             debug!("instantiate(a_ty={:?} dir={:?} b_vid={:?})",
                    a_ty,
@@ -219,7 +225,6 @@ impl<'a, 'tcx> CombineFields<'a, 'tcx> {
             // Check whether `vid` has been instantiated yet.  If not,
             // make a generalized form of `ty` and instantiate with
             // that.
-            let b_ty = self.infcx.type_variables.borrow().probe(b_vid);
             let b_ty = match b_ty {
                 Some(t) => t, // ...already instantiated.
                 None => {     // ...not yet instantiated:
@@ -307,12 +312,17 @@ impl<'cx, 'tcx> ty::fold::TypeFolder<'tcx> for Generalizer<'cx, 'tcx> {
         //  where `$1` has already been instantiated with `Box<$0>`)
         match t.sty {
             ty::TyInfer(ty::TyVar(vid)) => {
+                let mut variables = self.infcx.type_variables.borrow_mut();
+                let vid = variables.root_var(vid);
                 if vid == self.for_vid {
                     self.cycle_detected = true;
                     self.tcx().types.err
                 } else {
-                    match self.infcx.type_variables.borrow().probe(vid) {
-                        Some(u) => self.fold_ty(u),
+                    match variables.probe_root(vid) {
+                        Some(u) => {
+                            drop(variables);
+                            self.fold_ty(u)
+                        }
                         None => t,
                     }
                 }

src/librustc/middle/infer/equate.rs

@@ -50,8 +50,8 @@ impl<'a, 'tcx> TypeRelation<'a,'tcx> for Equate<'a, 'tcx> {
         if a == b { return Ok(a); }
 
         let infcx = self.fields.infcx;
-        let a = infcx.type_variables.borrow().replace_if_possible(a);
-        let b = infcx.type_variables.borrow().replace_if_possible(b);
+        let a = infcx.type_variables.borrow_mut().replace_if_possible(a);
+        let b = infcx.type_variables.borrow_mut().replace_if_possible(b);
         match (&a.sty, &b.sty) {
             (&ty::TyInfer(TyVar(a_id)), &ty::TyInfer(TyVar(b_id))) => {
                 infcx.type_variables.borrow_mut().relate_vars(a_id, EqTo, b_id);

src/librustc/middle/infer/freshen.rs

@@ -111,8 +111,9 @@ impl<'a, 'tcx> TypeFolder<'tcx> for TypeFreshener<'a, 'tcx> {
 
         match t.sty {
             ty::TyInfer(ty::TyVar(v)) => {
+                let opt_ty = self.infcx.type_variables.borrow_mut().probe(v);
                 self.freshen(
-                    self.infcx.type_variables.borrow().probe(v),
+                    opt_ty,
                     ty::TyVar(v),
                     ty::FreshTy)
             }

src/librustc/middle/infer/higher_ranked/mod.rs

@@ -434,7 +434,7 @@ impl<'a,'tcx> InferCtxtExt for InferCtxt<'a,'tcx> {
             self.region_vars.vars_created_since_snapshot(&snapshot.region_vars_snapshot);
 
         let escaping_types =
-            self.type_variables.borrow().types_escaping_snapshot(&snapshot.type_snapshot);
+            self.type_variables.borrow_mut().types_escaping_snapshot(&snapshot.type_snapshot);
 
         let mut escaping_region_vars = FnvHashSet();
         for ty in &escaping_types {

src/librustc/middle/infer/lattice.rs

@@ -60,8 +60,8 @@ pub fn super_lattice_tys<'a,'tcx,L:LatticeDir<'a,'tcx>>(this: &mut L,
     }
 
     let infcx = this.infcx();
-    let a = infcx.type_variables.borrow().replace_if_possible(a);
-    let b = infcx.type_variables.borrow().replace_if_possible(b);
+    let a = infcx.type_variables.borrow_mut().replace_if_possible(a);
+    let b = infcx.type_variables.borrow_mut().replace_if_possible(b);
     match (&a.sty, &b.sty) {
         (&ty::TyInfer(TyVar(..)), &ty::TyInfer(TyVar(..)))
             if infcx.type_var_diverges(a) && infcx.type_var_diverges(b) => {

src/librustc/middle/infer/mod.rs

@@ -637,7 +637,7 @@ impl<'a, 'tcx> InferCtxt<'a, 'tcx> {
         let mut variables = Vec::new();
 
         let unbound_ty_vars = self.type_variables
-                                  .borrow()
+                                  .borrow_mut()
                                   .unsolved_variables()
                                   .into_iter()
                                   .map(|t| self.tcx.mk_var(t));
@@ -1162,7 +1162,7 @@ impl<'a, 'tcx> InferCtxt<'a, 'tcx> {
                 // structurally), and we prevent cycles in any case,
                 // so this recursion should always be of very limited
                 // depth.
-                self.type_variables.borrow()
+                self.type_variables.borrow_mut()
                     .probe(v)
                     .map(|t| self.shallow_resolve(t))
                     .unwrap_or(typ)

src/librustc/middle/infer/sub.rs

@@ -65,8 +65,8 @@ impl<'a, 'tcx> TypeRelation<'a, 'tcx> for Sub<'a, 'tcx> {
         if a == b { return Ok(a); }
 
         let infcx = self.fields.infcx;
-        let a = infcx.type_variables.borrow().replace_if_possible(a);
-        let b = infcx.type_variables.borrow().replace_if_possible(b);
+        let a = infcx.type_variables.borrow_mut().replace_if_possible(a);
+        let b = infcx.type_variables.borrow_mut().replace_if_possible(b);
         match (&a.sty, &b.sty) {
             (&ty::TyInfer(TyVar(a_id)), &ty::TyInfer(TyVar(b_id))) => {
                 infcx.type_variables

src/librustc/middle/infer/type_variable.rs

@@ -20,9 +20,11 @@ use std::marker::PhantomData;
 use std::mem;
 use std::u32;
 use rustc_data_structures::snapshot_vec as sv;
+use rustc_data_structures::unify as ut;
 
 pub struct TypeVariableTable<'tcx> {
     values: sv::SnapshotVec<Delegate<'tcx>>,
+    eq_relations: ut::UnificationTable<ty::TyVid>,
 }
 
 struct TypeVariableData<'tcx> {
@@ -50,20 +52,22 @@ pub struct Default<'tcx> {
 }
 
 pub struct Snapshot {
-    snapshot: sv::Snapshot
+    snapshot: sv::Snapshot,
+    eq_snapshot: ut::Snapshot<ty::TyVid>,
 }
 
 enum UndoEntry<'tcx> {
     // The type of the var was specified.
     SpecifyVar(ty::TyVid, Vec<Relation>, Option<Default<'tcx>>),
     Relate(ty::TyVid, ty::TyVid),
+    RelateRange(ty::TyVid, usize),
 }
 
 struct Delegate<'tcx>(PhantomData<&'tcx ()>);
 
 type Relation = (RelationDir, ty::TyVid);
 
-#[derive(Copy, Clone, PartialEq, Debug)]
+#[derive(Copy, Clone, Eq, PartialEq, Hash, Debug)]
 pub enum RelationDir {
     SubtypeOf, SupertypeOf, EqTo, BiTo
 }
@@ -81,7 +85,10 @@ impl RelationDir {
 
 impl<'tcx> TypeVariableTable<'tcx> {
     pub fn new() -> TypeVariableTable<'tcx> {
-        TypeVariableTable { values: sv::SnapshotVec::new() }
+        TypeVariableTable {
+            values: sv::SnapshotVec::new(),
+            eq_relations: ut::UnificationTable::new(),
+        }
     }
 
     fn relations<'a>(&'a mut self, a: ty::TyVid) -> &'a mut Vec<Relation> {
@@ -103,22 +110,48 @@ impl<'tcx> TypeVariableTable<'tcx> {
     ///
     /// Precondition: neither `a` nor `b` are known.
     pub fn relate_vars(&mut self, a: ty::TyVid, dir: RelationDir, b: ty::TyVid) {
+        let a = self.root_var(a);
+        let b = self.root_var(b);
         if a != b {
-            self.relations(a).push((dir, b));
-            self.relations(b).push((dir.opposite(), a));
-            self.values.record(Relate(a, b));
+            if dir == EqTo {
+                // a and b must be equal which we mark in the unification table
+                let root = self.eq_relations.union(a, b);
+                // In addition to being equal, all relations from the variable which is no longer
+                // the root must be added to the root so they are not forgotten as the other
+                // variable should no longer be referenced (other than to get the root)
+                let other = if a == root { b } else { a };
+                let count = {
+                    let (relations, root_relations) = if other.index < root.index {
+                        let (pre, post) = self.values.split_at_mut(root.index as usize);
+                        (relations(&mut pre[other.index as usize]), relations(&mut post[0]))
+                    } else {
+                        let (pre, post) = self.values.split_at_mut(other.index as usize);
+                        (relations(&mut post[0]), relations(&mut pre[root.index as usize]))
+                    };
+                    root_relations.extend_from_slice(relations);
+                    relations.len()
+                };
+                self.values.record(RelateRange(root, count));
+            } else {
+                self.relations(a).push((dir, b));
+                self.relations(b).push((dir.opposite(), a));
+                self.values.record(Relate(a, b));
+            }
         }
     }
 
     /// Instantiates `vid` with the type `ty` and then pushes an entry onto `stack` for each of the
     /// relations of `vid` to other variables. The relations will have the form `(ty, dir, vid1)`
     /// where `vid1` is some other variable id.
+    ///
+    /// Precondition: `vid` must be a root in the unification table
     pub fn instantiate_and_push(
         &mut self,
         vid: ty::TyVid,
         ty: Ty<'tcx>,
         stack: &mut Vec<(Ty<'tcx>, RelationDir, ty::TyVid)>)
     {
+        debug_assert!(self.root_var(vid) == vid);
         let old_value = {
             let value_ptr = &mut self.values.get_mut(vid.index as usize).value;
             mem::replace(value_ptr, Known(ty))
@@ -140,21 +173,33 @@ impl<'tcx> TypeVariableTable<'tcx> {
     pub fn new_var(&mut self,
                    diverging: bool,
                    default: Option<Default<'tcx>>) -> ty::TyVid {
+        self.eq_relations.new_key(());
         let index = self.values.push(TypeVariableData {
             value: Bounded { relations: vec![], default: default },
             diverging: diverging
         });
         ty::TyVid { index: index as u32 }
     }
 
-    pub fn probe(&self, vid: ty::TyVid) -> Option<Ty<'tcx>> {
+    pub fn root_var(&mut self, vid: ty::TyVid) -> ty::TyVid {
+        self.eq_relations.find(vid)
+    }
+
+    pub fn probe(&mut self, vid: ty::TyVid) -> Option<Ty<'tcx>> {
+        let vid = self.root_var(vid);
+        self.probe_root(vid)
+    }
+
+    /// Retrieves the type of `vid` given that it is currently a root in the unification table
+    pub fn probe_root(&mut self, vid: ty::TyVid) -> Option<Ty<'tcx>> {
+        debug_assert!(self.root_var(vid) == vid);
         match self.values.get(vid.index as usize).value {
             Bounded { .. } => None,
             Known(t) => Some(t)
         }
     }
 
-    pub fn replace_if_possible(&self, t: Ty<'tcx>) -> Ty<'tcx> {
+    pub fn replace_if_possible(&mut self, t: Ty<'tcx>) -> Ty<'tcx> {
         match t.sty {
             ty::TyInfer(ty::TyVar(v)) => {
                 match self.probe(v) {
@@ -167,18 +212,23 @@ impl<'tcx> TypeVariableTable<'tcx> {
     }
 
     pub fn snapshot(&mut self) -> Snapshot {
-        Snapshot { snapshot: self.values.start_snapshot() }
+        Snapshot {
+            snapshot: self.values.start_snapshot(),
+            eq_snapshot: self.eq_relations.snapshot(),
+        }
     }
 
     pub fn rollback_to(&mut self, s: Snapshot) {
         self.values.rollback_to(s.snapshot);
+        self.eq_relations.rollback_to(s.eq_snapshot);
     }
 
     pub fn commit(&mut self, s: Snapshot) {
         self.values.commit(s.snapshot);
+        self.eq_relations.commit(s.eq_snapshot);
     }
 
-    pub fn types_escaping_snapshot(&self, s: &Snapshot) -> Vec<Ty<'tcx>> {
+    pub fn types_escaping_snapshot(&mut self, s: &Snapshot) -> Vec<Ty<'tcx>> {
         /*!
          * Find the set of type variables that existed *before* `s`
          * but which have only been unified since `s` started, and
@@ -208,7 +258,10 @@ impl<'tcx> TypeVariableTable<'tcx> {
                     if vid.index < new_elem_threshold {
                         // quick check to see if this variable was
                         // created since the snapshot started or not.
-                        let escaping_type = self.probe(vid).unwrap();
+                        let escaping_type = match self.values.get(vid.index as usize).value {
+                            Bounded { .. } => unreachable!(),
+                            Known(ty) => ty,
+                        };
                         escaping_types.push(escaping_type);
                     }
                     debug!("SpecifyVar({:?}) new_elem_threshold={}", vid, new_elem_threshold);
@@ -221,13 +274,15 @@ impl<'tcx> TypeVariableTable<'tcx> {
         escaping_types
     }
 
-    pub fn unsolved_variables(&self) -> Vec<ty::TyVid> {
-        self.values
-            .iter()
-            .enumerate()
-            .filter_map(|(i, value)| match &value.value {
-                &TypeVariableValue::Known(_) => None,
-                &TypeVariableValue::Bounded { .. } => Some(ty::TyVid { index: i as u32 })
+    pub fn unsolved_variables(&mut self) -> Vec<ty::TyVid> {
+        (0..self.values.len())
+            .filter_map(|i| {
+                let vid = ty::TyVid { index: i as u32 };
+                if self.probe(vid).is_some() {
+                    None
+                } else {
+                    Some(vid)
+                }
             })
             .collect()
     }
@@ -250,6 +305,13 @@ impl<'tcx> sv::SnapshotVecDelegate for Delegate<'tcx> {
                 relations(&mut (*values)[a.index as usize]).pop();
                 relations(&mut (*values)[b.index as usize]).pop();
             }
+
+            RelateRange(i, n) => {
+                let relations = relations(&mut (*values)[i.index as usize]);
+                for _ in 0..n {
+                    relations.pop();
+                }
+            }
         }
     }
 }

src/librustc/middle/infer/unify_key.rs

@@ -73,3 +73,10 @@ impl<'tcx> ToType<'tcx> for ast::FloatTy {
         tcx.mk_mach_float(*self)
     }
 }
+
+impl UnifyKey for ty::TyVid {
+    type Value = ();
+    fn index(&self) -> u32 { self.index }
+    fn from_index(i: u32) -> ty::TyVid { ty::TyVid { index: i } }
+    fn tag(_: Option<ty::TyVid>) -> &'static str { "TyVid" }
+}