use mir::Visitor when collecting alloc_ids in pretty printing

compiler/rustc_middle/src/mir/pretty.rs

@@ -17,9 +17,8 @@ use rustc_middle::mir::interpret::{
 use rustc_middle::mir::visit::Visitor;
 use rustc_middle::mir::MirSource;
 use rustc_middle::mir::*;
-use rustc_middle::ty::{self, TyCtxt, TypeFoldable, TypeVisitor};
+use rustc_middle::ty::{self, TyCtxt};
 use rustc_target::abi::Size;
-use std::ops::ControlFlow;
 
 const INDENT: &str = "    ";
 /// Alignment for lining up comments following MIR statements
@@ -669,6 +668,7 @@ pub fn write_allocations<'tcx>(
     fn alloc_ids_from_alloc(alloc: &Allocation) -> impl DoubleEndedIterator<Item = AllocId> + '_ {
         alloc.relocations().values().map(|id| *id)
     }
+
     fn alloc_ids_from_const(val: ConstValue<'_>) -> impl Iterator<Item = AllocId> + '_ {
         match val {
             ConstValue::Scalar(interpret::Scalar::Ptr(ptr, _size)) => {
@@ -682,17 +682,29 @@ pub fn write_allocations<'tcx>(
             }
         }
     }
+
     struct CollectAllocIds(BTreeSet<AllocId>);
-    impl<'tcx> TypeVisitor<'tcx> for CollectAllocIds {
-        fn visit_const(&mut self, c: ty::Const<'tcx>) -> ControlFlow<Self::BreakTy> {
-            if let ty::ConstKind::Value(val) = c.val() {
+
+    impl<'tcx> Visitor<'tcx> for CollectAllocIds {
+        fn visit_const(&mut self, c: &&'tcx ty::Const<'tcx>, _loc: Location) {
+            if let ty::ConstKind::Value(val) = c.val {
                 self.0.extend(alloc_ids_from_const(val));
             }
-            c.super_visit_with(self)
+        }
+
+        fn visit_constant(&mut self, c: &Constant<'tcx>, loc: Location) {
+            match c.literal {
+                ConstantKind::Ty(c) => self.visit_const(&c, loc),
+                ConstantKind::Val(val, _) => {
+                    self.0.extend(alloc_ids_from_const(val));
+                }
+            }
         }
     }
+
     let mut visitor = CollectAllocIds(Default::default());
-    body.visit_with(&mut visitor);
+    visitor.visit_body(body);
+
     // `seen` contains all seen allocations, including the ones we have *not* printed yet.
     // The protocol is to first `insert` into `seen`, and only if that returns `true`
     // then push to `todo`.