Arena-allocate slices to replace Vecs in SpirvType and SpirvConst.

Author: eddyb

crates/rustc_codegen_spirv/src/abi.rs

@@ -17,8 +17,8 @@ use rustc_middle::ty::{
 };
 use rustc_middle::{bug, span_bug};
 use rustc_span::def_id::DefId;
-use rustc_span::Span;
 use rustc_span::DUMMY_SP;
+use rustc_span::{Span, Symbol};
 use rustc_target::abi::call::{ArgAbi, ArgAttributes, FnAbi, PassMode};
 use rustc_target::abi::{
     Abi, Align, FieldsShape, LayoutS, Primitive, Scalar, Size, TagEncoding, VariantIdx, Variants,
@@ -300,6 +300,7 @@ impl<'tcx> ConvSpirvType<'tcx> for PointeeTy<'tcx> {
 
 impl<'tcx> ConvSpirvType<'tcx> for FnAbi<'tcx, Ty<'tcx>> {
     fn spirv_type(&self, span: Span, cx: &CodegenCx<'tcx>) -> Word {
+        // FIXME(eddyb) use `AccumulateVec`s just like `rustc` itself does.
         let mut argument_types = Vec::new();
 
         let return_type = match self.ret.mode {
@@ -332,7 +333,7 @@ impl<'tcx> ConvSpirvType<'tcx> for FnAbi<'tcx, Ty<'tcx>> {
 
         SpirvType::Function {
             return_type,
-            arguments: argument_types,
+            arguments: &argument_types,
         }
         .def(span, cx)
     }
@@ -364,8 +365,8 @@ impl<'tcx> ConvSpirvType<'tcx> for TyAndLayout<'tcx> {
                 def_id: def_id_for_spirv_type_adt(*self),
                 size: Some(Size::ZERO),
                 align: Align::from_bytes(0).unwrap(),
-                field_types: Vec::new(),
-                field_offsets: Vec::new(),
+                field_types: &[],
+                field_offsets: &[],
                 field_names: None,
             }
             .def_with_name(cx, span, TyLayoutNameKey::from(*self)),
@@ -416,23 +417,24 @@ impl<'tcx> ConvSpirvType<'tcx> for TyAndLayout<'tcx> {
                 } else {
                     Some(self.size)
                 };
+                // FIXME(eddyb) use `ArrayVec` here.
                 let mut field_names = Vec::new();
                 if let TyKind::Adt(adt, _) = self.ty.kind() {
                     if let Variants::Single { index } = self.variants {
                         for i in self.fields.index_by_increasing_offset() {
                             let field = &adt.variants()[index].fields[i];
-                            field_names.push(field.name.to_ident_string());
+                            field_names.push(field.name);
                         }
                     }
                 }
                 SpirvType::Adt {
                     def_id: def_id_for_spirv_type_adt(*self),
                     size,
                     align: self.align.abi,
-                    field_types: vec![a, b],
-                    field_offsets: vec![a_offset, b_offset],
+                    field_types: &[a, b],
+                    field_offsets: &[a_offset, b_offset],
                     field_names: if field_names.len() == 2 {
-                        Some(field_names)
+                        Some(&field_names)
                     } else {
                         None
                     },
@@ -598,8 +600,8 @@ fn trans_aggregate<'tcx>(cx: &CodegenCx<'tcx>, span: Span, ty: TyAndLayout<'tcx>
             def_id: def_id_for_spirv_type_adt(ty),
             size: Some(Size::ZERO),
             align: Align::from_bytes(0).unwrap(),
-            field_types: Vec::new(),
-            field_offsets: Vec::new(),
+            field_types: &[],
+            field_offsets: &[],
             field_names: None,
         }
         .def_with_name(cx, span, TyLayoutNameKey::from(ty))
@@ -664,6 +666,7 @@ pub fn auto_struct_layout<'tcx>(
     cx: &CodegenCx<'tcx>,
     field_types: &[Word],
 ) -> (Vec<Size>, Option<Size>, Align) {
+    // FIXME(eddyb) use `AccumulateVec`s just like `rustc` itself does.
     let mut field_offsets = Vec::with_capacity(field_types.len());
     let mut offset = Some(Size::ZERO);
     let mut max_align = Align::from_bytes(0).unwrap();
@@ -688,6 +691,7 @@ pub fn auto_struct_layout<'tcx>(
 fn trans_struct<'tcx>(cx: &CodegenCx<'tcx>, span: Span, ty: TyAndLayout<'tcx>) -> Word {
     let size = if ty.is_unsized() { None } else { Some(ty.size) };
     let align = ty.align.abi;
+    // FIXME(eddyb) use `AccumulateVec`s just like `rustc` itself does.
     let mut field_types = Vec::new();
     let mut field_offsets = Vec::new();
     let mut field_names = Vec::new();
@@ -699,17 +703,18 @@ fn trans_struct<'tcx>(cx: &CodegenCx<'tcx>, span: Span, ty: TyAndLayout<'tcx>) -
         if let Variants::Single { index } = ty.variants {
             if let TyKind::Adt(adt, _) = ty.ty.kind() {
                 let field = &adt.variants()[index].fields[i];
-                field_names.push(field.name.to_ident_string());
+                field_names.push(field.name);
             } else {
-                field_names.push(format!("{}", i));
+                // FIXME(eddyb) this looks like something that should exist in rustc.
+                field_names.push(Symbol::intern(&format!("{i}")));
             }
         } else {
             if let TyKind::Adt(_, _) = ty.ty.kind() {
             } else {
                 span_bug!(span, "Variants::Multiple not TyKind::Adt");
             }
             if i == 0 {
-                field_names.push("discriminant".to_string());
+                field_names.push(cx.sym.discriminant);
             } else {
                 cx.tcx.sess.fatal("Variants::Multiple has multiple fields")
             }
@@ -719,9 +724,9 @@ fn trans_struct<'tcx>(cx: &CodegenCx<'tcx>, span: Span, ty: TyAndLayout<'tcx>) -
         def_id: def_id_for_spirv_type_adt(ty),
         size,
         align,
-        field_types,
-        field_offsets,
-        field_names: Some(field_names),
+        field_types: &field_types,
+        field_offsets: &field_offsets,
+        field_names: Some(&field_names),
     }
     .def_with_name(cx, span, TyLayoutNameKey::from(ty))
 }

crates/rustc_codegen_spirv/src/builder/builder_methods.rs

@@ -176,7 +176,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
         semantics
     }
 
-    fn memset_const_pattern(&self, ty: &SpirvType, fill_byte: u8) -> Word {
+    fn memset_const_pattern(&self, ty: &SpirvType<'tcx>, fill_byte: u8) -> Word {
         match *ty {
             SpirvType::Void => self.fatal("memset invalid on void pattern"),
             SpirvType::Bool => self.fatal("memset invalid on bool pattern"),
@@ -212,7 +212,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
             SpirvType::Vector { element, count } | SpirvType::Matrix { element, count } => {
                 let elem_pat = self.memset_const_pattern(&self.lookup_type(element), fill_byte);
                 self.constant_composite(
-                    ty.clone().def(self.span(), self),
+                    ty.def(self.span(), self),
                     iter::repeat(elem_pat).take(count as usize),
                 )
                 .def(self)
@@ -221,7 +221,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                 let elem_pat = self.memset_const_pattern(&self.lookup_type(element), fill_byte);
                 let count = self.builder.lookup_const_u64(count).unwrap() as usize;
                 self.constant_composite(
-                    ty.clone().def(self.span(), self),
+                    ty.def(self.span(), self),
                     iter::repeat(elem_pat).take(count),
                 )
                 .def(self)
@@ -242,7 +242,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
         }
     }
 
-    fn memset_dynamic_pattern(&self, ty: &SpirvType, fill_var: Word) -> Word {
+    fn memset_dynamic_pattern(&self, ty: &SpirvType<'tcx>, fill_var: Word) -> Word {
         match *ty {
             SpirvType::Void => self.fatal("memset invalid on void pattern"),
             SpirvType::Bool => self.fatal("memset invalid on bool pattern"),
@@ -270,7 +270,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                 let count = self.builder.lookup_const_u64(count).unwrap() as usize;
                 self.emit()
                     .composite_construct(
-                        ty.clone().def(self.span(), self),
+                        ty.def(self.span(), self),
                         None,
                         iter::repeat(elem_pat).take(count),
                     )
@@ -280,7 +280,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                 let elem_pat = self.memset_dynamic_pattern(&self.lookup_type(element), fill_var);
                 self.emit()
                     .composite_construct(
-                        ty.clone().def(self.span(), self),
+                        ty.def(self.span(), self),
                         None,
                         iter::repeat(elem_pat).take(count as usize),
                     )
@@ -1260,9 +1260,7 @@ impl<'a, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'tcx> {
         };
         let pointee_kind = self.lookup_type(pointee);
         let result_pointee_type = match pointee_kind {
-            SpirvType::Adt {
-                ref field_types, ..
-            } => field_types[idx as usize],
+            SpirvType::Adt { field_types, .. } => field_types[idx as usize],
             SpirvType::Array { element, .. }
             | SpirvType::RuntimeArray { element, .. }
             | SpirvType::Vector { element, .. }
@@ -2345,7 +2343,7 @@ impl<'a, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'tcx> {
             ),
         };
 
-        for (argument, argument_type) in args.iter().zip(argument_types) {
+        for (argument, &argument_type) in args.iter().zip(argument_types) {
             assert_ty_eq!(self, argument.ty, argument_type);
         }
         let libm_intrinsic = self.libm_intrinsics.borrow().get(&callee_val).copied();

crates/rustc_codegen_spirv/src/builder/spirv_asm.rs

@@ -1081,7 +1081,7 @@ impl<'cx, 'tcx> Builder<'cx, 'tcx> {
             (OperandKind::LiteralContextDependentNumber, Some(word)) => {
                 assert!(matches!(inst.class.opcode, Op::Constant | Op::SpecConstant));
                 let ty = inst.result_type.unwrap();
-                fn parse(ty: SpirvType, w: &str) -> Result<dr::Operand, String> {
+                fn parse(ty: SpirvType<'_>, w: &str) -> Result<dr::Operand, String> {
                     fn fmt(x: impl ToString) -> String {
                         x.to_string()
                     }

crates/rustc_codegen_spirv/src/builder_spirv.rs

@@ -13,7 +13,6 @@ use rustc_span::symbol::Symbol;
 use rustc_span::{Span, DUMMY_SP};
 use std::assert_matches::assert_matches;
 use std::cell::{RefCell, RefMut};
-use std::rc::Rc;
 use std::{fs::File, io::Write, path::Path};
 
 #[derive(Copy, Clone, Debug, Ord, PartialOrd, Eq, PartialEq, Hash)]
@@ -68,7 +67,7 @@ impl SpirvValue {
     pub fn const_fold_load(self, cx: &CodegenCx<'_>) -> Option<Self> {
         match self.kind {
             SpirvValueKind::Def(id) | SpirvValueKind::IllegalConst(id) => {
-                let entry = cx.builder.id_to_const.borrow().get(&id)?.clone();
+                let &entry = cx.builder.id_to_const.borrow().get(&id)?;
                 match entry.val {
                     SpirvConst::PtrTo { pointee } => {
                         let ty = match cx.lookup_type(self.ty) {
@@ -213,8 +212,8 @@ impl SpirvValueExt for Word {
     }
 }
 
-#[derive(Debug, Clone, Ord, PartialOrd, Eq, PartialEq, Hash)]
-pub enum SpirvConst {
+#[derive(Debug, Copy, Clone, Ord, PartialOrd, Eq, PartialEq, Hash)]
+pub enum SpirvConst<'tcx> {
     U32(u32),
     U64(u64),
     /// f32 isn't hash, so store bits
@@ -232,8 +231,7 @@ pub enum SpirvConst {
     // different functions, but of the same type, don't overlap their zombies.
     ZombieUndefForFnAddr,
 
-    // FIXME(eddyb) use `tcx.arena.dropless` to get `&'tcx [_]`, instead of `Rc`.
-    Composite(Rc<[Word]>),
+    Composite(&'tcx [Word]),
 
     /// Pointer to constant data, i.e. `&pointee`, represented as an `OpVariable`
     /// in the `Private` storage class, and with `pointee` as its initializer.
@@ -242,6 +240,40 @@ pub enum SpirvConst {
     },
 }
 
+impl SpirvConst<'_> {
+    /// Replace `&[T]` fields with `&'tcx [T]` ones produced by calling
+    /// `tcx.arena.dropless.alloc_slice(...)` - this is done late for two reasons:
+    /// 1. it avoids allocating in the arena when the cache would be hit anyway,
+    ///    which would create "garbage" (as in, unreachable allocations)
+    ///    (ideally these would also be interned, but that's even more refactors)
+    /// 2. an empty slice is disallowed (as it's usually handled as a special
+    ///    case elsewhere, e.g. `rustc`'s `ty::List` - sadly we can't use that)
+    fn tcx_arena_alloc_slices<'tcx>(self, cx: &CodegenCx<'tcx>) -> SpirvConst<'tcx> {
+        fn arena_alloc_slice<'tcx, T: Copy>(cx: &CodegenCx<'tcx>, xs: &[T]) -> &'tcx [T] {
+            if xs.is_empty() {
+                &[]
+            } else {
+                cx.tcx.arena.dropless.alloc_slice(xs)
+            }
+        }
+
+        match self {
+            // FIXME(eddyb) these are all noop cases, could they be automated?
+            SpirvConst::U32(v) => SpirvConst::U32(v),
+            SpirvConst::U64(v) => SpirvConst::U64(v),
+            SpirvConst::F32(v) => SpirvConst::F32(v),
+            SpirvConst::F64(v) => SpirvConst::F64(v),
+            SpirvConst::Bool(v) => SpirvConst::Bool(v),
+            SpirvConst::Null => SpirvConst::Null,
+            SpirvConst::Undef => SpirvConst::Undef,
+            SpirvConst::ZombieUndefForFnAddr => SpirvConst::ZombieUndefForFnAddr,
+            SpirvConst::PtrTo { pointee } => SpirvConst::PtrTo { pointee },
+
+            SpirvConst::Composite(fields) => SpirvConst::Composite(arena_alloc_slice(cx, fields)),
+        }
+    }
+}
+
 #[derive(Debug, Clone, Ord, PartialOrd, Eq, PartialEq, Hash)]
 struct WithType<V> {
     ty: Word,
@@ -317,22 +349,22 @@ pub struct BuilderCursor {
     pub block: Option<usize>,
 }
 
-pub struct BuilderSpirv {
+pub struct BuilderSpirv<'tcx> {
     builder: RefCell<Builder>,
 
     // Bidirectional maps between `SpirvConst` and the ID of the defined global
     // (e.g. `OpConstant...`) instruction.
     // NOTE(eddyb) both maps have `WithConstLegality` around their keys, which
     // allows getting that legality information without additional lookups.
-    const_to_id: RefCell<FxHashMap<WithType<SpirvConst>, WithConstLegality<Word>>>,
-    id_to_const: RefCell<FxHashMap<Word, WithConstLegality<SpirvConst>>>,
+    const_to_id: RefCell<FxHashMap<WithType<SpirvConst<'tcx>>, WithConstLegality<Word>>>,
+    id_to_const: RefCell<FxHashMap<Word, WithConstLegality<SpirvConst<'tcx>>>>,
     string_cache: RefCell<FxHashMap<String, Word>>,
 
     enabled_capabilities: FxHashSet<Capability>,
     enabled_extensions: FxHashSet<Symbol>,
 }
 
-impl BuilderSpirv {
+impl<'tcx> BuilderSpirv<'tcx> {
     pub fn new(sym: &Symbols, target: &SpirvTarget, features: &[TargetFeature]) -> Self {
         let version = target.spirv_version();
         let memory_model = target.memory_model();
@@ -457,7 +489,12 @@ impl BuilderSpirv {
         bug!("Function not found: {}", id);
     }
 
-    pub fn def_constant(&self, ty: Word, val: SpirvConst) -> SpirvValue {
+    pub(crate) fn def_constant_cx(
+        &self,
+        ty: Word,
+        val: SpirvConst<'_>,
+        cx: &CodegenCx<'tcx>,
+    ) -> SpirvValue {
         let val_with_type = WithType { ty, val };
         let mut builder = self.builder(BuilderCursor::default());
         if let Some(entry) = self.const_to_id.borrow().get(&val_with_type) {
@@ -486,7 +523,7 @@ impl BuilderSpirv {
             SpirvConst::Null => builder.constant_null(ty),
             SpirvConst::Undef | SpirvConst::ZombieUndefForFnAddr => builder.undef(ty, None),
 
-            SpirvConst::Composite(ref v) => builder.constant_composite(ty, v.iter().copied()),
+            SpirvConst::Composite(v) => builder.constant_composite(ty, v.iter().copied()),
 
             SpirvConst::PtrTo { pointee } => {
                 builder.variable(ty, None, StorageClass::Private, Some(pointee))
@@ -517,7 +554,7 @@ impl BuilderSpirv {
                 Ok(())
             }
 
-            SpirvConst::Composite(ref v) => v.iter().fold(Ok(()), |composite_legal, field| {
+            SpirvConst::Composite(v) => v.iter().fold(Ok(()), |composite_legal, field| {
                 let field_entry = &self.id_to_const.borrow()[field];
                 let field_legal_in_composite = field_entry.legal.and(
                     // `field` is itself some legal `SpirvConst`, but can we have
@@ -556,14 +593,11 @@ impl BuilderSpirv {
                 }
             },
         };
+        let val = val.tcx_arena_alloc_slices(cx);
         assert_matches!(
-            self.const_to_id.borrow_mut().insert(
-                WithType {
-                    ty,
-                    val: val.clone()
-                },
-                WithConstLegality { val: id, legal }
-            ),
+            self.const_to_id
+                .borrow_mut()
+                .insert(WithType { ty, val }, WithConstLegality { val: id, legal }),
             None
         );
         assert_matches!(
@@ -581,10 +615,10 @@ impl BuilderSpirv {
         SpirvValue { kind, ty }
     }
 
-    pub fn lookup_const(&self, def: SpirvValue) -> Option<SpirvConst> {
+    pub fn lookup_const(&self, def: SpirvValue) -> Option<SpirvConst<'tcx>> {
         match def.kind {
             SpirvValueKind::Def(id) | SpirvValueKind::IllegalConst(id) => {
-                Some(self.id_to_const.borrow().get(&id)?.val.clone())
+                Some(self.id_to_const.borrow().get(&id)?.val)
             }
             _ => None,
         }