From efab0dcb2582fced369a3cc991660760f49a4712 Mon Sep 17 00:00:00 2001
From: "Celina G. Val" <celinval@amazon.com>
Date: Thu, 11 Jan 2024 17:15:21 -0800
Subject: [PATCH] Add more information to `visit_projection_elem`

Without the starting place, it's hard to retrieve any useful information
from visiting a projection.
---
 compiler/stable_mir/src/mir/body.rs  | 33 +++++++++++++++-------------
 compiler/stable_mir/src/mir/visit.rs | 24 ++++++++++++++++----
 2 files changed, 38 insertions(+), 19 deletions(-)

diff --git a/compiler/stable_mir/src/mir/body.rs b/compiler/stable_mir/src/mir/body.rs
index 72227a04bf189..38877f7a77fc4 100644
--- a/compiler/stable_mir/src/mir/body.rs
+++ b/compiler/stable_mir/src/mir/body.rs
@@ -963,7 +963,7 @@ pub enum PointerCoercion {
     /// Go from a safe fn pointer to an unsafe fn pointer.
     UnsafeFnPointer,
 
-    /// Go from a non-capturing closure to an fn pointer or an unsafe fn pointer.
+    /// Go from a non-capturing closure to a fn pointer or an unsafe fn pointer.
     /// It cannot convert a closure that requires unsafe.
     ClosureFnPointer(Safety),
 
@@ -1037,21 +1037,24 @@ impl Place {
     /// locals from the function body where this place originates from.
     pub fn ty(&self, locals: &[LocalDecl]) -> Result<Ty, Error> {
         let start_ty = locals[self.local].ty;
-        self.projection.iter().fold(Ok(start_ty), |place_ty, elem| {
-            let ty = place_ty?;
-            match elem {
-                ProjectionElem::Deref => Self::deref_ty(ty),
-                ProjectionElem::Field(_idx, fty) => Ok(*fty),
-                ProjectionElem::Index(_) | ProjectionElem::ConstantIndex { .. } => {
-                    Self::index_ty(ty)
-                }
-                ProjectionElem::Subslice { from, to, from_end } => {
-                    Self::subslice_ty(ty, from, to, from_end)
-                }
-                ProjectionElem::Downcast(_) => Ok(ty),
-                ProjectionElem::OpaqueCast(ty) | ProjectionElem::Subtype(ty) => Ok(*ty),
+        self.projection.iter().fold(Ok(start_ty), |place_ty, elem| elem.ty(place_ty?))
+    }
+}
+
+impl ProjectionElem {
+    /// Get the expected type after applying this projection to a given place type.
+    pub fn ty(&self, place_ty: Ty) -> Result<Ty, Error> {
+        let ty = place_ty;
+        match &self {
+            ProjectionElem::Deref => Self::deref_ty(ty),
+            ProjectionElem::Field(_idx, fty) => Ok(*fty),
+            ProjectionElem::Index(_) | ProjectionElem::ConstantIndex { .. } => Self::index_ty(ty),
+            ProjectionElem::Subslice { from, to, from_end } => {
+                Self::subslice_ty(ty, from, to, from_end)
             }
-        })
+            ProjectionElem::Downcast(_) => Ok(ty),
+            ProjectionElem::OpaqueCast(ty) | ProjectionElem::Subtype(ty) => Ok(*ty),
+        }
     }
 
     fn index_ty(ty: Ty) -> Result<Ty, Error> {
diff --git a/compiler/stable_mir/src/mir/visit.rs b/compiler/stable_mir/src/mir/visit.rs
index ab57ff0f8f5d7..24296e9e8778b 100644
--- a/compiler/stable_mir/src/mir/visit.rs
+++ b/compiler/stable_mir/src/mir/visit.rs
@@ -37,7 +37,7 @@
 
 use crate::mir::*;
 use crate::ty::{Const, GenericArgs, Region, Ty};
-use crate::{Opaque, Span};
+use crate::{Error, Opaque, Span};
 
 pub trait MirVisitor {
     fn visit_body(&mut self, body: &Body) {
@@ -76,12 +76,14 @@ pub trait MirVisitor {
         self.super_place(place, ptx, location)
     }
 
-    fn visit_projection_elem(
+    fn visit_projection_elem<'a>(
         &mut self,
+        place_ref: PlaceRef<'a>,
         elem: &ProjectionElem,
         ptx: PlaceContext,
         location: Location,
     ) {
+        let _ = place_ref;
         self.super_projection_elem(elem, ptx, location);
     }
 
@@ -284,8 +286,9 @@ pub trait MirVisitor {
         let _ = ptx;
         self.visit_local(&place.local, ptx, location);
 
-        for elem in &place.projection {
-            self.visit_projection_elem(elem, ptx, location);
+        for (idx, elem) in place.projection.iter().enumerate() {
+            let place_ref = PlaceRef { local: place.local, projection: &place.projection[..idx] };
+            self.visit_projection_elem(place_ref, elem, ptx, location);
         }
     }
 
@@ -453,6 +456,19 @@ impl Location {
     }
 }
 
+/// Reference to a place used to represent a partial projection.
+pub struct PlaceRef<'a> {
+    pub local: Local,
+    pub projection: &'a [ProjectionElem],
+}
+
+impl<'a> PlaceRef<'a> {
+    /// Get the type of this place.
+    pub fn ty(&self, locals: &[LocalDecl]) -> Result<Ty, Error> {
+        self.projection.iter().fold(Ok(locals[self.local].ty), |place_ty, elem| elem.ty(place_ty?))
+    }
+}
+
 /// Information about a place's usage.
 #[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
 pub struct PlaceContext {