Simplify condition graph.

Author: cjgillot

compiler/rustc_mir_transform/src/jump_threading.rs

@@ -134,7 +134,10 @@ impl<'tcx> crate::MirPass<'tcx> for JumpThreading {
             finder.entry_states[bb] = state;
         }
 
-        if let Some(opportunities) = OpportunitySet::new(body, finder.entry_states) {
+        let mut entry_states = finder.entry_states;
+        simplify_conditions(body, &mut entry_states);
+
+        if let Some(opportunities) = OpportunitySet::new(body, entry_states) {
             opportunities.apply();
         }
     }
@@ -779,6 +782,91 @@ impl<'a, 'tcx> TOFinder<'a, 'tcx> {
     }
 }
 
+/// Propagate fulfilled conditions forward in the CFG to reduce the amount of duplication.
+#[instrument(level = "debug", skip(body, entry_states))]
+fn simplify_conditions(body: &Body<'_>, entry_states: &mut IndexVec<BasicBlock, ConditionSet>) {
+    let basic_blocks = &body.basic_blocks;
+    let reverse_postorder = basic_blocks.reverse_postorder();
+
+    // Start by computing the number of *incoming edges* for each block.
+    // We do not use the cached `basic_blocks.predecessors` as we only want reachable predecessors.
+    let mut predecessors = IndexVec::from_elem(0, &entry_states);
+    predecessors[START_BLOCK] = 1; // Account for the implicit entry edge.
+    for &bb in reverse_postorder {
+        let term = basic_blocks[bb].terminator();
+        for s in term.successors() {
+            predecessors[s] += 1;
+        }
+    }
+
+    // Compute the number of edges into each block that carry each condition.
+    let mut fulfill_in_pred_count = IndexVec::from_fn_n(
+        |bb: BasicBlock| IndexVec::from_elem_n(0, entry_states[bb].targets.len()),
+        entry_states.len(),
+    );
+
+    // By traversing in RPO, we increase the likelihood to visit predecessors before successors.
+    for &bb in reverse_postorder {
+        let preds = predecessors[bb];
+        trace!(?bb, ?preds);
+
+        // We have removed all the input edges towards this block. Just skip visiting it.
+        if preds == 0 {
+            continue;
+        }
+
+        let state = &mut entry_states[bb];
+        trace!(?state);
+
+        // Conditions that are fulfilled in all the predecessors, are fulfilled in `bb`.
+        trace!(fulfilled_count = ?fulfill_in_pred_count[bb]);
+        for (condition, &cond_preds) in fulfill_in_pred_count[bb].iter_enumerated() {
+            if cond_preds == preds {
+                trace!(?condition);
+                state.fulfilled.push(condition);
+            }
+        }
+
+        // We want to count how many times each condition is fulfilled,
+        // so ensure we are not counting the same edge twice.
+        let mut targets: Vec<_> = state
+            .fulfilled
+            .iter()
+            .flat_map(|&index| state.targets[index].iter().copied())
+            .collect();
+        targets.sort();
+        targets.dedup();
+        trace!(?targets);
+
+        // We may modify the set of successors by applying edges, so track them here.
+        let mut successors = basic_blocks[bb].terminator().successors().collect::<Vec<_>>();
+
+        targets.reverse();
+        while let Some(target) = targets.pop() {
+            match target {
+                ConditionTarget::Goto(target) => {
+                    // We update the count of predecessors. If target or any successor has not been
+                    // processed yet, this increases the likelihood we find something relevant.
+                    predecessors[target] += 1;
+                    for &s in successors.iter() {
+                        predecessors[s] -= 1;
+                    }
+                    // Only process edges that still exist.
+                    targets.retain(|t| t.block() == target);
+                    successors.clear();
+                    successors.push(target);
+                }
+                ConditionTarget::Chain(target, condition) => {
+                    // `predecessors` is the number of incoming *edges* in each block.
+                    // Count the number of edges that apply `condition` into `target`.
+                    let count = successors.iter().filter(|&&s| s == target).count();
+                    fulfill_in_pred_count[target][condition] += count;
+                }
+            }
+        }
+    }
+}
+
 struct OpportunitySet<'a, 'tcx> {
     basic_blocks: &'a mut IndexVec<BasicBlock, BasicBlockData<'tcx>>,
     entry_states: IndexVec<BasicBlock, ConditionSet>,

tests/mir-opt/jump_threading.chained_conditions.JumpThreading.panic-abort.diff

@@ -61,8 +61,7 @@
           _22 = const true;
           _23 = const true;
           _4 = discriminant(_2);
--         switchInt(move _4) -> [0: bb3, 1: bb4, otherwise: bb2];
-+         switchInt(move _4) -> [0: bb21, 1: bb4, otherwise: bb2];
+          switchInt(move _4) -> [0: bb3, 1: bb4, otherwise: bb2];
       }
 
       bb2: {
@@ -71,7 +70,8 @@
 
       bb3: {
           _1 = chained_conditions::BacktraceStyle::Off;
-          goto -> bb18;
+-         goto -> bb18;
++         goto -> bb35;
       }
 
       bb4: {
@@ -101,8 +101,7 @@
           _23 = const false;
           _5 = move ((_2 as Some).0: std::string::String);
           _1 = chained_conditions::BacktraceStyle::Full;
--         drop(_5) -> [return: bb7, unwind unreachable];
-+         drop(_5) -> [return: bb23, unwind unreachable];
+          drop(_5) -> [return: bb7, unwind unreachable];
       }
 
       bb6: {
@@ -131,7 +130,8 @@
       bb7: {
           StorageDead(_5);
           StorageDead(_6);
-          goto -> bb18;
+-         goto -> bb18;
++         goto -> bb22;
       }
 
       bb8: {
@@ -143,8 +143,7 @@
           _23 = const false;
           _12 = move ((_2 as Some).0: std::string::String);
           _1 = chained_conditions::BacktraceStyle::Off;
--         drop(_12) -> [return: bb10, unwind unreachable];
-+         drop(_12) -> [return: bb30, unwind unreachable];
+          drop(_12) -> [return: bb10, unwind unreachable];
       }
 
       bb9: {
@@ -155,13 +154,14 @@
           StorageDead(_13);
           _1 = chained_conditions::BacktraceStyle::Short;
 -         goto -> bb18;
-+         goto -> bb36;
++         goto -> bb31;
       }
 
       bb10: {
           StorageDead(_12);
           StorageDead(_13);
-          goto -> bb18;
+-         goto -> bb18;
++         goto -> bb27;
       }
 
       bb11: {
@@ -218,118 +218,89 @@
 +     }
 + 
 +     bb21: {
-+         _1 = chained_conditions::BacktraceStyle::Off;
-+         goto -> bb40;
-+     }
-+ 
-+     bb22: {
-+         StorageDead(_5);
-+         StorageDead(_6);
-+         goto -> bb18;
-+     }
-+ 
-+     bb23: {
-+         StorageDead(_5);
-+         StorageDead(_6);
-+         goto -> bb25;
-+     }
-+ 
-+     bb24: {
 +         _24 = discriminant(_2);
 +         switchInt(move _24) -> [1: bb16, otherwise: bb15];
 +     }
 + 
-+     bb25: {
++     bb22: {
 +         _24 = discriminant(_2);
-+         switchInt(move _24) -> [1: bb28, otherwise: bb26];
++         switchInt(move _24) -> [1: bb25, otherwise: bb23];
 +     }
 + 
-+     bb26: {
++     bb23: {
 +         _22 = const false;
 +         _23 = const false;
 +         StorageDead(_2);
 +         _19 = discriminant(_1);
 +         goto -> bb11;
 +     }
 + 
-+     bb27: {
++     bb24: {
 +         switchInt(copy _23) -> [0: bb15, otherwise: bb17];
 +     }
 + 
-+     bb28: {
-+         goto -> bb26;
-+     }
-+ 
-+     bb29: {
-+         StorageDead(_12);
-+         StorageDead(_13);
-+         goto -> bb18;
-+     }
-+ 
-+     bb30: {
-+         StorageDead(_12);
-+         StorageDead(_13);
-+         goto -> bb32;
++     bb25: {
++         goto -> bb23;
 +     }
 + 
-+     bb31: {
++     bb26: {
 +         _24 = discriminant(_2);
 +         switchInt(move _24) -> [1: bb16, otherwise: bb15];
 +     }
 + 
-+     bb32: {
++     bb27: {
 +         _24 = discriminant(_2);
-+         switchInt(move _24) -> [1: bb35, otherwise: bb33];
++         switchInt(move _24) -> [1: bb30, otherwise: bb28];
 +     }
 + 
-+     bb33: {
++     bb28: {
 +         _22 = const false;
 +         _23 = const false;
 +         StorageDead(_2);
 +         _19 = discriminant(_1);
 +         goto -> bb13;
 +     }
 + 
-+     bb34: {
++     bb29: {
 +         switchInt(copy _23) -> [0: bb15, otherwise: bb17];
 +     }
 + 
-+     bb35: {
-+         goto -> bb33;
++     bb30: {
++         goto -> bb28;
 +     }
 + 
-+     bb36: {
++     bb31: {
 +         _24 = discriminant(_2);
-+         switchInt(move _24) -> [1: bb38, otherwise: bb37];
++         switchInt(move _24) -> [1: bb33, otherwise: bb32];
 +     }
 + 
-+     bb37: {
++     bb32: {
 +         _22 = const false;
 +         _23 = const false;
 +         StorageDead(_2);
 +         _19 = discriminant(_1);
 +         goto -> bb12;
 +     }
 + 
-+     bb38: {
-+         switchInt(copy _23) -> [0: bb37, otherwise: bb39];
++     bb33: {
++         switchInt(copy _23) -> [0: bb32, otherwise: bb34];
 +     }
 + 
-+     bb39: {
-+         drop(((_2 as Some).0: std::string::String)) -> [return: bb37, unwind unreachable];
++     bb34: {
++         drop(((_2 as Some).0: std::string::String)) -> [return: bb32, unwind unreachable];
 +     }
 + 
-+     bb40: {
++     bb35: {
 +         _24 = discriminant(_2);
-+         switchInt(move _24) -> [1: bb41, otherwise: bb33];
++         switchInt(move _24) -> [1: bb36, otherwise: bb28];
 +     }
 + 
-+     bb41: {
-+         goto -> bb42;
++     bb36: {
++         goto -> bb37;
 +     }
 + 
-+     bb42: {
-+         drop(((_2 as Some).0: std::string::String)) -> [return: bb33, unwind unreachable];
++     bb37: {
++         drop(((_2 as Some).0: std::string::String)) -> [return: bb28, unwind unreachable];
       }
   }