Rollup merge of #117095 - klinvill:smir-fn-arg-count, r=oli-obk

Add way to differentiate argument locals from other locals in Stable MIR

This PR resolves rust-lang/project-stable-mir#47 which request a way to differentiate argument locals in a SMIR `Body` from other locals.

Specifically, this PR exposes the `arg_count` field from the MIR `Body`. However, I'm opening this as a draft PR because I think there are a few outstanding questions on how this information should be exposed and described. Namely:

- Is exposing `arg_count` the best way to surface this information to SMIR users? Would it be better to leave `arg_count` as a private field and add public methods (e.g. `fn arguments(&self) -> Iter<'_, LocalDecls>`) that may use the underlying `arg_count` info from the MIR body, but expose this information to users in a more convenient form? Or is it best to stick close to the current MIR convention?
- If the answer to the above point is to stick with the current MIR convention (`arg_count`), is it reasonable to also commit to sticking to the current MIR convention that the first local is always the return local, while the next `arg_count` locals are always the (in-order) argument locals?
- Should `Body` in SMIR only represent function bodies (as implied by the comment I added)? That seems to be the current case in MIR, but should this restriction always be the case for SMIR?

r? `@celinval`
r? `@oli-obk`

compiler/stable_mir/src/mir/body.rs

@@ -2,10 +2,60 @@ use crate::ty::{AdtDef, ClosureDef, Const, CoroutineDef, GenericArgs, Movability
 use crate::Opaque;
 use crate::{ty::Ty, Span};
 
+/// The SMIR representation of a single function.
 #[derive(Clone, Debug)]
 pub struct Body {
     pub blocks: Vec<BasicBlock>,
-    pub locals: LocalDecls,
+
+    // Declarations of locals within the function.
+    //
+    // The first local is the return value pointer, followed by `arg_count`
+    // locals for the function arguments, followed by any user-declared
+    // variables and temporaries.
+    locals: LocalDecls,
+
+    // The number of arguments this function takes.
+    arg_count: usize,
+}
+
+impl Body {
+    /// Constructs a `Body`.
+    ///
+    /// A constructor is required to build a `Body` from outside the crate
+    /// because the `arg_count` and `locals` fields are private.
+    pub fn new(blocks: Vec<BasicBlock>, locals: LocalDecls, arg_count: usize) -> Self {
+        // If locals doesn't contain enough entries, it can lead to panics in
+        // `ret_local`, `arg_locals`, and `inner_locals`.
+        assert!(
+            locals.len() > arg_count,
+            "A Body must contain at least a local for the return value and each of the function's arguments"
+        );
+        Self { blocks, locals, arg_count }
+    }
+
+    /// Return local that holds this function's return value.
+    pub fn ret_local(&self) -> &LocalDecl {
+        &self.locals[0]
+    }
+
+    /// Locals in `self` that correspond to this function's arguments.
+    pub fn arg_locals(&self) -> &[LocalDecl] {
+        &self.locals[1..][..self.arg_count]
+    }
+
+    /// Inner locals for this function. These are the locals that are
+    /// neither the return local nor the argument locals.
+    pub fn inner_locals(&self) -> &[LocalDecl] {
+        &self.locals[self.arg_count + 1..]
+    }
+
+    /// Convenience function to get all the locals in this function.
+    ///
+    /// Locals are typically accessed via the more specific methods `ret_local`,
+    /// `arg_locals`, and `inner_locals`.
+    pub fn locals(&self) -> &[LocalDecl] {
+        &self.locals
+    }
 }
 
 type LocalDecls = Vec<LocalDecl>;
@@ -467,7 +517,7 @@ pub enum NullOp {
 }
 
 impl Operand {
-    pub fn ty(&self, locals: &LocalDecls) -> Ty {
+    pub fn ty(&self, locals: &[LocalDecl]) -> Ty {
         match self {
             Operand::Copy(place) | Operand::Move(place) => place.ty(locals),
             Operand::Constant(c) => c.ty(),
@@ -482,7 +532,7 @@ impl Constant {
 }
 
 impl Place {
-    pub fn ty(&self, locals: &LocalDecls) -> Ty {
+    pub fn ty(&self, locals: &[LocalDecl]) -> Ty {
         let _start_ty = locals[self.local].ty;
         todo!("Implement projection")
     }