Basic cleanup in rustc_transmute

compiler/rustc_transmute/src/layout/dfa.rs

@@ -5,6 +5,9 @@ use std::sync::atomic::{AtomicU32, Ordering};
 use super::{Byte, Reference, Region, Tree, Type, Uninhabited};
 use crate::{Map, Set};
 
+#[cfg(test)]
+mod tests;
+
 #[derive(PartialEq)]
 #[cfg_attr(test, derive(Clone))]
 pub(crate) struct Dfa<R, T>
@@ -37,7 +40,9 @@ where
     }
 }
 
-/// The states in a [`Dfa`] represent byte offsets.
+/// An identifier for a node in a [`Dfa`].
+///
+/// The numeric identifier does not encode a byte offset.
 #[derive(Hash, Eq, PartialEq, PartialOrd, Ord, Copy, Clone)]
 pub(crate) struct State(pub(crate) u32);
 
@@ -308,7 +313,8 @@ where
 {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         writeln!(f, "digraph {{")?;
-        writeln!(f, "    {:?} [shape = doublecircle]", self.start)?;
+        writeln!(f, "    start [shape = point, style = invis]")?;
+        writeln!(f, "    start -> {:?}", self.start)?;
         writeln!(f, "    {:?} [shape = doublecircle]", self.accept)?;
 
         for (src, transitions) in self.transitions.iter() {
@@ -380,15 +386,7 @@ mod edge_set {
         }
 
         pub(crate) fn map_states<SS>(self, mut f: impl FnMut(S) -> SS) -> EdgeSet<SS> {
-            EdgeSet {
-                // NOTE: It appears as through `<Vec<_> as
-                // IntoIterator>::IntoIter` and `std::iter::Map` both implement
-                // `TrustedLen`, which in turn means that this `.collect()`
-                // allocates the correct number of elements once up-front [1].
-                //
-                // [1] https://doc.rust-lang.org/1.85.0/src/alloc/vec/spec_from_iter_nested.rs.html#47
-                runs: self.runs.into_iter().map(|(b, s)| (b, f(s))).collect(),
-            }
+            EdgeSet { runs: self.runs.into_iter().map(|(b, s)| (b, f(s))).collect() }
         }
 
         /// Unions two edge sets together.
@@ -428,7 +426,13 @@ mod edge_set {
     }
 }
 
-/// Merges two sorted sequences into one sorted sequence.
+/// Partitions two sequences of byte edges into sorted, non-overlapping ranges.
+///
+/// Within each input, ranges must be non-empty, non-overlapping, and sorted by
+/// ascending start. Each output item contains a range and the destination of the
+/// edge from each input that covers it. An input with no edge covering the range
+/// contributes `None`. Ranges covered by neither input are omitted. Adjacent
+/// output ranges are not coalesced.
 pub(crate) fn union<S: Copy, X: Iterator<Item = (Byte, S)>, Y: Iterator<Item = (Byte, S)>>(
     xs: X,
     ys: Y,

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

@@ -0,0 +1,131 @@
+use std::ops::Range;
+
+use super::{Byte, Dfa, EdgeSet, union};
+
+fn bytes(range: Range<u16>) -> Byte {
+    Byte { start: range.start, end: range.end }
+}
+
+fn assert_union(
+    xs: &[(Byte, u8)],
+    ys: &[(Byte, u8)],
+    expected: &[(Byte, (Option<u8>, Option<u8>))],
+) {
+    let actual: Vec<_> = union(xs.iter().copied(), ys.iter().copied()).collect();
+    assert_eq!(actual, expected);
+
+    let actual: Vec<_> = union(ys.iter().copied(), xs.iter().copied()).collect();
+    let expected: Vec<_> = expected.iter().map(|&(range, (x, y))| (range, (y, x))).collect();
+    assert_eq!(actual, expected);
+}
+
+#[test]
+fn union_empty_inputs() {
+    assert_union(&[], &[], &[]);
+    assert_union(
+        &[(bytes(1..3), 1), (bytes(5..7), 2)],
+        &[],
+        &[(bytes(1..3), (Some(1), None)), (bytes(5..7), (Some(2), None))],
+    );
+}
+
+#[test]
+fn union_disjoint_and_adjacent_ranges() {
+    assert_union(
+        &[(bytes(1..3), 1), (bytes(6..8), 2)],
+        &[(bytes(3..5), 3), (bytes(8..9), 4)],
+        &[
+            (bytes(1..3), (Some(1), None)),
+            (bytes(3..5), (None, Some(3))),
+            (bytes(6..8), (Some(2), None)),
+            (bytes(8..9), (None, Some(4))),
+        ],
+    );
+}
+
+#[test]
+fn union_equal_ranges() {
+    assert_union(&[(bytes(1..4), 1)], &[(bytes(1..4), 2)], &[(bytes(1..4), (Some(1), Some(2)))]);
+}
+
+#[test]
+fn union_partially_overlapping_ranges() {
+    assert_union(
+        &[(bytes(1..4), 1)],
+        &[(bytes(3..6), 2)],
+        &[
+            (bytes(1..3), (Some(1), None)),
+            (bytes(3..4), (Some(1), Some(2))),
+            (bytes(4..6), (None, Some(2))),
+        ],
+    );
+}
+
+#[test]
+fn union_contained_ranges() {
+    assert_union(
+        &[(bytes(0..8), 1)],
+        &[(bytes(2..4), 2), (bytes(4..6), 3)],
+        &[
+            (bytes(0..2), (Some(1), None)),
+            (bytes(2..4), (Some(1), Some(2))),
+            (bytes(4..6), (Some(1), Some(3))),
+            (bytes(6..8), (Some(1), None)),
+        ],
+    );
+}
+
+#[test]
+fn union_matching_starts_and_ends() {
+    assert_union(
+        &[(bytes(0..4), 1), (bytes(6..10), 2)],
+        &[(bytes(0..2), 3), (bytes(8..10), 4)],
+        &[
+            (bytes(0..2), (Some(1), Some(3))),
+            (bytes(2..4), (Some(1), None)),
+            (bytes(6..8), (Some(2), None)),
+            (bytes(8..10), (Some(2), Some(4))),
+        ],
+    );
+}
+
+#[test]
+fn union_uninit_boundary() {
+    let uninit = Byte::UNINIT;
+    assert_union(
+        &[(Byte::uninit(), 1)],
+        &[(bytes(255..uninit), 2), (bytes(uninit..uninit + 1), 3)],
+        &[
+            (bytes(0..255), (Some(1), None)),
+            (bytes(255..uninit), (Some(1), Some(2))),
+            (bytes(uninit..uninit + 1), (Some(1), Some(3))),
+        ],
+    );
+}
+
+#[test]
+fn edge_set_union_coalesces_only_adjacent_ranges_with_equal_destinations() {
+    let xs = EdgeSet::from_edges(vec![(bytes(0..4), 1), (bytes(8..10), 1)]);
+    let ys = EdgeSet::from_edges(vec![(bytes(2..6), 2), (bytes(10..12), 3)]);
+    let merged = xs.union(&ys, |_, y| if y == Some(3) { 8 } else { 7 });
+
+    assert_eq!(
+        merged.iter().collect::<Vec<_>>(),
+        [(bytes(0..6), 7), (bytes(8..10), 7), (bytes(10..12), 8)],
+    );
+}
+
+#[test]
+fn dot_distinguishes_start_and_accept() {
+    for dfa in [Dfa::<!, !>::from_edges(0, 1, &[(0, 0u8, 1)]), Dfa::unit()] {
+        let dot = format!("{dfa:?}");
+        assert!(dot.lines().any(|line| line == "    start [shape = point, style = invis]"));
+        let start_edge = format!("    start -> {:?}", dfa.start);
+        assert!(dot.lines().any(|line| line == start_edge));
+
+        let accept_marker = format!("    {:?} [shape = doublecircle]", dfa.accept);
+        let accept_markers: Vec<_> =
+            dot.lines().filter(|line| line.contains("doublecircle")).collect();
+        assert_eq!(accept_markers, [accept_marker.as_str()]);
+    }
+}

compiler/rustc_transmute/src/layout/mod.rs

@@ -11,12 +11,14 @@ pub(crate) use dfa::{Dfa, union};
 #[derive(Debug)]
 pub(crate) struct Uninhabited;
 
-/// A range of byte values (including an uninit byte value).
+/// A half-open range of byte values, which may include an uninitialized byte.
+///
+/// The range is `[start, end)`. Values `0..=255` represent initialized bytes,
+/// and `256` represents an uninitialized byte. A range containing that value
+/// therefore has an exclusive upper bound of `257`.
 #[derive(Hash, Eq, PartialEq, Ord, PartialOrd, Clone, Copy)]
 pub(crate) struct Byte {
-    // An inclusive-exclusive range. We use this instead of `Range` because `Range: !Copy`.
-    //
-    // Uninit byte value is represented by 256.
+    // Store the endpoints separately because `Range` is not `Copy`.
     pub(crate) start: u16,
     pub(crate) end: u16,
 }
@@ -37,6 +39,7 @@ impl Byte {
         Byte { start: val, end: val + 1 }
     }
 
+    /// A byte that may be initialized to any value or uninitialized.
     #[inline]
     fn uninit() -> Byte {
         Byte { start: 0, end: Self::UNINIT + 1 }
@@ -136,7 +139,7 @@ pub mod rustc {
     use rustc_middle::ty::layout::{HasTyCtxt, LayoutCx, LayoutError};
     use rustc_middle::ty::{self, Region, Ty};
 
-    /// A visibility node in the layout.
+    /// A layout annotation used to check for safety invariants during pruning.
     #[derive(Debug, Hash, Eq, PartialEq, Clone, Copy)]
     pub enum Def<'tcx> {
         Adt(ty::AdtDef<'tcx>),
@@ -147,9 +150,9 @@ pub mod rustc {
 
     impl<'tcx> super::Def for Def<'tcx> {
         fn has_safety_invariants(&self) -> bool {
-            // Rust presently has no notion of 'unsafe fields', so for now we
-            // make the conservative assumption that everything besides
-            // primitive types carry safety invariants.
+            // Conservatively treat every ADT, variant, and field definition as
+            // potentially carrying safety invariants. This check does not
+            // inspect `unsafe` field annotations.
             self != &Self::Primitive
         }
     }

compiler/rustc_transmute/src/layout/tree.rs

@@ -7,6 +7,14 @@ mod tests;
 
 /// A tree-based representation of a type layout.
 ///
+/// A `Seq` concatenates layouts, while an `Alt` chooses among them. An empty
+/// `Seq` represents an inhabited, zero-sized layout; an empty `Alt` represents
+/// an uninhabited layout.
+///
+/// `Def` nodes are zero-width annotations used by [`Tree::prune`] to decide
+/// which branches to retain. Pruning removes these annotations and produces
+/// a `Tree<!, R, T>` for conversion with [`super::Dfa::from_tree`].
+///
 /// Invariants:
 /// 1. All paths through the layout have the same length (in bytes).
 ///
@@ -25,7 +33,7 @@ where
     Seq(Vec<Self>),
     /// A choice between alternative layouts.
     Alt(Vec<Self>),
-    /// A definition node.
+    /// A zero-width definition annotation used during pruning.
     Def(D),
     /// A reference node.
     Ref(Reference<R, T>),
@@ -70,7 +78,7 @@ where
         Self::Seq(Vec::new())
     }
 
-    /// A `Tree` containing a single, uninitialized byte.
+    /// A `Tree` containing one byte that may be initialized to any value or uninitialized.
     pub(crate) fn uninit() -> Self {
         Self::Byte(Byte::uninit())
     }
@@ -134,12 +142,18 @@ where
     }
 
     /// A `Tree` whose layout is entirely padding of the given width.
+    ///
+    /// Each padding byte may be initialized to any value or uninitialized.
     pub(crate) fn padding(width_in_bytes: usize) -> Self {
         Self::Seq(vec![Self::uninit(); width_in_bytes])
     }
 
-    /// Remove all `Def` nodes, and all branches of the layout for which `f`
-    /// produces `true`.
+    /// Removes all `Def` nodes and rejects branches whose definitions make `f` return `true`.
+    ///
+    /// A sequence becomes uninhabited if any of its elements becomes uninhabited;
+    /// alternatives retain their surviving branches. Retained definitions become
+    /// empty sequences, so the result contains no definition nodes, as expressed
+    /// by its `Tree<!, R, T>` type.
     pub(crate) fn prune<F>(self, f: &F) -> Tree<!, R, T>
     where
         F: Fn(D) -> bool,
@@ -473,7 +487,7 @@ pub(crate) mod rustc {
             };
 
             // When this function is invoked with enum variants,
-            // `ty_and_layout.size` does not encompass the entire size of the
+            // `layout.size` does not encompass the entire size of the
             // enum. We rely on `total_size` for this.
             assert!(layout.size <= total_size);