Abstract out locals manipulation

charon-ml/src/GAst.ml

@@ -96,21 +96,26 @@ and assertion = { cond : operand; expected : bool }
         nude = true (* Don't inherit VisitorsRuntime *);
       }]
 
-(** An expression body.
-    TODO: arg_count should be stored in GFunDecl below. But then,
-          the print is obfuscated and Aeneas may need some refactoring.
- *)
-type 'a0 gexpr_body = {
-  span : span;
+(** The local variables of a body. *)
+type locals = {
   arg_count : int;
       (** The number of local variables used for the input arguments. *)
-  locals : var list;
+  vars : var list;
       (** The local variables.
         We always have, in the following order:
         - the local used for the return value (index 0)
-        - the input arguments
+        - the `arg_count` input arguments
         - the remaining locals, used for the intermediate computations
      *)
+}
+
+(** An expression body.
+    TODO: arg_count should be stored in GFunDecl below. But then,
+          the print is obfuscated and Aeneas may need some refactoring.
+ *)
+and 'a0 gexpr_body = {
+  span : span;
+  locals : locals;  (** The local variables. *)
   body : 'a0;
 }
 

charon-ml/src/GAstOfJson.ml

@@ -349,6 +349,15 @@ and var_of_json (js : json) : (var, string) result =
         Ok ({ index; name; var_ty } : var)
     | _ -> Error "")
 
+and locals_of_json (js : json) : (locals, string) result =
+  combine_error_msgs js __FUNCTION__
+    (match js with
+    | `Assoc [ ("arg_count", arg_count); ("vars", vars) ] ->
+        let* arg_count = int_of_json arg_count in
+        let* vars = vector_of_json var_id_of_json var_of_json vars in
+        Ok ({ arg_count; vars } : locals)
+    | _ -> Error "")
+
 and gexpr_body_of_json :
       'a0.
       (json -> ('a0, string) result) ->
@@ -359,18 +368,12 @@ and gexpr_body_of_json :
   combine_error_msgs js __FUNCTION__
     (match js with
     | `Assoc
-        [
-          ("span", span);
-          ("arg_count", arg_count);
-          ("locals", locals);
-          ("comments", _);
-          ("body", body);
-        ] ->
+        [ ("span", span); ("locals", locals); ("comments", _); ("body", body) ]
+      ->
         let* span = span_of_json id_to_file span in
-        let* arg_count = int_of_json arg_count in
-        let* locals = vector_of_json var_id_of_json var_of_json locals in
+        let* locals = locals_of_json locals in
         let* body = arg0_of_json body in
-        Ok ({ span; arg_count; locals; body } : _ gexpr_body)
+        Ok ({ span; locals; body } : _ gexpr_body)
     | _ -> Error "")
 
 and item_kind_of_json (js : json) : (item_kind, string) result =

charon-ml/src/GAstUtils.ml

@@ -34,12 +34,12 @@ let list_ordered_ancestor_region_groups (regions_hierarchy : region_var_groups)
   let parents = List.map (fun (rg : region_var_group) -> rg.id) parents in
   parents
 
-let gexpr_body_get_input_vars (fbody : 'body gexpr_body) : var list =
-  let locals = List.tl fbody.locals in
-  Collections.List.prefix fbody.arg_count locals
+let locals_get_input_vars (locals : locals) : var list =
+  let args = List.tl locals.vars in
+  Collections.List.prefix locals.arg_count args
 
 let fun_body_get_input_vars (fbody : 'body gexpr_body) : var list =
-  gexpr_body_get_input_vars fbody
+  locals_get_input_vars fbody.locals
 
 let g_declaration_group_to_list (g : 'a g_declaration_group) : 'a list =
   match g with RecGroup ids -> ids | NonRecGroup id -> [ id ]

charon-ml/src/LlbcAstUtils.ml

@@ -13,10 +13,7 @@ let fun_decl_list_from_crate (crate : crate) : fun_decl list =
 *)
 let get_fun_args (fun_decl : fun_decl) : var list option =
   match fun_decl.body with
-  | Some body ->
-      let input_number = body.arg_count in
-      let input_list = List.tl body.locals in
-      Some (fst (List.split_at input_list input_number))
+  | Some body -> Some (GAstUtils.locals_get_input_vars body.locals)
   | None -> None
 
 (** Check if a {!type:Charon.LlbcAst.statement} contains loops *)

charon-ml/src/PrintGAst.ml

@@ -102,14 +102,11 @@ let gfun_decl_to_string (env : ('a, 'b) fmt_env) (indent : string)
         (Some name) None sg
   | Some body ->
       (* Locally update the environment *)
-      let locals = List.map (fun v -> (v.index, v.name)) body.locals in
+      let locals = List.map (fun v -> (v.index, v.name)) body.locals.vars in
       let env = { env with locals } in
 
       (* Arguments *)
-      let inputs = List.tl body.locals in
-      let inputs, _aux_locals =
-        Collections.List.split_at inputs body.arg_count
-      in
+      let inputs = GAstUtils.locals_get_input_vars body.locals in
 
       (* All the locals (with erased regions) *)
       let locals =
@@ -118,7 +115,7 @@ let gfun_decl_to_string (env : ('a, 'b) fmt_env) (indent : string)
             indent ^ indent_incr ^ var_to_string var ^ " : "
             ^ ty_to_string env var.var_ty
             ^ ";")
-          body.locals
+          body.locals.vars
       in
       let locals = String.concat "\n" locals in
 

charon-ml/src/Substitute.ml

@@ -338,7 +338,7 @@ let fun_body_substitute_in_body (subst : subst) (body : fun_body) :
   let locals =
     List.map
       (fun (v : var) -> { v with var_ty = ty_substitute subst v.var_ty })
-      body.locals
+      body.locals.vars
   in
   let body = statement_substitute subst body.body in
   (locals, body)

charon/src/ast/expressions_utils.rs

@@ -49,13 +49,9 @@ impl BorrowKind {
 
 impl Place {
     /// Compute the type of a place.
-    pub fn ty(
-        &self,
-        type_decls: &Vector<TypeDeclId, TypeDecl>,
-        locals: &Vector<VarId, Var>,
-    ) -> Result<Ty, ()> {
+    pub fn ty(&self, type_decls: &Vector<TypeDeclId, TypeDecl>, locals: &Locals) -> Result<Ty, ()> {
         Ok(match &self.kind {
-            PlaceKind::Base(var_id) => locals.get(*var_id).ok_or(())?.ty.clone(),
+            PlaceKind::Base(var_id) => locals.vars.get(*var_id).ok_or(())?.ty.clone(),
             PlaceKind::Projection(sub_place, proj) => {
                 let sub_ty = sub_place.ty(type_decls, locals)?;
                 proj.project_type(type_decls, &sub_ty)?

charon/src/ast/gast.rs

@@ -1,5 +1,4 @@
 //! Definitions common to [crate::ullbc_ast] and [crate::llbc_ast]
-pub use super::gast_utils::*;
 use crate::expressions::*;
 use crate::generate_index_type;
 use crate::ids::Vector;
@@ -35,21 +34,28 @@ pub struct Var {
 #[derive(Debug, Copy, Clone, Serialize, Deserialize, Drive, DriveMut)]
 pub struct Opaque;
 
+/// The local variables of a body.
+#[derive(Debug, Default, Clone, Serialize, Deserialize, Drive, DriveMut)]
+pub struct Locals {
+    /// The number of local variables used for the input arguments.
+    pub arg_count: usize,
+    /// The local variables.
+    /// We always have, in the following order:
+    /// - the local used for the return value (index 0)
+    /// - the `arg_count` input arguments
+    /// - the remaining locals, used for the intermediate computations
+    pub vars: Vector<VarId, Var>,
+}
+
 /// An expression body.
 /// TODO: arg_count should be stored in GFunDecl below. But then,
 ///       the print is obfuscated and Aeneas may need some refactoring.
 #[derive(Debug, Clone, Serialize, Deserialize)]
 #[charon::rename("GexprBody")]
 pub struct GExprBody<T> {
     pub span: Span,
-    /// The number of local variables used for the input arguments.
-    pub arg_count: usize,
     /// The local variables.
-    /// We always have, in the following order:
-    /// - the local used for the return value (index 0)
-    /// - the input arguments
-    /// - the remaining locals, used for the intermediate computations
-    pub locals: Vector<VarId, Var>,
+    pub locals: Locals,
     /// For each line inside the body, we record any whole-line `//` comments found before it. They
     /// are added to statements in the late `recover_body_comments` pass.
     #[charon::opaque]
@@ -62,7 +68,6 @@ impl<T: Drive> Drive for GExprBody<T> {
     fn drive<V: Visitor>(&self, visitor: &mut V) {
         visitor.visit(self, Event::Enter);
         self.span.drive(visitor);
-        self.arg_count.drive(visitor);
         self.locals.drive(visitor);
         self.body.drive(visitor);
         visitor.visit(self, Event::Exit);
@@ -72,7 +77,6 @@ impl<T: DriveMut> DriveMut for GExprBody<T> {
     fn drive_mut<V: VisitorMut>(&mut self, visitor: &mut V) {
         visitor.visit(self, Event::Enter);
         self.span.drive_mut(visitor);
-        self.arg_count.drive_mut(visitor);
         self.locals.drive_mut(visitor);
         self.body.drive_mut(visitor);
         visitor.visit(self, Event::Exit);

charon/src/ast/gast_utils.rs

@@ -1,22 +1,9 @@
 //! Implementations for [crate::gast]
 
 use crate::ast::*;
-use crate::ids::Vector;
 use crate::llbc_ast;
 use crate::ullbc_ast;
 
-/// Makes a lambda that generates a new variable id, pushes a new variable in
-/// the body locals with the given type and returns its id.
-pub fn make_locals_generator(locals: &mut Vector<VarId, Var>) -> impl FnMut(Ty) -> VarId + '_ {
-    move |ty| {
-        locals.push_with(|index| Var {
-            index,
-            name: None,
-            ty,
-        })
-    }
-}
-
 impl FunIdOrTraitMethodRef {
     pub fn mk_builtin(aid: BuiltinFunId) -> Self {
         Self::Fun(FunId::Builtin(aid))
@@ -54,3 +41,28 @@ impl Body {
         }
     }
 }
+
+impl Locals {
+    /// Creates a new variable and returns a place pointing to it.
+    pub fn new_var(&mut self, name: Option<String>, ty: Ty) -> Place {
+        let var_id = self.vars.push_with(|index| Var { index, name, ty });
+        Place::new(var_id)
+    }
+
+    /// Gets a place pointing to the corresponding variable.
+    pub fn place_for_var(&self, var_id: VarId) -> Place {
+        Place::new(var_id)
+    }
+
+    /// The place where we write the return value.
+    pub fn return_place(&self) -> Place {
+        self.place_for_var(VarId::new(0))
+    }
+}
+
+impl std::ops::Index<VarId> for Locals {
+    type Output = Var;
+    fn index(&self, var_id: VarId) -> &Self::Output {
+        &self.vars[var_id]
+    }
+}

charon/src/ast/ullbc_ast_utils.rs

@@ -106,7 +106,7 @@ impl BlockData {
 impl ExprBody {
     pub fn transform_sequences<F>(&mut self, f: &mut F)
     where
-        F: FnMut(&mut Vector<VarId, Var>, &mut [Statement]) -> Vec<(usize, Vec<Statement>)>,
+        F: FnMut(&mut Locals, &mut [Statement]) -> Vec<(usize, Vec<Statement>)>,
     {
         for block in &mut self.body {
             block.transform_sequences(&mut |seq| f(&mut self.locals, seq));

charon/src/bin/charon-driver/translate/translate_ctx.rs

@@ -256,8 +256,8 @@ pub(crate) struct BodyTransCtx<'tcx, 'ctx, 'ctx1> {
     /// Cache the translation of types. This harnesses the deduplication of `TyKind` that hax does.
     pub type_trans_cache: HashMap<HashByAddr<Arc<hax::TyKind>>, Ty>,
 
-    /// The "regular" variables
-    pub vars: Vector<VarId, ast::Var>,
+    /// The (regular) variables in the current function body.
+    pub locals: Locals,
     /// The map from rust variable indices to translated variables indices.
     pub vars_map: HashMap<usize, VarId>,
     /// The translated blocks. We can't use `ast::Vector<BlockId, ast::BlockData>`
@@ -905,7 +905,7 @@ impl<'tcx, 'ctx, 'ctx1> BodyTransCtx<'tcx, 'ctx, 'ctx1> {
             parent_trait_clauses: Default::default(),
             item_trait_clauses: Default::default(),
             type_trans_cache: Default::default(),
-            vars: Default::default(),
+            locals: Default::default(),
             vars_map: Default::default(),
             blocks: Default::default(),
             blocks_map: Default::default(),
@@ -1089,7 +1089,7 @@ impl<'tcx, 'ctx, 'ctx1> BodyTransCtx<'tcx, 'ctx, 'ctx1> {
     }
 
     pub(crate) fn push_var(&mut self, rid: usize, ty: Ty, name: Option<String>) {
-        let var_id = self.vars.push_with(|index| Var { index, name, ty });
+        let var_id = self.locals.vars.push_with(|index| Var { index, name, ty });
         self.vars_map.insert(rid, var_id);
     }
 
@@ -1151,7 +1151,7 @@ impl<'tcx, 'ctx, 'ctx1, 'a> IntoFormatter for &'a BodyTransCtx<'tcx, 'ctx, 'ctx1
         FmtCtx {
             translated: Some(&self.t_ctx.translated),
             generics,
-            locals: Some(&self.vars),
+            locals: Some(&self.locals),
         }
     }
 }

charon/src/bin/charon-driver/translate/translate_functions_to_ullbc.rs

@@ -1278,6 +1278,7 @@ impl<'tcx, 'ctx, 'ctx1> BodyTransCtx<'tcx, 'ctx, 'ctx1> {
             return Ok(Err(Opaque));
         };
 
+        self.locals.arg_count = arg_count;
         // Here, we have to create a MIR state, which contains the body
         // Yes, we have to clone, this is annoying: we end up cloning the body twice
         let state = self
@@ -1311,8 +1312,7 @@ impl<'tcx, 'ctx, 'ctx1> BodyTransCtx<'tcx, 'ctx, 'ctx1> {
         // Create the body
         Ok(Ok(Body::Unstructured(ExprBody {
             span,
-            arg_count,
-            locals: mem::take(&mut self.vars),
+            locals: mem::take(&mut self.locals),
             comments: self.translate_body_comments(def, span),
             body: blocks,
         })))

charon/src/bin/generate-ml/main.rs

@@ -1296,6 +1296,7 @@ fn generate_ml(
                     "ClosureInfo",
                     "FunSig",
                     "ItemKind",
+                    "Locals",
                     "GExprBody",
                     "TraitDecl",
                     "TraitImpl",

charon/src/pretty/fmt_with_ctx.rs

@@ -506,11 +506,11 @@ where
 
         // Format the local variables
         let mut locals: Vec<String> = Vec::new();
-        for v in &self.locals {
+        for v in &self.locals.vars {
             let index = v.index.index();
             let comment = if index == 0 {
                 "// return".to_string()
-            } else if index <= self.arg_count {
+            } else if index <= self.locals.arg_count {
                 format!("// arg #{index}").to_string()
             } else {
                 match &v.name {

charon/src/pretty/formatter.rs

@@ -91,13 +91,13 @@ impl<'a, 'b> SetGenerics<'a> for FmtCtx<'b> {
 pub trait SetLocals<'a> {
     type C: 'a + AstFormatter;
 
-    fn set_locals(&'a self, locals: &'a Vector<VarId, ast::Var>) -> Self::C;
+    fn set_locals(&'a self, locals: &'a Locals) -> Self::C;
 }
 
 impl<'a, 'b> SetLocals<'a> for FmtCtx<'b> {
     type C = FmtCtx<'a>;
 
-    fn set_locals(&'a self, locals: &'a Vector<VarId, ast::Var>) -> Self::C {
+    fn set_locals(&'a self, locals: &'a Locals) -> Self::C {
         FmtCtx {
             translated: self.translated.as_deref(),
             generics: self.generics.clone(),
@@ -168,7 +168,7 @@ pub struct FmtCtx<'a> {
     /// Generics form a stack, where each binder introduces a new level. For DeBruijn indices to
     /// work, we keep the innermost parameters at the start of the vector.
     pub generics: VecDeque<Cow<'a, GenericParams>>,
-    pub locals: Option<&'a Vector<VarId, ast::Var>>,
+    pub locals: Option<&'a Locals>,
 }
 
 impl<'a> FmtCtx<'a> {
@@ -465,7 +465,7 @@ impl<'a> Formatter<VarId> for FmtCtx<'a> {
     fn format_object(&self, id: VarId) -> String {
         match &self.locals {
             None => id.to_pretty_string(),
-            Some(vars) => match vars.get(id) {
+            Some(locals) => match locals.vars.get(id) {
                 None => id.to_pretty_string(),
                 Some(v) => v.to_string(),
             },