From 0db06bf0046b70dee03de07fd6b455173a90c117 Mon Sep 17 00:00:00 2001
From: Michael Goulet <michael@errs.io>
Date: Mon, 18 Mar 2024 13:41:08 -0400
Subject: [PATCH] Detect allocator for box in must_not_suspend lint

---
 compiler/rustc_mir_transform/src/coroutine.rs | 14 ++++++---
 tests/ui/lint/must_not_suspend/allocator.rs   | 30 +++++++++++++++++++
 .../ui/lint/must_not_suspend/allocator.stderr | 22 ++++++++++++++
 3 files changed, 62 insertions(+), 4 deletions(-)
 create mode 100644 tests/ui/lint/must_not_suspend/allocator.rs
 create mode 100644 tests/ui/lint/must_not_suspend/allocator.stderr

diff --git a/compiler/rustc_mir_transform/src/coroutine.rs b/compiler/rustc_mir_transform/src/coroutine.rs
index 54b13a40e92dd..0d18d4fd69e6f 100644
--- a/compiler/rustc_mir_transform/src/coroutine.rs
+++ b/compiler/rustc_mir_transform/src/coroutine.rs
@@ -1964,15 +1964,21 @@ fn check_must_not_suspend_ty<'tcx>(
     debug!("Checking must_not_suspend for {}", ty);
 
     match *ty.kind() {
-        ty::Adt(..) if ty.is_box() => {
-            let boxed_ty = ty.boxed_ty();
-            let descr_pre = &format!("{}boxed ", data.descr_pre);
+        ty::Adt(_, args) if ty.is_box() => {
+            let boxed_ty = args.type_at(0);
+            let allocator_ty = args.type_at(1);
             check_must_not_suspend_ty(
                 tcx,
                 boxed_ty,
                 hir_id,
                 param_env,
-                SuspendCheckData { descr_pre, ..data },
+                SuspendCheckData { descr_pre: &format!("{}boxed ", data.descr_pre), ..data },
+            ) || check_must_not_suspend_ty(
+                tcx,
+                allocator_ty,
+                hir_id,
+                param_env,
+                SuspendCheckData { descr_pre: &format!("{}allocator ", data.descr_pre), ..data },
             )
         }
         ty::Adt(def, _) => check_must_not_suspend_def(tcx, def.did(), hir_id, data),
diff --git a/tests/ui/lint/must_not_suspend/allocator.rs b/tests/ui/lint/must_not_suspend/allocator.rs
new file mode 100644
index 0000000000000..c2ceb3297f37f
--- /dev/null
+++ b/tests/ui/lint/must_not_suspend/allocator.rs
@@ -0,0 +1,30 @@
+//@ edition: 2021
+
+#![feature(must_not_suspend, allocator_api)]
+#![deny(must_not_suspend)]
+
+use std::alloc::*;
+use std::ptr::NonNull;
+
+#[must_not_suspend]
+struct MyAllocatorWhichMustNotSuspend;
+
+unsafe impl Allocator for MyAllocatorWhichMustNotSuspend {
+    fn allocate(&self, l: Layout) -> Result<NonNull<[u8]>, AllocError> {
+        Global.allocate(l)
+    }
+    unsafe fn deallocate(&self, p: NonNull<u8>, l: Layout) {
+        Global.deallocate(p, l)
+    }
+}
+
+async fn suspend() {}
+
+async fn foo() {
+    let x = Box::new_in(1i32, MyAllocatorWhichMustNotSuspend);
+    //~^ ERROR allocator `MyAllocatorWhichMustNotSuspend` held across a suspend point, but should not be
+    suspend().await;
+    drop(x);
+}
+
+fn main() {}
diff --git a/tests/ui/lint/must_not_suspend/allocator.stderr b/tests/ui/lint/must_not_suspend/allocator.stderr
new file mode 100644
index 0000000000000..959945f2626cd
--- /dev/null
+++ b/tests/ui/lint/must_not_suspend/allocator.stderr
@@ -0,0 +1,22 @@
+error: allocator `MyAllocatorWhichMustNotSuspend` held across a suspend point, but should not be
+  --> $DIR/allocator.rs:24:9
+   |
+LL |     let x = Box::new_in(1i32, MyAllocatorWhichMustNotSuspend);
+   |         ^
+LL |
+LL |     suspend().await;
+   |               ----- the value is held across this suspend point
+   |
+help: consider using a block (`{ ... }`) to shrink the value's scope, ending before the suspend point
+  --> $DIR/allocator.rs:24:9
+   |
+LL |     let x = Box::new_in(1i32, MyAllocatorWhichMustNotSuspend);
+   |         ^
+note: the lint level is defined here
+  --> $DIR/allocator.rs:4:9
+   |
+LL | #![deny(must_not_suspend)]
+   |         ^^^^^^^^^^^^^^^^
+
+error: aborting due to 1 previous error
+