Pretty print records using field names (#927)

* get GlobalEnv inside PrettyCx * done pretty record * only drop struct name for top-level index * add comment * remove collect
flux-rs · Dec 6, 2024 · 184296b · 184296b
1 parent a6fed47
commit 184296b
Show file tree

Hide file tree

Showing 8 changed files with 116 additions and 31 deletions.
diff --git a/crates/flux-common/src/dbg.rs b/crates/flux-common/src/dbg.rs
@@ -70,10 +70,11 @@ 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) => {{
+    ($pos:literal, $stmt:expr, $rcx:expr, $env:expr, $span:expr, $genv:expr) => {{
         if config::dump_checker_trace() {
-            let rcx_json = RefineCtxtTrace::new($rcx);
-            let env_json = TypeEnvTrace::new($env);
+            let genv = $genv;
+            let rcx_json = RefineCtxtTrace::new(genv, $rcx);
+            let env_json = TypeEnvTrace::new(genv, $env);
             tracing::info!(event = concat!("statement_", $pos), stmt = ?$stmt, stmt_span = ?$span, rcx = ?$rcx, env = ?$env, rcx_json = ?rcx_json, env_json = ?env_json)
         }
     }};

diff --git a/crates/flux-infer/src/refine_tree.rs b/crates/flux-infer/src/refine_tree.rs
@@ -8,6 +8,7 @@ use flux_common::{index::IndexVec, iter::IterExt};
 use flux_macros::DebugAsJson;
 use flux_middle::{
     global_env::GlobalEnv,
+    pretty::{PrettyCx, WithCx},
     queries::QueryResult,
     rty::{
         canonicalize::{Hoister, HoisterDelegate},
@@ -799,10 +800,12 @@ struct RcxBind {
     sort: String,
 }
 impl RefineCtxtTrace {
-    pub fn new(rcx: &RefineCtxt) -> Self {
+    pub fn new(genv: GlobalEnv, rcx: &RefineCtxt) -> Self {
         let parents = ParentsIter::new(NodePtr::clone(&rcx.ptr)).collect_vec();
         let mut bindings = vec![];
         let mut exprs = vec![];
+        let cx = PrettyCx::default_with_genv(genv);
+
         parents.into_iter().rev().for_each(|ptr| {
             let node = ptr.borrow();
             match &node.kind {
@@ -811,10 +814,11 @@ impl RefineCtxtTrace {
                     bindings.push(bind);
                 }
                 NodeKind::Assumption(e) if !e.simplify().is_trivially_true() => {
+                    let e = WithCx::new(&cx, e);
                     exprs.push(format!("{e:?}"));
                 }
-                NodeKind::Root(bs) => {
-                    for (name, sort) in bs {
+                NodeKind::Root(binds) => {
+                    for (name, sort) in binds {
                         let bind = RcxBind { name: format!("{name:?}"), sort: format!("{sort:?}") };
                         bindings.push(bind);
                     }

diff --git a/crates/flux-middle/src/pretty.rs b/crates/flux-middle/src/pretty.rs
@@ -149,7 +149,10 @@ pub fn pprint_with_default_cx<T: Pretty>(
 }
 
 pub use crate::_impl_debug_with_default_cx as impl_debug_with_default_cx;
-use crate::rty::{BoundReft, BoundReftKind, BoundVariableKind};
+use crate::{
+    global_env::GlobalEnv,
+    rty::{BoundReft, BoundReftKind, BoundVariableKind},
+};
 
 #[derive(Copy, Clone)]
 pub enum KVarArgs {
@@ -158,8 +161,9 @@ pub enum KVarArgs {
     Hide,
 }
 
-pub struct PrettyCx<'tcx> {
+pub struct PrettyCx<'genv, 'tcx> {
     pub tcx: TyCtxt<'tcx>,
+    pub genv: Option<GlobalEnv<'genv, 'tcx>>,
     pub kvar_args: KVarArgs,
     pub fully_qualified_paths: bool,
     pub simplify_exprs: bool,
@@ -210,9 +214,9 @@ impl Env {
     }
 }
 
-pub struct WithCx<'a, 'tcx, T> {
+pub struct WithCx<'a, 'genv, 'tcx, T> {
     data: T,
-    cx: &'a PrettyCx<'tcx>,
+    cx: &'a PrettyCx<'genv, 'tcx>,
 }
 
 pub struct Join<'a, I> {
@@ -255,10 +259,11 @@ macro_rules! set_opts {
     };
 }
 
-impl PrettyCx<'_> {
-    pub fn default(tcx: TyCtxt) -> PrettyCx {
+impl<'genv, 'tcx> PrettyCx<'genv, 'tcx> {
+    pub fn default(tcx: TyCtxt<'tcx>) -> Self {
         PrettyCx {
             tcx,
+            genv: None,
             kvar_args: KVarArgs::SelfOnly,
             fully_qualified_paths: false,
             simplify_exprs: true,
@@ -274,6 +279,11 @@ impl PrettyCx<'_> {
         }
     }
 
+    pub fn default_with_genv(genv: GlobalEnv<'genv, 'tcx>) -> Self {
+        let def = Self::default(genv.tcx());
+        Self { genv: Some(genv), ..def }
+    }
+
     pub fn merge(&mut self, opts: &config::Value) {
         set_opts!(
             self,
@@ -380,8 +390,8 @@ impl PrettyCx<'_> {
     }
 }
 
-impl<'a, 'tcx, T> WithCx<'a, 'tcx, T> {
-    pub fn new(cx: &'a PrettyCx<'tcx>, data: T) -> Self {
+impl<'a, 'genv, 'tcx, T> WithCx<'a, 'genv, 'tcx, T> {
+    pub fn new(cx: &'a PrettyCx<'genv, 'tcx>, data: T) -> Self {
         Self { data, cx }
     }
 }
@@ -451,7 +461,7 @@ where
     }
 }
 
-impl<T: Pretty> fmt::Debug for WithCx<'_, '_, T> {
+impl<T: Pretty> fmt::Debug for WithCx<'_, '_, '_, T> {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         <T as Pretty>::fmt(&self.data, self.cx, f)
     }

diff --git a/crates/flux-middle/src/rty/expr.rs b/crates/flux-middle/src/rty/expr.rs
@@ -28,6 +28,7 @@ use crate::{
     big_int::BigInt,
     fhir::SpecFuncKind,
     global_env::GlobalEnv,
+    pretty::*,
     queries::QueryResult,
     rty::{
         fold::{TypeFoldable, TypeFolder, TypeSuperFoldable},
@@ -1077,11 +1078,22 @@ impl From<char> for Constant {
 impl_internable!(ExprKind);
 impl_slice_internable!(Expr, KVar);
 
+#[derive(Debug)]
+pub struct FieldBind<T> {
+    pub name: Symbol,
+    pub value: T,
+}
+
+impl<T: Pretty> Pretty for FieldBind<T> {
+    fn fmt(&self, cx: &PrettyCx, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        define_scoped!(cx, f);
+        w!("{}: {:?}", ^self.name, &self.value)
+    }
+}
 mod pretty {
     use flux_rustc_bridge::def_id_to_string;
 
     use super::*;
-    use crate::pretty::*;
 
     #[derive(PartialEq, Eq, PartialOrd, Ord)]
     enum Precedence {
@@ -1163,10 +1175,20 @@ mod pretty {
                     }
                 }
                 ExprKind::FieldProj(e, proj) => {
+                    // base
                     if e.is_atom() {
-                        w!("{:?}.{:?}", e, ^proj.field_idx())
+                        w!("{:?}", e)?;
                     } else {
-                        w!("({:?}).{:?}", e, ^proj.field_idx())
+                        w!("({:?})", e)?;
+                    };
+                    // proj
+                    if let Some(genv) = cx.genv
+                        && let FieldProj::Adt { def_id, field } = proj
+                        && let Ok(adt_sort_def) = genv.adt_sort_def_of(def_id)
+                    {
+                        w!(".{}", ^adt_sort_def.field_names()[*field as usize])
+                    } else {
+                        w!(".{:?}", ^proj.field_idx())
                     }
                 }
                 ExprKind::Aggregate(AggregateKind::Tuple(_), flds) => {
@@ -1177,7 +1199,19 @@ mod pretty {
                     }
                 }
                 ExprKind::Aggregate(AggregateKind::Adt(def_id), flds) => {
-                    w!("{:?} {{ {:?} }}", def_id, join!(", ", flds))
+                    if let Some(genv) = cx.genv
+                        && let Ok(adt_sort_def) = genv.adt_sort_def_of(def_id)
+                    {
+                        let field_binds = adt_sort_def
+                            .field_names()
+                            .iter()
+                            .zip(flds)
+                            .map(|(name, value)| FieldBind { name: *name, value: value.clone() })
+                            .collect_vec();
+                        w!("{:?} {{ {:?} }}", def_id, join!(", ", field_binds))
+                    } else {
+                        w!("{:?} {{ {:?} }}", def_id, join!(", ", flds))
+                    }
                 }
                 ExprKind::PathProj(e, field) => {
                     if e.is_atom() {

diff --git a/crates/flux-middle/src/rty/pretty.rs b/crates/flux-middle/src/rty/pretty.rs
@@ -1,5 +1,6 @@
-use std::fmt;
+use std::{fmt, iter};
 
+use expr::FieldBind;
 use flux_rustc_bridge::ty::region_to_string;
 use rustc_type_ir::DebruijnIndex;
 use ty::{UnevaluatedConst, ValTree};
@@ -235,6 +236,28 @@ impl Pretty for SubsetTy {
     }
 }
 
+// This is a trick to avoid pretty printing `S [S { x: 10, y: 20}]`
+// and instead just print `S[{x: 10, y: 20}]` for struct-valued indices.
+struct IdxFmt(Expr);
+
+impl Pretty for IdxFmt {
+    fn fmt(&self, cx: &PrettyCx, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        define_scoped!(cx, f);
+        let e = &self.0;
+        let e = if cx.simplify_exprs { e.simplify() } else { e.clone() };
+        if let ExprKind::Aggregate(AggregateKind::Adt(def_id), flds) = e.kind()
+            && let Some(genv) = cx.genv
+            && let Ok(adt_sort_def) = genv.adt_sort_def_of(def_id)
+        {
+            let field_binds = iter::zip(adt_sort_def.field_names(), flds)
+                .map(|(name, value)| FieldBind { name: *name, value: value.clone() });
+            w!("{{ {:?} }}", join!(", ", field_binds))
+        } else {
+            w!("{:?}", &self.0)
+        }
+    }
+}
+
 impl Pretty for Ty {
     fn fmt(&self, cx: &PrettyCx, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         define_scoped!(cx, f);
@@ -249,7 +272,7 @@ impl Pretty for Ty {
                         w!("[]")?;
                     }
                 } else {
-                    w!("[{:?}]", idx)?;
+                    w!("[{:?}]", IdxFmt(idx.clone()))?;
                 }
                 Ok(())
             }
@@ -277,8 +300,19 @@ impl Pretty for Ty {
             }
             TyKind::Param(param_ty) => w!("{}", ^param_ty),
             TyKind::Downcast(adt, .., variant_idx, fields) => {
-                w!("{:?}::{}", adt.did(), ^adt.variant(*variant_idx).name)?;
-                if !fields.is_empty() {
+                let is_struct = adt.is_struct();
+                // base-name
+                w!("{:?}", adt.did())?;
+                // variant-name: if it is not a struct
+                if !is_struct {
+                    w!("::{}", ^adt.variant(*variant_idx).name)?;
+                }
+                // fields: use curly-braces + names for structs, otherwise use parens
+                if is_struct {
+                    let field_binds = iter::zip(&adt.variant(*variant_idx).fields, fields)
+                        .map(|(field_def, value)| FieldBind { name: field_def.name, value });
+                    w!(" {{ {:?} }}", join!(", ", field_binds))?;
+                } else if !fields.is_empty() {
                     w!("({:?})", join!(", ", fields))?;
                 }
                 Ok(())

diff --git a/crates/flux-refineck/src/checker.rs b/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);
+                dbg::statement!("start", stmt, &infcx, &env, span, self.genv);
                 self.check_statement(&mut infcx, &mut env, stmt)?;
-                dbg::statement!("end", stmt, &infcx, &env, span);
+                dbg::statement!("end", stmt, &infcx, &env, span, self.genv);
                 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);
+        dbg::statement!("start", stmt, infcx, env, span, self.genv);
         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);
+        dbg::statement!("end", stmt, infcx, env, span, self.genv);
         Ok(())
     }
 

diff --git a/crates/flux-refineck/src/type_env.rs b/crates/flux-refineck/src/type_env.rs
@@ -11,6 +11,7 @@ use flux_infer::{
 use flux_macros::DebugAsJson;
 use flux_middle::{
     global_env::GlobalEnv,
+    pretty::{PrettyCx, WithCx},
     queries::QueryResult,
     rty::{
         canonicalize::{Hoister, LocalHoister},
@@ -886,16 +887,18 @@ struct TypeEnvBind {
 }
 
 impl TypeEnvTrace {
-    pub fn new(env: &TypeEnv) -> Self {
+    pub fn new(genv: GlobalEnv, env: &TypeEnv) -> Self {
         let mut bindings = vec![];
+        let cx = PrettyCx::default_with_genv(genv);
         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 kind = format!("{:?}", binding.kind);
-                let ty = format!("{:?}", binding.ty);
+                let ty = WithCx::new(&cx, binding.ty.clone());
+                let ty = format!("{:?}", ty);
                 bindings.push(TypeEnvBind { loc, kind, ty });
             });
 

diff --git a/tools/vscode/src/extension.ts b/tools/vscode/src/extension.ts
@@ -257,7 +257,7 @@ function collapseBindings(bindings: RcxBind[]): RcxBind[] {
         }
     }
     if (names.length > 0) { binds.push({name: names, sort: sort}) };
-    console.log("collapsed bindings", binds);
+    // console.log("collapsed bindings", binds);
     return binds;
 }
 
@@ -366,7 +366,6 @@ class FluxViewProvider implements vscode.WebviewViewProvider {
     private _getHtmlForInfo() {
         const rcxBindings = this._currentRcx?.bindings.map(bind => {
           const name = typeof bind.name === 'string' ? bind.name : bind.name.join(' ');
-          console.log("bind", bind, name);
           return `
             <tr>
                 <td><b style="color: #F26123">${name}</b> : ${bind.sort} </td>