1
Fork 0

intern valtrees

This commit is contained in:
Lukas Markeffsky 2025-02-07 19:33:58 +01:00
parent ef148cd7eb
commit 885e0f1b96
15 changed files with 149 additions and 117 deletions

View file

@ -90,6 +90,7 @@ macro_rules! arena_types {
[] autodiff_item: rustc_ast::expand::autodiff_attrs::AutoDiffItem,
[] ordered_name_set: rustc_data_structures::fx::FxIndexSet<rustc_span::Symbol>,
[] pats: rustc_middle::ty::PatternKind<'tcx>,
[] valtree: rustc_middle::ty::ValTreeKind<'tcx>,
// Note that this deliberately duplicates items in the `rustc_hir::arena`,
// since we need to allocate this type on both the `rustc_hir` arena

View file

@ -146,6 +146,12 @@ impl<'tcx, E: TyEncoder<I = TyCtxt<'tcx>>> Encodable<E> for ty::Pattern<'tcx> {
}
}
impl<'tcx, E: TyEncoder<I = TyCtxt<'tcx>>> Encodable<E> for ty::ValTree<'tcx> {
fn encode(&self, e: &mut E) {
self.0.0.encode(e);
}
}
impl<'tcx, E: TyEncoder<I = TyCtxt<'tcx>>> Encodable<E> for ConstAllocation<'tcx> {
fn encode(&self, e: &mut E) {
self.inner().encode(e)
@ -355,12 +361,9 @@ impl<'tcx, D: TyDecoder<I = TyCtxt<'tcx>>> Decodable<D> for ty::Pattern<'tcx> {
}
}
impl<'tcx, D: TyDecoder<I = TyCtxt<'tcx>>> RefDecodable<'tcx, D> for [ty::ValTree<'tcx>] {
fn decode(decoder: &mut D) -> &'tcx Self {
decoder
.interner()
.arena
.alloc_from_iter((0..decoder.read_usize()).map(|_| Decodable::decode(decoder)))
impl<'tcx, D: TyDecoder<I = TyCtxt<'tcx>>> Decodable<D> for ty::ValTree<'tcx> {
fn decode(decoder: &mut D) -> Self {
decoder.interner().intern_valtree(Decodable::decode(decoder))
}
}

View file

@ -20,7 +20,7 @@ pub type ConstKind<'tcx> = ir::ConstKind<TyCtxt<'tcx>>;
pub type UnevaluatedConst<'tcx> = ir::UnevaluatedConst<TyCtxt<'tcx>>;
#[cfg(target_pointer_width = "64")]
rustc_data_structures::static_assert_size!(ConstKind<'_>, 32);
rustc_data_structures::static_assert_size!(ConstKind<'_>, 24);
#[derive(Copy, Clone, PartialEq, Eq, Hash, HashStable)]
#[rustc_pass_by_value]
@ -190,7 +190,7 @@ impl<'tcx> Const<'tcx> {
.size;
ty::Const::new_value(
tcx,
ty::ValTree::from_scalar_int(ScalarInt::try_from_uint(bits, size).unwrap()),
ty::ValTree::from_scalar_int(tcx, ScalarInt::try_from_uint(bits, size).unwrap()),
ty,
)
}
@ -198,7 +198,7 @@ impl<'tcx> Const<'tcx> {
#[inline]
/// Creates an interned zst constant.
pub fn zero_sized(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>) -> Self {
ty::Const::new_value(tcx, ty::ValTree::zst(), ty)
ty::Const::new_value(tcx, ty::ValTree::zst(tcx), ty)
}
#[inline]

View file

@ -1,4 +1,8 @@
use rustc_macros::{HashStable, TyDecodable, TyEncodable, TypeFoldable, TypeVisitable};
use std::fmt;
use std::ops::Deref;
use rustc_data_structures::intern::Interned;
use rustc_macros::{HashStable, Lift, TyDecodable, TyEncodable, TypeFoldable, TypeVisitable};
use super::ScalarInt;
use crate::mir::interpret::Scalar;
@ -16,9 +20,9 @@ use crate::ty::{self, Ty, TyCtxt};
///
/// `ValTree` does not have this problem with representation, as it only contains integers or
/// lists of (nested) `ValTree`.
#[derive(Copy, Clone, Debug, Hash, Eq, PartialEq)]
#[derive(Clone, Debug, Hash, Eq, PartialEq)]
#[derive(HashStable, TyEncodable, TyDecodable)]
pub enum ValTree<'tcx> {
pub enum ValTreeKind<'tcx> {
/// integers, `bool`, `char` are represented as scalars.
/// See the `ScalarInt` documentation for how `ScalarInt` guarantees that equal values
/// of these types have the same representation.
@ -33,58 +37,98 @@ pub enum ValTree<'tcx> {
/// the fields of the variant.
///
/// ZST types are represented as an empty slice.
Branch(&'tcx [ValTree<'tcx>]),
Branch(Box<[ValTree<'tcx>]>),
}
impl<'tcx> ValTree<'tcx> {
pub fn zst() -> Self {
Self::Branch(&[])
}
impl<'tcx> ValTreeKind<'tcx> {
#[inline]
pub fn unwrap_leaf(self) -> ScalarInt {
pub fn unwrap_leaf(&self) -> ScalarInt {
match self {
Self::Leaf(s) => s,
Self::Leaf(s) => *s,
_ => bug!("expected leaf, got {:?}", self),
}
}
#[inline]
pub fn unwrap_branch(self) -> &'tcx [Self] {
pub fn unwrap_branch(&self) -> &[ValTree<'tcx>] {
match self {
Self::Branch(branch) => branch,
Self::Branch(branch) => &**branch,
_ => bug!("expected branch, got {:?}", self),
}
}
pub fn from_raw_bytes<'a>(tcx: TyCtxt<'tcx>, bytes: &'a [u8]) -> Self {
let branches = bytes.iter().map(|b| Self::Leaf(ScalarInt::from(*b)));
let interned = tcx.arena.alloc_from_iter(branches);
Self::Branch(interned)
}
pub fn from_scalar_int(i: ScalarInt) -> Self {
Self::Leaf(i)
}
pub fn try_to_scalar(self) -> Option<Scalar> {
pub fn try_to_scalar(&self) -> Option<Scalar> {
self.try_to_scalar_int().map(Scalar::Int)
}
pub fn try_to_scalar_int(self) -> Option<ScalarInt> {
pub fn try_to_scalar_int(&self) -> Option<ScalarInt> {
match self {
Self::Leaf(s) => Some(s),
Self::Leaf(s) => Some(*s),
Self::Branch(_) => None,
}
}
pub fn try_to_branch(&self) -> Option<&[ValTree<'tcx>]> {
match self {
Self::Branch(branch) => Some(&**branch),
Self::Leaf(_) => None,
}
}
}
/// An interned valtree. Use this rather than `ValTreeKind`, whenever possible.
///
/// See the docs of [`ValTreeKind`] or the [dev guide] for an explanation of this type.
///
/// [dev guide]: https://rustc-dev-guide.rust-lang.org/mir/index.html#valtrees
#[derive(Copy, Clone, Hash, Eq, PartialEq)]
#[derive(HashStable)]
pub struct ValTree<'tcx>(pub(crate) Interned<'tcx, ValTreeKind<'tcx>>);
impl<'tcx> ValTree<'tcx> {
/// Returns the zero-sized valtree: `Branch([])`.
pub fn zst(tcx: TyCtxt<'tcx>) -> Self {
tcx.consts.valtree_zst
}
pub fn is_zst(self) -> bool {
matches!(*self, ValTreeKind::Branch(box []))
}
pub fn from_raw_bytes(tcx: TyCtxt<'tcx>, bytes: &[u8]) -> Self {
let branches = bytes.iter().map(|&b| Self::from_scalar_int(tcx, b.into()));
Self::from_branches(tcx, branches)
}
pub fn from_branches(tcx: TyCtxt<'tcx>, branches: impl IntoIterator<Item = Self>) -> Self {
tcx.intern_valtree(ValTreeKind::Branch(branches.into_iter().collect()))
}
pub fn from_scalar_int(tcx: TyCtxt<'tcx>, i: ScalarInt) -> Self {
tcx.intern_valtree(ValTreeKind::Leaf(i))
}
}
impl<'tcx> Deref for ValTree<'tcx> {
type Target = &'tcx ValTreeKind<'tcx>;
#[inline]
fn deref(&self) -> &&'tcx ValTreeKind<'tcx> {
&self.0.0
}
}
impl fmt::Debug for ValTree<'_> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
(**self).fmt(f)
}
}
/// A type-level constant value.
///
/// Represents a typed, fully evaluated constant.
#[derive(Copy, Clone, Debug, Hash, Eq, PartialEq)]
#[derive(HashStable, TyEncodable, TyDecodable, TypeFoldable, TypeVisitable)]
#[derive(HashStable, TyEncodable, TyDecodable, TypeFoldable, TypeVisitable, Lift)]
pub struct Value<'tcx> {
pub ty: Ty<'tcx>,
pub valtree: ValTree<'tcx>,

View file

@ -80,7 +80,7 @@ use crate::ty::{
GenericArgsRef, GenericParamDefKind, List, ListWithCachedTypeInfo, ParamConst, ParamTy,
Pattern, PatternKind, PolyExistentialPredicate, PolyFnSig, Predicate, PredicateKind,
PredicatePolarity, Region, RegionKind, ReprOptions, TraitObjectVisitor, Ty, TyKind, TyVid,
Visibility,
ValTree, ValTreeKind, Visibility,
};
#[allow(rustc::usage_of_ty_tykind)]
@ -806,6 +806,7 @@ pub struct CtxtInterners<'tcx> {
local_def_ids: InternedSet<'tcx, List<LocalDefId>>,
captures: InternedSet<'tcx, List<&'tcx ty::CapturedPlace<'tcx>>>,
offset_of: InternedSet<'tcx, List<(VariantIdx, FieldIdx)>>,
valtree: InternedSet<'tcx, ty::ValTreeKind<'tcx>>,
}
impl<'tcx> CtxtInterners<'tcx> {
@ -835,6 +836,7 @@ impl<'tcx> CtxtInterners<'tcx> {
local_def_ids: Default::default(),
captures: Default::default(),
offset_of: Default::default(),
valtree: Default::default(),
}
}
@ -1026,6 +1028,8 @@ pub struct CommonConsts<'tcx> {
pub unit: Const<'tcx>,
pub true_: Const<'tcx>,
pub false_: Const<'tcx>,
/// Use [`ty::ValTree::zst`] instead.
pub(crate) valtree_zst: ValTree<'tcx>,
}
impl<'tcx> CommonTypes<'tcx> {
@ -1129,19 +1133,30 @@ impl<'tcx> CommonConsts<'tcx> {
)
};
let mk_valtree = |v| {
ty::ValTree(Interned::new_unchecked(
interners.valtree.intern(v, |v| InternedInSet(interners.arena.alloc(v))).0,
))
};
let valtree_zst = mk_valtree(ty::ValTreeKind::Branch(Box::default()));
let valtree_true = mk_valtree(ty::ValTreeKind::Leaf(ty::ScalarInt::TRUE));
let valtree_false = mk_valtree(ty::ValTreeKind::Leaf(ty::ScalarInt::FALSE));
CommonConsts {
unit: mk_const(ty::ConstKind::Value(ty::Value {
ty: types.unit,
valtree: ty::ValTree::zst(),
valtree: valtree_zst,
})),
true_: mk_const(ty::ConstKind::Value(ty::Value {
ty: types.bool,
valtree: ty::ValTree::Leaf(ty::ScalarInt::TRUE),
valtree: valtree_true,
})),
false_: mk_const(ty::ConstKind::Value(ty::Value {
ty: types.bool,
valtree: ty::ValTree::Leaf(ty::ScalarInt::FALSE),
valtree: valtree_false,
})),
valtree_zst,
}
}
}
@ -2259,6 +2274,7 @@ nop_lift! { const_allocation; ConstAllocation<'a> => ConstAllocation<'tcx> }
nop_lift! { predicate; Predicate<'a> => Predicate<'tcx> }
nop_lift! { predicate; Clause<'a> => Clause<'tcx> }
nop_lift! { layout; Layout<'a> => Layout<'tcx> }
nop_lift! { valtree; ValTree<'a> => ValTree<'tcx> }
nop_list_lift! { type_lists; Ty<'a> => Ty<'tcx> }
nop_list_lift! {
@ -2269,26 +2285,6 @@ nop_list_lift! { bound_variable_kinds; ty::BoundVariableKind => ty::BoundVariabl
// This is the impl for `&'a GenericArgs<'a>`.
nop_list_lift! { args; GenericArg<'a> => GenericArg<'tcx> }
macro_rules! nop_slice_lift {
($ty:ty => $lifted:ty) => {
impl<'a, 'tcx> Lift<TyCtxt<'tcx>> for &'a [$ty] {
type Lifted = &'tcx [$lifted];
fn lift_to_interner(self, tcx: TyCtxt<'tcx>) -> Option<Self::Lifted> {
if self.is_empty() {
return Some(&[]);
}
tcx.interners
.arena
.dropless
.contains_slice(self)
.then(|| unsafe { mem::transmute(self) })
}
}
};
}
nop_slice_lift! { ty::ValTree<'a> => ty::ValTree<'tcx> }
macro_rules! sty_debug_print {
($fmt: expr, $ctxt: expr, $($variant: ident),*) => {{
// Curious inner module to allow variant names to be used as
@ -2533,6 +2529,7 @@ macro_rules! direct_interners {
// crate only, and have a corresponding `mk_` function.
direct_interners! {
region: pub(crate) intern_region(RegionKind<'tcx>): Region -> Region<'tcx>,
valtree: pub(crate) intern_valtree(ValTreeKind<'tcx>): ValTree -> ValTree<'tcx>,
pat: pub mk_pat(PatternKind<'tcx>): Pattern -> Pattern<'tcx>,
const_allocation: pub mk_const_alloc(Allocation): ConstAllocation -> ConstAllocation<'tcx>,
layout: pub mk_layout(LayoutData<FieldIdx, VariantIdx>): Layout -> Layout<'tcx>,

View file

@ -60,7 +60,8 @@ pub use self::closure::{
place_to_string_for_capture,
};
pub use self::consts::{
Const, ConstInt, ConstKind, Expr, ExprKind, ScalarInt, UnevaluatedConst, ValTree, Value,
Const, ConstInt, ConstKind, Expr, ExprKind, ScalarInt, UnevaluatedConst, ValTree, ValTreeKind,
Value,
};
pub use self::context::{
CtxtInterners, CurrentGcx, DeducedParamAttrs, Feed, FreeRegionInfo, GlobalCtxt, Lift, TyCtxt,

View file

@ -1639,14 +1639,12 @@ pub trait PrettyPrinter<'tcx>: Printer<'tcx> + fmt::Write {
match ty.kind() {
// Byte strings (&[u8; N])
ty::Ref(_, inner, _) => {
if let ty::Array(elem, len) = inner.kind()
if let ty::Array(elem, ct_len) = inner.kind()
&& let ty::Uint(ty::UintTy::U8) = elem.kind()
&& let ty::ConstKind::Value(cv) = len.kind()
&& let ty::ValTree::Leaf(int) = cv.valtree
&& let Some(len) = ct_len.try_to_target_usize(self.tcx())
{
match self.tcx().try_get_global_alloc(prov.alloc_id()) {
Some(GlobalAlloc::Memory(alloc)) => {
let len = int.to_bits(self.tcx().data_layout.pointer_size);
let range = AllocRange { start: offset, size: Size::from_bytes(len) };
if let Ok(byte_str) =
alloc.inner().get_bytes_strip_provenance(&self.tcx(), range)
@ -1800,8 +1798,8 @@ pub trait PrettyPrinter<'tcx>: Printer<'tcx> + fmt::Write {
}
let u8_type = self.tcx().types.u8;
match (cv.valtree, *cv.ty.kind()) {
(ty::ValTree::Branch(_), ty::Ref(_, inner_ty, _)) => match inner_ty.kind() {
match (*cv.valtree, *cv.ty.kind()) {
(ty::ValTreeKind::Branch(_), ty::Ref(_, inner_ty, _)) => match inner_ty.kind() {
ty::Slice(t) if *t == u8_type => {
let bytes = cv.try_to_raw_bytes(self.tcx()).unwrap_or_else(|| {
bug!(
@ -1826,7 +1824,7 @@ pub trait PrettyPrinter<'tcx>: Printer<'tcx> + fmt::Write {
return Ok(());
}
},
(ty::ValTree::Branch(_), ty::Array(t, _)) if t == u8_type => {
(ty::ValTreeKind::Branch(_), ty::Array(t, _)) if t == u8_type => {
let bytes = cv.try_to_raw_bytes(self.tcx()).unwrap_or_else(|| {
bug!("expected to convert valtree to raw bytes for type {:?}", t)
});
@ -1835,7 +1833,7 @@ pub trait PrettyPrinter<'tcx>: Printer<'tcx> + fmt::Write {
return Ok(());
}
// Aggregates, printed as array/tuple/struct/variant construction syntax.
(ty::ValTree::Branch(_), ty::Array(..) | ty::Tuple(..) | ty::Adt(..)) => {
(ty::ValTreeKind::Branch(_), ty::Array(..) | ty::Tuple(..) | ty::Adt(..)) => {
let contents = self.tcx().destructure_const(ty::Const::new_value(
self.tcx(),
cv.valtree,
@ -1891,12 +1889,12 @@ pub trait PrettyPrinter<'tcx>: Printer<'tcx> + fmt::Write {
}
return Ok(());
}
(ty::ValTree::Leaf(leaf), ty::Ref(_, inner_ty, _)) => {
(ty::ValTreeKind::Leaf(leaf), ty::Ref(_, inner_ty, _)) => {
p!(write("&"));
return self.pretty_print_const_scalar_int(leaf, inner_ty, print_ty);
return self.pretty_print_const_scalar_int(*leaf, inner_ty, print_ty);
}
(ty::ValTree::Leaf(leaf), _) => {
return self.pretty_print_const_scalar_int(leaf, cv.ty, print_ty);
(ty::ValTreeKind::Leaf(leaf), _) => {
return self.pretty_print_const_scalar_int(*leaf, cv.ty, print_ty);
}
(_, ty::FnDef(def_id, args)) => {
// Never allowed today, but we still encounter them in invalid const args.
@ -1909,7 +1907,7 @@ pub trait PrettyPrinter<'tcx>: Printer<'tcx> + fmt::Write {
}
// fallback
if cv.valtree == ty::ValTree::zst() {
if cv.valtree.is_zst() {
p!(write("<ZST>"));
} else {
p!(write("{:?}", cv.valtree));

View file

@ -165,13 +165,9 @@ impl<'tcx> fmt::Debug for ty::consts::Expr<'tcx> {
impl<'tcx> fmt::Debug for ty::Const<'tcx> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
// If this is a value, we spend some effort to make it look nice.
if let ConstKind::Value(_) = self.kind() {
if let ConstKind::Value(cv) = self.kind() {
return ty::tls::with(move |tcx| {
// ValTrees aren't interned, so we lift the entire constant.
let lifted = tcx.lift(*self).unwrap();
let ConstKind::Value(cv) = lifted.kind() else {
bug!("we checked that this is a valtree")
};
let cv = tcx.lift(cv).unwrap();
let mut cx = FmtPrinter::new(tcx, Namespace::ValueNS);
cx.pretty_print_const_valtree(cv, /*print_ty*/ true)?;
f.write_str(&cx.into_buffer())