Safe Transmute: Enable handling references, including recursive types

This patch enables support for references in Safe Transmute, by generating
nested obligations during trait selection. Specifically, when we call
`confirm_transmutability_candidate(...)`, we now recursively traverse the
`rustc_transmute::Answer` tree and create obligations for all the `Answer`
variants, some of which include multiple nested `Answer`s.

Also, to handle recursive types, enable support for coinduction for the Safe
Transmute trait (`BikeshedIntrinsicFrom`) by adding the `#[rustc_coinduction]`
annotation.

Also fix some small logic issues when reducing the `or` and `and` combinations
in `rustc_transmute`, so that we don't end up with additional redundant
`Answer`s in the tree.

Co-authored-by: Jack Wrenn <jack@wrenn.fyi>

Author: bryangarza

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

@@ -2783,6 +2783,14 @@ impl<'tcx> InferCtxtPrivExt<'tcx> for TypeErrCtxt<'_, 'tcx> {
                     rustc_transmute::Reason::DstIsTooBig => {
                         format!("The size of `{src}` is smaller than the size of `{dst}`")
                     }
+                    rustc_transmute::Reason::DstHasStricterAlignment => {
+                        format!(
+                            "The alignment of `{src}` should be stricter than that of `{dst}`, but it is not"
+                        )
+                    }
+                    rustc_transmute::Reason::DstIsMoreUnique => {
+                        format!("`{src}` is a shared reference, but `{dst}` is a unique reference")
+                    }
                 };
                 (custom_err_msg, Some(reason_msg))
             }
@@ -2791,7 +2799,7 @@ impl<'tcx> InferCtxtPrivExt<'tcx> for TypeErrCtxt<'_, 'tcx> {
                 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

@@ -13,7 +13,7 @@ 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,10 +279,61 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> {
         ImplSourceBuiltinData { nested: obligations }
     }
 
+    #[instrument(skip(self))]
     fn confirm_transmutability_candidate(
         &mut self,
         obligation: &TraitObligation<'tcx>,
     ) -> Result<ImplSourceBuiltinData<PredicateObligation<'tcx>>, SelectionError<'tcx>> {
+        fn flatten_answer_tree<'tcx>(
+            tcx: TyCtxt<'tcx>,
+            obligation: &TraitObligation<'tcx>,
+            predicate: TraitPredicate<'tcx>,
+            answer: rustc_transmute::Answer<rustc_transmute::layout::rustc::Ref<'tcx>>,
+        ) -> Result<Vec<PredicateObligation<'tcx>>, SelectionError<'tcx>> {
+            match answer {
+                rustc_transmute::Answer::Yes => Ok(vec![]),
+                rustc_transmute::Answer::No(_) => Err(Unimplemented),
+                // FIXME(bryangarza): Add separate `IfAny` case, instead of treating as `IfAll`
+                rustc_transmute::Answer::IfAll(answers)
+                | rustc_transmute::Answer::IfAny(answers) => {
+                    let mut nested = vec![];
+                    for flattened in answers
+                        .into_iter()
+                        .map(|answer| flatten_answer_tree(tcx, obligation, predicate, answer))
+                    {
+                        nested.extend(flattened?);
+                    }
+                    Ok(nested)
+                }
+                rustc_transmute::Answer::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 = tcx.mk_trait_ref(
+                            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,
+                        )
+                    };
+
+                    // // FIXME(bryangarza): Check src.mutability or dst.mutability to know whether dst -> src obligation is needed
+                    Ok(vec![make_obl(src.ty, dst.ty), make_obl(dst.ty, src.ty)])
+                }
+            }
+        }
+
         debug!(?obligation, "confirm_transmutability_candidate");
 
         // We erase regions here because transmutability calls layout queries,
@@ -312,10 +363,10 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> {
             assume,
         );
 
-        match maybe_transmutable {
-            rustc_transmute::Answer::Yes => Ok(ImplSourceBuiltinData { nested: vec![] }),
-            _ => Err(Unimplemented),
-        }
+        info!(?maybe_transmutable);
+        let nested = flatten_answer_tree(self.tcx(), obligation, predicate, maybe_transmutable)?;
+        info!(?nested);
+        Ok(ImplSourceBuiltinData { nested })
     }
 
     /// This handles the case where an `auto trait Foo` impl is being used.

compiler/rustc_transmute/src/layout/mod.rs

@@ -30,33 +30,46 @@ 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 {
+        1
+    }
+
+    fn is_mutable(&self) -> bool {
+        false
+    }
+}
 
 impl Def for ! {}
 impl Ref for ! {}
 
 #[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

@@ -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),
             }
         }

compiler/rustc_transmute/src/lib.rs

@@ -8,7 +8,7 @@ extern crate tracing;
 
 pub(crate) use rustc_data_structures::fx::{FxIndexMap as Map, FxIndexSet as Set};
 
-pub(crate) mod layout;
+pub mod layout;
 pub(crate) mod maybe_transmutable;
 
 #[derive(Default)]
@@ -21,10 +21,7 @@ pub struct Assume {
 
 /// The type encodes answers to the question: "Are these types transmutable?"
 #[derive(Debug, Hash, Eq, PartialEq, PartialOrd, Ord, Clone)]
-pub enum Answer<R>
-where
-    R: layout::Ref,
-{
+pub enum Answer<R> {
     /// `Src` is transmutable into `Dst`.
     Yes,
 
@@ -54,6 +51,10 @@ pub enum Reason {
     DstIsPrivate,
     /// `Dst` is larger than `Src`, and the excess bytes were not exclusively uninitialized.
     DstIsTooBig,
+    /// Src should have a stricter alignment than Dst, but it does not.
+    DstHasStricterAlignment,
+    /// Can't go from shared pointer to unique pointer
+    DstIsMoreUnique,
 }
 
 #[cfg(feature = "rustc")]