JIT: Optimize data flow used in assertion prop/CSE

src/coreclr/jit/assertionprop.cpp

@@ -5488,8 +5488,15 @@ class AssertionPropFlowCallback
     //   lastTryBlock  - the last block of the try for "block" handler;.
     //
     // Notes:
-    //   We can jump to the handler from any instruction in the try region.
-    //   It means we can propagate only assertions that are valid for the whole try region.
+    //   We can jump to the handler from any instruction in the try region. It
+    //   means we can propagate only assertions that are valid for the whole
+    //   try region.
+    //
+    //   It suffices to intersect with only the head 'try' block's assertions,
+    //   since that block dominates all other blocks in the try and we never
+    //   kill assertions in global AP. Note that if we did kill assertions we
+    //   would need to be more careful about our mid-block handling when in a
+    //   try region.
     void MergeHandler(BasicBlock* block, BasicBlock* firstTryBlock, BasicBlock* lastTryBlock)
     {
         if (VerboseDataflow())
@@ -5498,11 +5505,8 @@ class AssertionPropFlowCallback
             Compiler::optDumpAssertionIndices("in -> ", block->bbAssertionIn, "; ");
             JITDUMP("firstTryBlock " FMT_BB " ", firstTryBlock->bbNum);
             Compiler::optDumpAssertionIndices("in -> ", firstTryBlock->bbAssertionIn, "; ");
-            JITDUMP("lastTryBlock " FMT_BB " ", lastTryBlock->bbNum);
-            Compiler::optDumpAssertionIndices("out -> ", lastTryBlock->bbAssertionOut, "\n");
         }
         BitVecOps::IntersectionD(apTraits, block->bbAssertionIn, firstTryBlock->bbAssertionIn);
-        BitVecOps::IntersectionD(apTraits, block->bbAssertionIn, lastTryBlock->bbAssertionOut);
     }
 
     // At the end of the merge store results of the dataflow equations, in a postmerge state.

src/coreclr/jit/dataflow.h

@@ -58,19 +58,47 @@ void DataFlow::ForwardAnalysis(TCallback& callback)
         }
         else
         {
-            FlowEdge* preds = m_pCompiler->BlockPredsWithEH(block);
-            for (FlowEdge* pred = preds; pred; pred = pred->getNextPredEdge())
+            for (FlowEdge* pred = block->bbPreds; pred; pred = pred->getNextPredEdge())
             {
                 callback.Merge(block, pred->getSourceBlock(), pred->getDupCount());
             }
         }
 
         if (callback.EndMerge(block))
         {
-            block->VisitAllSuccs(m_pCompiler, [&worklist](BasicBlock* succ) {
+            // The clients using DataFlow (CSE, assertion prop) currently do
+            // not need EH successors here:
+            //
+            // 1. CSE does not CSE into handlers, so it considers no
+            // expressions available at the beginning of handlers;
+            //
+            // 2. Assertion prop never kills any assertions, so it is
+            // sufficient to propagate facts from the 'try' head block, since
+            // that block dominates all other blocks in the 'try'. That happens
+            // implicitly below.
+            //
+            block->VisitRegularSuccs(m_pCompiler, [&worklist](BasicBlock* succ) {
                 worklist.insert(worklist.end(), succ);
                 return BasicBlockVisit::Continue;
             });
         }
+
+        if (m_pCompiler->bbIsTryBeg(block))
+        {
+            // Handlers of the try are reachable (and may require special
+            // handling compared to the normal "at-the-end" propagation above).
+            EHblkDsc* eh = m_pCompiler->ehGetDsc(block->getTryIndex());
+            do
+            {
+                worklist.insert(worklist.end(), eh->ExFlowBlock());
+
+                if (eh->ebdEnclosingTryIndex == EHblkDsc::NO_ENCLOSING_INDEX)
+                {
+                    break;
+                }
+
+                eh = m_pCompiler->ehGetDsc(eh->ebdEnclosingTryIndex);
+            } while (eh->ebdTryBeg == block);
+        }
     }
 }