fixed fields chain, added new test cases

Author: Kivooeo

compiler/rustc_mir_build/src/check_unsafety.rs

@@ -45,6 +45,8 @@ struct UnsafetyVisitor<'a, 'tcx> {
     /// Flag to ensure that we only suggest wrapping the entire function body in
     /// an unsafe block once.
     suggest_unsafe_block: bool,
+    /// Controls how union field accesses are checked
+    union_field_access_mode: UnionFieldAccessMode,
 }
 
 impl<'tcx> UnsafetyVisitor<'_, 'tcx> {
@@ -218,6 +220,7 @@ impl<'tcx> UnsafetyVisitor<'_, 'tcx> {
                 inside_adt: false,
                 warnings: self.warnings,
                 suggest_unsafe_block: self.suggest_unsafe_block,
+                union_field_access_mode: UnionFieldAccessMode::Normal,
             };
             // params in THIR may be unsafe, e.g. a union pattern.
             for param in &inner_thir.params {
@@ -658,18 +661,25 @@ impl<'a, 'tcx> Visitor<'a, 'tcx> for UnsafetyVisitor<'a, 'tcx> {
                     } else if adt_def.is_union() {
                         // Check if this field access is part of a raw borrow operation
                         // If so, we've already handled it above and shouldn't reach here
-                        if let Some(assigned_ty) = self.assignment_info {
-                            if assigned_ty.needs_drop(self.tcx, self.typing_env) {
-                                // This would be unsafe, but should be outright impossible since we
-                                // reject such unions.
-                                assert!(
-                                    self.tcx.dcx().has_errors().is_some(),
-                                    "union fields that need dropping should be impossible: {assigned_ty}"
-                                );
+                        match self.union_field_access_mode {
+                            UnionFieldAccessMode::SuppressUnionFieldAccessError => {
+                                // Suppress AccessToUnionField error for union fields chains
+                            }
+                            UnionFieldAccessMode::Normal => {
+                                if let Some(assigned_ty) = self.assignment_info {
+                                    if assigned_ty.needs_drop(self.tcx, self.typing_env) {
+                                        // This would be unsafe, but should be outright impossible since we
+                                        // reject such unions.
+                                        assert!(
+                                            self.tcx.dcx().has_errors().is_some(),
+                                            "union fields that need dropping should be impossible: {assigned_ty}"
+                                        );
+                                    }
+                                } else {
+                                    // Only require unsafe if this is not a raw borrow operation
+                                    self.requires_unsafe(expr.span, AccessToUnionField);
+                                }
                             }
-                        } else {
-                            // Only require unsafe if this is not a raw borrow operation
-                            self.requires_unsafe(expr.span, AccessToUnionField);
                         }
                     }
                 }
@@ -728,7 +738,7 @@ impl<'a, 'tcx> UnsafetyVisitor<'a, 'tcx> {
         match self.thir[expr_id].kind {
             ExprKind::Field { lhs, .. } => {
                 let lhs = &self.thir[lhs];
-                if let ty::Adt(adt_def, _) = lhs.ty.kind() { adt_def.is_union() } else { false }
+                matches!(lhs.ty.kind(), ty::Adt(adt_def, _) if adt_def.is_union())
             }
             _ => false,
         }
@@ -737,28 +747,28 @@ impl<'a, 'tcx> UnsafetyVisitor<'a, 'tcx> {
     /// Visit a union field access in the context of a raw borrow operation
     /// This ensures we still check safety of nested operations while allowing
     /// the raw pointer creation itself
-    fn visit_union_field_for_raw_borrow(&mut self, expr_id: ExprId) {
-        match self.thir[expr_id].kind {
-            ExprKind::Field { lhs, variant_index, name } => {
-                let lhs_expr = &self.thir[lhs];
-                if let ty::Adt(adt_def, _) = lhs_expr.ty.kind() {
-                    // Check for unsafe fields but skip the union access check
-                    if adt_def.variant(variant_index).fields[name].safety.is_unsafe() {
-                        self.requires_unsafe(self.thir[expr_id].span, UseOfUnsafeField);
-                    }
-                    // For unions, we don't require unsafe for raw pointer creation
-                    // But we still need to check the LHS for safety
-                    self.visit_expr(lhs_expr);
-                } else {
-                    // Not a union, use normal visiting
-                    visit::walk_expr(self, &self.thir[expr_id]);
+    fn visit_union_field_for_raw_borrow(&mut self, mut expr_id: ExprId) {
+        let prev = self.union_field_access_mode;
+        self.union_field_access_mode = UnionFieldAccessMode::SuppressUnionFieldAccessError;
+        // Walk through the chain of union field accesses using while let
+        while let ExprKind::Field { lhs, variant_index, name } = self.thir[expr_id].kind {
+            let lhs_expr = &self.thir[lhs];
+            if let ty::Adt(adt_def, _) = lhs_expr.ty.kind() {
+                // Check for unsafe fields but skip the union access check
+                if adt_def.variant(variant_index).fields[name].safety.is_unsafe() {
+                    self.requires_unsafe(self.thir[expr_id].span, UseOfUnsafeField);
                 }
-            }
-            _ => {
-                // Not a field access, use normal visiting
+                // If the LHS is also a union field access, keep walking
+                expr_id = lhs;
+            } else {
+                // Not a union, use normal visiting
                 visit::walk_expr(self, &self.thir[expr_id]);
+                return;
             }
         }
+        // Visit the base expression for any nested safety checks
+        self.visit_expr(&self.thir[expr_id]);
+        self.union_field_access_mode = prev;
     }
 }
 
@@ -770,6 +780,13 @@ enum SafetyContext {
     UnsafeBlock { span: Span, hir_id: HirId, used: bool, nested_used_blocks: Vec<NestedUsedBlock> },
 }
 
+/// Controls how union field accesses are checked
+#[derive(Clone, Copy)]
+enum UnionFieldAccessMode {
+    Normal,
+    SuppressUnionFieldAccessError,
+}
+
 #[derive(Clone, Copy)]
 struct NestedUsedBlock {
     hir_id: HirId,
@@ -1244,6 +1261,7 @@ pub(crate) fn check_unsafety(tcx: TyCtxt<'_>, def: LocalDefId) {
         inside_adt: false,
         warnings: &mut warnings,
         suggest_unsafe_block: true,
+        union_field_access_mode: UnionFieldAccessMode::Normal,
     };
     // params in THIR may be unsafe, e.g. a union pattern.
     for param in &thir.params {

tests/ui/union/union-unsafe.rs

@@ -96,4 +96,38 @@ fn main() {
     let mut u3 = U3 { a: ManuallyDrop::new(String::from("old")) }; // OK
     u3.a = ManuallyDrop::new(String::from("new")); // OK (assignment does not drop)
     *u3.a = String::from("new"); //~ ERROR access to union field is unsafe
+
+    let mut unions = [U1 { a: 1 }, U1 { a: 2 }];
+
+    // Array indexing + union field raw borrow - should be OK
+    let ptr = &raw mut unions[0].a; // OK
+    let ptr2 = &raw const unions[1].a; // OK
+
+    // Test for union fields chain, this should be allowed
+    #[derive(Copy, Clone)]
+    union Inner {
+        a: u8,
+    }
+
+    union MoreInner {
+        moreinner: ManuallyDrop<Inner>,
+    }
+
+    union LessOuter {
+        lessouter: ManuallyDrop<MoreInner>,
+    }
+
+    union Outer {
+        outer: ManuallyDrop<LessOuter>,
+    }
+
+    let super_outer = Outer {
+        outer: ManuallyDrop::new(LessOuter {
+            lessouter: ManuallyDrop::new(MoreInner {
+                moreinner: ManuallyDrop::new(Inner { a: 42 }),
+            }),
+        }),
+    };
+
+    let ptr = &raw const super_outer.outer.lessouter.moreinner.a;
 }