Make checker trace more JSON-ish (#911)

* rename dbg::statement_span

* make checker trace more json-ish

* remove dependency on display json; restore compatibility with logreader.py

Cargo.toml

@@ -37,6 +37,7 @@ home = "0.5.9"
 itertools = "0.13.0"
 pad-adapter = "0.1.1"
 serde = { version = "1.0", features = ["derive"] }
+serde_json = "1.0"
 toml = "0.8"
 toposort-scc = "0.5.4"
 

crates/flux-common/src/dbg.rs

@@ -70,19 +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) => {{
-        tracing::debug!(event = concat!("statement_", $pos), stmt = ?$stmt, rcx = ?$rcx, env = ?$env)
-    }};
-}
-pub use crate::_statement as statement;
-
-#[macro_export]
-macro_rules! _statement_span{
     ($pos:literal, $stmt:expr, $rcx:expr, $env:expr, $span:expr) => {{
-        tracing::debug!(event = concat!("statement_", $pos), stmt = ?$stmt, stmt_span = ?$span, rcx = ?$rcx, env = ?$env)
+        let rcx_json = RefineCtxtTrace::new($rcx);
+        let env_json = TypeEnvTrace::new($env);
+        tracing::debug!(event = concat!("statement_", $pos), stmt = ?$stmt, stmt_span = ?$span, rcx = ?$rcx, env = ?$env, rcx_json = ?rcx_json, env_json = ?env_json)
     }};
 }
-pub use crate::_statement_span as statement_span;
+pub use crate::_statement as statement;
 
 #[macro_export]
 macro_rules! _terminator{

crates/flux-infer/Cargo.toml

@@ -8,11 +8,13 @@ flux-common.workspace = true
 flux-config.workspace = true
 flux-errors.workspace = true
 flux-middle.workspace = true
+flux-macros.workspace = true
 
 liquid-fixpoint.workspace = true
-
 itertools.workspace = true
 pad-adapter.workspace = true
+serde.workspace = true
+serde_json.workspace = true
 
 [lints]
 workspace = true

crates/flux-infer/src/refine_tree.rs

@@ -5,6 +5,7 @@ use std::{
 };
 
 use flux_common::{index::IndexVec, iter::IterExt};
+use flux_macros::DebugAsJson;
 use flux_middle::{
     global_env::GlobalEnv,
     queries::QueryResult,
@@ -18,6 +19,7 @@ use flux_middle::{
 };
 use itertools::Itertools;
 use rustc_hir::def_id::DefId;
+use serde::Serialize;
 
 use crate::{
     fixpoint_encoding::{fixpoint, FixpointCtxt},
@@ -776,3 +778,37 @@ mod pretty {
         Scope,
     );
 }
+
+/// A very explicit representation of [`RefineCtxt`] for debugging/tracing/serialization ONLY.
+#[derive(Serialize, DebugAsJson)]
+pub struct RefineCtxtTrace {
+    bindings: Vec<RcxBind>,
+    exprs: Vec<String>,
+}
+
+#[derive(Serialize)]
+struct RcxBind {
+    name: String,
+    sort: String,
+}
+impl RefineCtxtTrace {
+    pub fn new(rcx: &RefineCtxt) -> Self {
+        let parents = ParentsIter::new(NodePtr::clone(&rcx.ptr)).collect_vec();
+        let mut bindings = vec![];
+        let mut exprs = vec![];
+        parents.into_iter().rev().for_each(|ptr| {
+            let node = ptr.borrow();
+            match &node.kind {
+                NodeKind::ForAll(name, sort) => {
+                    let bind = RcxBind { name: format!("{name:?}"), sort: format!("{sort:?}") };
+                    bindings.push(bind);
+                }
+                NodeKind::Assumption(e) if !e.simplify().is_trivially_true() => {
+                    exprs.push(format!("{e:?}"))
+                }
+                _ => (),
+            }
+        });
+        Self { bindings, exprs }
+    }
+}

crates/flux-macros/src/lib.rs

@@ -4,6 +4,7 @@ mod diagnostics;
 mod fold;
 mod primops;
 
+use quote::quote;
 use synstructure::decl_derive;
 
 decl_derive!(
@@ -63,3 +64,18 @@ pub fn fluent_messages(input: proc_macro::TokenStream) -> proc_macro::TokenStrea
 pub fn primop_rules(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
     primops::primop_rules(input)
 }
+
+decl_derive!(
+    [DebugAsJson] => debug_as_json
+);
+
+fn debug_as_json(s: synstructure::Structure<'_>) -> proc_macro2::TokenStream {
+    s.gen_impl(quote! {
+        gen impl std::fmt::Debug for @Self {
+            fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+                let s = serde_json::to_string(&self).map_err(|_| std::fmt::Error::default())?;
+                write!(f, "{}", s)
+            }
+        }
+    })
+}

crates/flux-refineck/Cargo.toml

@@ -14,7 +14,8 @@ flux-infer.workspace = true
 flux-macros.workspace = true
 flux-middle.workspace = true
 flux-rustc-bridge.workspace = true
-
+serde.workspace = true
+serde_json.workspace = true
 liquid-fixpoint.workspace = true
 
 itertools.workspace = true

crates/flux-refineck/src/checker.rs

@@ -5,7 +5,7 @@ use flux_config as config;
 use flux_infer::{
     fixpoint_encoding::{self, KVarGen},
     infer::{ConstrReason, InferCtxt, InferCtxtRoot, SubtypeReason},
-    refine_tree::{AssumeInvariants, RefineTree, Snapshot},
+    refine_tree::{AssumeInvariants, RefineCtxtTrace, RefineTree, Snapshot},
 };
 use flux_middle::{
     global_env::GlobalEnv,
@@ -51,7 +51,7 @@ use crate::{
     ghost_statements::{GhostStatement, GhostStatements, Point},
     primops,
     queue::WorkQueue,
-    type_env::{BasicBlockEnv, BasicBlockEnvShape, PtrToRefBound, TypeEnv},
+    type_env::{BasicBlockEnv, BasicBlockEnvShape, PtrToRefBound, TypeEnv, TypeEnvTrace},
 };
 
 type Result<T = ()> = std::result::Result<T, CheckerError>;
@@ -509,9 +509,9 @@ impl<'ck, 'genv, 'tcx, M: Mode> Checker<'ck, 'genv, 'tcx, M> {
                 span,
             )?;
             bug::track_span(span, || {
-                dbg::statement_span!("start", stmt, infcx, env, span);
+                dbg::statement!("start", stmt, &infcx, &env, span);
                 self.check_statement(&mut infcx, &mut env, stmt)?;
-                dbg::statement_span!("end", stmt, infcx, env, span);
+                dbg::statement!("end", stmt, &infcx, &env, span);
                 Ok(())
             })?;
             if !stmt.is_nop() {
@@ -1534,7 +1534,7 @@ impl<'ck, 'genv, 'tcx, M: Mode> Checker<'ck, 'genv, 'tcx, M> {
         stmt: &GhostStatement,
         span: Span,
     ) -> Result {
-        dbg::statement!("start", stmt, infcx, env);
+        dbg::statement!("start", stmt, infcx, env, span);
         match stmt {
             GhostStatement::Fold(place) => {
                 env.fold(&mut infcx.at(span), place).with_span(span)?;
@@ -1549,7 +1549,7 @@ impl<'ck, 'genv, 'tcx, M: Mode> Checker<'ck, 'genv, 'tcx, M> {
                     .with_span(span)?;
             }
         }
-        dbg::statement!("end", stmt, infcx, env);
+        dbg::statement!("end", stmt, infcx, env, span);
         Ok(())
     }
 

crates/flux-refineck/src/type_env.rs

@@ -8,6 +8,7 @@ use flux_infer::{
     infer::{ConstrReason, InferCtxt, InferCtxtAt, InferResult},
     refine_tree::{AssumeInvariants, RefineCtxt, Scope},
 };
+use flux_macros::DebugAsJson;
 use flux_middle::{
     global_env::GlobalEnv,
     queries::QueryResult,
@@ -31,6 +32,7 @@ use itertools::{izip, Itertools};
 use rustc_index::IndexSlice;
 use rustc_middle::{mir::RETURN_PLACE, ty::TyCtxt};
 use rustc_type_ir::BoundVar;
+use serde::Serialize;
 
 use self::place_ty::{LocKind, PlacesTree};
 use super::rty::Sort;
@@ -873,3 +875,32 @@ mod pretty {
         BasicBlockEnv => "basic_block_env"
     }
 }
+
+/// A very explicit representation of [`TypeEnv`] for debugging/tracing/serialization ONLY.
+#[derive(Serialize, DebugAsJson)]
+pub struct TypeEnvTrace(Vec<TypeEnvBind>);
+
+#[derive(Serialize)]
+struct TypeEnvBind {
+    loc: String,
+    kind: String,
+    ty: String,
+}
+
+impl TypeEnvTrace {
+    pub fn new<'a>(env: &TypeEnv<'a>) -> Self {
+        let mut bindings = vec![];
+        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);
+                bindings.push(TypeEnvBind { loc, kind, ty });
+            });
+
+        TypeEnvTrace(bindings)
+    }
+}