Auto merge of #143182 - xdoardo:more-addrspace, r=<try>

Allow custom default address spaces and parse `p-` specifications in the datalayout string

Some targets, such as CHERI, use as default an address space different from the "normal" default address space `0` (in the case of CHERI, [200 is used](https://www.cl.cam.ac.uk/techreports/UCAM-CL-TR-877.pdf)). Currently, `rustc` does not allow to specify custom address spaces and does not take into consideration [`p-` specifications in the datalayout string](https://llvm.org/docs/LangRef.html#langref-datalayout).

This patch tries to mitigate these problems by allowing targets to define a custom default address space (while keeping the default value to address space `0`) and adding the code to parse the `p-` specifications in `rustc_abi`. The main changes are that `TargetDataLayout` now uses functions to refer to pointer-related informations, instead of having specific fields for the size and alignment of pointers in the default address space; furthermore, the two `pointer_size` and `pointer_align` fields in `TargetDataLayout` are replaced with an `FxHashMap` that holds info for all the possible address spaces, as parsed by the `p-` specifications.

The potential performance drawbacks of not having ad-hoc fields for the default address space will be tested in this PR's CI run.

r? workingjubilee

Author: rust-bors[bot]

compiler/rustc_abi/src/layout/ty.rs

@@ -118,11 +118,11 @@ impl<'a> Layout<'a> {
 
     /// Whether the layout is from a type that implements [`std::marker::PointerLike`].
     ///
-    /// Currently, that means that the type is pointer-sized, pointer-aligned,
-    /// and has a initialized (non-union), scalar ABI.
+    /// Currently, that means that the type is pointer-sized, pointer-aligned, and has a initialized
+    /// (non-union), scalar ABI; all of this with respect with the default address space.
     pub fn is_pointer_like(self, data_layout: &TargetDataLayout) -> bool {
-        self.size() == data_layout.pointer_size
-            && self.align().abi == data_layout.pointer_align.abi
+        self.size() == data_layout.pointer_size()
+            && self.align().abi == data_layout.pointer_align().abi
             && matches!(self.backend_repr(), BackendRepr::Scalar(Scalar::Initialized { .. }))
     }
 }

compiler/rustc_abi/src/lib.rs

@@ -221,6 +221,17 @@ impl ReprOptions {
 /// * Cranelift stores the base-2 log of the lane count in a 4 bit integer.
 pub const MAX_SIMD_LANES: u64 = 1 << 0xF;
 
+/// Informations relative to a specific address space.
+#[derive(Copy, Clone, Debug, PartialEq, Eq)]
+pub struct AddressSpaceInfo {
+    /// The size of the bitwise representation of the pointer.
+    pointer_size: Size,
+    /// The alignment requirements for pointers in this address space.
+    pointer_align: AbiAlign,
+    /// The size of the index that used for address calculations on pointers in this address space.
+    pointer_index: Size,
+}
+
 /// Parsed [Data layout](https://llvm.org/docs/LangRef.html#data-layout)
 /// for a target, which contains everything needed to compute layouts.
 #[derive(Debug, PartialEq, Eq)]
@@ -236,13 +247,14 @@ pub struct TargetDataLayout {
     pub f32_align: AbiAlign,
     pub f64_align: AbiAlign,
     pub f128_align: AbiAlign,
-    pub pointer_size: Size,
-    pub pointer_align: AbiAlign,
     pub aggregate_align: AbiAlign,
 
     /// Alignments for vector types.
     pub vector_align: Vec<(Size, AbiAlign)>,
 
+    pub default_address_space: AddressSpace,
+    pub address_space_info: Vec<(AddressSpace, AddressSpaceInfo)>,
+
     pub instruction_address_space: AddressSpace,
 
     /// Minimum size of #[repr(C)] enums (default c_int::BITS, usually 32)
@@ -267,14 +279,21 @@ impl Default for TargetDataLayout {
             f32_align: AbiAlign::new(align(32)),
             f64_align: AbiAlign::new(align(64)),
             f128_align: AbiAlign::new(align(128)),
-            pointer_size: Size::from_bits(64),
-            pointer_align: AbiAlign::new(align(64)),
             aggregate_align: AbiAlign { abi: align(8) },
             vector_align: vec![
                 (Size::from_bits(64), AbiAlign::new(align(64))),
                 (Size::from_bits(128), AbiAlign::new(align(128))),
             ],
-            instruction_address_space: AddressSpace::DATA,
+            default_address_space: AddressSpace::ZERO,
+            address_space_info: vec![(
+                AddressSpace::ZERO,
+                AddressSpaceInfo {
+                    pointer_size: Size::from_bits(64),
+                    pointer_align: AbiAlign::new(align(64)),
+                    pointer_index: Size::from_bits(64),
+                },
+            )],
+            instruction_address_space: AddressSpace::ZERO,
             c_enum_min_size: Integer::I32,
         }
     }
@@ -288,6 +307,7 @@ pub enum TargetDataLayoutErrors<'a> {
     InconsistentTargetArchitecture { dl: &'a str, target: &'a str },
     InconsistentTargetPointerWidth { pointer_size: u64, target: u32 },
     InvalidBitsSize { err: String },
+    MissingAddressSpaceInfo { addr_space: AddressSpace },
 }
 
 impl TargetDataLayout {
@@ -298,6 +318,7 @@ impl TargetDataLayout {
     /// determined from llvm string.
     pub fn parse_from_llvm_datalayout_string<'a>(
         input: &'a str,
+        default_address_space: AddressSpace,
     ) -> Result<TargetDataLayout, TargetDataLayoutErrors<'a>> {
         // Parse an address space index from a string.
         let parse_address_space = |s: &'a str, cause: &'a str| {
@@ -334,6 +355,8 @@ impl TargetDataLayout {
         };
 
         let mut dl = TargetDataLayout::default();
+        dl.default_address_space = default_address_space;
+
         let mut i128_align_src = 64;
         for spec in input.split('-') {
             let spec_parts = spec.split(':').collect::<Vec<_>>();
@@ -349,13 +372,47 @@ impl TargetDataLayout {
                 ["f32", a @ ..] => dl.f32_align = parse_align(a, "f32")?,
                 ["f64", a @ ..] => dl.f64_align = parse_align(a, "f64")?,
                 ["f128", a @ ..] => dl.f128_align = parse_align(a, "f128")?,
-                // FIXME(erikdesjardins): we should be parsing nonzero address spaces
-                // this will require replacing TargetDataLayout::{pointer_size,pointer_align}
-                // with e.g. `fn pointer_size_in(AddressSpace)`
-                [p @ "p", s, a @ ..] | [p @ "p0", s, a @ ..] => {
-                    dl.pointer_size = parse_size(s, p)?;
-                    dl.pointer_align = parse_align(a, p)?;
+                [p, s, a @ ..] if p.starts_with("p") => {
+                    let p = p.strip_prefix(char::is_alphabetic).unwrap_or_default();
+
+                    let addr_space = if !p.is_empty() {
+                        parse_address_space(p, "p")?
+                    } else {
+                        AddressSpace::ZERO
+                    };
+
+                    let pointer_size = parse_size(s, p)?;
+                    let info = AddressSpaceInfo {
+                        pointer_index: pointer_size,
+                        pointer_size,
+                        pointer_align: parse_align(a, p)?,
+                    };
+                    match dl.address_space_info.iter_mut().find(|(a, _)| *a == addr_space) {
+                        Some(e) => e.1 = info,
+                        None => dl.address_space_info.push((addr_space, info)),
+                    }
+                }
+                [p, s, _pr, i, a @ ..] if p.starts_with("p") => {
+                    let p = p.strip_prefix(char::is_alphabetic).unwrap_or_default();
+
+                    let addr_space = if !p.is_empty() {
+                        parse_address_space(p, "p")?
+                    } else {
+                        AddressSpace::ZERO
+                    };
+
+                    let info = AddressSpaceInfo {
+                        pointer_align: parse_align(a, p)?,
+                        pointer_size: parse_size(s, p)?,
+                        pointer_index: parse_size(i, p)?,
+                    };
+
+                    match dl.address_space_info.iter_mut().find(|(a, _)| *a == addr_space) {
+                        Some(e) => e.1 = info,
+                        None => dl.address_space_info.push((addr_space, info)),
+                    }
                 }
+
                 [s, a @ ..] if s.starts_with('i') => {
                     let Ok(bits) = s[1..].parse::<u64>() else {
                         parse_size(&s[1..], "i")?; // For the user error.
@@ -390,10 +447,34 @@ impl TargetDataLayout {
                 _ => {} // Ignore everything else.
             }
         }
+
+        if dl.address_space_info.iter().find(|(a, _)| *a == default_address_space).is_none() {
+            return Err(TargetDataLayoutErrors::MissingAddressSpaceInfo {
+                addr_space: default_address_space,
+            });
+        }
+
+        // Inherit, if not given, address space informations for specific LLVM elements from the
+        // default data address space.
+
+        if dl.address_space_info.iter().find(|(a, _)| *a == dl.instruction_address_space).is_none()
+        {
+            dl.address_space_info.push((
+                dl.instruction_address_space,
+                dl.address_space_info
+                    .iter()
+                    .find(|(a, _)| *a == default_address_space)
+                    .unwrap()
+                    .1
+                    .clone(),
+            ));
+        }
+
         Ok(dl)
     }
 
-    /// Returns **exclusive** upper bound on object size in bytes.
+    /// Returns **exclusive** upper bound on object size in bytes, in the default data address
+    /// space.
     ///
     /// The theoretical maximum object size is defined as the maximum positive `isize` value.
     /// This ensures that the `offset` semantics remain well-defined by allowing it to correctly
@@ -404,7 +485,21 @@ impl TargetDataLayout {
     /// so we adopt such a more-constrained size bound due to its technical limitations.
     #[inline]
     pub fn obj_size_bound(&self) -> u64 {
-        match self.pointer_size.bits() {
+        self.obj_size_bound_in(self.default_address_space)
+    }
+
+    /// Returns **exclusive** upper bound on object size in bytes.
+    ///
+    /// The theoretical maximum object size is defined as the maximum positive `isize` value.
+    /// This ensures that the `offset` semantics remain well-defined by allowing it to correctly
+    /// index every address within an object along with one byte past the end, along with allowing
+    /// `isize` to store the difference between any two pointers into an object.
+    ///
+    /// LLVM uses a 64-bit integer to represent object size in *bits*, but we care only for bytes,
+    /// so we adopt such a more-constrained size bound due to its technical limitations.
+    #[inline]
+    pub fn obj_size_bound_in(&self, address_space: AddressSpace) -> u64 {
+        match self.pointer_size_in(address_space).bits() {
             16 => 1 << 15,
             32 => 1 << 31,
             64 => 1 << 61,
@@ -414,8 +509,13 @@ impl TargetDataLayout {
 
     #[inline]
     pub fn ptr_sized_integer(&self) -> Integer {
+        self.ptr_sized_integer_in(self.default_address_space)
+    }
+
+    #[inline]
+    pub fn ptr_sized_integer_in(&self, address_space: AddressSpace) -> Integer {
         use Integer::*;
-        match self.pointer_size.bits() {
+        match self.pointer_index_in(address_space).bits() {
             16 => I16,
             32 => I32,
             64 => I64,
@@ -439,6 +539,54 @@ impl TargetDataLayout {
             Align::from_bytes(vec_size.bytes().next_power_of_two()).unwrap(),
         ))
     }
+
+    /// Get the pointer size in the default data address space.
+    #[inline]
+    pub fn pointer_size(&self) -> Size {
+        self.pointer_size_in(self.default_address_space)
+    }
+
+    /// Get the pointer size in a specific address space.
+    #[inline]
+    pub fn pointer_size_in(&self, c: AddressSpace) -> Size {
+        if let Some(e) = self.address_space_info.iter().find(|(a, _)| a == &c) {
+            e.1.pointer_size
+        } else {
+            panic!("Use of unknown address space {c:?}");
+        }
+    }
+
+    /// Get the pointer index in the default data address space.
+    #[inline]
+    pub fn pointer_index(&self) -> Size {
+        self.pointer_index_in(self.default_address_space)
+    }
+
+    /// Get the pointer index in a specific address space.
+    #[inline]
+    pub fn pointer_index_in(&self, c: AddressSpace) -> Size {
+        if let Some(e) = self.address_space_info.iter().find(|(a, _)| a == &c) {
+            e.1.pointer_index
+        } else {
+            panic!("Use of unknown address space {c:?}");
+        }
+    }
+
+    /// Get the pointer alignment in the default data address space.
+    #[inline]
+    pub fn pointer_align(&self) -> AbiAlign {
+        self.pointer_align_in(self.default_address_space)
+    }
+
+    /// Get the pointer alignment in a specific address space.
+    #[inline]
+    pub fn pointer_align_in(&self, c: AddressSpace) -> AbiAlign {
+        if let Some(e) = self.address_space_info.iter().find(|(a, _)| a == &c) {
+            e.1.pointer_align
+        } else {
+            panic!("Use of unknown address space {c:?}");
+        }
+    }
 }
 
 pub trait HasDataLayout {
@@ -1104,7 +1252,7 @@ impl Primitive {
             // FIXME(erikdesjardins): ignoring address space is technically wrong, pointers in
             // different address spaces can have different sizes
             // (but TargetDataLayout doesn't currently parse that part of the DL string)
-            Pointer(_) => dl.pointer_size,
+            Pointer(a) => dl.pointer_size_in(a),
         }
     }
 
@@ -1118,7 +1266,7 @@ impl Primitive {
             // FIXME(erikdesjardins): ignoring address space is technically wrong, pointers in
             // different address spaces can have different alignments
             // (but TargetDataLayout doesn't currently parse that part of the DL string)
-            Pointer(_) => dl.pointer_align,
+            Pointer(a) => dl.pointer_align_in(a),
         }
     }
 }
@@ -1422,8 +1570,8 @@ impl<FieldIdx: Idx> FieldsShape<FieldIdx> {
 pub struct AddressSpace(pub u32);
 
 impl AddressSpace {
-    /// The default address space, corresponding to data space.
-    pub const DATA: Self = AddressSpace(0);
+    /// LLVM's `0` address space.
+    pub const ZERO: Self = AddressSpace(0);
 }
 
 /// The way we represent values to the backend

compiler/rustc_ast_lowering/src/format.rs

@@ -55,7 +55,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
     /// Get the maximum value of int_ty. It is platform-dependent due to the byte size of isize
     fn int_ty_max(&self, int_ty: IntTy) -> u128 {
         match int_ty {
-            IntTy::Isize => self.tcx.data_layout.pointer_size.signed_int_max() as u128,
+            IntTy::Isize => self.tcx.data_layout.pointer_size().signed_int_max() as u128,
             IntTy::I8 => i8::MAX as u128,
             IntTy::I16 => i16::MAX as u128,
             IntTy::I32 => i32::MAX as u128,
@@ -67,7 +67,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
     /// Get the maximum value of uint_ty. It is platform-dependent due to the byte size of usize
     fn uint_ty_max(&self, uint_ty: UintTy) -> u128 {
         match uint_ty {
-            UintTy::Usize => self.tcx.data_layout.pointer_size.unsigned_int_max(),
+            UintTy::Usize => self.tcx.data_layout.pointer_size().unsigned_int_max(),
             UintTy::U8 => u8::MAX as u128,
             UintTy::U16 => u16::MAX as u128,
             UintTy::U32 => u32::MAX as u128,

compiler/rustc_codegen_cranelift/src/abi/mod.rs

@@ -831,7 +831,7 @@ pub(crate) fn codegen_drop<'tcx>(
 
 pub(crate) fn lib_call_arg_param(tcx: TyCtxt<'_>, ty: Type, is_signed: bool) -> AbiParam {
     let param = AbiParam::new(ty);
-    if ty.is_int() && u64::from(ty.bits()) < tcx.data_layout.pointer_size.bits() {
+    if ty.is_int() && u64::from(ty.bits()) < tcx.data_layout.pointer_size().bits() {
         match (&*tcx.sess.target.arch, &*tcx.sess.target.vendor) {
             ("x86_64", _) | ("aarch64", "apple") => match (ty, is_signed) {
                 (types::I8 | types::I16, true) => param.sext(),

compiler/rustc_codegen_cranelift/src/abi/pass_mode.rs

@@ -127,7 +127,7 @@ impl<'tcx> ArgAbiExt<'tcx> for ArgAbi<'tcx, Ty<'tcx>> {
             PassMode::Indirect { attrs, meta_attrs: None, on_stack } => {
                 if on_stack {
                     // Abi requires aligning struct size to pointer size
-                    let size = self.layout.size.align_to(tcx.data_layout.pointer_align.abi);
+                    let size = self.layout.size.align_to(tcx.data_layout.pointer_align().abi);
                     let size = u32::try_from(size.bytes()).unwrap();
                     smallvec![apply_attrs_to_abi_param(
                         AbiParam::special(pointer_ty(tcx), ArgumentPurpose::StructArgument(size),),

compiler/rustc_codegen_cranelift/src/common.rs

@@ -15,7 +15,7 @@ use crate::debuginfo::FunctionDebugContext;
 use crate::prelude::*;
 
 pub(crate) fn pointer_ty(tcx: TyCtxt<'_>) -> types::Type {
-    match tcx.data_layout.pointer_size.bits() {
+    match tcx.data_layout.pointer_size().bits() {
         16 => types::I16,
         32 => types::I32,
         64 => types::I64,

compiler/rustc_codegen_cranelift/src/constant.rs

@@ -443,7 +443,7 @@ fn define_all_allocs(tcx: TyCtxt<'_>, module: &mut dyn Module, cx: &mut Constant
             let addend = {
                 let endianness = tcx.data_layout.endian;
                 let offset = offset.bytes() as usize;
-                let ptr_size = tcx.data_layout.pointer_size;
+                let ptr_size = tcx.data_layout.pointer_size();
                 let bytes = &alloc.inspect_with_uninit_and_ptr_outside_interpreter(
                     offset..offset + ptr_size.bytes() as usize,
                 );

compiler/rustc_codegen_gcc/src/common.rs

@@ -170,7 +170,7 @@ impl<'gcc, 'tcx> ConstCodegenMethods for CodegenCx<'gcc, 'tcx> {
     }
 
     fn const_usize(&self, i: u64) -> RValue<'gcc> {
-        let bit_size = self.data_layout().pointer_size.bits();
+        let bit_size = self.data_layout().pointer_size().bits();
         if bit_size < 64 {
             // make sure it doesn't overflow
             assert!(i < (1 << bit_size));

compiler/rustc_codegen_gcc/src/consts.rs

@@ -294,7 +294,7 @@ pub(crate) fn const_alloc_to_gcc_uncached<'gcc>(
     let alloc = alloc.inner();
     let mut llvals = Vec::with_capacity(alloc.provenance().ptrs().len() + 1);
     let dl = cx.data_layout();
-    let pointer_size = dl.pointer_size.bytes() as usize;
+    let pointer_size = dl.pointer_size().bytes() as usize;
 
     let mut next_offset = 0;
     for &(offset, prov) in alloc.provenance().ptrs().iter() {
@@ -331,7 +331,7 @@ pub(crate) fn const_alloc_to_gcc_uncached<'gcc>(
             ),
             abi::Scalar::Initialized {
                 value: Primitive::Pointer(address_space),
-                valid_range: WrappingRange::full(dl.pointer_size),
+                valid_range: WrappingRange::full(dl.pointer_size()),
             },
             cx.type_i8p_ext(address_space),
         ));