Transmute cleanup

compiler/rustc_transmute/src/layout/dfa.rs

@@ -124,6 +124,11 @@ where
     }
 
     /// Concatenate two `Dfa`s.
+    ///
+    /// A unit (`start == accept`) acts as an identity. For other inputs, state
+    /// identifiers must be disjoint, and `self.accept` must not be a key in
+    /// `self.transitions`. The method joins the inputs by replacing `other.start`
+    /// with `self.accept`.
     pub(crate) fn concat(self, other: Self) -> Self {
         if self.start == self.accept {
             return other;
@@ -139,11 +144,10 @@ where
         for (source, transition) in other.transitions {
             let fix_state = |state| if state == other.start { self.accept } else { state };
             let byte_transitions = transition.byte_transitions.map_states(&fix_state);
-            let ref_transitions = transition
-                .ref_transitions
-                .into_iter()
-                .map(|(r, state)| (r, fix_state(state)))
-                .collect();
+            let mut ref_transitions = transition.ref_transitions;
+            for state in ref_transitions.values_mut() {
+                *state = fix_state(*state);
+            }
 
             let old = transitions
                 .insert(fix_state(source), Transitions { byte_transitions, ref_transitions });
@@ -154,6 +158,10 @@ where
     }
 
     /// Compute the union of two `Dfa`s.
+    ///
+    /// `new_state` must return a distinct identifier on each call. The output has
+    /// its own state namespace: its identifiers may also occur in either input,
+    /// and the two inputs may share identifiers with each other.
     pub(crate) fn union(self, other: Self, mut new_state: impl FnMut() -> State) -> Self {
         // We implement `union` by lazily initializing a set of states
         // corresponding to the product of states in `self` and `other`, and
@@ -230,18 +238,23 @@ where
                 },
             );
 
-            let ref_transitions =
-                a_transitions.ref_transitions.keys().chain(b_transitions.ref_transitions.keys());
-
-            let ref_transitions = ref_transitions
-                .map(|ref_transition| {
-                    let a_dst = a_transitions.ref_transitions.get(ref_transition).copied();
-                    let b_dst = b_transitions.ref_transitions.get(ref_transition).copied();
-
-                    assert!(a_dst.is_some() || b_dst.is_some());
+            let a_refs = a_transitions.ref_transitions.iter().map(|(r, &a_dst)| {
+                let b_dst = b_transitions.ref_transitions.get(r).copied();
+                (*r, Some(a_dst), b_dst)
+            });
+            let b_refs = b_transitions.ref_transitions.iter().filter_map(|(r, &b_dst)| {
+                if a_transitions.ref_transitions.contains_key(r) {
+                    None
+                } else {
+                    Some((*r, None, Some(b_dst)))
+                }
+            });
 
+            let ref_transitions = a_refs
+                .chain(b_refs)
+                .map(|(r, a_dst, b_dst)| {
                     queue.enqueue(a_dst, b_dst);
-                    (*ref_transition, mapped((a_dst, b_dst)))
+                    (r, mapped((a_dst, b_dst)))
                 })
                 .collect();
 
@@ -366,6 +379,15 @@ mod edge_set {
             S: Ord,
         {
             edges.sort();
+            for (range, _) in &edges {
+                assert!(
+                    range.start < range.end && range.end <= Byte::UNINIT + 1,
+                    "invalid byte edge range: {range:?}",
+                );
+            }
+            for pair in edges.windows(2) {
+                assert!(pair[0].0.end <= pair[1].0.start, "byte edge ranges overlap");
+            }
             Self { runs: edges.into() }
         }
 

compiler/rustc_transmute/src/layout/dfa/tests.rs

@@ -1,6 +1,6 @@
 use std::ops::Range;
 
-use super::{Byte, Dfa, EdgeSet, union};
+use super::{Byte, Dfa, EdgeSet, Reference, State, union};
 
 fn bytes(range: Range<u16>) -> Byte {
     Byte { start: range.start, end: range.end }
@@ -115,6 +115,115 @@ fn edge_set_union_coalesces_only_adjacent_ranges_with_equal_destinations() {
     );
 }
 
+#[test]
+fn edge_set_from_edges_accepts_empty_input() {
+    assert_eq!(EdgeSet::<u8>::from_edges(vec![]), EdgeSet::empty());
+}
+
+#[test]
+fn edge_set_from_edges_sorts_valid_ranges() {
+    let uninit = Byte::UNINIT;
+    let edges = EdgeSet::from_edges(vec![
+        (bytes(uninit..uninit + 1), 3),
+        (bytes(4..6), 2),
+        (bytes(0..4), 1),
+    ]);
+    assert_eq!(
+        edges.iter().collect::<Vec<_>>(),
+        [(bytes(0..4), 1), (bytes(4..6), 2), (bytes(uninit..uninit + 1), 3)],
+    );
+}
+
+#[test]
+#[should_panic(expected = "invalid byte edge range")]
+fn edge_set_from_edges_rejects_empty_range() {
+    EdgeSet::from_edges(vec![(bytes(2..2), 0)]);
+}
+
+#[test]
+#[should_panic(expected = "invalid byte edge range")]
+fn edge_set_from_edges_rejects_reversed_range() {
+    EdgeSet::from_edges(vec![(bytes(4..2), 0)]);
+}
+
+#[test]
+#[should_panic(expected = "invalid byte edge range")]
+fn edge_set_from_edges_rejects_range_past_uninit() {
+    EdgeSet::from_edges(vec![(bytes(0..Byte::UNINIT + 2), 0)]);
+}
+
+#[test]
+#[should_panic(expected = "byte edge ranges overlap")]
+fn edge_set_from_edges_rejects_overlapping_ranges() {
+    EdgeSet::from_edges(vec![(bytes(3..5), 0), (bytes(1..4), 1)]);
+}
+
+fn reference(region: usize) -> Reference<usize, ()> {
+    Reference { region, is_mut: false, referent: (), referent_size: 0, referent_align: 1 }
+}
+
+#[test]
+fn concat_preserves_reference_edges() {
+    let first_ref = reference(1);
+    let second_ref = reference(2);
+    let first = Dfa::from_ref(first_ref);
+    let second = Dfa::from_ref(second_ref).concat(Dfa::from_byte(7u8.into()));
+    let start = first.start;
+    let boundary = first.accept;
+    let second_start = second.start;
+    let after_second_ref = second.refs_from(second.start).next().unwrap().1;
+    let accept = second.accept;
+
+    let concatenated = first.concat(second);
+
+    assert_eq!(concatenated.start, start);
+    assert_eq!(concatenated.accept, accept);
+    assert_eq!(concatenated.refs_from(start).collect::<Vec<_>>(), [(first_ref, boundary)]);
+    assert_eq!(
+        concatenated.refs_from(boundary).collect::<Vec<_>>(),
+        [(second_ref, after_second_ref)],
+    );
+    assert_eq!(
+        concatenated.bytes_from(after_second_ref).collect::<Vec<_>>(),
+        [(7u8.into(), accept)]
+    );
+    assert!(!concatenated.transitions.contains_key(&second_start));
+}
+
+#[test]
+fn unit_is_concat_identity_for_reference_graph() {
+    let dfa = Dfa::from_ref(reference(1)).concat(Dfa::from_byte(7u8.into()));
+    assert_eq!(Dfa::unit().concat(dfa.clone()), dfa);
+    assert_eq!(dfa.clone().concat(Dfa::unit()), dfa);
+}
+
+#[test]
+fn union_merges_shared_and_distinct_reference_edges_in_order() {
+    let shared_ref = reference(1);
+    let a_only_ref = reference(2);
+    let b_only_ref = reference(3);
+    let mut a = Dfa::from_ref(shared_ref).concat(Dfa::from_byte(1u8.into()));
+    let mut b = Dfa::from_ref(shared_ref).concat(Dfa::from_byte(5u8.into()));
+    let a_continuation = a.refs_from(a.start).next().unwrap().1;
+    let b_continuation = b.refs_from(b.start).next().unwrap().1;
+    a.transitions.get_mut(&a.start).unwrap().ref_transitions.insert(a_only_ref, a_continuation);
+    b.transitions.get_mut(&b.start).unwrap().ref_transitions.insert(b_only_ref, b_continuation);
+
+    let merged = a.union(b, State::new);
+    let refs: Vec<_> = merged.refs_from(merged.start).collect();
+
+    assert_eq!(
+        refs.iter().map(|&(r, _)| r).collect::<Vec<_>>(),
+        [shared_ref, a_only_ref, b_only_ref]
+    );
+    assert_eq!(
+        merged.bytes_from(refs[0].1).collect::<Vec<_>>(),
+        [(1u8.into(), merged.accept), (5u8.into(), merged.accept)],
+    );
+    assert_eq!(merged.bytes_from(refs[1].1).collect::<Vec<_>>(), [(1u8.into(), merged.accept)]);
+    assert_eq!(merged.bytes_from(refs[2].1).collect::<Vec<_>>(), [(5u8.into(), merged.accept)]);
+}
+
 #[test]
 fn dot_distinguishes_start_and_accept() {
     for dfa in [Dfa::<!, !>::from_edges(0, 1, &[(0, 0u8, 1)]), Dfa::unit()] {

compiler/rustc_transmute/src/layout/tree.rs

@@ -302,33 +302,24 @@ pub(crate) mod rustc {
 
     impl<'tcx> Tree<Def<'tcx>, Region<'tcx>, Ty<'tcx>> {
         pub(crate) fn from_ty(ty: Ty<'tcx>, cx: LayoutCx<'tcx>) -> Result<Self, Err> {
-            use rustc_abi::HasDataLayout;
             let layout = layout_of(cx, ty)?;
+            Self::from_ty_and_layout((ty, layout), cx)
+        }
 
+        /// Constructs a tree using a layout already queried for `ty` in `cx`.
+        /// Callers pass the original type so reference nodes retain its regions.
+        fn from_ty_and_layout(
+            (ty, layout): (Ty<'tcx>, Layout<'tcx>),
+            cx: LayoutCx<'tcx>,
+        ) -> Result<Self, Err> {
             if let Err(e) = ty.error_reported() {
                 return Err(Err::TypeError(e));
             }
 
-            let target = cx.data_layout();
-            let pointer_size = target.pointer_size();
-
             match ty.kind() {
                 ty::Bool => Ok(Self::bool()),
 
-                ty::Float(nty) => {
-                    let width = nty.bit_width() / 8;
-                    Ok(Self::number(width.try_into().unwrap()))
-                }
-
-                ty::Int(nty) => {
-                    let width = nty.normalize(pointer_size.bits() as _).bit_width().unwrap() / 8;
-                    Ok(Self::number(width.try_into().unwrap()))
-                }
-
-                ty::Uint(nty) => {
-                    let width = nty.normalize(pointer_size.bits() as _).bit_width().unwrap() / 8;
-                    Ok(Self::number(width.try_into().unwrap()))
-                }
+                ty::Float(_) | ty::Int(_) | ty::Uint(_) => Ok(Self::number(layout.size.bytes())),
 
                 ty::Tuple(members) => Self::from_tuple((ty, layout), members, cx),
 
@@ -338,7 +329,7 @@ pub(crate) mod rustc {
                     };
                     let inner_layout = layout_of(cx, *inner_ty)?;
                     assert_eq!(*stride, inner_layout.size);
-                    let elt = Tree::from_ty(*inner_ty, cx)?;
+                    let elt = Self::from_ty_and_layout((*inner_ty, inner_layout), cx)?;
                     Ok(std::iter::repeat_n(elt, *count as usize)
                         .fold(Tree::unit(), |tree, elt| tree.then(elt)))
                 }
@@ -451,13 +442,11 @@ pub(crate) mod rustc {
                     // currently always the first field of the layout.
                     assert_eq!(tag_field, FieldIdx::ZERO);
 
-                    let variants = def.discriminants(cx.tcx()).try_fold(
-                        Self::uninhabited(),
-                        |variants, (idx, _discriminant)| {
+                    let variants =
+                        def.variant_range().try_fold(Self::uninhabited(), |variants, idx| {
                             let variant = layout_of_variant(idx, Some(tag_encoding))?;
                             Result::<Self, Err>::Ok(variants.or(variant))
-                        },
-                    )?;
+                        })?;
 
                     Ok(Self::def(Def::Adt(def)).then(variants))
                 }
@@ -517,7 +506,7 @@ pub(crate) mod rustc {
 
                 let field_ty = ty_field(cx, (ty, layout), field_idx);
                 let field_layout = layout_of(cx, field_ty)?;
-                let field_tree = Self::from_ty(field_ty, cx)?;
+                let field_tree = Self::from_ty_and_layout((field_ty, field_layout), cx)?;
 
                 struct_tree = struct_tree.then(padding).then(field_tree);
 
@@ -574,7 +563,7 @@ pub(crate) mod rustc {
                 |fields, (idx, _field_def)| {
                     let field_ty = ty_field(cx, (ty, layout), idx);
                     let field_layout = layout_of(cx, field_ty)?;
-                    let field = Self::from_ty(field_ty, cx)?;
+                    let field = Self::from_ty_and_layout((field_ty, field_layout), cx)?;
                     let trailing_padding_needed = layout.size - field_layout.size;
                     let trailing_padding = Self::padding(trailing_padding_needed.bytes_usize());
                     let field_and_padding = field.then(trailing_padding);