Replace arg_count in public API with return/arg getters

This commit hides the arg_count field in Body and instead exposes more
stable and user-friendly methods to get the return and argument locals.
As a result, Body instances must now be constructed using the `new`
function.

Author: klinvill

compiler/rustc_smir/src/rustc_smir/mod.rs

@@ -287,9 +287,8 @@ impl<'tcx> Stable<'tcx> for mir::Body<'tcx> {
     type T = stable_mir::mir::Body;
 
     fn stable(&self, tables: &mut Tables<'tcx>) -> Self::T {
-        stable_mir::mir::Body {
-            blocks: self
-                .basic_blocks
+        stable_mir::mir::Body::new(
+            self.basic_blocks
                 .iter()
                 .map(|block| stable_mir::mir::BasicBlock {
                     terminator: block.terminator().stable(tables),
@@ -300,16 +299,15 @@ impl<'tcx> Stable<'tcx> for mir::Body<'tcx> {
                         .collect(),
                 })
                 .collect(),
-            locals: self
-                .local_decls
+            self.local_decls
                 .iter()
                 .map(|decl| stable_mir::mir::LocalDecl {
                     ty: decl.ty.stable(tables),
                     span: decl.source_info.span.stable(tables),
                 })
                 .collect(),
-            arg_count: self.arg_count,
-        }
+            self.arg_count,
+        )
     }
 }
 

compiler/stable_mir/src/mir/body.rs

@@ -8,14 +8,40 @@ pub struct Body {
     pub blocks: Vec<BasicBlock>,
 
     /// Declarations of locals.
-    ///
-    /// 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.
+    //
+    // 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.
     pub locals: LocalDecls,
 
-    /// The number of arguments this function takes.
-    pub arg_count: usize,
+    // 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` field is 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` and `arg_locals`.
+        assert!(
+            locals.len() >= arg_count + 1,
+            "A Body must contain at least a local for the return value and each of the function's arguments"
+        );
+        Self { blocks, locals, arg_count }
+    }
+
+    /// Gets the function's return local.
+    pub fn ret_local(&self) -> &LocalDecl {
+        &self.locals[0]
+    }
+
+    /// Gets the locals in `self` that correspond to the function's arguments.
+    pub fn arg_locals(&self) -> &[LocalDecl] {
+        &self.locals[1..self.arg_count + 1]
+    }
 }
 
 type LocalDecls = Vec<LocalDecl>;