cleanup pprint

crates/flux-infer/src/refine_tree.rs

@@ -752,7 +752,7 @@ impl RefineCtxtTrace {
         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);
+        let cx = &PrettyCx::default(genv);
 
         parents.into_iter().rev().for_each(|ptr| {
             let node = ptr.borrow();

crates/flux-middle/src/big_int.rs

@@ -4,12 +4,12 @@ use rustc_macros::{Decodable, Encodable};
 
 /// A signed integer in the range [-2^128, 2^128], represented by a `u128` and an explicit sign.
 ///
-/// In the logic, we work with natural numbers so we could represent them with arbitrary precision
+/// In the logic, we work mathematical integers so we could represent them with arbitrary precision
 /// integers. We instead take the simpler approach of using a fixed size representation that allows
-/// us to store any Rust literal, i.e., we can represent both `i128::MIN` and `u128::MAX`. Since we
-/// never do arithmetic, this representation doesn't present any problems. We may choose to change
-/// the representation in the future (and use arbitrary precision integers) if this ever become a
-/// problem, e.g., if we want to do (precise) arithmetic during constant folding.
+/// us to store any Rust literal, i.e., we can represent both `i128::MIN` and `u128::MAX`. This works
+/// because we never do arithmetic. We can change the representation in the future (and use arbitrary
+/// precision integers) if this ever becomes a problem, e.g., if we want to do (precise) arithmetic
+/// during constant folding.
 #[derive(Clone, Debug, Copy, PartialEq, Eq, Hash, Encodable, Decodable, PartialOrd, Ord)]
 pub struct BigInt {
     sign: Sign,

crates/flux-middle/src/pretty.rs

@@ -146,98 +146,54 @@ pub enum KVarArgs {
     Hide,
 }
 
-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,
-    pub tags: bool,
-    pub bindings_chain: bool,
-    pub preds_chain: bool,
-    pub full_spans: bool,
-    pub hide_uninit: bool,
-    pub hide_refinements: bool,
-    pub hide_regions: bool,
-    pub hide_sorts: bool,
-    env: BoundVarEnv,
+#[derive(Clone, Copy)]
+pub enum GenvOrTcx<'genv, 'tcx> {
+    Genv(GlobalEnv<'genv, 'tcx>),
+    Tcx(TyCtxt<'tcx>),
 }
 
-newtype_index! {
-    /// Name used during pretty printing to format anonymous bound variables
-    #[debug_format = "b{}"]
-    struct BoundVarName {}
-}
-
-#[derive(Default)]
-struct BoundVarEnv {
-    name_gen: IndexGen<BoundVarName>,
-    layers: RefCell<Vec<UnordMap<BoundVar, BoundVarName>>>,
-}
-
-impl BoundVarEnv {
-    fn lookup(&self, debruijn: DebruijnIndex, var: BoundVar) -> Option<BoundVarName> {
-        let layers = self.layers.borrow();
-        layers
-            .get(layers.len().checked_sub(debruijn.as_usize() + 1)?)?
-            .get(&var)
-            .copied()
-    }
-
-    fn push_layer(&self, vars: &[BoundVariableKind]) {
-        let mut layer = UnordMap::default();
-        for (idx, var) in vars.iter().enumerate() {
-            if let BoundVariableKind::Refine(_, _, BoundReftKind::Annon) = var {
-                layer.insert(BoundVar::from_usize(idx), self.name_gen.fresh());
-            }
+impl<'genv, 'tcx> GenvOrTcx<'genv, 'tcx> {
+    fn tcx(self) -> TyCtxt<'tcx> {
+        match self {
+            GenvOrTcx::Genv(genv) => genv.tcx(),
+            GenvOrTcx::Tcx(tcx) => tcx,
         }
-        self.layers.borrow_mut().push(layer);
     }
 
-    fn pop_layer(&self) {
-        self.layers.borrow_mut().pop();
-    }
-}
-
-pub struct WithCx<'a, 'genv, 'tcx, T> {
-    data: T,
-    cx: &'a PrettyCx<'genv, 'tcx>,
-}
-
-pub struct Join<'a, I> {
-    sep: &'a str,
-    iter: RefCell<Option<I>>,
-}
-
-pub struct Parens<'a, T> {
-    val: &'a T,
-    parenthesize: bool,
-}
-
-pub trait Pretty {
-    fn fmt(&self, cx: &PrettyCx, f: &mut fmt::Formatter<'_>) -> fmt::Result;
-
-    fn default_cx(tcx: TyCtxt) -> PrettyCx {
-        PrettyCx::default(tcx)
+    fn genv(self) -> Option<GlobalEnv<'genv, 'tcx>> {
+        match self {
+            GenvOrTcx::Genv(genv) => Some(genv),
+            GenvOrTcx::Tcx(_) => None,
+        }
     }
 }
 
-impl Pretty for String {
-    fn fmt(&self, _cx: &PrettyCx, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        write!(f, "{}", self)
+impl<'tcx> From<TyCtxt<'tcx>> for GenvOrTcx<'_, 'tcx> {
+    fn from(v: TyCtxt<'tcx>) -> Self {
+        Self::Tcx(v)
     }
 }
 
-impl<'a, I> Join<'a, I> {
-    pub fn new<T: IntoIterator<IntoIter = I>>(sep: &'a str, iter: T) -> Self {
-        Self { sep, iter: RefCell::new(Some(iter.into_iter())) }
+impl<'genv, 'tcx> From<GlobalEnv<'genv, 'tcx>> for GenvOrTcx<'genv, 'tcx> {
+    fn from(v: GlobalEnv<'genv, 'tcx>) -> Self {
+        Self::Genv(v)
     }
 }
 
-impl<'a, T> Parens<'a, T> {
-    pub fn new(val: &'a T, parenthesize: bool) -> Self {
-        Self { val, parenthesize }
-    }
+pub struct PrettyCx<'genv, 'tcx> {
+    pub cx: GenvOrTcx<'genv, 'tcx>,
+    pub kvar_args: KVarArgs,
+    pub fully_qualified_paths: bool,
+    pub simplify_exprs: bool,
+    pub tags: bool,
+    pub bindings_chain: bool,
+    pub preds_chain: bool,
+    pub full_spans: bool,
+    pub hide_uninit: bool,
+    pub hide_refinements: bool,
+    pub hide_regions: bool,
+    pub hide_sorts: bool,
+    env: BoundVarEnv,
 }
 
 macro_rules! set_opts {
@@ -251,10 +207,9 @@ macro_rules! set_opts {
 }
 
 impl<'genv, 'tcx> PrettyCx<'genv, 'tcx> {
-    pub fn default(tcx: TyCtxt<'tcx>) -> Self {
+    pub fn default(cx: impl Into<GenvOrTcx<'genv, 'tcx>>) -> Self {
         PrettyCx {
-            tcx,
-            genv: None,
+            cx: cx.into(),
             kvar_args: KVarArgs::SelfOnly,
             fully_qualified_paths: false,
             simplify_exprs: true,
@@ -270,9 +225,12 @@ impl<'genv, 'tcx> PrettyCx<'genv, 'tcx> {
         }
     }
 
-    pub fn default_with_genv(genv: GlobalEnv<'genv, 'tcx>) -> Self {
-        let def = Self::default(genv.tcx());
-        Self { genv: Some(genv), ..def }
+    pub fn tcx(&self) -> TyCtxt<'tcx> {
+        self.cx.tcx()
+    }
+
+    pub fn genv(&self) -> Option<GlobalEnv<'genv, 'tcx>> {
+        self.cx.genv()
     }
 
     pub fn merge(&mut self, opts: &config::Value) {
@@ -379,6 +337,83 @@ impl<'genv, 'tcx> PrettyCx<'genv, 'tcx> {
     }
 }
 
+newtype_index! {
+    /// Name used during pretty printing to format anonymous bound variables
+    #[debug_format = "b{}"]
+    struct BoundVarName {}
+}
+
+#[derive(Default)]
+struct BoundVarEnv {
+    name_gen: IndexGen<BoundVarName>,
+    layers: RefCell<Vec<UnordMap<BoundVar, BoundVarName>>>,
+}
+
+impl BoundVarEnv {
+    fn lookup(&self, debruijn: DebruijnIndex, var: BoundVar) -> Option<BoundVarName> {
+        let layers = self.layers.borrow();
+        layers
+            .get(layers.len().checked_sub(debruijn.as_usize() + 1)?)?
+            .get(&var)
+            .copied()
+    }
+
+    fn push_layer(&self, vars: &[BoundVariableKind]) {
+        let mut layer = UnordMap::default();
+        for (idx, var) in vars.iter().enumerate() {
+            if let BoundVariableKind::Refine(_, _, BoundReftKind::Annon) = var {
+                layer.insert(BoundVar::from_usize(idx), self.name_gen.fresh());
+            }
+        }
+        self.layers.borrow_mut().push(layer);
+    }
+
+    fn pop_layer(&self) {
+        self.layers.borrow_mut().pop();
+    }
+}
+
+pub struct WithCx<'a, 'genv, 'tcx, T> {
+    data: T,
+    cx: &'a PrettyCx<'genv, 'tcx>,
+}
+
+pub struct Join<'a, I> {
+    sep: &'a str,
+    iter: RefCell<Option<I>>,
+}
+
+pub struct Parens<'a, T> {
+    val: &'a T,
+    parenthesize: bool,
+}
+
+pub trait Pretty {
+    fn fmt(&self, cx: &PrettyCx, f: &mut fmt::Formatter<'_>) -> fmt::Result;
+
+    fn default_cx(tcx: TyCtxt) -> PrettyCx {
+        PrettyCx::default(tcx)
+    }
+}
+
+impl Pretty for String {
+    fn fmt(&self, _cx: &PrettyCx, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "{}", self)
+    }
+}
+
+impl<'a, I> Join<'a, I> {
+    pub fn new<T: IntoIterator<IntoIter = I>>(sep: &'a str, iter: T) -> Self {
+        Self { sep, iter: RefCell::new(Some(iter.into_iter())) }
+    }
+}
+
+impl<'a, T> Parens<'a, T> {
+    pub fn new(val: &'a T, parenthesize: bool) -> Self {
+        Self { val, parenthesize }
+    }
+}
+
 impl<'a, 'genv, 'tcx, T> WithCx<'a, 'genv, 'tcx, T> {
     pub fn new(cx: &'a PrettyCx<'genv, 'tcx>, data: T) -> Self {
         Self { data, cx }
@@ -458,9 +493,9 @@ impl<T: Pretty> fmt::Debug for WithCx<'_, '_, '_, T> {
 
 impl Pretty for DefId {
     fn fmt(&self, cx: &PrettyCx, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        let path = cx.tcx.def_path(*self);
+        let path = cx.tcx().def_path(*self);
         if cx.fully_qualified_paths {
-            let krate = cx.tcx.crate_name(self.krate);
+            let krate = cx.tcx().crate_name(self.krate);
             w!(cx, f, "{}{}", ^krate, ^path.to_string_no_crate_verbose())
         } else {
             w!(cx, f, "{}", ^path.data.last().unwrap())
@@ -479,7 +514,7 @@ impl Pretty for Span {
         if cx.full_spans {
             write!(f, "{self:?}")
         } else {
-            let src_map = cx.tcx.sess.source_map();
+            let src_map = cx.tcx().sess.source_map();
             let lo = src_map.lookup_char_pos(self.lo());
             let hi = src_map.lookup_char_pos(self.hi());
             // use rustc_span::FileName;

crates/flux-middle/src/rty/expr.rs

@@ -1186,7 +1186,7 @@ pub(crate) mod pretty {
                         w!(cx, f, "({:?})", e)?;
                     };
                     // proj
-                    if let Some(genv) = cx.genv
+                    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)
                     {
@@ -1203,7 +1203,7 @@ pub(crate) mod pretty {
                     }
                 }
                 ExprKind::Aggregate(AggregateKind::Adt(def_id), flds) => {
-                    if let Some(genv) = cx.genv
+                    if let Some(genv) = cx.genv()
                         && let Ok(adt_sort_def) = genv.adt_sort_def_of(def_id)
                     {
                         let field_binds = adt_sort_def
@@ -1394,7 +1394,7 @@ pub(crate) mod pretty {
         flds: &[Expr],
         is_named: bool,
     ) -> Result<NestedString, fmt::Error> {
-        let keys = if let Some(genv) = cx.genv
+        let keys = if let Some(genv) = cx.genv()
             && let Ok(adt_sort_def) = genv.adt_sort_def_of(def_id)
         {
             adt_sort_def
@@ -1471,7 +1471,7 @@ pub(crate) mod pretty {
                 ExprKind::FieldProj(e, proj) => {
                     let e_d = e.fmt_nested(cx)?;
                     let proj_text =   // proj
-                if let Some(genv) = cx.genv
+                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)
                 {

crates/flux-middle/src/rty/mod.rs

@@ -1565,10 +1565,6 @@ impl BaseTy {
         matches!(self, BaseTy::Slice(..))
     }
 
-    fn is_adt(&self) -> bool {
-        matches!(self, BaseTy::Adt(..))
-    }
-
     pub fn is_box(&self) -> bool {
         matches!(self, BaseTy::Adt(adt_def, _) if adt_def.is_box())
     }

crates/flux-middle/src/rty/pretty.rs

@@ -235,10 +235,9 @@ impl PrettyNested for IdxFmt {
 
 impl Pretty for IdxFmt {
     fn fmt(&self, cx: &PrettyCx, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        let e = &self.0;
-        let e = if cx.simplify_exprs { e.simplify() } else { e.clone() };
+        let e = if cx.simplify_exprs { self.0.simplify() } else { self.0.clone() };
         if let ExprKind::Aggregate(AggregateKind::Adt(def_id), flds) = e.kind()
-            && let Some(genv) = cx.genv
+            && 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)
@@ -254,16 +253,14 @@ impl Pretty for Ty {
     fn fmt(&self, cx: &PrettyCx, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         match self.kind() {
             TyKind::Indexed(bty, idx) => {
-                w!(cx, f, "{:?}", parens!(bty, !bty.is_atom()))?;
                 if cx.hide_refinements {
+                    w!(cx, f, "{:?}", bty)?;
                     return Ok(());
                 }
                 if idx.is_unit() {
-                    if bty.is_adt() {
-                        w!(cx, f, "[]")?;
-                    }
+                    w!(cx, f, "{:?}", bty)?;
                 } else {
-                    w!(cx, f, "[{:?}]", IdxFmt(idx.clone()))?;
+                    w!(cx, f, "{:?}[{:?}]", parens!(bty, !bty.is_atom()), IdxFmt(idx.clone()))?;
                 }
                 Ok(())
             }
@@ -455,9 +452,9 @@ fn fmt_alias_ty(
             }
         }
         AliasKind::Projection => {
-            let assoc_name = cx.tcx.item_name(alias_ty.def_id);
-            let trait_ref = cx.tcx.parent(alias_ty.def_id);
-            let trait_generic_count = cx.tcx.generics_of(trait_ref).count() - 1;
+            let assoc_name = cx.tcx().item_name(alias_ty.def_id);
+            let trait_ref = cx.tcx().parent(alias_ty.def_id);
+            let trait_generic_count = cx.tcx().generics_of(trait_ref).count() - 1;
 
             let [self_ty, args @ ..] = &alias_ty.args[..] else {
                 return w!(cx, f, "<alias_ty>");

crates/flux-refineck/src/type_env.rs

@@ -919,7 +919,7 @@ impl TypeEnvTrace {
         env: &TypeEnv,
     ) -> Self {
         let mut bindings = vec![];
-        let cx = PrettyCx::default_with_genv(genv).hide_regions(true);
+        let cx = PrettyCx::default(genv).hide_regions(true);
         env.bindings
             .iter()
             .filter(|(_, binding)| !binding.ty.is_uninit())