Merge pull request #3708 from cfallin/cold-blocks-clif

Add syntax for cold blocks to CLIF.

cranelift/codegen/src/write.rs

@@ -216,12 +216,18 @@ pub fn write_block_header(
     block: Block,
     indent: usize,
 ) -> fmt::Result {
+    let cold = if func.layout.is_cold(block) {
+        " cold"
+    } else {
+        ""
+    };
+
     // The `indent` is the instruction indentation. block headers are 4 spaces out from that.
     write!(w, "{1:0$}{2}", indent - 4, "", block)?;
 
     let mut args = func.dfg.block_params(block).iter().cloned();
     match args.next() {
-        None => return writeln!(w, ":"),
+        None => return writeln!(w, "{}:", cold),
         Some(arg) => {
             write!(w, "(")?;
             write_arg(w, func, arg)?;
@@ -232,7 +238,7 @@ pub fn write_block_header(
         write!(w, ", ")?;
         write_arg(w, func, arg)?;
     }
-    writeln!(w, "):")
+    writeln!(w, "){}:", cold)
 }
 
 fn decorate_block<FW: FuncWriter>(
@@ -666,4 +672,26 @@ mod tests {
             "function u0:0() fast {\nblock0(v3: i32):\n    v0 -> v3\n    v2 -> v0\n    v4 = iconst.i32 42\n    v5 = iadd v0, v0\n    v1 -> v5\n    v6 = iconst.i32 23\n    v7 = iadd v1, v1\n}\n"
         );
     }
+
+    #[test]
+    fn cold_blocks() {
+        let mut func = Function::new();
+        {
+            let mut pos = FuncCursor::new(&mut func);
+
+            let block0 = pos.func.dfg.make_block();
+            pos.insert_block(block0);
+            pos.func.layout.set_cold(block0);
+
+            let block1 = pos.func.dfg.make_block();
+            pos.insert_block(block1);
+            pos.func.dfg.append_block_param(block1, types::I32);
+            pos.func.layout.set_cold(block1);
+        }
+
+        assert_eq!(
+            func.to_string(),
+            "function u0:0() fast {\nblock0 cold:\n\nblock1(v0: i32) cold:\n}\n"
+        );
+    }
 }

cranelift/reader/src/lexer.rs

@@ -34,6 +34,7 @@ pub enum Token<'a> {
     Type(types::Type),    // i32, f32, b32x4, ...
     Value(Value),         // v12, v7
     Block(Block),         // block3
+    Cold,                 // cold (flag on block)
     StackSlot(u32),       // ss3
     GlobalValue(u32),     // gv3
     Heap(u32),            // heap2
@@ -326,6 +327,7 @@ impl<'a> Lexer<'a> {
                 .unwrap_or_else(|| match text {
                     "iflags" => Token::Type(types::IFLAGS),
                     "fflags" => Token::Type(types::FFLAGS),
+                    "cold" => Token::Cold,
                     _ => Token::Identifier(text),
                 }),
             loc,

cranelift/reader/src/parser.rs

@@ -447,6 +447,11 @@ impl Context {
         self.function.layout.append_block(block);
         Ok(block)
     }
+
+    /// Set a block as cold.
+    fn set_cold_block(&mut self, block: Block) {
+        self.function.layout.set_cold(block);
+    }
 }
 
 impl<'a> Parser<'a> {
@@ -1856,7 +1861,8 @@ impl<'a> Parser<'a> {
     // Parse a basic block, add contents to `ctx`.
     //
     // extended-basic-block ::= * block-header { instruction }
-    // block-header           ::= Block(block) [block-params] ":"
+    // block-header         ::= Block(block) [block-params] [block-flags] ":"
+    // block-flags          ::= [Cold]
     //
     fn parse_basic_block(&mut self, ctx: &mut Context) -> ParseResult<()> {
         // Collect comments for the next block.
@@ -1869,12 +1875,16 @@ impl<'a> Parser<'a> {
             return Err(self.error("too many blocks"));
         }
 
-        if !self.optional(Token::Colon) {
-            // block-header ::= Block(block) [ * block-params ] ":"
+        if self.token() == Some(Token::LPar) {
             self.parse_block_params(ctx, block)?;
-            self.match_token(Token::Colon, "expected ':' after block parameters")?;
         }
 
+        if self.optional(Token::Cold) {
+            ctx.set_cold_block(block);
+        }
+
+        self.match_token(Token::Colon, "expected ':' after block parameters")?;
+
         // Collect any trailing comments.
         self.token();
         self.claim_gathered_comments(block);
@@ -3728,4 +3738,23 @@ mod tests {
             "0x00000003000000020000000100000000"
         );
     }
+
+    #[test]
+    fn parse_cold_blocks() {
+        let code = "function %test() {
+        block0 cold:
+            return
+        block1(v0: i32) cold:
+            return
+        block2(v1: i32):
+            return
+        }";
+
+        let mut parser = Parser::new(code);
+        let func = parser.parse_function().unwrap().0;
+        assert_eq!(func.layout.blocks().count(), 3);
+        assert!(func.layout.is_cold(Block::from_u32(0)));
+        assert!(func.layout.is_cold(Block::from_u32(1)));
+        assert!(!func.layout.is_cold(Block::from_u32(2)));
+    }
 }