[jit] Refine Type of CondBranchCheckType When Its Branch Leads to Unr…

…eachable

Summary:
For example, in this program:
```
bb 0 {
  v0:Object = LoadArg<0>
  CondBranchCheckType<1, 2, TupleExact> v0
}
bb 1 {
  ...
}
bb 2 {
  Unreachable
}
```
We know that because `bb 2` is unreachable, `v0` must be a `TupleExact`. This diff inserts a `RefineType<TupleExact> v0` before the `Branch` that replaces the `CondBranchCheckType`. And if `bb 1` had been unreachable, the `RefineType `would be for `Object - TupleExact`.

Test Plan: unit tests

Reviewers: bsimmers, emacs, #cinderjit

Reviewed By: bsimmers

Differential Revision: https://phabricator.intern.facebook.com/D41747741

Tasks: T138210636

Tags: cinder-310-exclusive

LaMa Project: L1125187

Jit/hir/optimization.cpp

@@ -530,6 +530,8 @@ static bool absorbDstBlock(BasicBlock* block) {
 bool CleanCFG::RemoveUnreachableInstructions(CFG* cfg) {
   bool modified = false;
   std::vector<BasicBlock*> blocks = cfg->GetPostOrderTraversal();
+  DominatorAnalysis dom(*cfg->func);
+  RegUses reg_uses = collectDirectRegUses(*cfg->func);
 
   for (BasicBlock* block : blocks) {
     auto it = block->begin();
@@ -597,9 +599,31 @@ bool CleanCFG::RemoveUnreachableInstructions(CFG* cfg) {
                 "true branch must be unreachable");
             target = cond_branch->false_bb();
           }
-          // TODO(T138210636): When replacing CondBranchCheckType with Branch
-          // due to an unreachable target, leave an appropriate AssertType in
-          // its place
+
+          if (branch->IsCondBranchCheckType()) {
+            auto check_type_branch = static_cast<CondBranchCheckType*>(branch);
+            Register* refined_value = cfg->func->env.AllocateRegister();
+            Type check_type = check_type_branch->type();
+            if (target == cond_branch->false_bb()) {
+              check_type = TTop - check_type_branch->type();
+            }
+
+            Register* operand = check_type_branch->GetOperand(0);
+            RefineType::create(refined_value, check_type, operand)
+                ->InsertBefore(*cond_branch);
+            auto uses = reg_uses.find(operand);
+            if (uses == reg_uses.end()) {
+              break;
+            }
+            std::unordered_set<Instr*>& instrs_using_reg = uses->second;
+            const std::unordered_set<const BasicBlock*>& dom_set =
+                dom.getBlocksDominatedBy(target);
+            for (Instr* instr : instrs_using_reg) {
+              if (dom_set.count(instr->block())) {
+                instr->ReplaceUsesOf(operand, refined_value);
+              }
+            }
+          }
           cond_branch->ReplaceWith(*Branch::create(target));
         } else {
           JIT_CHECK(false, "Unexpected branch instruction %s", *branch);

RuntimeTests/hir_tests/clean_cfg_test.txt

@@ -204,7 +204,7 @@ CleanUpCondBranchToTwoUnreachableEdges
 # HIR
 fun test {
   bb 0 {
-    v0 = LoadArg<0; CInt32>
+    v0 = LoadArg<0>
     v1 = GuardType<Tuple> v0
     CondBranch<1, 2> v0
   }
@@ -237,7 +237,7 @@ CleanUpExistingUnreachable
 # HIR
 fun test {
   bb 0 {
-    v0 = LoadArg<0; CInt32>
+    v0 = LoadArg<0>
     v1 = GuardType<Tuple> v0
     CondBranch<1, 2> v0
   }
@@ -274,7 +274,7 @@ CleanUpLoopWithUnreachable
 # HIR
 fun test {
   bb 0 {
-    v0 = LoadArg<0; CInt32>
+    v0 = LoadArg<0>
     v1 = GuardType<Tuple> v0
     Branch<1>
   }
@@ -329,3 +329,96 @@ fun test {
   }
 }
 ---
+CleanUpTruthyBranchOfCondBranchCheckType
+---
+# HIR
+fun test {
+  bb 0 {
+    v0 = LoadArg<0>
+    CondBranchCheckType<1, 2, Tuple> v0
+  }
+
+  bb 1 (preds 0) {
+    v1:Object = BinaryOp<Add> v0 v0
+    CondBranch<3, 4> v1
+  }
+
+  bb 2 (preds 0) {
+    v2:Bottom = LoadConst<Bottom>
+    Branch<3>
+  }
+
+  bb 3 (preds 1, 2) {
+    v4 = Phi<1, 2> v1 v2
+    Return v4
+  }
+
+  bb 4 (preds 1) {
+    v5:Object = BinaryOp<Add> v1 v1
+    CondBranch<1, 3> v5
+  }
+}
+---
+fun test {
+  bb 0 {
+    v0:Object = LoadArg<0>
+    v6:Tuple = RefineType<Tuple> v0
+    Branch<1>
+  }
+
+  bb 1 (preds 0, 4) {
+    v1:Object = BinaryOp<Add> v6 v6 {
+      FrameState {
+        NextInstrOffset 0
+      }
+    }
+    CondBranch<3, 4> v1
+  }
+
+  bb 4 (preds 1) {
+    v5:Object = BinaryOp<Add> v1 v1 {
+      FrameState {
+        NextInstrOffset 0
+      }
+    }
+    CondBranch<1, 3> v5
+  }
+
+  bb 3 (preds 1, 4) {
+    Return v1
+  }
+}
+---
+CleanUpFalseBranchOfCondBranchCheckType
+---
+# HIR
+fun test {
+  bb 0 {
+    v0 = LoadArg<0>
+    CondBranchCheckType<2, 1, Tuple> v0
+  }
+
+  bb 1 (preds 0) {
+    v1:Object = BinaryOp<Add> v0 v0
+    Return v1
+  }
+
+  bb 2 (preds 0) {
+    v2:Bottom = LoadConst<Bottom>
+    Return v2
+  }
+}
+---
+fun test {
+  bb 0 {
+    v0:Object = LoadArg<0>
+    v3:{Array|BaseException|Bytes|Cell|Code|Dict|Float|Frame|Func|Gen|List|Long|NoneType|ObjectExact|ObjectUser|Set|Slice|Type|Unicode|WaitHandle} = RefineType<{Array|BaseException|Bytes|CBool|CDouble|CInt|CPtr|Cell|Code|Dict|Float|Frame|Func|Gen|List|Long|NoneType|Nullptr|ObjectExact|ObjectUser|Set|Slice|Type|Unicode|WaitHandle}> v0
+    v1:Object = BinaryOp<Add> v3 v3 {
+      FrameState {
+        NextInstrOffset 0
+      }
+    }
+    Return v1
+  }
+}
+---