From c64aea5e377d46ebb09848e9e295b8a96477a74a Mon Sep 17 00:00:00 2001
From: Zyad Hassan <zyadh@amazon.com>
Date: Tue, 1 Oct 2024 16:03:51 -0700
Subject: [PATCH] [aeneas] Preserve variable names

---
 .../codegen_aeneas_llbc/mir_to_ullbc/mod.rs   | 30 +++++++++++++------
 tests/expected/llbc/basic0/expected           |  4 +--
 tests/expected/llbc/basic1/expected           | 12 ++++----
 3 files changed, 29 insertions(+), 17 deletions(-)

diff --git a/kani-compiler/src/codegen_aeneas_llbc/mir_to_ullbc/mod.rs b/kani-compiler/src/codegen_aeneas_llbc/mir_to_ullbc/mod.rs
index b8a7c2357a83..6973e5014365 100644
--- a/kani-compiler/src/codegen_aeneas_llbc/mir_to_ullbc/mod.rs
+++ b/kani-compiler/src/codegen_aeneas_llbc/mir_to_ullbc/mod.rs
@@ -16,7 +16,7 @@ use charon_lib::ast::Rvalue as CharonRvalue;
 use charon_lib::ast::Span as CharonSpan;
 use charon_lib::ast::meta::{AttrInfo, Loc, RawSpan};
 use charon_lib::ast::types::Ty as CharonTy;
-use charon_lib::ast::{AbortKind, Body as CharonBody, Var, VarId, make_locals_generator};
+use charon_lib::ast::{AbortKind, Body as CharonBody, Var, VarId};
 use charon_lib::ast::{
     AnyTransId, Assert, BodyId, BuiltinTy, Disambiguator, FileName, FunDecl, FunSig, GenericArgs,
     GenericParams, IntegerTy, ItemKind, ItemMeta, ItemOpacity, Literal, LiteralTy, Name, Opaque,
@@ -39,14 +39,16 @@ use charon_lib::ullbc_ast::{
     Terminator as CharonTerminator,
 };
 use charon_lib::{error_assert, error_or_panic};
+use rustc_data_structures::fx::FxHashMap;
 use rustc_errors::MultiSpan;
 use rustc_middle::ty::TyCtxt;
 use rustc_smir::rustc_internal;
 use rustc_span::def_id::DefId as InternalDefId;
 use stable_mir::abi::PassMode;
+use stable_mir::mir::VarDebugInfoContents;
 use stable_mir::mir::mono::Instance;
 use stable_mir::mir::{
-    BasicBlock, BinOp, Body, BorrowKind, CastKind, ConstOperand, Mutability, Operand, Place,
+    BasicBlock, BinOp, Body, BorrowKind, CastKind, ConstOperand, Local, Mutability, Operand, Place,
     ProjectionElem, Rvalue, Statement, StatementKind, SwitchTargets, Terminator, TerminatorKind,
 };
 use stable_mir::ty::{
@@ -64,6 +66,7 @@ pub struct Context<'a, 'tcx> {
     instance: Instance,
     translated: &'a mut TranslatedCrate,
     errors: &'a mut ErrorCtx<'tcx>,
+    local_names: FxHashMap<Local, String>,
 }
 
 impl<'a, 'tcx> Context<'a, 'tcx> {
@@ -75,7 +78,17 @@ impl<'a, 'tcx> Context<'a, 'tcx> {
         translated: &'a mut TranslatedCrate,
         errors: &'a mut ErrorCtx<'tcx>,
     ) -> Self {
-        Self { tcx, instance, translated, errors }
+        let mut local_names = FxHashMap::default();
+        // populate names of locals
+        for info in instance.body().unwrap().var_debug_info {
+            if let VarDebugInfoContents::Place(p) = info.value {
+                if p.projection.is_empty() {
+                    local_names.insert(p.local, info.name);
+                }
+            }
+        }
+
+        Self { tcx, instance, translated, errors, local_names }
     }
 
     fn tcx(&self) -> TyCtxt<'tcx> {
@@ -463,12 +476,11 @@ impl<'a, 'tcx> Context<'a, 'tcx> {
         // - the input arguments
         // - the remaining locals, used for the intermediate computations
         let mut locals = Vector::new();
-        {
-            let mut add_variable = make_locals_generator(&mut locals);
-            mir_body.local_decls().for_each(|(_, local)| {
-                add_variable(self.translate_ty(local.ty));
-            });
-        }
+        mir_body.local_decls().for_each(|(local, local_decl)| {
+            let ty = self.translate_ty(local_decl.ty);
+            let name = self.local_names.get(&local);
+            locals.push_with(|index| Var { index, name: name.cloned(), ty });
+        });
         locals
     }
 
diff --git a/tests/expected/llbc/basic0/expected b/tests/expected/llbc/basic0/expected
index 112a67a21548..45d073737f3b 100644
--- a/tests/expected/llbc/basic0/expected
+++ b/tests/expected/llbc/basic0/expected
@@ -1,8 +1,8 @@
 fn test::is_zero(@1: i32) -> bool\
 {\
     let @0: bool; // return\
-    let @1: i32; // arg #1\
+    let i@1: i32; // arg #1\
 
-    @0 := copy (@1) == const (0 : i32)\
+    @0 := copy (i@1) == const (0 : i32)\
     return\
 }
diff --git a/tests/expected/llbc/basic1/expected b/tests/expected/llbc/basic1/expected
index 9cb0bef6f7c6..233940af9752 100644
--- a/tests/expected/llbc/basic1/expected
+++ b/tests/expected/llbc/basic1/expected
@@ -1,15 +1,15 @@
 fn test::select(@1: bool, @2: i32, @3: i32) -> i32
 {
     let @0: i32; // return
-    let @1: bool; // arg #1
-    let @2: i32; // arg #2
-    let @3: i32; // arg #3
+    let s@1: bool; // arg #1
+    let x@2: i32; // arg #2
+    let y@3: i32; // arg #3
 
-    if copy (@1) {
-        @0 := copy (@2)
+    if copy (s@1) {
+        @0 := copy (x@2)
     }
     else {
-        @0 := copy (@3)
+        @0 := copy (y@3)
     }
     return
 }