Support arbitrary let statements in custom mir

Author: JakobDegen

library/core/src/intrinsics/mir.rs

@@ -90,10 +90,14 @@ pub macro mir {
     (
         $(let $local_decl:ident $(: $local_decl_ty:ty)? ;)*
 
-        $entry_block:block
+        {
+            $($entry:tt)*
+        }
 
         $(
-            $block_name:ident = $block:block
+            $block_name:ident = {
+                $($block:tt)*
+            }
         )*
     ) => {{
         // First, we declare all basic blocks.
@@ -109,15 +113,146 @@ pub macro mir {
                 let $local_decl $(: $local_decl_ty)? ;
             )*
 
+            ::core::intrinsics::mir::__internal_extract_let!($($entry)*);
+            $(
+                ::core::intrinsics::mir::__internal_extract_let!($($block)*);
+            )*
+
             {
                 // Finally, the contents of the basic blocks
-                $entry_block;
+                ::core::intrinsics::mir::__internal_remove_let!({
+                    {}
+                    { $($entry)* }
+                });
                 $(
-                    $block;
+                    ::core::intrinsics::mir::__internal_remove_let!({
+                        {}
+                        { $($block)* }
+                    });
                 )*
 
                 RET
             }
         }
     }}
 }
+
+/// Helper macro that extracts the `let` declarations out of a bunch of statements.
+///
+/// This macro is written using the "statement muncher" strategy. Each invocation parses the first
+/// statement out of the input, does the appropriate thing with it, and then recursively calls the
+/// same macro on the remainder of the input.
+#[doc(hidden)]
+pub macro __internal_extract_let {
+    // If it's a `let` like statement, keep the `let`
+    (
+        let $var:ident $(: $ty:ty)? = $expr:expr; $($rest:tt)*
+    ) => {
+        let $var $(: $ty)?;
+        ::core::intrinsics::mir::__internal_extract_let!($($rest)*);
+    },
+    // Otherwise, output nothing
+    (
+        $stmt:stmt; $($rest:tt)*
+    ) => {
+        ::core::intrinsics::mir::__internal_extract_let!($($rest)*);
+    },
+    (
+        $expr:expr
+    ) => {}
+}
+
+/// Helper macro that removes the `let` declarations from a bunch of statements.
+///
+/// Because expression position macros cannot expand to statements + expressions, we need to be
+/// slightly creative here. The general strategy is also statement munching as above, but the output
+/// of the macro is "stored" in the subsequent macro invocation. Easiest understood via example:
+/// ```text
+/// invoke!(
+///     {
+///         {
+///             x = 5;
+///         }
+///         {
+///             let d = e;
+///             Call()
+///         }
+///     }
+/// )
+/// ```
+/// becomes
+/// ```text
+/// invoke!(
+///     {
+///         {
+///             x = 5;
+///             d = e;
+///         }
+///         {
+///             Call()
+///         }
+///     }
+/// )
+/// ```
+#[doc(hidden)]
+pub macro __internal_remove_let {
+    // If it's a `let` like statement, remove the `let`
+    (
+        {
+            {
+                $($already_parsed:tt)*
+            }
+            {
+                let $var:ident $(: $ty:ty)? = $expr:expr;
+                $($rest:tt)*
+            }
+        }
+    ) => { ::core::intrinsics::mir::__internal_remove_let!(
+        {
+            {
+                $($already_parsed)*
+                $var = $expr;
+            }
+            {
+                $($rest)*
+            }
+        }
+    )},
+    // Otherwise, keep going
+    (
+        {
+            {
+                $($already_parsed:tt)*
+            }
+            {
+                $stmt:stmt;
+                $($rest:tt)*
+            }
+        }
+    ) => { ::core::intrinsics::mir::__internal_remove_let!(
+        {
+            {
+                $($already_parsed)*
+                $stmt;
+            }
+            {
+                $($rest)*
+            }
+        }
+    )},
+    (
+        {
+            {
+                $($already_parsed:tt)*
+            }
+            {
+                $expr:expr
+            }
+        }
+    ) => {
+        {
+            $($already_parsed)*
+            $expr
+        }
+    },
+}

src/test/mir-opt/building/custom/arbitrary_let.arbitrary_let.built.after.mir

@@ -0,0 +1,22 @@
+// MIR for `arbitrary_let` after built
+
+fn arbitrary_let(_1: i32) -> i32 {
+    let mut _0: i32;                     // return place in scope 0 at $DIR/arbitrary_let.rs:+0:29: +0:32
+    let mut _2: i32;                     // in scope 0 at $SRC_DIR/core/src/intrinsics/mir.rs:LL:COL
+    let mut _3: i32;                     // in scope 0 at $SRC_DIR/core/src/intrinsics/mir.rs:LL:COL
+
+    bb0: {
+        _2 = _1;                         // scope 0 at $DIR/arbitrary_let.rs:+0:1: +0:32
+        goto -> bb2;                     // scope 0 at $DIR/arbitrary_let.rs:+0:1: +0:32
+    }
+
+    bb1: {
+        _0 = _3;                         // scope 0 at $DIR/arbitrary_let.rs:+0:1: +0:32
+        return;                          // scope 0 at $DIR/arbitrary_let.rs:+0:1: +0:32
+    }
+
+    bb2: {
+        _3 = _2;                         // scope 0 at $DIR/arbitrary_let.rs:+0:1: +0:32
+        goto -> bb1;                     // scope 0 at $DIR/arbitrary_let.rs:+0:1: +0:32
+    }
+}

src/test/mir-opt/building/custom/arbitrary_let.rs

@@ -0,0 +1,28 @@
+#![feature(custom_mir, core_intrinsics)]
+
+extern crate core;
+use core::intrinsics::mir::*;
+use core::ptr::{addr_of, addr_of_mut};
+
+// EMIT_MIR arbitrary_let.arbitrary_let.built.after.mir
+#[custom_mir(dialect = "built")]
+fn arbitrary_let(x: i32) -> i32 {
+    mir!(
+        {
+            let y = x;
+            Goto(second)
+        }
+        third = {
+            RET = z;
+            Return()
+        }
+        second = {
+            let z = y;
+            Goto(third)
+        }
+    )
+}
+
+fn main() {
+    assert_eq!(arbitrary_let(5), 5);
+}