Safe Transmute: Enable handling references

compiler/rustc_trait_selection/src/solve/eval_ctxt.rs

@@ -668,18 +668,16 @@ impl<'tcx> EvalCtxt<'_, 'tcx> {
         scope: Ty<'tcx>,
         assume: rustc_transmute::Assume,
     ) -> Result<Certainty, NoSolution> {
+        use rustc_transmute::Answer;
         // FIXME(transmutability): This really should be returning nested goals for `Answer::If*`
         match rustc_transmute::TransmuteTypeEnv::new(self.infcx).is_transmutable(
             ObligationCause::dummy(),
             src_and_dst,
             scope,
             assume,
         ) {
-            rustc_transmute::Answer::Yes => Ok(Certainty::Yes),
-            rustc_transmute::Answer::No(_)
-            | rustc_transmute::Answer::IfTransmutable { .. }
-            | rustc_transmute::Answer::IfAll(_)
-            | rustc_transmute::Answer::IfAny(_) => Err(NoSolution),
+            Answer::Yes => Ok(Certainty::Yes),
+            Answer::No(_) | Answer::If(_) => Err(NoSolution),
         }
     }
 }

compiler/rustc_trait_selection/src/traits/error_reporting/mod.rs

@@ -65,6 +65,11 @@ pub struct ImplCandidate<'tcx> {
     pub similarity: CandidateSimilarity,
 }
 
+enum GetSafeTransmuteErrorAndReason {
+    Silent,
+    Error { err_msg: String, safe_transmute_explanation: String },
+}
+
 pub trait InferCtxtExt<'tcx> {
     /// Given some node representing a fn-like thing in the HIR map,
     /// returns a span and `ArgKind` information that describes the
@@ -724,11 +729,17 @@ impl<'tcx> TypeErrCtxtExt<'tcx> for TypeErrCtxt<'_, 'tcx> {
                             == self.tcx.lang_items().transmute_trait()
                         {
                             // Recompute the safe transmute reason and use that for the error reporting
-                            self.get_safe_transmute_error_and_reason(
+                            match self.get_safe_transmute_error_and_reason(
                                 obligation.clone(),
                                 trait_ref,
                                 span,
-                            )
+                            ) {
+                                GetSafeTransmuteErrorAndReason::Silent => return,
+                                GetSafeTransmuteErrorAndReason::Error {
+                                    err_msg,
+                                    safe_transmute_explanation,
+                                } => (err_msg, Some(safe_transmute_explanation)),
+                            }
                         } else {
                             (err_msg, None)
                         };
@@ -1292,7 +1303,7 @@ trait InferCtxtPrivExt<'tcx> {
         obligation: PredicateObligation<'tcx>,
         trait_ref: ty::PolyTraitRef<'tcx>,
         span: Span,
-    ) -> (String, Option<String>);
+    ) -> GetSafeTransmuteErrorAndReason;
 
     fn add_tuple_trait_message(
         &self,
@@ -2738,7 +2749,9 @@ impl<'tcx> InferCtxtPrivExt<'tcx> for TypeErrCtxt<'_, 'tcx> {
         obligation: PredicateObligation<'tcx>,
         trait_ref: ty::PolyTraitRef<'tcx>,
         span: Span,
-    ) -> (String, Option<String>) {
+    ) -> GetSafeTransmuteErrorAndReason {
+        use rustc_transmute::Answer;
+
         // Erase regions because layout code doesn't particularly care about regions.
         let trait_ref = self.tcx.erase_regions(self.tcx.erase_late_bound_regions(trait_ref));
 
@@ -2751,19 +2764,20 @@ impl<'tcx> InferCtxtPrivExt<'tcx> for TypeErrCtxt<'_, 'tcx> {
             rustc_transmute::Assume::from_const(self.infcx.tcx, obligation.param_env, trait_ref.substs.const_at(3)) else {
                 span_bug!(span, "Unable to construct rustc_transmute::Assume where it was previously possible");
             };
+
         match rustc_transmute::TransmuteTypeEnv::new(self.infcx).is_transmutable(
             obligation.cause,
             src_and_dst,
             scope,
             assume,
         ) {
-            rustc_transmute::Answer::No(reason) => {
+            Answer::No(reason) => {
                 let dst = trait_ref.substs.type_at(0);
                 let src = trait_ref.substs.type_at(1);
-                let custom_err_msg = format!(
+                let err_msg = format!(
                     "`{src}` cannot be safely transmuted into `{dst}` in the defining scope of `{scope}`"
                 );
-                let reason_msg = match reason {
+                let safe_transmute_explanation = match reason {
                     rustc_transmute::Reason::SrcIsUnspecified => {
                         format!("`{src}` does not have a well-specified layout")
                     }
@@ -2779,19 +2793,39 @@ impl<'tcx> InferCtxtPrivExt<'tcx> for TypeErrCtxt<'_, 'tcx> {
                     rustc_transmute::Reason::DstIsPrivate => format!(
                         "`{dst}` is or contains a type or field that is not visible in that scope"
                     ),
-                    // FIXME(bryangarza): Include the number of bytes of src and dst
                     rustc_transmute::Reason::DstIsTooBig => {
                         format!("The size of `{src}` is smaller than the size of `{dst}`")
                     }
+                    rustc_transmute::Reason::DstHasStricterAlignment {
+                        src_min_align,
+                        dst_min_align,
+                    } => {
+                        format!(
+                            "The minimum alignment of `{src}` ({src_min_align}) should be greater than that of `{dst}` ({dst_min_align})"
+                        )
+                    }
+                    rustc_transmute::Reason::DstIsMoreUnique => {
+                        format!("`{src}` is a shared reference, but `{dst}` is a unique reference")
+                    }
+                    // Already reported by rustc
+                    rustc_transmute::Reason::TypeError => {
+                        return GetSafeTransmuteErrorAndReason::Silent;
+                    }
+                    rustc_transmute::Reason::SrcLayoutUnknown => {
+                        format!("`{src}` has an unknown layout")
+                    }
+                    rustc_transmute::Reason::DstLayoutUnknown => {
+                        format!("`{dst}` has an unknown layout")
+                    }
                 };
-                (custom_err_msg, Some(reason_msg))
+                GetSafeTransmuteErrorAndReason::Error { err_msg, safe_transmute_explanation }
             }
             // Should never get a Yes at this point! We already ran it before, and did not get a Yes.
-            rustc_transmute::Answer::Yes => span_bug!(
+            Answer::Yes => span_bug!(
                 span,
                 "Inconsistent rustc_transmute::is_transmutable(...) result, got Yes",
             ),
-            _ => span_bug!(span, "Unsupported rustc_transmute::Reason variant"),
+            other => span_bug!(span, "Unsupported rustc_transmute::Answer variant: `{other:?}`"),
         }
     }
 

compiler/rustc_trait_selection/src/traits/select/confirmation.rs

@@ -6,14 +6,15 @@
 //!
 //! [rustc dev guide]:
 //! https://rustc-dev-guide.rust-lang.org/traits/resolution.html#confirmation
+use rustc_ast::Mutability;
 use rustc_data_structures::stack::ensure_sufficient_stack;
 use rustc_hir::lang_items::LangItem;
 use rustc_infer::infer::LateBoundRegionConversionTime::HigherRankedType;
 use rustc_infer::infer::{DefineOpaqueTypes, InferOk};
 use rustc_middle::traits::SelectionOutputTypeParameterMismatch;
 use rustc_middle::ty::{
     self, Binder, GenericParamDefKind, InternalSubsts, SubstsRef, ToPolyTraitRef, ToPredicate,
-    TraitRef, Ty, TyCtxt, TypeVisitableExt,
+    TraitPredicate, TraitRef, Ty, TyCtxt, TypeVisitableExt,
 };
 use rustc_session::config::TraitSolver;
 use rustc_span::def_id::DefId;
@@ -279,11 +280,60 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> {
         ImplSourceBuiltinData { nested: obligations }
     }
 
+    #[instrument(level = "debug", skip(self))]
     fn confirm_transmutability_candidate(
         &mut self,
         obligation: &TraitObligation<'tcx>,
     ) -> Result<ImplSourceBuiltinData<PredicateObligation<'tcx>>, SelectionError<'tcx>> {
-        debug!(?obligation, "confirm_transmutability_candidate");
+        use rustc_transmute::{Answer, Condition};
+        #[instrument(level = "debug", skip(tcx, obligation, predicate))]
+        fn flatten_answer_tree<'tcx>(
+            tcx: TyCtxt<'tcx>,
+            obligation: &TraitObligation<'tcx>,
+            predicate: TraitPredicate<'tcx>,
+            cond: Condition<rustc_transmute::layout::rustc::Ref<'tcx>>,
+        ) -> Vec<PredicateObligation<'tcx>> {
+            match cond {
+                // FIXME(bryangarza): Add separate `IfAny` case, instead of treating as `IfAll`
+                // Not possible until the trait solver supports disjunctions of obligations
+                Condition::IfAll(conds) | Condition::IfAny(conds) => conds
+                    .into_iter()
+                    .flat_map(|cond| flatten_answer_tree(tcx, obligation, predicate, cond))
+                    .collect(),
+                Condition::IfTransmutable { src, dst } => {
+                    let trait_def_id = obligation.predicate.def_id();
+                    let scope = predicate.trait_ref.substs.type_at(2);
+                    let assume_const = predicate.trait_ref.substs.const_at(3);
+                    let make_obl = |from_ty, to_ty| {
+                        let trait_ref1 = ty::TraitRef::new(
+                            tcx,
+                            trait_def_id,
+                            [
+                                ty::GenericArg::from(to_ty),
+                                ty::GenericArg::from(from_ty),
+                                ty::GenericArg::from(scope),
+                                ty::GenericArg::from(assume_const),
+                            ],
+                        );
+                        Obligation::with_depth(
+                            tcx,
+                            obligation.cause.clone(),
+                            obligation.recursion_depth + 1,
+                            obligation.param_env,
+                            trait_ref1,
+                        )
+                    };
+
+                    // If Dst is mutable, check bidirectionally.
+                    // For example, transmuting bool -> u8 is OK as long as you can't update that u8
+                    // to be > 1, because you could later transmute the u8 back to a bool and get UB.
+                    match dst.mutability {
+                        Mutability::Not => vec![make_obl(src.ty, dst.ty)],
+                        Mutability::Mut => vec![make_obl(src.ty, dst.ty), make_obl(dst.ty, src.ty)],
+                    }
+                }
+            }
+        }
 
         // We erase regions here because transmutability calls layout queries,
         // which does not handle inference regions and doesn't particularly
@@ -301,21 +351,25 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> {
             return Err(Unimplemented);
         };
 
+        let dst = predicate.trait_ref.substs.type_at(0);
+        let src = predicate.trait_ref.substs.type_at(1);
+        debug!(?src, ?dst);
         let mut transmute_env = rustc_transmute::TransmuteTypeEnv::new(self.infcx);
         let maybe_transmutable = transmute_env.is_transmutable(
             obligation.cause.clone(),
-            rustc_transmute::Types {
-                dst: predicate.trait_ref.substs.type_at(0),
-                src: predicate.trait_ref.substs.type_at(1),
-            },
+            rustc_transmute::Types { dst, src },
             predicate.trait_ref.substs.type_at(2),
             assume,
         );
 
-        match maybe_transmutable {
-            rustc_transmute::Answer::Yes => Ok(ImplSourceBuiltinData { nested: vec![] }),
-            _ => Err(Unimplemented),
-        }
+        let fully_flattened = match maybe_transmutable {
+            Answer::No(_) => Err(Unimplemented)?,
+            Answer::If(cond) => flatten_answer_tree(self.tcx(), obligation, predicate, cond),
+            Answer::Yes => vec![],
+        };
+
+        debug!(?fully_flattened);
+        Ok(ImplSourceBuiltinData { nested: fully_flattened })
     }
 
     /// This handles the case where an `auto trait Foo` impl is being used.

compiler/rustc_transmute/src/layout/mod.rs

@@ -30,33 +30,49 @@ impl fmt::Debug for Byte {
 }
 
 pub(crate) trait Def: Debug + Hash + Eq + PartialEq + Copy + Clone {}
-pub trait Ref: Debug + Hash + Eq + PartialEq + Copy + Clone {}
+pub trait Ref: Debug + Hash + Eq + PartialEq + Copy + Clone {
+    fn min_align(&self) -> usize;
+
+    fn is_mutable(&self) -> bool;
+}
 
 impl Def for ! {}
-impl Ref for ! {}
+impl Ref for ! {
+    fn min_align(&self) -> usize {
+        unreachable!()
+    }
+    fn is_mutable(&self) -> bool {
+        unreachable!()
+    }
+}
 
 #[cfg(feature = "rustc")]
-pub(crate) mod rustc {
+pub mod rustc {
     use rustc_middle::mir::Mutability;
-    use rustc_middle::ty;
-    use rustc_middle::ty::Region;
-    use rustc_middle::ty::Ty;
+    use rustc_middle::ty::{self, Ty};
 
     /// A reference in the layout.
     #[derive(Debug, Hash, Eq, PartialEq, PartialOrd, Ord, Clone, Copy)]
     pub struct Ref<'tcx> {
-        lifetime: Region<'tcx>,
-        ty: Ty<'tcx>,
-        mutability: Mutability,
+        pub lifetime: ty::Region<'tcx>,
+        pub ty: Ty<'tcx>,
+        pub mutability: Mutability,
+        pub align: usize,
     }
 
-    impl<'tcx> super::Ref for Ref<'tcx> {}
+    impl<'tcx> super::Ref for Ref<'tcx> {
+        fn min_align(&self) -> usize {
+            self.align
+        }
 
-    impl<'tcx> Ref<'tcx> {
-        pub fn min_align(&self) -> usize {
-            todo!()
+        fn is_mutable(&self) -> bool {
+            match self.mutability {
+                Mutability::Mut => true,
+                Mutability::Not => false,
+            }
         }
     }
+    impl<'tcx> Ref<'tcx> {}
 
     /// A visibility node in the layout.
     #[derive(Debug, Hash, Eq, PartialEq, Clone, Copy)]

compiler/rustc_transmute/src/layout/tree.rs

@@ -188,14 +188,14 @@ pub(crate) mod rustc {
         /// The layout of the type is unspecified.
         Unspecified,
         /// This error will be surfaced elsewhere by rustc, so don't surface it.
-        Unknown,
+        UnknownLayout,
         TypeError(ErrorGuaranteed),
     }
 
     impl<'tcx> From<LayoutError<'tcx>> for Err {
         fn from(err: LayoutError<'tcx>) -> Self {
             match err {
-                LayoutError::Unknown(..) => Self::Unknown,
+                LayoutError::Unknown(..) => Self::UnknownLayout,
                 err => unimplemented!("{:?}", err),
             }
         }
@@ -365,6 +365,17 @@ pub(crate) mod rustc {
                         }
                     }))
                 }
+
+                ty::Ref(lifetime, ty, mutability) => {
+                    let align = layout_of(tcx, *ty)?.align();
+                    Ok(Tree::Ref(Ref {
+                        lifetime: *lifetime,
+                        ty: *ty,
+                        mutability: *mutability,
+                        align,
+                    }))
+                }
+
                 _ => Err(Err::Unspecified),
             }
         }