Fix substitution

charon/src/ast/mod.rs

@@ -27,4 +27,5 @@ pub use krate::*;
 pub use meta::*;
 pub use names::*;
 pub use types::*;
+pub use types_utils::TyVisitable;
 pub use values::*;

charon/src/ast/types/vars.rs

@@ -124,6 +124,10 @@ impl DeBruijnId {
         DeBruijnId { index: 0 }
     }
 
+    pub fn one() -> Self {
+        DeBruijnId { index: 1 }
+    }
+
     pub fn new(index: usize) -> Self {
         DeBruijnId { index }
     }
@@ -143,6 +147,18 @@ impl DeBruijnId {
             index: self.index - 1,
         }
     }
+
+    pub fn plus(&self, delta: Self) -> Self {
+        DeBruijnId {
+            index: self.index + delta.index,
+        }
+    }
+
+    pub fn sub(&self, delta: Self) -> Option<Self> {
+        Some(DeBruijnId {
+            index: self.index.checked_sub(delta.index)?,
+        })
+    }
 }
 
 impl<Id> DeBruijnVar<Id>
@@ -161,19 +177,44 @@ where
         DeBruijnVar::Bound(index, id)
     }
 
+    pub fn incr(&self) -> Self {
+        match *self {
+            DeBruijnVar::Bound(dbid, varid) => DeBruijnVar::Bound(dbid.incr(), varid),
+            DeBruijnVar::Free(varid) => DeBruijnVar::Free(varid),
+        }
+    }
+
     pub fn decr(&self) -> Self {
         match *self {
             DeBruijnVar::Bound(dbid, varid) => DeBruijnVar::Bound(dbid.decr(), varid),
             DeBruijnVar::Free(varid) => DeBruijnVar::Free(varid),
         }
     }
 
+    /// Returns the variable id if it is bound as the given depth.
     pub fn bound_at_depth(&self, depth: DeBruijnId) -> Option<Id> {
         match *self {
             DeBruijnVar::Bound(dbid, varid) if dbid == depth => Some(varid),
             _ => None,
         }
     }
+
+    /// Move the variable out of `depth` binders. Returns `None` if the variable is bound in one of
+    /// these `depth` binders.
+    pub fn move_out_from_depth(&self, depth: DeBruijnId) -> Option<Self> {
+        Some(match *self {
+            DeBruijnVar::Bound(dbid, varid) => DeBruijnVar::Bound(dbid.sub(depth)?, varid),
+            DeBruijnVar::Free(_) => *self,
+        })
+    }
+
+    /// Move under `depth` binders.
+    pub fn move_under_binders(&self, depth: DeBruijnId) -> Self {
+        match *self {
+            DeBruijnVar::Bound(dbid, varid) => DeBruijnVar::Bound(dbid.plus(depth), varid),
+            DeBruijnVar::Free(_) => *self,
+        }
+    }
 }
 
 impl TypeVar {
@@ -182,6 +223,12 @@ impl TypeVar {
     }
 }
 
+impl Default for DeBruijnId {
+    fn default() -> Self {
+        Self::zero()
+    }
+}
+
 // The derive macro doesn't handle generics.
 impl<Id: Drive> Drive for DeBruijnVar<Id> {
     fn drive<V: Visitor>(&self, visitor: &mut V) {

charon/src/ast/types_utils.rs

@@ -1,10 +1,9 @@
 //! This file groups everything which is linked to implementations about [crate::types]
 use crate::types::*;
 use crate::{common::visitor_event::VisitEvent, ids::Vector};
-use derive_visitor::{Drive, DriveMut, Event, Visitor, VisitorMut};
+use derive_visitor::{visitor_enter_fn_mut, Drive, DriveMut, Event, Visitor, VisitorMut};
 use std::collections::HashSet;
 use std::{collections::HashMap, iter::Iterator};
-use DeBruijnVar::Free;
 
 impl GenericParams {
     pub fn empty() -> Self {
@@ -553,6 +552,7 @@ impl<V> std::ops::DerefMut for VisitInsideTy<V> {
 type FnTys = (Vec<Ty>, Ty);
 
 /// Visitor for the [Ty::substitute] function.
+/// This substitutes only free variables; this does not work for non-top-level binders.
 #[derive(VisitorMut)]
 #[visitor(
     PolyTraitDeclRef(enter, exit),
@@ -580,6 +580,13 @@ impl<'a> SubstVisitor<'a> {
         })
     }
 
+    fn should_subst<Id: Copy>(&self, var: DeBruijnVar<Id>) -> Option<Id> {
+        match var {
+            DeBruijnVar::Free(id) => Some(id),
+            DeBruijnVar::Bound(..) => None,
+        }
+    }
+
     fn enter_poly_trait_decl_ref(&mut self, _: &mut PolyTraitDeclRef) {
         self.binder_depth = self.binder_depth.incr();
     }
@@ -599,8 +606,8 @@ impl<'a> SubstVisitor<'a> {
     fn exit_region(&mut self, r: &mut Region) {
         match r {
             Region::Var(var) => {
-                if let Some(varid) = var.bound_at_depth(self.binder_depth) {
-                    *r = self.generics.regions.get(varid).unwrap().clone()
+                if let Some(varid) = self.should_subst(*var) {
+                    *r = self.generics.regions[varid].move_under_binders(self.binder_depth)
                 }
             }
             _ => (),
@@ -609,39 +616,78 @@ impl<'a> SubstVisitor<'a> {
 
     fn exit_ty(&mut self, ty: &mut Ty) {
         match ty.kind() {
-            TyKind::TypeVar(Free(id)) => {
-                {}
-                *ty = self.generics.types.get(*id).unwrap().clone()
+            TyKind::TypeVar(var) => {
+                if let Some(id) = self.should_subst(*var) {
+                    *ty = self.generics.types[id]
+                        .clone()
+                        .move_under_binders(self.binder_depth)
+                }
             }
             _ => (),
         }
     }
 
     fn exit_const_generic(&mut self, cg: &mut ConstGeneric) {
         match cg {
-            ConstGeneric::Var(Free(id)) => {
-                *cg = self.generics.const_generics.get(*id).unwrap().clone()
+            ConstGeneric::Var(var) => {
+                if let Some(id) = self.should_subst(*var) {
+                    *cg = self.generics.const_generics[id]
+                        .clone()
+                        .move_under_binders(self.binder_depth)
+                }
             }
             _ => (),
         }
     }
 
     fn exit_trait_ref(&mut self, tr: &mut TraitRef) {
         match &mut tr.kind {
-            TraitRefKind::Clause(Free(id)) => {
-                *tr = self.generics.trait_refs.get(*id).unwrap().clone()
+            TraitRefKind::Clause(var) => {
+                if let Some(id) = self.should_subst(*var) {
+                    *tr = self.generics.trait_refs[id]
+                        .clone()
+                        .move_under_binders(self.binder_depth)
+                }
             }
             _ => (),
         }
     }
 }
 
-impl Ty {
-    pub fn substitute(&mut self, generics: &GenericArgs) {
-        self.drive_inner_mut(&mut SubstVisitor::new(generics));
+/// Types that are involved at the type-level and may be substituted around.
+pub trait TyVisitable: Sized + Drive + DriveMut {
+    fn substitute(&mut self, generics: &GenericArgs) {
+        self.drive_mut(&mut SubstVisitor::new(generics));
+    }
+
+    /// Move under `depth` binders.
+    fn move_under_binders(mut self, depth: DeBruijnId) -> Self {
+        self.visit_db_id(|id| *id = id.plus(depth));
+        self
+    }
+
+    /// Move the region out of `depth` binders. Returns `None` if the variable is bound in one of
+    /// these `depth` binders.
+    fn move_from_under_binders(mut self, depth: DeBruijnId) -> Option<Self> {
+        let mut ok = true;
+        self.visit_db_id(|id| match id.sub(depth) {
+            Some(sub) => *id = sub,
+            None => ok = false,
+        });
+        ok.then_some(self)
+    }
+
+    /// Helper function.
+    fn visit_db_id(&mut self, f: impl FnMut(&mut DeBruijnId)) {
+        self.drive_mut(&mut Ty::visit_inside(visitor_enter_fn_mut(f)));
     }
 }
 
+impl TyVisitable for ConstGeneric {}
+impl TyVisitable for Region {}
+impl TyVisitable for TraitRef {}
+impl TyVisitable for Ty {}
+
 impl PartialEq for TraitClause {
     fn eq(&self, other: &Self) -> bool {
         // Skip `span` and `origin`