Add ability to parse likely flags for ifs in Cmm.

Adding the ability to parse likely flags in Cmm allows better codegen
for cmm files.

Test Plan: ci

Reviewers: bgamari, simonmar

Reviewed By: bgamari

Subscribers: rwbarton, thomie, carter

GHC Trac Issues: #14672

Differential Revision: https://phabricator.haskell.org/D4316

2 files changed, 41 insertions, 21 deletions

diff --git a/compiler/cmm/CmmLex.x b/compiler/cmm/CmmLex.x
index a68f155883..691ca5eb28 100644
--- a/compiler/cmm/CmmLex.x
+++ b/compiler/cmm/CmmLex.x
@@ -99,6 +99,10 @@ $white_no_nl+           ;
   "&&"                  { kw CmmT_BoolAnd }
   "||"                  { kw CmmT_BoolOr }
 
+  "True"                { kw CmmT_True  }
+  "False"               { kw CmmT_False }
+  "likely"              { kw CmmT_likely}
+
   P@decimal             { global_regN (\n -> VanillaReg n VGcPtr) }
   R@decimal             { global_regN (\n -> VanillaReg n VNonGcPtr) }
   F@decimal             { global_regN FloatReg }
@@ -180,6 +184,9 @@ data CmmToken
   | CmmT_Int       Integer
   | CmmT_Float     Rational
   | CmmT_EOF
+  | CmmT_False
+  | CmmT_True
+  | CmmT_likely
   deriving (Show)
 
 -- -----------------------------------------------------------------------------
@@ -266,7 +273,10 @@ reservedWordsFM = listToUFM $
         ( "b512",               CmmT_bits512 ),
         ( "f32",                CmmT_float32 ),
         ( "f64",                CmmT_float64 ),
-        ( "gcptr",              CmmT_gcptr )
+        ( "gcptr",              CmmT_gcptr ),
+        ( "likely",             CmmT_likely),
+        ( "True",               CmmT_True  ),
+        ( "False",              CmmT_False )
         ]
 
 tok_decimal span buf len
diff --git a/compiler/cmm/CmmParse.y b/compiler/cmm/CmmParse.y
index 8afbd2f9d9..cf660d274f 100644
--- a/compiler/cmm/CmmParse.y
+++ b/compiler/cmm/CmmParse.y
@@ -299,6 +299,10 @@ import qualified Data.Map as M
         '&&'    { L _ (CmmT_BoolAnd) }
         '||'    { L _ (CmmT_BoolOr) }
 
+        'True'  { L _ (CmmT_True ) }
+        'False' { L _ (CmmT_False) }
+        'likely'{ L _ (CmmT_likely)}
+
         'CLOSURE'       { L _ (CmmT_CLOSURE) }
         'INFO_TABLE'    { L _ (CmmT_INFO_TABLE) }
         'INFO_TABLE_RET'{ L _ (CmmT_INFO_TABLE_RET) }
@@ -629,10 +633,10 @@ stmt    :: { CmmParse () }
                 { doCall $2 [] $4 }
         | '(' formals ')' '=' 'call' expr '(' exprs0 ')' ';'
                 { doCall $6 $2 $8 }
-        | 'if' bool_expr 'goto' NAME
-                { do l <- lookupLabel $4; cmmRawIf $2 l }
-        | 'if' bool_expr '{' body '}' else      
-                { cmmIfThenElse $2 (withSourceNote $3 $5 $4) $6 }
+        | 'if' bool_expr cond_likely 'goto' NAME
+                { do l <- lookupLabel $5; cmmRawIf $2 l $3 }
+        | 'if' bool_expr cond_likely '{' body '}' else
+                { cmmIfThenElse $2 (withSourceNote $4 $6 $5) $7 $3 }
         | 'push' '(' exprs0 ')' maybe_body
                 { pushStackFrame $3 $5 }
         | 'reserve' expr '=' lreg maybe_body
@@ -721,6 +725,12 @@ else    :: { CmmParse () }
         : {- empty -}                   { return () }
         | 'else' '{' body '}'           { withSourceNote $2 $4 $3 }
 
+cond_likely :: { Maybe Bool }
+        : '(' 'likely' ':' 'True'  ')'  { Just True  }
+        | '(' 'likely' ':' 'False' ')'  { Just False }
+        | {- empty -}                   { Nothing }
+
+
 -- we have to write this out longhand so that Happy's precedence rules
 -- can kick in.
 expr    :: { CmmParse CmmExpr }
@@ -1289,11 +1299,11 @@ data BoolExpr
 
 -- ToDo: smart constructors which simplify the boolean expression.
 
-cmmIfThenElse cond then_part else_part = do
+cmmIfThenElse cond then_part else_part likely = do
      then_id <- newBlockId
      join_id <- newBlockId
      c <- cond
-     emitCond c then_id
+     emitCond c then_id likely
      else_part
      emit (mkBranch join_id)
      emitLabel then_id
@@ -1301,38 +1311,38 @@ cmmIfThenElse cond then_part else_part = do
      -- fall through to join
      emitLabel join_id
 
-cmmRawIf cond then_id = do
+cmmRawIf cond then_id likely = do
     c <- cond
-    emitCond c then_id
+    emitCond c then_id likely
 
 -- 'emitCond cond true_id'  emits code to test whether the cond is true,
 -- branching to true_id if so, and falling through otherwise.
-emitCond (BoolTest e) then_id = do
+emitCond (BoolTest e) then_id likely = do
   else_id <- newBlockId
-  emit (mkCbranch e then_id else_id Nothing)
+  emit (mkCbranch e then_id else_id likely)
   emitLabel else_id
-emitCond (BoolNot (BoolTest (CmmMachOp op args))) then_id
+emitCond (BoolNot (BoolTest (CmmMachOp op args))) then_id likely
   | Just op' <- maybeInvertComparison op
-  = emitCond (BoolTest (CmmMachOp op' args)) then_id
-emitCond (BoolNot e) then_id = do
+  = emitCond (BoolTest (CmmMachOp op' args)) then_id (not <$> likely)
+emitCond (BoolNot e) then_id likely = do
   else_id <- newBlockId
-  emitCond e else_id
+  emitCond e else_id likely
   emit (mkBranch then_id)
   emitLabel else_id
-emitCond (e1 `BoolOr` e2) then_id = do
-  emitCond e1 then_id
-  emitCond e2 then_id
-emitCond (e1 `BoolAnd` e2) then_id = do
+emitCond (e1 `BoolOr` e2) then_id likely = do
+  emitCond e1 then_id likely
+  emitCond e2 then_id likely
+emitCond (e1 `BoolAnd` e2) then_id likely = do
         -- we'd like to invert one of the conditionals here to avoid an
         -- extra branch instruction, but we can't use maybeInvertComparison
         -- here because we can't look too closely at the expression since
         -- we're in a loop.
   and_id <- newBlockId
   else_id <- newBlockId
-  emitCond e1 and_id
+  emitCond e1 and_id likely
   emit (mkBranch else_id)
   emitLabel and_id
-  emitCond e2 then_id
+  emitCond e2 then_id likely
   emitLabel else_id
 
 -- -----------------------------------------------------------------------------