Source-level names in JSON trace (#930)

* get GlobalEnv inside PrettyCx

* done pretty record

* only drop struct name for top-level index

* add comment

* remove lifetimes from json trace

* clippy

* move local_names into Body::new

crates/flux-common/src/dbg.rs

@@ -70,11 +70,13 @@ pub use crate::_basic_block_start as basic_block_start;
 
 #[macro_export]
 macro_rules! _statement{
-    ($pos:literal, $stmt:expr, $rcx:expr, $env:expr, $span:expr, $genv:expr) => {{
+    ($pos:literal, $stmt:expr, $rcx:expr, $env:expr, $span:expr, $checker:expr) => {{
         if config::dump_checker_trace() {
-            let genv = $genv;
+            let ck = $checker;
+            let genv = ck.genv;
+            let local_names = &ck.body.local_names;
             let rcx_json = RefineCtxtTrace::new(genv, $rcx);
-            let env_json = TypeEnvTrace::new(genv, $env);
+            let env_json = TypeEnvTrace::new(genv, local_names, $env);
             tracing::info!(event = concat!("statement_", $pos), stmt = ?$stmt, stmt_span = ?$span, rcx = ?$rcx, env = ?$env, rcx_json = ?rcx_json, env_json = ?env_json)
         }
     }};

crates/flux-refineck/src/checker.rs

@@ -515,9 +515,9 @@ impl<'ck, 'genv, 'tcx, M: Mode> Checker<'ck, 'genv, 'tcx, M> {
                 span,
             )?;
             bug::track_span(span, || {
-                dbg::statement!("start", stmt, &infcx, &env, span, self.genv);
+                dbg::statement!("start", stmt, &infcx, &env, span, &self);
                 self.check_statement(&mut infcx, &mut env, stmt)?;
-                dbg::statement!("end", stmt, &infcx, &env, span, self.genv);
+                dbg::statement!("end", stmt, &infcx, &env, span, &self);
                 Ok(())
             })?;
             if !stmt.is_nop() {
@@ -1538,7 +1538,7 @@ impl<'ck, 'genv, 'tcx, M: Mode> Checker<'ck, 'genv, 'tcx, M> {
         stmt: &GhostStatement,
         span: Span,
     ) -> Result {
-        dbg::statement!("start", stmt, infcx, env, span, self.genv);
+        dbg::statement!("start", stmt, infcx, env, span, &self);
         match stmt {
             GhostStatement::Fold(place) => {
                 env.fold(&mut infcx.at(span), place).with_span(span)?;
@@ -1553,7 +1553,7 @@ impl<'ck, 'genv, 'tcx, M: Mode> Checker<'ck, 'genv, 'tcx, M> {
                     .with_span(span)?;
             }
         }
-        dbg::statement!("end", stmt, infcx, env, span, self.genv);
+        dbg::statement!("end", stmt, infcx, env, span, &self);
         Ok(())
     }
 

crates/flux-refineck/src/type_env.rs

@@ -30,8 +30,10 @@ use flux_rustc_bridge::{
     ty,
 };
 use itertools::{izip, Itertools};
+use rustc_data_structures::unord::UnordMap;
 use rustc_index::IndexSlice;
 use rustc_middle::{mir::RETURN_PLACE, ty::TyCtxt};
+use rustc_span::Symbol;
 use rustc_type_ir::BoundVar;
 use serde::Serialize;
 
@@ -881,25 +883,48 @@ pub struct TypeEnvTrace(Vec<TypeEnvBind>);
 
 #[derive(Serialize)]
 struct TypeEnvBind {
-    loc: String,
+    local: LocInfo,
+    name: Option<String>,
     kind: String,
     ty: String,
 }
 
+#[derive(Serialize)]
+enum LocInfo {
+    Local(String),
+    Var(String),
+}
+
+fn loc_info(loc: &Loc) -> LocInfo {
+    match loc {
+        Loc::Local(local) => LocInfo::Local(format!("{local:?}")),
+        Loc::Var(var) => LocInfo::Var(format!("{var:?}")),
+    }
+}
+
+fn loc_name(local_names: &UnordMap<Local, Symbol>, loc: &Loc) -> Option<String> {
+    if let Loc::Local(local) = loc {
+        let name = local_names.get(local)?;
+        return Some(format!("{}", name));
+    }
+    None
+}
+
 impl TypeEnvTrace {
-    pub fn new(genv: GlobalEnv, env: &TypeEnv) -> Self {
+    pub fn new(genv: GlobalEnv, local_names: &UnordMap<Local, Symbol>, env: &TypeEnv) -> Self {
         let mut bindings = vec![];
-        let cx = PrettyCx::default_with_genv(genv);
+        let cx = PrettyCx::default_with_genv(genv).hide_regions(true);
         env.bindings
             .iter()
             .filter(|(_, binding)| !binding.ty.is_uninit())
             .sorted_by(|(loc1, _), (loc2, _)| loc1.cmp(loc2))
             .for_each(|(loc, binding)| {
-                let loc = format!("{loc:?}");
+                let name = loc_name(local_names, loc);
+                let local = loc_info(loc);
                 let kind = format!("{:?}", binding.kind);
                 let ty = WithCx::new(&cx, binding.ty.clone());
                 let ty = format!("{:?}", ty);
-                bindings.push(TypeEnvBind { loc, kind, ty });
+                bindings.push(TypeEnvBind { name, local, kind, ty });
             });
 
         TypeEnvTrace(bindings)

crates/flux-rustc-bridge/src/mir.rs

@@ -11,13 +11,14 @@ use rustc_borrowck::consumers::{BodyWithBorrowckFacts, BorrowIndex};
 use rustc_data_structures::{
     fx::FxIndexMap,
     graph::{self, dominators::Dominators, DirectedGraph, StartNode},
+    unord::UnordMap,
 };
 use rustc_hir::def_id::{DefId, LocalDefId};
 use rustc_index::IndexSlice;
 use rustc_infer::infer::TyCtxtInferExt;
 use rustc_macros::{TyDecodable, TyEncodable};
 use rustc_middle::{
-    mir,
+    mir::{self, VarDebugInfoContents},
     ty::{FloatTy, IntTy, ParamConst, TyCtxt, TypingMode, UintTy},
 };
 pub use rustc_middle::{
@@ -75,6 +76,7 @@ pub struct Body<'tcx> {
     /// See [`mk_fake_predecessors`]
     fake_predecessors: IndexVec<BasicBlock, usize>,
     body_with_facts: BodyWithBorrowckFacts<'tcx>,
+    pub local_names: UnordMap<Local, Symbol>,
 }
 
 #[derive(Debug)]
@@ -382,14 +384,27 @@ impl<'tcx> Body<'tcx> {
         for (rank, bb) in (0u32..).zip(reverse_post_order) {
             dominator_order_rank[bb] = rank;
         }
-
+        let local_names = body_with_facts
+            .body
+            .var_debug_info
+            .iter()
+            .flat_map(|var_debug_info| {
+                if let VarDebugInfoContents::Place(place) = var_debug_info.value {
+                    let local = place.as_local()?;
+                    Some((local, var_debug_info.name))
+                } else {
+                    None
+                }
+            })
+            .collect();
         Self {
             basic_blocks,
             local_decls,
             infcx,
             fake_predecessors,
             body_with_facts,
             dominator_order_rank,
+            local_names,
         }
     }
 

tools/vscode/src/extension.ts

@@ -257,7 +257,6 @@ function collapseBindings(bindings: RcxBind[]): RcxBind[] {
         }
     }
     if (names.length > 0) { binds.push({name: names, sort: sort}) };
-    // console.log("collapsed bindings", binds);
     return binds;
 }
 
@@ -269,7 +268,7 @@ function parseRcx(rcx: string): Rcx {
 }
 
 function parseEnv(env: string): TypeEnv {
-    return JSON.parse(env);
+    return JSON.parse(env).filter((bind: TypeEnvBind) => bind.name)
 }
 
 class FluxViewProvider implements vscode.WebviewViewProvider {
@@ -381,7 +380,8 @@ class FluxViewProvider implements vscode.WebviewViewProvider {
 
         const envBindings = this._currentEnv?.map(bind => `
             <tr>
-                <td><b style="color: #F26123">${bind.loc}</b> : ${bind.ty} </td>
+                <td><b style="color: #F26123">${bind.name}</b></td>
+                <td>: ${bind.ty}</td>
             </tr>
           `).join('');
 
@@ -509,10 +509,11 @@ async function readFluxCheckerTrace(): Promise<Map<string, LineInfo[]>> {
     }
 }
 
+// TODO: add local-info to TypeEnvBind
 type TypeEnvBind = {
-  loc: String,
-  kind: String,
-  ty: String,
+  name: string | null,
+  kind: string,
+  ty: string,
 }
 type TypeEnv = TypeEnvBind[];