Improve pretty printing of valtrees for references

2 files changed, 86 insertions, 0 deletions

diff --git a/compiler/rustc_ty_utils/src/consts.rs b/compiler/rustc_ty_utils/src/consts.rs
new file mode 100644
index 00000000000..22cfd7e1985
--- /dev/null
+++ b/compiler/rustc_ty_utils/src/consts.rs
@@ -0,0 +1,84 @@
+use rustc_middle::ty::{self, TyCtxt};
+use rustc_target::abi::VariantIdx;
+
+/// Tries to destructure constants of type Array or Adt into the constants
+/// of its fields.
+pub(crate) fn destructure_const<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    const_: ty::Const<'tcx>,
+) -> ty::DestructuredConst<'tcx> {
+    if let ty::ConstKind::Value(valtree) = const_.kind() {
+        let branches = match valtree {
+            ty::ValTree::Branch(b) => b,
+            _ => bug!("cannot destructure constant {:?}", const_),
+        };
+
+        let (fields, variant) = match const_.ty().kind() {
+            ty::Array(inner_ty, _) | ty::Slice(inner_ty) => {
+                // construct the consts for the elements of the array/slice
+                let field_consts = branches
+                    .iter()
+                    .map(|b| {
+                        tcx.mk_const(ty::ConstS { kind: ty::ConstKind::Value(*b), ty: *inner_ty })
+                    })
+                    .collect::<Vec<_>>();
+                debug!(?field_consts);
+
+                (field_consts, None)
+            }
+            ty::Adt(def, _) if def.variants().is_empty() => bug!("unreachable"),
+            ty::Adt(def, substs) => {
+                let variant_idx = if def.is_enum() {
+                    VariantIdx::from_u32(branches[0].unwrap_leaf().try_to_u32().unwrap())
+                } else {
+                    VariantIdx::from_u32(0)
+                };
+                let fields = &def.variant(variant_idx).fields;
+                let mut field_consts = Vec::with_capacity(fields.len());
+
+                // Note: First element inValTree corresponds to variant of enum
+                let mut valtree_idx = if def.is_enum() { 1 } else { 0 };
+                for field in fields {
+                    let field_ty = field.ty(tcx, substs);
+                    let field_valtree = branches[valtree_idx]; // first element of branches is variant
+                    let field_const = tcx.mk_const(ty::ConstS {
+                        kind: ty::ConstKind::Value(field_valtree),
+                        ty: field_ty,
+                    });
+                    field_consts.push(field_const);
+                    valtree_idx += 1;
+                }
+                debug!(?field_consts);
+
+                (field_consts, Some(variant_idx))
+            }
+            ty::Tuple(elem_tys) => {
+                let fields = elem_tys
+                    .iter()
+                    .enumerate()
+                    .map(|(i, elem_ty)| {
+                        let elem_valtree = branches[i];
+                        tcx.mk_const(ty::ConstS {
+                            kind: ty::ConstKind::Value(elem_valtree),
+                            ty: elem_ty,
+                        })
+                    })
+                    .collect::<Vec<_>>();
+
+                (fields, None)
+            }
+            _ => bug!("cannot destructure constant {:?}", const_),
+        };
+
+        let fields = tcx.arena.alloc_from_iter(fields.into_iter());
+
+        ty::DestructuredConst { variant, fields }
+    } else {
+        bug!("cannot destructure constant {:?}", const_)
+    }
+}
+
+pub fn provide(providers: &mut ty::query::Providers) {
+    *providers =
+        ty::query::Providers { destructure_const, ..*providers };
+}
diff --git a/compiler/rustc_ty_utils/src/lib.rs b/compiler/rustc_ty_utils/src/lib.rs
index 484967bbef8..a27cb7cf07e 100644
--- a/compiler/rustc_ty_utils/src/lib.rs
+++ b/compiler/rustc_ty_utils/src/lib.rs
@@ -18,6 +18,7 @@ use rustc_middle::ty::query::Providers;
 
 mod assoc;
 mod common_traits;
+pub mod consts;
 pub mod instance;
 mod needs_drop;
 pub mod representability;
@@ -26,6 +27,7 @@ mod ty;
 pub fn provide(providers: &mut Providers) {
     assoc::provide(providers);
     common_traits::provide(providers);
+    consts::provide(providers);
     needs_drop::provide(providers);
     ty::provide(providers);
     instance::provide(providers);