Auto merge of #85263 - Smittyvb:thir-unsafeck-union-field, r=oli-obk

Check for union field accesses in THIR unsafeck

see also #85259, #83129, https://github.com/rust-lang/project-thir-unsafeck/issues/7

r? `@LeSeulArtichaut`

2 files changed, 130 insertions, 6 deletions

diff --git a/compiler/rustc_mir_build/src/check_unsafety.rs b/compiler/rustc_mir_build/src/check_unsafety.rs
index 9bef8ac2413..c3af98fae9d 100644
--- a/compiler/rustc_mir_build/src/check_unsafety.rs
+++ b/compiler/rustc_mir_build/src/check_unsafety.rs
@@ -26,6 +26,8 @@ struct UnsafetyVisitor<'a, 'tcx> {
     /// calls to functions with `#[target_feature]` (RFC 2396).
     body_target_features: &'tcx Vec<Symbol>,
     is_const: bool,
+    in_possible_lhs_union_assign: bool,
+    in_union_destructure: bool,
 }
 
 impl<'tcx> UnsafetyVisitor<'_, 'tcx> {
@@ -158,14 +160,115 @@ impl<'a, 'tcx> Visitor<'a, 'tcx> for UnsafetyVisitor<'a, 'tcx> {
         }
     }
 
+    fn visit_pat(&mut self, pat: &Pat<'tcx>) {
+        use PatKind::*;
+
+        if self.in_union_destructure {
+            match *pat.kind {
+                // binding to a variable allows getting stuff out of variable
+                Binding { .. }
+                // match is conditional on having this value
+                | Constant { .. }
+                | Variant { .. }
+                | Leaf { .. }
+                | Deref { .. }
+                | Range { .. }
+                | Slice { .. }
+                | Array { .. } => {
+                    self.requires_unsafe(pat.span, AccessToUnionField);
+                    return; // don't walk pattern
+                }
+                // wildcard doesn't take anything
+                Wild |
+                // these just wrap other patterns
+                Or { .. } |
+                AscribeUserType { .. } => {}
+            }
+        };
+
+        if let ty::Adt(adt_def, _) = pat.ty.kind() {
+            // check for extracting values from union via destructuring
+            if adt_def.is_union() {
+                match *pat.kind {
+                    // assigning the whole union is okay
+                    // let x = Union { ... };
+                    // let y = x; // safe
+                    Binding { .. } |
+                    // binding to wildcard is okay since that never reads anything and stops double errors
+                    // with implict wildcard branches from `if let`s
+                    Wild |
+                    // doesn't have any effect on semantics
+                    AscribeUserType { .. } |
+                    // creating a union literal
+                    Constant { .. } => {},
+                    Leaf { .. } | Or { .. } => {
+                        // pattern matching with a union and not doing something like v = Union { bar: 5 }
+                        self.in_union_destructure = true;
+                        visit::walk_pat(self, pat);
+                        self.in_union_destructure = false;
+                        return; // don't walk pattern
+                    }
+                    Variant { .. } | Deref { .. } | Range { .. } | Slice { .. } | Array { .. } =>
+                        unreachable!("impossible union destructuring type"),
+                }
+            }
+        }
+
+        visit::walk_pat(self, pat);
+    }
+
     fn visit_expr(&mut self, expr: &Expr<'tcx>) {
+        // could we be in a the LHS of an assignment of a union?
+        match expr.kind {
+            ExprKind::Field { .. }
+            | ExprKind::VarRef { .. }
+            | ExprKind::UpvarRef { .. }
+            | ExprKind::Scope { .. }
+            | ExprKind::Cast { .. } => {}
+
+            ExprKind::AddressOf { .. }
+            | ExprKind::Adt { .. }
+            | ExprKind::Array { .. }
+            | ExprKind::Binary { .. }
+            | ExprKind::Block { .. }
+            | ExprKind::Borrow { .. }
+            | ExprKind::Literal { .. }
+            | ExprKind::ConstBlock { .. }
+            | ExprKind::Deref { .. }
+            | ExprKind::Index { .. }
+            | ExprKind::NeverToAny { .. }
+            | ExprKind::PlaceTypeAscription { .. }
+            | ExprKind::ValueTypeAscription { .. }
+            | ExprKind::Pointer { .. }
+            | ExprKind::Repeat { .. }
+            | ExprKind::StaticRef { .. }
+            | ExprKind::ThreadLocalRef { .. }
+            | ExprKind::Tuple { .. }
+            | ExprKind::Unary { .. }
+            | ExprKind::Call { .. }
+            | ExprKind::Assign { .. }
+            | ExprKind::AssignOp { .. }
+            | ExprKind::Break { .. }
+            | ExprKind::Closure { .. }
+            | ExprKind::Continue { .. }
+            | ExprKind::Return { .. }
+            | ExprKind::Yield { .. }
+            | ExprKind::Loop { .. }
+            | ExprKind::Match { .. }
+            | ExprKind::Box { .. }
+            | ExprKind::If { .. }
+            | ExprKind::InlineAsm { .. }
+            | ExprKind::LlvmInlineAsm { .. }
+            | ExprKind::LogicalOp { .. }
+            | ExprKind::Use { .. } => self.in_possible_lhs_union_assign = false,
+        };
         match expr.kind {
             ExprKind::Scope { value, lint_level: LintLevel::Explicit(hir_id), region_scope: _ } => {
                 let prev_id = self.hir_context;
                 self.hir_context = hir_id;
                 self.visit_expr(&self.thir[value]);
                 self.hir_context = prev_id;
-                return;
+                return; // don't visit the whole expression
             }
             ExprKind::Call { fun, ty: _, args: _, from_hir_call: _, fn_span: _ } => {
                 if self.thir[fun].ty.fn_sig(self.tcx).unsafety() == hir::Unsafety::Unsafe {
@@ -246,9 +349,29 @@ impl<'a, 'tcx> Visitor<'a, 'tcx> for UnsafetyVisitor<'a, 'tcx> {
                 // Unsafe blocks can be used in closures, make sure to take it into account
                 self.safety_context = closure_visitor.safety_context;
             }
+            ExprKind::Field { lhs, .. } => {
+                // assigning to union field is okay for AccessToUnionField
+                if let ty::Adt(adt_def, _) = &self.thir[lhs].ty.kind() {
+                    if adt_def.is_union() {
+                        if self.in_possible_lhs_union_assign {
+                            // FIXME: trigger AssignToDroppingUnionField unsafety if needed
+                        } else {
+                            self.requires_unsafe(expr.span, AccessToUnionField);
+                        }
+                    }
+                }
+            }
+            // don't have any special handling for AssignOp since it causes a read *and* write to lhs
+            ExprKind::Assign { lhs, rhs } => {
+                // assigning to a union is safe, check here so it doesn't get treated as a read later
+                self.in_possible_lhs_union_assign = true;
+                visit::walk_expr(self, &self.thir()[lhs]);
+                self.in_possible_lhs_union_assign = false;
+                visit::walk_expr(self, &self.thir()[rhs]);
+                return; // don't visit the whole expression
+            }
             _ => {}
         }
-
         visit::walk_expr(self, expr);
     }
 }
@@ -296,7 +419,6 @@ enum UnsafeOpKind {
     DerefOfRawPointer,
     #[allow(dead_code)] // FIXME
     AssignToDroppingUnionField,
-    #[allow(dead_code)] // FIXME
     AccessToUnionField,
     #[allow(dead_code)] // FIXME
     MutationOfLayoutConstrainedField,
@@ -417,6 +539,8 @@ pub fn check_unsafety<'tcx>(tcx: TyCtxt<'tcx>, def: ty::WithOptConstParam<LocalD
         body_unsafety,
         body_target_features,
         is_const,
+        in_possible_lhs_union_assign: false,
+        in_union_destructure: false,
     };
     visitor.visit_expr(&thir[expr]);
 }
diff --git a/compiler/rustc_mir_build/src/thir/visit.rs b/compiler/rustc_mir_build/src/thir/visit.rs
index f611bb6eb43..ce5d4362c08 100644
--- a/compiler/rustc_mir_build/src/thir/visit.rs
+++ b/compiler/rustc_mir_build/src/thir/visit.rs
@@ -153,8 +153,8 @@ pub fn walk_expr<'a, 'tcx: 'a, V: Visitor<'a, 'tcx>>(visitor: &mut V, expr: &Exp
 }
 
 pub fn walk_stmt<'a, 'tcx: 'a, V: Visitor<'a, 'tcx>>(visitor: &mut V, stmt: &Stmt<'tcx>) {
-    match stmt.kind {
-        StmtKind::Expr { expr, scope: _ } => visitor.visit_expr(&visitor.thir()[expr]),
+    match &stmt.kind {
+        StmtKind::Expr { expr, scope: _ } => visitor.visit_expr(&visitor.thir()[*expr]),
         StmtKind::Let {
             initializer,
             remainder_scope: _,
@@ -163,7 +163,7 @@ pub fn walk_stmt<'a, 'tcx: 'a, V: Visitor<'a, 'tcx>>(visitor: &mut V, stmt: &Stm
             lint_level: _,
         } => {
             if let Some(init) = initializer {
-                visitor.visit_expr(&visitor.thir()[init]);
+                visitor.visit_expr(&visitor.thir()[*init]);
             }
             visitor.visit_pat(pattern);
         }