Rollup merge of #158170 - usamoi:loop-unroll-stmt, r=JonathanBrouwer,saethlin

parse attrs if a loop's parent hir node is a stmt

fix this case:

```rust
// unroll doesn't work
#[unsafe(no_mangle)]
fn f(n: usize) -> usize {
    let mut sum = 0;
    let mut i = 0;
    #[unroll(4)]
    loop { if i < n { sum += i; i += 1; } else { break; } }
    sum
}

// unroll works
#[unsafe(no_mangle)]
fn g(n: usize) {
    let mut i = 0;
    #[unroll(4)]
    loop { if i < n { i += 1; } else { break; } }
}
```

cc #156874

r? saethlin

Author: jhpratt

compiler/rustc_mir_build/src/thir/cx/expr.rs

@@ -22,7 +22,7 @@ use rustc_middle::ty::{
     UpvarArgs,
 };
 use rustc_middle::{bug, span_bug};
-use rustc_span::Span;
+use rustc_span::{DesugaringKind, Span};
 use tracing::{debug, info, instrument, trace};
 
 use crate::diagnostics::*;
@@ -72,18 +72,24 @@ impl<'tcx> ThirBuildCx<'tcx> {
 
         let mut attrs = ThinVec::new();
 
-        if let ExprKind::Loop { .. } = expr.kind {
-            // For loops defined with loop and while, the expr already has the attrs
-            if let hir::Node::Block(_) = self.tcx.parent_hir_node(hir_expr.hir_id) {
-                attrs = parsed_attrs(hir_expr.hir_id, self.tcx);
-            }
-
-            // For loop desugaring puts us pretty deep down the HIR tree
-            if let hir::Node::Arm(arm) = self.tcx.parent_hir_node(hir_expr.hir_id)
-                && let hir::Node::Expr(expr) = self.tcx.parent_hir_node(arm.hir_id)
-                && let hir::Node::Expr(expr) = self.tcx.parent_hir_node(expr.hir_id)
-            {
-                attrs = parsed_attrs(expr.hir_id, self.tcx);
+        if let hir::ExprKind::Loop(_, _, _, span) = hir_expr.kind {
+            match span.desugaring_kind() {
+                // `for` loop desugaring puts us pretty deep down the HIR tree
+                Some(DesugaringKind::ForLoop) => {
+                    let arm = self.tcx.parent_hir_node(hir_expr.hir_id).expect_arm();
+                    let expr = self.tcx.parent_hir_node(arm.hir_id).expect_expr();
+                    std::assert_matches!(expr.kind, hir::ExprKind::Match(..));
+                    // ignore async for loops
+                    if let hir::Node::Expr(expr) = self.tcx.parent_hir_node(expr.hir_id) {
+                        std::assert_matches!(expr.kind, hir::ExprKind::DropTemps(..));
+                        attrs = parsed_attrs(expr.hir_id, self.tcx)
+                    }
+                }
+                // For loops defined with `loop` and `while`, the expr already has the attrs
+                Some(DesugaringKind::WhileLoop) | None => {
+                    attrs = parsed_attrs(hir_expr.hir_id, self.tcx);
+                }
+                _ => (),
             }
         }
 

tests/codegen-llvm/loop-attrs/unroll-loop-metadata.rs

@@ -2,6 +2,7 @@
 
 #![crate_type = "lib"]
 #![feature(loop_hints)]
+#![feature(stmt_expr_attributes)]
 
 // This test ensures that we emit the expected LLVM metadata for loop hint attributes.
 // It does not test that loops are optimized as expected, because successful unrolling removes the
@@ -26,5 +27,63 @@ pub fn unroll_hint() {
     }
 }
 
+// HIR for a `loop` statement is a bit different, make sure we still apply the metadata in that
+// case.
+
+#[no_mangle]
+pub fn unroll_full() {
+    // CHECK-LABEL: @unroll_full
+    // CHECK: !llvm.loop ![[FULL:[0-9]+]]
+    let mut i = 0;
+    let _return = (#[unroll(full)]
+    loop {
+        unsafe { maybe_has_side_effect() }
+        i += 1;
+        if i >= 10 {
+            break 1;
+        }
+    });
+}
+
+#[no_mangle]
+pub fn unroll_never() {
+    // CHECK-LABEL: @unroll_never
+    // CHECK: !llvm.loop ![[DISABLE:[0-9]+]]
+    let mut i = 0;
+    let _return = (1 + #[unroll(never)]
+    loop {
+        unsafe { maybe_has_side_effect() }
+        i += 1;
+        if i >= 10 {
+            break 1;
+        }
+    });
+}
+
+#[no_mangle]
+pub fn unroll_count() {
+    // CHECK-LABEL: @unroll_count
+    // CHECK: !llvm.loop ![[COUNT:[0-9]+]]
+    let mut i = 0;
+    #[unroll(5)]
+    loop {
+        unsafe { maybe_has_side_effect() }
+        i += 1;
+        if i >= 10 {
+            break;
+        }
+    }
+    unsafe { maybe_has_side_effect() }
+}
+
 // CHECK: ![[HINT]] = distinct !{![[HINT]], ![[INNER_HINT:[0-9]+]]}
 // CHECK: ![[INNER_HINT]] = !{!"llvm.loop.unroll.enable"}
+
+// CHECK: ![[FULL]] = distinct !{![[FULL]], ![[INNER_FULL:[0-9]+]]}
+// CHECK: ![[INNER_FULL]] = !{!"llvm.loop.unroll.full"}
+
+// CHECK: ![[DISABLE]] = distinct !{![[DISABLE]], ![[INNER_DISABLE:[0-9]+]]}
+// CHECK: ![[INNER_DISABLE]] = !{!"llvm.loop.unroll.disable"}
+
+// CHECK: ![[COUNT]] = distinct !{![[COUNT]], ![[INNER_COUNT:[0-9]+]]}
+// CHECK: ![[INNER_COUNT]] = !{!"llvm.loop.unroll.count", i32 5}