[AVR] Ensure that function pointers stored within aggregates are annotated with the correct space

Before this patch, a function pointer stored within an aggregate like a
struct or an enum would always have the default address space `0`.

This patch removes this assumption and instead, introspects the inner
type being pointed at, storing the target address space in the PointeeInfo
struct so that downstream users may query it.

3 files changed, 40 insertions, 13 deletions

diff --git a/src/librustc_codegen_llvm/type_of.rs b/src/librustc_codegen_llvm/type_of.rs
index 5a0da6be598..1d0adc5783f 100644
--- a/src/librustc_codegen_llvm/type_of.rs
+++ b/src/librustc_codegen_llvm/type_of.rs
@@ -7,7 +7,7 @@ use rustc_middle::bug;
 use rustc_middle::ty::layout::{FnAbiExt, TyAndLayout};
 use rustc_middle::ty::print::obsolete::DefPathBasedNames;
 use rustc_middle::ty::{self, Ty, TypeFoldable};
-use rustc_target::abi::{Abi, Align, FieldsShape};
+use rustc_target::abi::{Abi, AddressSpace, Align, FieldsShape};
 use rustc_target::abi::{Int, Pointer, F32, F64};
 use rustc_target::abi::{LayoutOf, PointeeInfo, Scalar, Size, TyAndLayoutMethods, Variants};
 
@@ -310,12 +310,13 @@ impl<'tcx> LayoutLlvmExt<'tcx> for TyAndLayout<'tcx> {
             F64 => cx.type_f64(),
             Pointer => {
                 // If we know the alignment, pick something better than i8.
-                let pointee = if let Some(pointee) = self.pointee_info_at(cx, offset) {
-                    cx.type_pointee_for_align(pointee.align)
-                } else {
-                    cx.type_i8()
-                };
-                cx.type_ptr_to(pointee)
+                let (pointee, address_space) =
+                    if let Some(pointee) = self.pointee_info_at(cx, offset) {
+                        (cx.type_pointee_for_align(pointee.align), pointee.address_space)
+                    } else {
+                        (cx.type_i8(), AddressSpace::DATA)
+                    };
+                cx.type_ptr_to_ext(pointee, address_space)
             }
         }
     }
diff --git a/src/librustc_middle/ty/layout.rs b/src/librustc_middle/ty/layout.rs
index cb937bf0112..dc775b15927 100644
--- a/src/librustc_middle/ty/layout.rs
+++ b/src/librustc_middle/ty/layout.rs
@@ -2166,16 +2166,31 @@ where
     }
 
     fn pointee_info_at(this: TyAndLayout<'tcx>, cx: &C, offset: Size) -> Option<PointeeInfo> {
-        match this.ty.kind {
+        let addr_space_of_ty = |ty: Ty<'tcx>| {
+            if ty.is_fn() { cx.data_layout().instruction_address_space } else { AddressSpace::DATA }
+        };
+
+        let pointee_info = match this.ty.kind {
             ty::RawPtr(mt) if offset.bytes() == 0 => {
                 cx.layout_of(mt.ty).to_result().ok().map(|layout| PointeeInfo {
                     size: layout.size,
                     align: layout.align.abi,
                     safe: None,
+                    address_space: addr_space_of_ty(mt.ty),
+                })
+            }
+            ty::FnPtr(fn_sig) if offset.bytes() == 0 => {
+                cx.layout_of(cx.tcx().mk_fn_ptr(fn_sig)).to_result().ok().map(|layout| {
+                    PointeeInfo {
+                        size: layout.size,
+                        align: layout.align.abi,
+                        safe: None,
+                        address_space: cx.data_layout().instruction_address_space,
+                    }
                 })
             }
-
             ty::Ref(_, ty, mt) if offset.bytes() == 0 => {
+                let address_space = addr_space_of_ty(ty);
                 let tcx = cx.tcx();
                 let is_freeze = ty.is_freeze(tcx.at(DUMMY_SP), cx.param_env());
                 let kind = match mt {
@@ -2210,6 +2225,7 @@ where
                     size: layout.size,
                     align: layout.align.abi,
                     safe: Some(kind),
+                    address_space,
                 })
             }
 
@@ -2254,7 +2270,9 @@ where
                             result = field.to_result().ok().and_then(|field| {
                                 if ptr_end <= field_start + field.size {
                                     // We found the right field, look inside it.
-                                    field.pointee_info_at(cx, offset - field_start)
+                                    let field_info =
+                                        field.pointee_info_at(cx, offset - field_start);
+                                    field_info
                                 } else {
                                     None
                                 }
@@ -2277,7 +2295,14 @@ where
 
                 result
             }
-        }
+        };
+
+        debug!(
+            "pointee_info_at (offset={:?}, type kind: {:?}) => {:?}",
+            offset, this.ty.kind, pointee_info
+        );
+
+        pointee_info
     }
 }
 
diff --git a/src/librustc_target/abi/mod.rs b/src/librustc_target/abi/mod.rs
index a570bd914ae..b3e5f5c0c74 100644
--- a/src/librustc_target/abi/mod.rs
+++ b/src/librustc_target/abi/mod.rs
@@ -1024,7 +1024,7 @@ impl<T, E> MaybeResult<T> for Result<T, E> {
     }
 }
 
-#[derive(Copy, Clone, PartialEq, Eq)]
+#[derive(Copy, Clone, PartialEq, Eq, Debug)]
 pub enum PointerKind {
     /// Most general case, we know no restrictions to tell LLVM.
     Shared,
@@ -1039,11 +1039,12 @@ pub enum PointerKind {
     UniqueOwned,
 }
 
-#[derive(Copy, Clone)]
+#[derive(Copy, Clone, Debug)]
 pub struct PointeeInfo {
     pub size: Size,
     pub align: Align,
     pub safe: Option<PointerKind>,
+    pub address_space: AddressSpace,
 }
 
 pub trait TyAndLayoutMethods<'a, C: LayoutOf<Ty = Self>>: Sized {