count bound vars

2 files changed, 92 insertions, 6 deletions

diff --git a/compiler/rustc_middle/src/ty/fold.rs b/compiler/rustc_middle/src/ty/fold.rs
index a6a1d1f73bb..15ce0ad3970 100644
--- a/compiler/rustc_middle/src/ty/fold.rs
+++ b/compiler/rustc_middle/src/ty/fold.rs
@@ -888,6 +888,54 @@ impl<'tcx> TypeVisitor<'tcx> for HasEscapingVarsVisitor {
 #[derive(Debug, PartialEq, Eq, Copy, Clone)]
 struct FoundFlags;
 
+crate struct CountBoundVars {
+    crate outer_index: ty::DebruijnIndex,
+    crate bound_tys: FxHashSet<ty::BoundTy>,
+    crate bound_regions: FxHashSet<ty::BoundRegion>,
+    crate bound_consts: FxHashSet<ty::BoundVar>,
+}
+
+impl<'tcx> TypeVisitor<'tcx> for CountBoundVars {
+    type BreakTy = ();
+
+    fn visit_binder<T: TypeFoldable<'tcx>>(&mut self, t: &Binder<T>) -> ControlFlow<Self::BreakTy> {
+        self.outer_index.shift_in(1);
+        let result = t.super_visit_with(self);
+        self.outer_index.shift_out(1);
+        result
+    }
+
+    fn visit_ty(&mut self, t: Ty<'tcx>) -> ControlFlow<Self::BreakTy> {
+        match t.kind {
+            ty::Bound(debruijn, ty) if debruijn == self.outer_index => {
+                self.bound_tys.insert(ty);
+                ControlFlow::CONTINUE
+            }
+            _ => t.super_visit_with(self),
+        }
+    }
+
+    fn visit_region(&mut self, r: ty::Region<'tcx>) -> ControlFlow<Self::BreakTy> {
+        match r {
+            ty::ReLateBound(debruijn, re) if *debruijn == self.outer_index => {
+                self.bound_regions.insert(*re);
+                ControlFlow::CONTINUE
+            }
+            _ => r.super_visit_with(self),
+        }
+    }
+
+    fn visit_const(&mut self, ct: &'tcx ty::Const<'tcx>) -> ControlFlow<Self::BreakTy> {
+        match ct.val {
+            ty::ConstKind::Bound(debruijn, c) if debruijn == self.outer_index => {
+                self.bound_consts.insert(c);
+                ControlFlow::CONTINUE
+            }
+            _ => ct.super_visit_with(self),
+        }
+    }
+}
+
 // FIXME: Optimize for checking for infer flags
 struct HasTypeFlagsVisitor {
     flags: ty::TypeFlags,
diff --git a/compiler/rustc_middle/src/ty/sty.rs b/compiler/rustc_middle/src/ty/sty.rs
index c7896d2b1e8..7c7afc0b296 100644
--- a/compiler/rustc_middle/src/ty/sty.rs
+++ b/compiler/rustc_middle/src/ty/sty.rs
@@ -959,24 +959,58 @@ impl<'tcx> PolyExistentialTraitRef<'tcx> {
 #[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Debug)]
 pub struct Binder<T>(T, u32);
 
-impl<T> Binder<T> {
+impl<'tcx, T> Binder<T>
+where
+    T: TypeFoldable<'tcx>,
+{
     /// Wraps `value` in a binder, asserting that `value` does not
     /// contain any bound vars that would be bound by the
     /// binder. This is commonly used to 'inject' a value T into a
     /// different binding level.
-    pub fn dummy<'tcx>(value: T) -> Binder<T>
-    where
-        T: TypeFoldable<'tcx>,
-    {
+    pub fn dummy(value: T) -> Binder<T> {
         debug_assert!(!value.has_escaping_bound_vars());
         Binder(value, 0)
     }
 
     /// Wraps `value` in a binder, binding higher-ranked vars (if any).
     pub fn bind(value: T) -> Binder<T> {
-        Binder(value, 0)
+        use crate::ty::fold::CountBoundVars;
+        use rustc_data_structures::fx::FxHashSet;
+        let mut counter = CountBoundVars {
+            outer_index: ty::INNERMOST,
+            bound_tys: FxHashSet::default(),
+            bound_regions: FxHashSet::default(),
+            bound_consts: FxHashSet::default(),
+        };
+        value.visit_with(&mut counter);
+        let bound_tys = counter.bound_tys.len();
+        let bound_regions = if !counter.bound_regions.is_empty() {
+            let mut env = false;
+            let mut anons = FxHashSet::default();
+            let mut named = FxHashSet::default();
+            for br in counter.bound_regions {
+                match br.kind {
+                    ty::BrAnon(idx) => {
+                        anons.insert(idx);
+                    }
+                    ty::BrNamed(def_id, _) => {
+                        named.insert(def_id);
+                    }
+                    ty::BrEnv => env = true,
+                }
+            }
+            (if env { 1 } else { 0 }) + anons.len() + named.len()
+        } else {
+            0
+        };
+        let bound_consts = counter.bound_consts.len();
+
+        let bound_vars = bound_tys + bound_regions + bound_consts;
+        Binder(value, bound_vars as u32)
     }
+}
 
+impl<T> Binder<T> {
     /// Skips the binder and returns the "bound" value. This is a
     /// risky thing to do because it's easy to get confused about
     /// De Bruijn indices and the like. It is usually better to
@@ -997,6 +1031,10 @@ impl<T> Binder<T> {
         self.0
     }
 
+    pub fn bound_vars(&self) -> u32 {
+        self.1
+    }
+
     pub fn as_ref(&self) -> Binder<&T> {
         Binder(&self.0, self.1)
     }