From 99066aa86feda0f7e8e077bbb4a98201d7525d34 Mon Sep 17 00:00:00 2001
From: Nico Lehmann <nico.lehmannm@gmail.com>
Date: Fri, 6 Dec 2024 11:34:18 -0800
Subject: [PATCH] fix formatting of RefineCtxt

---
 crates/flux-infer/src/refine_tree.rs | 50 ++++++++++------------------
 crates/flux-middle/src/pretty.rs     |  8 +++++
 2 files changed, 26 insertions(+), 32 deletions(-)

diff --git a/crates/flux-infer/src/refine_tree.rs b/crates/flux-infer/src/refine_tree.rs
index f6295ff065..462ac0e4a6 100644
--- a/crates/flux-infer/src/refine_tree.rs
+++ b/crates/flux-infer/src/refine_tree.rs
@@ -720,39 +720,25 @@ mod pretty {
     impl Pretty for RefineCtxt<'_> {
         fn fmt(&self, cx: &PrettyCx, f: &mut fmt::Formatter<'_>) -> fmt::Result {
             define_scoped!(cx, f);
-            let parents = ParentsIter::new(NodePtr::clone(&self.ptr)).collect_vec();
-            write!(
-                f,
-                "{{{}}}",
-                parents
-                    .into_iter()
-                    .rev()
-                    .filter(|ptr| {
-                        let node = ptr.borrow();
-                        match &node.kind {
-                            NodeKind::ForAll(..) => true,
-                            NodeKind::Assumption(e) => !e.simplify().is_trivially_true(),
-                            NodeKind::Root(_) => true,
-                            _ => false,
-                        }
-                    })
-                    .format_with(", ", |n, f| {
-                        let n = n.borrow();
-                        match &n.kind {
-                            NodeKind::Root(bindings) => {
-                                for (name, sort) in bindings {
-                                    f(&format_args_cx!("{:?}: {:?}", ^name, sort))?;
-                                }
-                                Ok(())
-                            }
-                            NodeKind::ForAll(name, sort) => {
-                                f(&format_args_cx!("{:?}: {:?}", ^name, sort))
-                            }
-                            NodeKind::Assumption(pred) => f(&format_args_cx!("{:?}", pred)),
-                            _ => unreachable!(),
+            let mut elements = vec![];
+            for node in ParentsIter::new(NodePtr::clone(&self.ptr)) {
+                let n = node.borrow();
+                match &n.kind {
+                    NodeKind::Root(bindings) => {
+                        for (name, sort) in bindings {
+                            elements.push(format_cx!("{:?} {:?}", ^name, sort));
                         }
-                    })
-            )
+                    }
+                    NodeKind::ForAll(name, sort) => {
+                        elements.push(format_cx!("{:?}: {:?}", ^name, sort));
+                    }
+                    NodeKind::Assumption(pred) => {
+                        elements.push(format_cx!("{:?}", pred));
+                    }
+                    _ => {}
+                }
+            }
+            write!(f, "{{{}}}", elements.into_iter().rev().format(", "))
         }
     }
 
diff --git a/crates/flux-middle/src/pretty.rs b/crates/flux-middle/src/pretty.rs
index 51ec855ce2..feaecb5059 100644
--- a/crates/flux-middle/src/pretty.rs
+++ b/crates/flux-middle/src/pretty.rs
@@ -63,6 +63,14 @@ macro_rules! _format_args_cx {
 }
 pub use crate::_format_args_cx as format_args_cx;
 
+#[macro_export]
+macro_rules! _format_cx {
+    ($($arg:tt)*) => {
+        std::fmt::format($crate::_format_args_cx!($($arg)*))
+    }
+}
+pub use crate::_format_cx as format_cx;
+
 #[macro_export]
 macro_rules! _w {
     ($fmt:literal, $($args:tt)*) => {