From 4f8cccff68aeb57d8a348d29e8fcc7c4053801e7 Mon Sep 17 00:00:00 2001
From: Erik Eckstein <eeckstein@apple.com>
Date: Thu, 4 Jan 2018 15:33:01 -0800
Subject: [PATCH] SimplifyCFG: do constant folding also in SimplifyCFG.

This enables to optimize chains of constant-foldable terminal instructions, for which we would have needed multiple ConstantPropagaion-SimplifyCFG pass pairs in the pipeline otherwise.
See the test file for an example.
---
 lib/SILOptimizer/Transforms/SimplifyCFG.cpp | 42 +++++++++++++++++----
 test/SILOptimizer/simplify_cfg.sil          | 35 +++++++++++++++++
 2 files changed, 69 insertions(+), 8 deletions(-)

diff --git a/lib/SILOptimizer/Transforms/SimplifyCFG.cpp b/lib/SILOptimizer/Transforms/SimplifyCFG.cpp
index 89d84943020fa..16d0269da4a1a 100644
--- a/lib/SILOptimizer/Transforms/SimplifyCFG.cpp
+++ b/lib/SILOptimizer/Transforms/SimplifyCFG.cpp
@@ -26,6 +26,7 @@
 #include "swift/SILOptimizer/Utils/CFG.h"
 #include "swift/SILOptimizer/Utils/CastOptimizer.h"
 #include "swift/SILOptimizer/Utils/Local.h"
+#include "swift/SILOptimizer/Utils/ConstantFolding.h"
 #include "swift/SILOptimizer/Utils/SILInliner.h"
 #include "swift/SILOptimizer/Utils/SILSSAUpdater.h"
 #include "llvm/ADT/SmallPtrSet.h"
@@ -75,13 +76,25 @@ namespace {
     // Dominance and post-dominance info for the current function
     DominanceInfo *DT = nullptr;
 
+    ConstantFolder ConstFolder;
+
+    void constFoldingCallback(SILInstruction *I) {
+      // If a terminal instruction gets constant folded (like cond_br), it
+      // enables further simplify-CFG optimizations.
+      if (isa<TermInst>(I))
+        addToWorklist(I->getParent());
+    }
+
     bool ShouldVerify;
     bool EnableJumpThread;
   public:
     SimplifyCFG(SILFunction &Fn, SILPassManager *PM, bool Verify,
                 bool EnableJumpThread)
-        : Fn(Fn), PM(PM), ShouldVerify(Verify),
-          EnableJumpThread(EnableJumpThread) {}
+        : Fn(Fn), PM(PM),
+          ConstFolder(PM->getOptions().AssertConfig,
+                      /* EnableDiagnostics */false,
+                      [&](SILInstruction *I) { constFoldingCallback(I); }),
+          ShouldVerify(Verify), EnableJumpThread(EnableJumpThread) {}
 
     bool run();
     
@@ -1092,11 +1105,15 @@ bool SimplifyCFG::tryJumpThreading(BranchInst *BI) {
 /// result in exposing opportunities for CFG simplification.
 bool SimplifyCFG::simplifyBranchOperands(OperandValueArrayRef Operands) {
   bool Simplified = false;
-  for (auto O = Operands.begin(), E = Operands.end(); O != E; ++O)
+  for (auto O = Operands.begin(), E = Operands.end(); O != E; ++O) {
     // All of our interesting simplifications are on single-value instructions
     // for now.
-    if (auto *I = dyn_cast<SingleValueInstruction>(*O))
-      if (SILValue Result = simplifyInstruction(I)) {
+    if (auto *I = dyn_cast<SingleValueInstruction>(*O)) {
+      SILValue Result = simplifyInstruction(I);
+
+      // The Result can be the same instruction I in case it is in an
+      // unreachable block. In this case it can reference itself as operand.
+      if (Result && Result != I) {
         DEBUG(llvm::dbgs() << "simplify branch operand " << *I);
         I->replaceAllUsesWith(Result);
         if (isInstructionTriviallyDead(I)) {
@@ -1104,6 +1121,8 @@ bool SimplifyCFG::simplifyBranchOperands(OperandValueArrayRef Operands) {
           Simplified = true;
         }
       }
+    }
+  }
   return Simplified;
 }
 
@@ -1228,9 +1247,16 @@ bool SimplifyCFG::simplifyBranchBlock(BranchInst *BI) {
     // If there are any BB arguments in the destination, replace them with the
     // branch operands, since they must dominate the dest block.
     for (unsigned i = 0, e = BI->getArgs().size(); i != e; ++i) {
-      if (DestBB->getArgument(i) != BI->getArg(i))
-        DestBB->getArgument(i)->replaceAllUsesWith(BI->getArg(i));
-      else {
+      if (DestBB->getArgument(i) != BI->getArg(i)) {
+        SILValue Val = BI->getArg(i);
+        DestBB->getArgument(i)->replaceAllUsesWith(Val);
+        if (auto *I = dyn_cast<SingleValueInstruction>(Val)) {
+          // Replacing operands may trigger constant folding which then could
+          // trigger other simplify-CFG optimizations.
+          ConstFolder.addToWorklist(I);
+          ConstFolder.processWorkList();
+        }
+      } else {
         // We must be processing an unreachable part of the cfg with a cycle.
         // bb1(arg1): // preds: bb3
         //   br bb2
diff --git a/test/SILOptimizer/simplify_cfg.sil b/test/SILOptimizer/simplify_cfg.sil
index 01dd1aab3e270..41422c0f9b0e7 100644
--- a/test/SILOptimizer/simplify_cfg.sil
+++ b/test/SILOptimizer/simplify_cfg.sil
@@ -2982,3 +2982,38 @@ bb7:
   br bb5
 }
 
+// CHECK-LABEL: sil @test_constant_folding
+// CHECK:  [[R:%[0-9]+]] = integer_literal $Builtin.Int32, 30
+// CHECK:  return [[R]] : $Builtin.Int32
+// CHECK-NEXT: }
+sil @test_constant_folding : $@convention(thin) () -> Builtin.Int32 {
+bb0:
+  %0 = integer_literal $Builtin.Int1, 0
+  %20 = integer_literal $Builtin.Int32, 20
+  %30 = integer_literal $Builtin.Int32, 30
+  cond_br %0, bb1, bb2
+bb1:
+  br bb3(%20 : $Builtin.Int32)
+bb2:
+  br bb3(%30 : $Builtin.Int32)
+
+bb3(%2 : $Builtin.Int32):
+  %3 = builtin "cmp_slt_Int32"(%2 : $Builtin.Int32, %30 : $Builtin.Int32) : $Builtin.Int1
+  cond_br %3, bb4, bb5
+bb4:
+  br bb6(%20 : $Builtin.Int32)
+bb5:
+  br bb6(%30 : $Builtin.Int32)
+
+bb6(%4 : $Builtin.Int32):
+  %5 = builtin "cmp_slt_Int32"(%4 : $Builtin.Int32, %30 : $Builtin.Int32) : $Builtin.Int1
+  cond_br %5, bb7, bb8
+bb7:
+  br bb9(%20 : $Builtin.Int32)
+bb8:
+  br bb9(%30 : $Builtin.Int32)
+
+bb9(%6 : $Builtin.Int32):
+  return %6 : $Builtin.Int32
+}
+