From 3311283f90dd62dcd322d6730f0dd09b3ee6d7fd Mon Sep 17 00:00:00 2001
From: Alex Crichton <alex@alexcrichton.com>
Date: Tue, 12 Mar 2024 15:10:01 -0700
Subject: [PATCH 1/7] Get wasm32-wasip2 compiling with its custom pal
 implementation

The ordering of targets in `pal/mod.rs` did not end up using the wasip2
implementation, so after reordering that I've edited the implementation
to compile correctly.
---
 library/std/src/sys/pal/mod.rs        | 6 +++---
 library/std/src/sys/pal/wasip2/mod.rs | 6 ------
 2 files changed, 3 insertions(+), 9 deletions(-)

diff --git a/library/std/src/sys/pal/mod.rs b/library/std/src/sys/pal/mod.rs
index 7c87deed371..8c75ac65299 100644
--- a/library/std/src/sys/pal/mod.rs
+++ b/library/std/src/sys/pal/mod.rs
@@ -37,12 +37,12 @@
     } else if #[cfg(target_os = "hermit")] {
         mod hermit;
         pub use self::hermit::*;
-    } else if #[cfg(target_os = "wasi")] {
-        mod wasi;
-        pub use self::wasi::*;
     } else if #[cfg(all(target_os = "wasi", target_env = "p2"))] {
         mod wasip2;
         pub use self::wasip2::*;
+    } else if #[cfg(target_os = "wasi")] {
+        mod wasi;
+        pub use self::wasi::*;
     } else if #[cfg(target_family = "wasm")] {
         mod wasm;
         pub use self::wasm::*;
diff --git a/library/std/src/sys/pal/wasip2/mod.rs b/library/std/src/sys/pal/wasip2/mod.rs
index d1d444d7b79..4f164a3ec18 100644
--- a/library/std/src/sys/pal/wasip2/mod.rs
+++ b/library/std/src/sys/pal/wasip2/mod.rs
@@ -10,8 +10,6 @@
 pub mod alloc;
 #[path = "../wasi/args.rs"]
 pub mod args;
-#[path = "../unix/cmath.rs"]
-pub mod cmath;
 #[path = "../wasi/env.rs"]
 pub mod env;
 #[path = "../wasi/fd.rs"]
@@ -28,10 +26,6 @@
 pub mod net;
 #[path = "../wasi/os.rs"]
 pub mod os;
-#[path = "../unix/os_str.rs"]
-pub mod os_str;
-#[path = "../unix/path.rs"]
-pub mod path;
 #[path = "../unsupported/pipe.rs"]
 pub mod pipe;
 #[path = "../unsupported/process.rs"]

From 5af8187325210e3006874b1457fa6ec69293f0e9 Mon Sep 17 00:00:00 2001
From: Alex Crichton <alex@alexcrichton.com>
Date: Tue, 12 Mar 2024 15:10:40 -0700
Subject: [PATCH 2/7] Provide `cabi_realloc` on `wasm32-wasip2` by default

This commit provides a component model intrinsic in the standard library
by default on the `wasm32-wasip2` target. This intrinsic is not
required by the component model itself but is quite common to use, for
example it's needed if a wasm module receives a string or a list.

The intention of this commit is to provide an overridable definition in
the standard library through a weak definition of this function. That
means that downstream crates can provide their own customized and more
specific versions if they'd like, but the standard library's version
should suffice for general-purpose use.
---
 .../std/src/sys/pal/wasip2/cabi_realloc.rs    | 65 +++++++++++++++++++
 library/std/src/sys/pal/wasip2/mod.rs         |  2 +
 2 files changed, 67 insertions(+)
 create mode 100644 library/std/src/sys/pal/wasip2/cabi_realloc.rs

diff --git a/library/std/src/sys/pal/wasip2/cabi_realloc.rs b/library/std/src/sys/pal/wasip2/cabi_realloc.rs
new file mode 100644
index 00000000000..820063173d6
--- /dev/null
+++ b/library/std/src/sys/pal/wasip2/cabi_realloc.rs
@@ -0,0 +1,65 @@
+//! This module contains a canonical definition of the `cabi_realloc` function
+//! for the component model.
+//!
+//! The component model's canonical ABI for representing datatypes in memory
+//! makes use of this function when transferring lists and strings, for example.
+//! This function behaves like C's `realloc` but also takes alignment into
+//! account.
+//!
+//! Components are notably not required to export this function, but nearly
+//! all components end up doing so currently. This definition in the standard
+//! library removes the need for all compilations to define this themselves.
+//!
+//! More information about the canonical ABI can be found at
+//! <https://github.com/WebAssembly/component-model/blob/main/design/mvp/CanonicalABI.md>
+//!
+//! Note that the name of this function is not standardized in the canonical ABI
+//! at this time. Instead it's a convention of the "componentization process"
+//! where a core wasm module is converted to a component to use this name.
+//! Additionally this is not the only possible definition of this function, so
+//! this is defined as a "weak" symbol. This means that other definitions are
+//! allowed to overwrite it if they are present in a compilation.
+
+use crate::alloc::{self, Layout};
+use crate::ptr;
+
+#[used]
+static FORCE_CODEGEN_OF_CABI_REALLOC: unsafe extern "C" fn(
+    *mut u8,
+    usize,
+    usize,
+    usize,
+) -> *mut u8 = cabi_realloc;
+
+#[linkage = "weak"]
+#[no_mangle]
+pub unsafe extern "C" fn cabi_realloc(
+    old_ptr: *mut u8,
+    old_len: usize,
+    align: usize,
+    new_len: usize,
+) -> *mut u8 {
+    let layout;
+    let ptr = if old_len == 0 {
+        if new_len == 0 {
+            return ptr::without_provenance_mut(align);
+        }
+        layout = Layout::from_size_align_unchecked(new_len, align);
+        alloc::alloc(layout)
+    } else {
+        debug_assert_ne!(new_len, 0, "non-zero old_len requires non-zero new_len!");
+        layout = Layout::from_size_align_unchecked(old_len, align);
+        alloc::realloc(old_ptr, layout, new_len)
+    };
+    if ptr.is_null() {
+        // Print a nice message in debug mode, but in release mode don't
+        // pull in so many dependencies related to printing so just emit an
+        // `unreachable` instruction.
+        if cfg!(debug_assertions) {
+            alloc::handle_alloc_error(layout);
+        } else {
+            super::abort_internal();
+        }
+    }
+    return ptr;
+}
diff --git a/library/std/src/sys/pal/wasip2/mod.rs b/library/std/src/sys/pal/wasip2/mod.rs
index 4f164a3ec18..94aa458d2f9 100644
--- a/library/std/src/sys/pal/wasip2/mod.rs
+++ b/library/std/src/sys/pal/wasip2/mod.rs
@@ -66,3 +66,5 @@
 use helpers::err2io;
 pub use helpers::hashmap_random_keys;
 pub use helpers::is_interrupted;
+
+mod cabi_realloc;

From 1f2d1420cb1c5c5cb52712955f8ed47dcfda8c67 Mon Sep 17 00:00:00 2001
From: Kai Luo <lkail@cn.ibm.com>
Date: Tue, 2 Apr 2024 17:17:13 +0800
Subject: [PATCH 3/7] Fix linking c++ runtimes on AIX

---
 compiler/rustc_llvm/build.rs | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/compiler/rustc_llvm/build.rs b/compiler/rustc_llvm/build.rs
index 4b0c1229da1..66fa4bde785 100644
--- a/compiler/rustc_llvm/build.rs
+++ b/compiler/rustc_llvm/build.rs
@@ -391,6 +391,11 @@ fn main() {
         }
     }
 
+    // libc++abi have to be specified explicitly on AIX.
+    if target.contains("aix") {
+        println!("cargo:rustc-link-lib=c++abi");
+    }
+
     // Libstdc++ depends on pthread which Rust doesn't link on MinGW
     // since nothing else requires it.
     if target.ends_with("windows-gnu") {

From 00f7f57159aac2e1df49db2eb87353338bef2bb0 Mon Sep 17 00:00:00 2001
From: Kai Luo <lkail@cn.ibm.com>
Date: Tue, 2 Apr 2024 17:25:22 +0800
Subject: [PATCH 4/7] Fix build on AIX

---
 compiler/rustc_llvm/build.rs | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/compiler/rustc_llvm/build.rs b/compiler/rustc_llvm/build.rs
index 66fa4bde785..e2c0ec90c7c 100644
--- a/compiler/rustc_llvm/build.rs
+++ b/compiler/rustc_llvm/build.rs
@@ -391,9 +391,10 @@ fn main() {
         }
     }
 
-    // libc++abi have to be specified explicitly on AIX.
+    // libc++abi and libunwind have to be specified explicitly on AIX.
     if target.contains("aix") {
         println!("cargo:rustc-link-lib=c++abi");
+        println!("cargo:rustc-link-lib=unwind");
     }
 
     // Libstdc++ depends on pthread which Rust doesn't link on MinGW

From a6b2d12c921552082d6849c75531da666a9c9c82 Mon Sep 17 00:00:00 2001
From: Tshepang Mbambo <tshepang@gmail.com>
Date: Tue, 2 Apr 2024 21:41:16 +0200
Subject: [PATCH 5/7] use a consistent style for links

---
 library/std/src/fs.rs | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/library/std/src/fs.rs b/library/std/src/fs.rs
index 9017ba79714..b1102b440e0 100644
--- a/library/std/src/fs.rs
+++ b/library/std/src/fs.rs
@@ -37,7 +37,7 @@
 ///
 /// # Examples
 ///
-/// Creates a new file and write bytes to it (you can also use [`write()`]):
+/// Creates a new file and write bytes to it (you can also use [`write`]):
 ///
 /// ```no_run
 /// use std::fs::File;
@@ -2018,7 +2018,7 @@ pub fn rename<P: AsRef<Path>, Q: AsRef<Path>>(from: P, to: Q) -> io::Result<()>
 /// the length of the `to` file as reported by `metadata`.
 ///
 /// If you want to copy the contents of one file to another and you’re
-/// working with [`File`]s, see the [`io::copy()`] function.
+/// working with [`File`]s, see the [`io::copy`](io::copy()) function.
 ///
 /// # Platform-specific behavior
 ///

From a1a1f41027c16a940c1de9e445064692110637c8 Mon Sep 17 00:00:00 2001
From: Michael Goulet <michael@errs.io>
Date: Mon, 1 Apr 2024 14:42:00 -0400
Subject: [PATCH 6/7] Fix capture analysis for by-move closure bodies

---
 .../src/coroutine/by_move_body.rs             | 68 ++++++++++----
 .../miri/tests/pass/async-closure-captures.rs | 89 +++++++++++++++++++
 .../tests/pass/async-closure-captures.stdout  | 10 +++
 .../ui/async-await/async-closures/captures.rs | 80 +++++++++++++++++
 .../async-closures/captures.run.stdout        | 10 +++
 5 files changed, 240 insertions(+), 17 deletions(-)
 create mode 100644 src/tools/miri/tests/pass/async-closure-captures.rs
 create mode 100644 src/tools/miri/tests/pass/async-closure-captures.stdout
 create mode 100644 tests/ui/async-await/async-closures/captures.rs
 create mode 100644 tests/ui/async-await/async-closures/captures.run.stdout

diff --git a/compiler/rustc_mir_transform/src/coroutine/by_move_body.rs b/compiler/rustc_mir_transform/src/coroutine/by_move_body.rs
index 99dbb342268..60b52fba219 100644
--- a/compiler/rustc_mir_transform/src/coroutine/by_move_body.rs
+++ b/compiler/rustc_mir_transform/src/coroutine/by_move_body.rs
@@ -3,6 +3,8 @@
 //! be a coroutine body that takes all of its upvars by-move, and which we stash
 //! into the `CoroutineInfo` for all coroutines returned by coroutine-closures.
 
+use itertools::Itertools;
+
 use rustc_data_structures::unord::UnordSet;
 use rustc_hir as hir;
 use rustc_middle::mir::visit::MutVisitor;
@@ -26,36 +28,68 @@ fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut mir::Body<'tcx>) {
         if coroutine_ty.references_error() {
             return;
         }
-        let ty::Coroutine(_, args) = *coroutine_ty.kind() else { bug!("{body:#?}") };
 
-        let coroutine_kind = args.as_coroutine().kind_ty().to_opt_closure_kind().unwrap();
+        let ty::Coroutine(_, args) = *coroutine_ty.kind() else { bug!("{body:#?}") };
+        let args = args.as_coroutine();
+
+        let coroutine_kind = args.kind_ty().to_opt_closure_kind().unwrap();
         if coroutine_kind == ty::ClosureKind::FnOnce {
             return;
         }
 
-        let mut by_ref_fields = UnordSet::default();
-        let by_move_upvars = Ty::new_tup_from_iter(
-            tcx,
-            tcx.closure_captures(coroutine_def_id).iter().enumerate().map(|(idx, capture)| {
-                if capture.is_by_ref() {
-                    by_ref_fields.insert(FieldIdx::from_usize(idx));
-                }
-                capture.place.ty()
-            }),
+        let parent_def_id = tcx.local_parent(coroutine_def_id);
+        let ty::CoroutineClosure(_, parent_args) =
+            *tcx.type_of(parent_def_id).instantiate_identity().kind()
+        else {
+            bug!();
+        };
+        let parent_args = parent_args.as_coroutine_closure();
+        let parent_upvars_ty = parent_args.tupled_upvars_ty();
+        let tupled_inputs_ty = tcx.instantiate_bound_regions_with_erased(
+            parent_args.coroutine_closure_sig().map_bound(|sig| sig.tupled_inputs_ty),
         );
+        let num_args = tupled_inputs_ty.tuple_fields().len();
+
+        let mut by_ref_fields = UnordSet::default();
+        for (idx, (coroutine_capture, parent_capture)) in tcx
+            .closure_captures(coroutine_def_id)
+            .iter()
+            // By construction we capture all the args first.
+            .skip(num_args)
+            .zip_eq(tcx.closure_captures(parent_def_id))
+            .enumerate()
+        {
+            // This argument is captured by-move from the parent closure, but by-ref
+            // from the inner async block. That means that it's being borrowed from
+            // the closure body -- we need to change the coroutine take it by move.
+            if coroutine_capture.is_by_ref() && !parent_capture.is_by_ref() {
+                by_ref_fields.insert(FieldIdx::from_usize(num_args + idx));
+            }
+
+            // Make sure we're actually talking about the same capture.
+            assert_eq!(coroutine_capture.place.ty(), parent_capture.place.ty());
+        }
+
         let by_move_coroutine_ty = Ty::new_coroutine(
             tcx,
             coroutine_def_id.to_def_id(),
             ty::CoroutineArgs::new(
                 tcx,
                 ty::CoroutineArgsParts {
-                    parent_args: args.as_coroutine().parent_args(),
+                    parent_args: args.parent_args(),
                     kind_ty: Ty::from_closure_kind(tcx, ty::ClosureKind::FnOnce),
-                    resume_ty: args.as_coroutine().resume_ty(),
-                    yield_ty: args.as_coroutine().yield_ty(),
-                    return_ty: args.as_coroutine().return_ty(),
-                    witness: args.as_coroutine().witness(),
-                    tupled_upvars_ty: by_move_upvars,
+                    resume_ty: args.resume_ty(),
+                    yield_ty: args.yield_ty(),
+                    return_ty: args.return_ty(),
+                    witness: args.witness(),
+                    // Concatenate the args + closure's captures (since they're all by move).
+                    tupled_upvars_ty: Ty::new_tup_from_iter(
+                        tcx,
+                        tupled_inputs_ty
+                            .tuple_fields()
+                            .iter()
+                            .chain(parent_upvars_ty.tuple_fields()),
+                    ),
                 },
             )
             .args,
diff --git a/src/tools/miri/tests/pass/async-closure-captures.rs b/src/tools/miri/tests/pass/async-closure-captures.rs
new file mode 100644
index 00000000000..acff4a38338
--- /dev/null
+++ b/src/tools/miri/tests/pass/async-closure-captures.rs
@@ -0,0 +1,89 @@
+#![feature(async_closure, noop_waker)]
+
+use std::future::Future;
+use std::pin::pin;
+use std::task::*;
+
+pub fn block_on<T>(fut: impl Future<Output = T>) -> T {
+    let mut fut = pin!(fut);
+    let ctx = &mut Context::from_waker(Waker::noop());
+
+    loop {
+        match fut.as_mut().poll(ctx) {
+            Poll::Pending => {}
+            Poll::Ready(t) => break t,
+        }
+    }
+}
+
+fn main() {
+    block_on(async_main());
+}
+
+async fn call<T>(f: &impl async Fn() -> T) -> T {
+    f().await
+}
+
+async fn call_once<T>(f: impl async FnOnce() -> T) -> T {
+    f().await
+}
+
+#[derive(Debug)]
+#[allow(unused)]
+struct Hello(i32);
+
+async fn async_main() {
+    // Capture something by-ref
+    {
+        let x = Hello(0);
+        let c = async || {
+            println!("{x:?}");
+        };
+        call(&c).await;
+        call_once(c).await;
+
+        let x = &Hello(1);
+        let c = async || {
+            println!("{x:?}");
+        };
+        call(&c).await;
+        call_once(c).await;
+    }
+
+    // Capture something and consume it (force to `AsyncFnOnce`)
+    {
+        let x = Hello(2);
+        let c = async || {
+            println!("{x:?}");
+            drop(x);
+        };
+        call_once(c).await;
+    }
+
+    // Capture something with `move`, don't consume it
+    {
+        let x = Hello(3);
+        let c = async move || {
+            println!("{x:?}");
+        };
+        call(&c).await;
+        call_once(c).await;
+
+        let x = &Hello(4);
+        let c = async move || {
+            println!("{x:?}");
+        };
+        call(&c).await;
+        call_once(c).await;
+    }
+
+    // Capture something with `move`, also consume it (so `AsyncFnOnce`)
+    {
+        let x = Hello(5);
+        let c = async move || {
+            println!("{x:?}");
+            drop(x);
+        };
+        call_once(c).await;
+    }
+}
diff --git a/src/tools/miri/tests/pass/async-closure-captures.stdout b/src/tools/miri/tests/pass/async-closure-captures.stdout
new file mode 100644
index 00000000000..a0db6d236fe
--- /dev/null
+++ b/src/tools/miri/tests/pass/async-closure-captures.stdout
@@ -0,0 +1,10 @@
+Hello(0)
+Hello(0)
+Hello(1)
+Hello(1)
+Hello(2)
+Hello(3)
+Hello(3)
+Hello(4)
+Hello(4)
+Hello(5)
diff --git a/tests/ui/async-await/async-closures/captures.rs b/tests/ui/async-await/async-closures/captures.rs
new file mode 100644
index 00000000000..46bbf53f0a7
--- /dev/null
+++ b/tests/ui/async-await/async-closures/captures.rs
@@ -0,0 +1,80 @@
+//@ aux-build:block-on.rs
+//@ edition:2021
+//@ run-pass
+//@ check-run-results
+
+#![feature(async_closure)]
+
+extern crate block_on;
+
+fn main() {
+    block_on::block_on(async_main());
+}
+
+async fn call<T>(f: &impl async Fn() -> T) -> T {
+    f().await
+}
+
+async fn call_once<T>(f: impl async FnOnce() -> T) -> T {
+    f().await
+}
+
+#[derive(Debug)]
+#[allow(unused)]
+struct Hello(i32);
+
+async fn async_main() {
+    // Capture something by-ref
+    {
+        let x = Hello(0);
+        let c = async || {
+            println!("{x:?}");
+        };
+        call(&c).await;
+        call_once(c).await;
+
+        let x = &Hello(1);
+        let c = async || {
+            println!("{x:?}");
+        };
+        call(&c).await;
+        call_once(c).await;
+    }
+
+    // Capture something and consume it (force to `AsyncFnOnce`)
+    {
+        let x = Hello(2);
+        let c = async || {
+            println!("{x:?}");
+            drop(x);
+        };
+        call_once(c).await;
+    }
+
+    // Capture something with `move`, don't consume it
+    {
+        let x = Hello(3);
+        let c = async move || {
+            println!("{x:?}");
+        };
+        call(&c).await;
+        call_once(c).await;
+
+        let x = &Hello(4);
+        let c = async move || {
+            println!("{x:?}");
+        };
+        call(&c).await;
+        call_once(c).await;
+    }
+
+    // Capture something with `move`, also consume it (so `AsyncFnOnce`)
+    {
+        let x = Hello(5);
+        let c = async move || {
+            println!("{x:?}");
+            drop(x);
+        };
+        call_once(c).await;
+    }
+}
diff --git a/tests/ui/async-await/async-closures/captures.run.stdout b/tests/ui/async-await/async-closures/captures.run.stdout
new file mode 100644
index 00000000000..a0db6d236fe
--- /dev/null
+++ b/tests/ui/async-await/async-closures/captures.run.stdout
@@ -0,0 +1,10 @@
+Hello(0)
+Hello(0)
+Hello(1)
+Hello(1)
+Hello(2)
+Hello(3)
+Hello(3)
+Hello(4)
+Hello(4)
+Hello(5)

From ec74a304bb737f24d6e7f2fb8c3a2b3cf3575f0f Mon Sep 17 00:00:00 2001
From: Michael Goulet <michael@errs.io>
Date: Tue, 2 Apr 2024 11:26:57 -0400
Subject: [PATCH 7/7] Comments, comments, comments

---
 .../src/coroutine/by_move_body.rs             | 131 ++++++++++++------
 .../miri/tests/pass/async-closure-captures.rs |   2 +
 .../ui/async-await/async-closures/captures.rs |   2 +
 3 files changed, 96 insertions(+), 39 deletions(-)

diff --git a/compiler/rustc_mir_transform/src/coroutine/by_move_body.rs b/compiler/rustc_mir_transform/src/coroutine/by_move_body.rs
index 60b52fba219..0866205dfd0 100644
--- a/compiler/rustc_mir_transform/src/coroutine/by_move_body.rs
+++ b/compiler/rustc_mir_transform/src/coroutine/by_move_body.rs
@@ -1,7 +1,64 @@
-//! A MIR pass which duplicates a coroutine's body and removes any derefs which
-//! would be present for upvars that are taken by-ref. The result of which will
-//! be a coroutine body that takes all of its upvars by-move, and which we stash
-//! into the `CoroutineInfo` for all coroutines returned by coroutine-closures.
+//! This pass constructs a second coroutine body sufficient for return from
+//! `FnOnce`/`AsyncFnOnce` implementations for coroutine-closures (e.g. async closures).
+//!
+//! Consider an async closure like:
+//! ```rust
+//! #![feature(async_closure)]
+//!
+//! let x = vec![1, 2, 3];
+//!
+//! let closure = async move || {
+//!     println!("{x:#?}");
+//! };
+//! ```
+//!
+//! This desugars to something like:
+//! ```rust,ignore (invalid-borrowck)
+//! let x = vec![1, 2, 3];
+//!
+//! let closure = move || {
+//!     async {
+//!         println!("{x:#?}");
+//!     }
+//! };
+//! ```
+//!
+//! Important to note here is that while the outer closure *moves* `x: Vec<i32>`
+//! into its upvars, the inner `async` coroutine simply captures a ref of `x`.
+//! This is the "magic" of async closures -- the futures that they return are
+//! allowed to borrow from their parent closure's upvars.
+//!
+//! However, what happens when we call `closure` with `AsyncFnOnce` (or `FnOnce`,
+//! since all async closures implement that too)? Well, recall the signature:
+//! ```
+//! use std::future::Future;
+//! pub trait AsyncFnOnce<Args>
+//! {
+//!     type CallOnceFuture: Future<Output = Self::Output>;
+//!     type Output;
+//!     fn async_call_once(
+//!         self,
+//!         args: Args
+//!     ) -> Self::CallOnceFuture;
+//! }
+//! ```
+//!
+//! This signature *consumes* the async closure (`self`) and returns a `CallOnceFuture`.
+//! How do we deal with the fact that the coroutine is supposed to take a reference
+//! to the captured `x` from the parent closure, when that parent closure has been
+//! destroyed?
+//!
+//! This is the second piece of magic of async closures. We can simply create a
+//! *second* `async` coroutine body where that `x` that was previously captured
+//! by reference is now captured by value. This means that we consume the outer
+//! closure and return a new coroutine that will hold onto all of these captures,
+//! and drop them when it is finished (i.e. after it has been `.await`ed).
+//!
+//! We do this with the analysis below, which detects the captures that come from
+//! borrowing from the outer closure, and we simply peel off a `deref` projection
+//! from them. This second body is stored alongside the first body, and optimized
+//! with it in lockstep. When we need to resolve a body for `FnOnce` or `AsyncFnOnce`,
+//! we use this "by move" body instead.
 
 use itertools::Itertools;
 
@@ -16,6 +73,8 @@
 
 impl<'tcx> MirPass<'tcx> for ByMoveBody {
     fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut mir::Body<'tcx>) {
+        // We only need to generate by-move coroutine bodies for coroutines that come
+        // from coroutine-closures.
         let Some(coroutine_def_id) = body.source.def_id().as_local() else {
             return;
         };
@@ -24,15 +83,19 @@ fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut mir::Body<'tcx>) {
         else {
             return;
         };
+
+        // Also, let's skip processing any bodies with errors, since there's no guarantee
+        // the MIR body will be constructed well.
         let coroutine_ty = body.local_decls[ty::CAPTURE_STRUCT_LOCAL].ty;
         if coroutine_ty.references_error() {
             return;
         }
 
-        let ty::Coroutine(_, args) = *coroutine_ty.kind() else { bug!("{body:#?}") };
-        let args = args.as_coroutine();
-
-        let coroutine_kind = args.kind_ty().to_opt_closure_kind().unwrap();
+        let ty::Coroutine(_, coroutine_args) = *coroutine_ty.kind() else { bug!("{body:#?}") };
+        // We don't need to generate a by-move coroutine if the kind of the coroutine is
+        // already `FnOnce` -- that means that any upvars that the closure consumes have
+        // already been taken by-value.
+        let coroutine_kind = coroutine_args.as_coroutine().kind_ty().to_opt_closure_kind().unwrap();
         if coroutine_kind == ty::ClosureKind::FnOnce {
             return;
         }
@@ -43,12 +106,13 @@ fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut mir::Body<'tcx>) {
         else {
             bug!();
         };
-        let parent_args = parent_args.as_coroutine_closure();
-        let parent_upvars_ty = parent_args.tupled_upvars_ty();
-        let tupled_inputs_ty = tcx.instantiate_bound_regions_with_erased(
-            parent_args.coroutine_closure_sig().map_bound(|sig| sig.tupled_inputs_ty),
-        );
-        let num_args = tupled_inputs_ty.tuple_fields().len();
+        let parent_closure_args = parent_args.as_coroutine_closure();
+        let num_args = parent_closure_args
+            .coroutine_closure_sig()
+            .skip_binder()
+            .tupled_inputs_ty
+            .tuple_fields()
+            .len();
 
         let mut by_ref_fields = UnordSet::default();
         for (idx, (coroutine_capture, parent_capture)) in tcx
@@ -59,41 +123,30 @@ fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut mir::Body<'tcx>) {
             .zip_eq(tcx.closure_captures(parent_def_id))
             .enumerate()
         {
-            // This argument is captured by-move from the parent closure, but by-ref
+            // This upvar is captured by-move from the parent closure, but by-ref
             // from the inner async block. That means that it's being borrowed from
-            // the closure body -- we need to change the coroutine take it by move.
+            // the outer closure body -- we need to change the coroutine to take the
+            // upvar by value.
             if coroutine_capture.is_by_ref() && !parent_capture.is_by_ref() {
                 by_ref_fields.insert(FieldIdx::from_usize(num_args + idx));
             }
 
             // Make sure we're actually talking about the same capture.
+            // FIXME(async_closures): We could look at the `hir::Upvar` instead?
             assert_eq!(coroutine_capture.place.ty(), parent_capture.place.ty());
         }
 
-        let by_move_coroutine_ty = Ty::new_coroutine(
-            tcx,
-            coroutine_def_id.to_def_id(),
-            ty::CoroutineArgs::new(
+        let by_move_coroutine_ty = tcx
+            .instantiate_bound_regions_with_erased(parent_closure_args.coroutine_closure_sig())
+            .to_coroutine_given_kind_and_upvars(
                 tcx,
-                ty::CoroutineArgsParts {
-                    parent_args: args.parent_args(),
-                    kind_ty: Ty::from_closure_kind(tcx, ty::ClosureKind::FnOnce),
-                    resume_ty: args.resume_ty(),
-                    yield_ty: args.yield_ty(),
-                    return_ty: args.return_ty(),
-                    witness: args.witness(),
-                    // Concatenate the args + closure's captures (since they're all by move).
-                    tupled_upvars_ty: Ty::new_tup_from_iter(
-                        tcx,
-                        tupled_inputs_ty
-                            .tuple_fields()
-                            .iter()
-                            .chain(parent_upvars_ty.tuple_fields()),
-                    ),
-                },
-            )
-            .args,
-        );
+                parent_closure_args.parent_args(),
+                coroutine_def_id.to_def_id(),
+                ty::ClosureKind::FnOnce,
+                tcx.lifetimes.re_erased,
+                parent_closure_args.tupled_upvars_ty(),
+                parent_closure_args.coroutine_captures_by_ref_ty(),
+            );
 
         let mut by_move_body = body.clone();
         MakeByMoveBody { tcx, by_ref_fields, by_move_coroutine_ty }.visit_body(&mut by_move_body);
diff --git a/src/tools/miri/tests/pass/async-closure-captures.rs b/src/tools/miri/tests/pass/async-closure-captures.rs
index acff4a38338..3e33de32efb 100644
--- a/src/tools/miri/tests/pass/async-closure-captures.rs
+++ b/src/tools/miri/tests/pass/async-closure-captures.rs
@@ -1,3 +1,5 @@
+// Same as rustc's `tests/ui/async-await/async-closures/captures.rs`, keep in sync
+
 #![feature(async_closure, noop_waker)]
 
 use std::future::Future;
diff --git a/tests/ui/async-await/async-closures/captures.rs b/tests/ui/async-await/async-closures/captures.rs
index 46bbf53f0a7..e3ab8713709 100644
--- a/tests/ui/async-await/async-closures/captures.rs
+++ b/tests/ui/async-await/async-closures/captures.rs
@@ -3,6 +3,8 @@
 //@ run-pass
 //@ check-run-results
 
+// Same as miri's `tests/pass/async-closure-captures.rs`, keep in sync
+
 #![feature(async_closure)]
 
 extern crate block_on;