Don't visit inside types by default

AeneasVerif · Oct 18, 2024 · 92e9f8b · 92e9f8b
1 parent 2950a29
commit 92e9f8b
Show file tree

Hide file tree

Showing 8 changed files with 122 additions and 41 deletions.
diff --git a/charon/src/ast/types.rs b/charon/src/ast/types.rs
@@ -648,7 +648,11 @@ pub enum ConstGeneric {
 }
 
 /// A type.
-#[derive(Debug, Clone, Hash, PartialEq, Eq, Serialize, Deserialize, Drive, DriveMut)]
+///
+/// Warning: for performance reasons, the `Drive` and `DriveMut` impls of `Ty` don't explore the
+/// contents of the type, they only yield a pointer to the type itself. To recurse into the type,
+/// use `drive_inner{_mut}` or `visit_inside`.
+#[derive(Debug, Clone, Hash, PartialEq, Eq, Serialize, Deserialize)]
 pub struct Ty(HashConsed<TyKind>);
 
 impl Ty {
@@ -659,6 +663,13 @@ impl Ty {
     pub fn kind(&self) -> &TyKind {
         self.0.inner()
     }
+
+    pub fn drive_inner<V: Visitor>(&self, visitor: &mut V) {
+        self.0.drive(visitor)
+    }
+    pub fn drive_inner_mut<V: VisitorMut>(&mut self, visitor: &mut V) {
+        self.0.drive_mut(visitor)
+    }
 }
 
 #[derive(

diff --git a/charon/src/ast/types_utils.rs b/charon/src/ast/types_utils.rs
@@ -286,6 +286,11 @@ impl Ty {
             _ => None,
         }
     }
+
+    /// Wrap a visitor to make it visit the contents of types it encounters.
+    pub fn visit_inside<V>(visitor: V) -> VisitInsideTy<V> {
+        VisitInsideTy(visitor)
+    }
 }
 
 impl TyKind {
@@ -373,3 +378,52 @@ impl<T: DriveMut> DriveMut for RegionBinder<T> {
         visitor.visit(self, Event::Exit);
     }
 }
+
+/// See comment on `Ty`: this impl doesn't recurse into the contents of the type.
+impl Drive for Ty {
+    fn drive<V: Visitor>(&self, visitor: &mut V) {
+        visitor.visit(self, Event::Enter);
+        visitor.visit(self, Event::Exit);
+    }
+}
+
+/// See comment on `Ty`: this impl doesn't recurse into the contents of the type.
+impl DriveMut for Ty {
+    fn drive_mut<V: VisitorMut>(&mut self, visitor: &mut V) {
+        visitor.visit(self, Event::Enter);
+        visitor.visit(self, Event::Exit);
+    }
+}
+
+pub struct VisitInsideTy<V>(pub V);
+
+impl<V: Visitor> Visitor for VisitInsideTy<V> {
+    fn visit(&mut self, item: &dyn std::any::Any, event: Event) {
+        let is_enter = matches!(event, Event::Enter);
+        self.0.visit(item, event);
+        if is_enter && let Some(ty) = item.downcast_ref::<Ty>() {
+            ty.drive_inner(self)
+        }
+    }
+}
+impl<V: VisitorMut> VisitorMut for VisitInsideTy<V> {
+    fn visit(&mut self, item: &mut dyn std::any::Any, event: Event) {
+        let is_enter = matches!(event, Event::Enter);
+        self.0.visit(item, event);
+        if is_enter && let Some(ty) = item.downcast_mut::<Ty>() {
+            ty.drive_inner_mut(self)
+        }
+    }
+}
+
+impl<V> std::ops::Deref for VisitInsideTy<V> {
+    type Target = V;
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+impl<V> std::ops::DerefMut for VisitInsideTy<V> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.0
+    }
+}
diff --git a/charon/src/bin/generate-ml/main.rs b/charon/src/bin/generate-ml/main.rs
@@ -84,17 +84,19 @@ fn contains_id(crate_data: &TranslatedCrate, haystack: TypeDeclId) -> HashMap<Ty
         let exploring_def_id = ty.def_id;
         let mut contains = false;
         let mut requires_parent = None;
-        ty.drive(&mut derive_visitor::visitor_enter_fn(|id: &TypeDeclId| {
-            if let Some(ty) = crate_data.type_decls.get(*id) {
-                match traverse_ty(crate_data, ty, stack, map) {
-                    Ok(true) => contains = true,
-                    Err(loop_id) if loop_id != exploring_def_id && stack.contains(&loop_id) => {
-                        requires_parent = Some(loop_id)
+        ty.drive(&mut Ty::visit_inside(derive_visitor::visitor_enter_fn(
+            |id: &TypeDeclId| {
+                if let Some(ty) = crate_data.type_decls.get(*id) {
+                    match traverse_ty(crate_data, ty, stack, map) {
+                        Ok(true) => contains = true,
+                        Err(loop_id) if loop_id != exploring_def_id && stack.contains(&loop_id) => {
+                            requires_parent = Some(loop_id)
+                        }
+                        _ => {}
                     }
-                    _ => {}
                 }
-            }
-        }));
+            },
+        )));
         stack.pop();
 
         if contains {
@@ -151,9 +153,11 @@ impl<'a> GenerateCtx<'a> {
             name_to_type.insert(long_name, ty);
 
             let mut contained = HashSet::new();
-            ty.drive(&mut visitor_enter_fn(|id: &TypeDeclId| {
-                contained.insert(*id);
-            }));
+            ty.drive(&mut Ty::visit_inside(visitor_enter_fn(
+                |id: &TypeDeclId| {
+                    contained.insert(*id);
+                },
+            )));
             type_tree.insert(ty.def_id, contained);
         }
         let contains_raw_span = {

diff --git a/charon/src/transform/check_generics.rs b/charon/src/transform/check_generics.rs
@@ -161,12 +161,12 @@ pub struct Check;
 impl TransformPass for Check {
     fn transform_ctx(&self, ctx: &mut TransformCtx<'_>) {
         for item in ctx.translated.all_items() {
-            let mut visitor = CheckGenericsVisitor {
+            let mut visitor = Ty::visit_inside(CheckGenericsVisitor {
                 translated: &ctx.translated,
                 error_ctx: &mut ctx.errors,
                 discharged_args: 0,
                 item_span: item.item_meta().span,
-            };
+            });
             item.drive(&mut visitor);
             assert_eq!(
                 visitor.discharged_args, 0,

diff --git a/charon/src/transform/hide_marker_traits.rs b/charon/src/transform/hide_marker_traits.rs
@@ -102,6 +102,7 @@ impl TransformPass for Transform {
             ctx.translated.trait_decls.remove(id);
         }
 
-        ctx.translated.drive_mut(&mut Visitor { exclude });
+        ctx.translated
+            .drive_mut(&mut Ty::visit_inside(Visitor { exclude }));
     }
 }
diff --git a/charon/src/transform/lift_associated_item_clauses.rs b/charon/src/transform/lift_associated_item_clauses.rs
@@ -46,27 +46,32 @@ impl TransformPass for Transform {
 
         // Update trait refs.
         ctx.translated
-            .drive_mut(&mut visitor_enter_fn_mut(|trkind: &mut TraitRefKind| {
-                use TraitRefKind::*;
-                if let ItemClause(..) = trkind {
-                    take_mut::take(trkind, |trkind| {
-                        let ItemClause(trait_ref, trait_decl, item_name, item_clause_id) = trkind
-                        else {
-                            unreachable!()
-                        };
-                        let new_id = (|| {
-                            let new_id = *trait_item_clause_ids
-                                .get(trait_decl)?
-                                .get(&item_name)?
-                                .get(item_clause_id)?;
-                            Some(new_id)
-                        })();
-                        match new_id {
-                            Some(new_id) => ParentClause(trait_ref, trait_decl, new_id),
-                            None => ItemClause(trait_ref, trait_decl, item_name, item_clause_id),
-                        }
-                    })
-                }
-            }));
+            .drive_mut(&mut Ty::visit_inside(visitor_enter_fn_mut(
+                |trkind: &mut TraitRefKind| {
+                    use TraitRefKind::*;
+                    if let ItemClause(..) = trkind {
+                        take_mut::take(trkind, |trkind| {
+                            let ItemClause(trait_ref, trait_decl, item_name, item_clause_id) =
+                                trkind
+                            else {
+                                unreachable!()
+                            };
+                            let new_id = (|| {
+                                let new_id = *trait_item_clause_ids
+                                    .get(trait_decl)?
+                                    .get(&item_name)?
+                                    .get(item_clause_id)?;
+                                Some(new_id)
+                            })();
+                            match new_id {
+                                Some(new_id) => ParentClause(trait_ref, trait_decl, new_id),
+                                None => {
+                                    ItemClause(trait_ref, trait_decl, item_name, item_clause_id)
+                                }
+                            }
+                        })
+                    }
+                },
+            )));
     }
 }
diff --git a/charon/src/transform/reorder_decls.rs b/charon/src/transform/reorder_decls.rs
@@ -126,7 +126,8 @@ impl Display for DeclarationGroup {
     GlobalDeclId(enter),
     TraitImplId(enter),
     TraitDeclId(enter),
-    BodyId(enter)
+    BodyId(enter),
+    Ty(enter)
 )]
 pub struct Deps<'tcx, 'ctx> {
     ctx: &'tcx TransformCtx<'ctx>,
@@ -280,6 +281,11 @@ impl Deps<'_, '_> {
             body.drive(self);
         }
     }
+
+    fn enter_ty(&mut self, ty: &Ty) {
+        // Recurse into the type, which doesn't happen by default.
+        ty.drive_inner(self);
+    }
 }
 
 impl AnyTransId {

diff --git a/charon/src/transform/update_closure_signatures.rs b/charon/src/transform/update_closure_signatures.rs
@@ -90,10 +90,10 @@ fn transform_function(
 
         // Explore the state and introduce fresh regions for the erased
         // regions we find.
-        let mut visitor = InsertRegions {
+        let mut visitor = Ty::visit_inside(InsertRegions {
             regions: &mut generics.regions,
             depth: 0,
-        };
+        });
         state.drive_mut(&mut visitor);
 
         // Update the inputs (slightly annoying to push to the front of