Auto merge of #61922 - tmandry:moar-generator-optimization, r=matthewjasper

Don't store locals that have been moved from in generators

This avoids reserving storage in generators for locals that are moved
out of (and not re-initialized) prior to yield points. Fixes #59123.

This adds a new dataflow analysis, `RequiresStorage`, to determine whether the storage of a local can be destroyed without being observed by the program. The rules are:

1. StorageLive(x) => mark x live
2. StorageDead(x) => mark x dead
3. If a local is moved from, _and has never had its address taken_, mark it dead
4. If (any part of) a local is initialized, mark it live'

This is used to determine whether to save a local in the generator object at all, as well as which locals can be overlapped in the generator layout.

Here's the size in bytes of all testcases included in the change, before and after the change:

async fn test    |Size before |Size after
-----------------|------------|----------
single           | 1028       | 1028
single_with_noop | 2056       | 1032
joined           | 5132       | 3084
joined_with_noop | 8208       | 3084

generator test              |Size before |Size after
----------------------------|------------|----------
move_before_yield           | 1028       | 1028
move_before_yield_with_noop | 2056       | 1032
overlap_move_points         | 3080       | 2056

## Future work

Note that there is a possible extension to this optimization, which modifies rule 3 to read: "If a local is moved from, _**and either has never had its address taken, or is Freeze and has never been mutably borrowed**_, mark it dead." This was discussed at length in #59123 and then #61849. Because this would cause some behavior to be UB which was not UB before, it's a step that needs to be taken carefully.

A more immediate priority for me is inlining `std::mem::size_of_val(&x)` so it becomes apparent that the address of `x` is not taken. This way, using `size_of_val` to look at the size of your inner futures does not affect the size of your outer future.

cc @cramertj @eddyb @Matthias247 @nikomatsakis @RalfJung @Zoxc

2 files changed, 160 insertions, 0 deletions

diff --git a/src/test/run-pass/async-await/async-fn-size-moved-locals.rs b/src/test/run-pass/async-await/async-fn-size-moved-locals.rs
new file mode 100644
index 00000000000..139be7fe013
--- /dev/null
+++ b/src/test/run-pass/async-await/async-fn-size-moved-locals.rs
@@ -0,0 +1,98 @@
+// Test that we don't duplicate storage for futures moved around in .await, and
+// for futures moved into other futures.
+//
+// The exact sizes can change by a few bytes (we'd like to know when they do).
+// What we don't want to see is the wrong multiple of 1024 (the size of BigFut)
+// being reflected in the size.
+//
+// See issue #59123 for a full explanation.
+
+// edition:2018
+
+#![feature(async_await)]
+
+use std::future::Future;
+use std::pin::Pin;
+use std::task::{Context, Poll};
+
+const BIG_FUT_SIZE: usize = 1024;
+struct BigFut([u8; BIG_FUT_SIZE]);
+
+impl BigFut {
+    fn new() -> Self {
+        BigFut([0; BIG_FUT_SIZE])
+    } }
+
+impl Drop for BigFut {
+    fn drop(&mut self) {}
+}
+
+impl Future for BigFut {
+    type Output = ();
+
+    fn poll(self: Pin<&mut Self>, _ctx: &mut Context<'_>) -> Poll<Self::Output> {
+        Poll::Ready(())
+    }
+}
+
+#[allow(dead_code)]
+struct Joiner {
+    a: Option<BigFut>,
+    b: Option<BigFut>,
+    c: Option<BigFut>,
+}
+
+impl Future for Joiner {
+    type Output = ();
+
+    fn poll(self: Pin<&mut Self>, _ctx: &mut Context<'_>) -> Poll<Self::Output> {
+        Poll::Ready(())
+    }
+}
+
+fn noop() {}
+
+async fn single() {
+    let x = BigFut::new();
+    x.await;
+}
+
+async fn single_with_noop() {
+    let x = BigFut::new();
+    noop();
+    x.await;
+}
+
+async fn joined() {
+    let a = BigFut::new();
+    let b = BigFut::new();
+    let c = BigFut::new();
+
+    let joiner = Joiner {
+        a: Some(a),
+        b: Some(b),
+        c: Some(c),
+    };
+    joiner.await
+}
+
+async fn joined_with_noop() {
+    let a = BigFut::new();
+    let b = BigFut::new();
+    let c = BigFut::new();
+
+    let joiner = Joiner {
+        a: Some(a),
+        b: Some(b),
+        c: Some(c),
+    };
+    noop();
+    joiner.await
+}
+
+fn main() {
+    assert_eq!(1028, std::mem::size_of_val(&single()));
+    assert_eq!(1032, std::mem::size_of_val(&single_with_noop()));
+    assert_eq!(3084, std::mem::size_of_val(&joined()));
+    assert_eq!(3084, std::mem::size_of_val(&joined_with_noop()));
+}
diff --git a/src/test/run-pass/generator/size-moved-locals.rs b/src/test/run-pass/generator/size-moved-locals.rs
new file mode 100644
index 00000000000..37e2e0cfdcc
--- /dev/null
+++ b/src/test/run-pass/generator/size-moved-locals.rs
@@ -0,0 +1,62 @@
+// Test that we don't duplicate storage for a variable that is moved to another
+// binding. This used to happen in the presence of unwind and drop edges (see
+// `complex` below.)
+//
+// The exact sizes here can change (we'd like to know when they do). What we
+// don't want to see is the `complex` generator size being upwards of 2048 bytes
+// (which would indicate it is reserving space for two copies of Foo.)
+//
+// See issue #59123 for a full explanation.
+
+// edition:2018
+
+#![feature(generators, generator_trait)]
+
+use std::ops::Generator;
+
+const FOO_SIZE: usize = 1024;
+struct Foo([u8; FOO_SIZE]);
+
+impl Drop for Foo {
+    fn drop(&mut self) {}
+}
+
+fn move_before_yield() -> impl Generator<Yield = (), Return = ()> {
+    static || {
+        let first = Foo([0; FOO_SIZE]);
+        let _second = first;
+        yield;
+        // _second dropped here
+    }
+}
+
+fn noop() {}
+
+fn move_before_yield_with_noop() -> impl Generator<Yield = (), Return = ()> {
+    static || {
+        let first = Foo([0; FOO_SIZE]);
+        noop();
+        let _second = first;
+        yield;
+        // _second dropped here
+    }
+}
+
+// Today we don't have NRVO (we allocate space for both `first` and `second`,)
+// but we can overlap `first` with `_third`.
+fn overlap_move_points() -> impl Generator<Yield = (), Return = ()> {
+    static || {
+        let first = Foo([0; FOO_SIZE]);
+        yield;
+        let second = first;
+        yield;
+        let _third = second;
+        yield;
+    }
+}
+
+fn main() {
+    assert_eq!(1028, std::mem::size_of_val(&move_before_yield()));
+    assert_eq!(1032, std::mem::size_of_val(&move_before_yield_with_noop()));
+    assert_eq!(2056, std::mem::size_of_val(&overlap_move_points()));
+}