From d517a1cbda49309c8364daaedfee040371e616fb Mon Sep 17 00:00:00 2001
From: Kirby Linvill <kjlinvill@gmail.com>
Date: Fri, 10 Nov 2023 11:21:46 -0700
Subject: [PATCH] Add SMIR visitor for Places and projections

---
 compiler/stable_mir/src/mir/body.rs  |  2 +-
 compiler/stable_mir/src/mir/visit.rs | 33 +++++++++++++++++++++++++++-
 2 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/compiler/stable_mir/src/mir/body.rs b/compiler/stable_mir/src/mir/body.rs
index 625e52c983e..48d4a3d37e2 100644
--- a/compiler/stable_mir/src/mir/body.rs
+++ b/compiler/stable_mir/src/mir/body.rs
@@ -406,7 +406,7 @@ pub struct Place {
 // ProjectionElem<Local, Ty>) and user-provided type annotations (for which the projection elements
 // are of type ProjectionElem<(), ()>). In SMIR we don't need this generality, so we just use
 // ProjectionElem for Places.
-#[derive(Clone, Debug)]
+#[derive(Clone, Debug, Eq, PartialEq)]
 pub enum ProjectionElem {
     /// Dereference projections (e.g. `*_1`) project to the address referenced by the base place.
     Deref,
diff --git a/compiler/stable_mir/src/mir/visit.rs b/compiler/stable_mir/src/mir/visit.rs
index 806dced71ff..475d6e9763d 100644
--- a/compiler/stable_mir/src/mir/visit.rs
+++ b/compiler/stable_mir/src/mir/visit.rs
@@ -76,6 +76,15 @@ fn visit_place(&mut self, place: &Place, ptx: PlaceContext, location: Location)
         self.super_place(place, ptx, location)
     }
 
+    fn visit_projection_elem(
+        &mut self,
+        elem: &ProjectionElem,
+        ptx: PlaceContext,
+        location: Location,
+    ) {
+        self.super_projection_elem(elem, ptx, location);
+    }
+
     fn visit_local(&mut self, local: &Local, ptx: PlaceContext, location: Location) {
         let _ = (local, ptx, location);
     }
@@ -264,7 +273,29 @@ fn super_span(&mut self, span: &Span) {
     fn super_place(&mut self, place: &Place, ptx: PlaceContext, location: Location) {
         let _ = location;
         let _ = ptx;
-        visit_opaque(&Opaque(place.projection.clone()));
+        self.visit_local(&place.local, ptx, location);
+
+        for elem in &place.projection {
+            self.visit_projection_elem(elem, ptx, location);
+        }
+    }
+
+    fn super_projection_elem(
+        &mut self,
+        elem: &ProjectionElem,
+        ptx: PlaceContext,
+        location: Location,
+    ) {
+        match elem {
+            ProjectionElem::Deref => {}
+            ProjectionElem::Field(_idx, ty) => self.visit_ty(ty, location),
+            ProjectionElem::Index(local) => self.visit_local(local, ptx, location),
+            ProjectionElem::ConstantIndex { offset: _, min_length: _, from_end: _ } => {}
+            ProjectionElem::Subslice { from: _, to: _, from_end: _ } => {}
+            ProjectionElem::Downcast(_idx) => {}
+            ProjectionElem::OpaqueCast(ty) => self.visit_ty(ty, location),
+            ProjectionElem::Subtype(ty) => self.visit_ty(ty, location),
+        }
     }
 
     fn super_rvalue(&mut self, rvalue: &Rvalue, location: Location) {