Simplify condition graph.

Author: cjgillot

compiler/rustc_mir_transform/src/jump_threading.rs

@@ -134,7 +134,8 @@ impl<'tcx> crate::MirPass<'tcx> for JumpThreading {
             finder.entry_states[bb] = state;
         }
 
-        if let Some(opportunities) = OpportunitySet::new(body, finder.entry_states) {
+        if let Some(mut opportunities) = OpportunitySet::new(body, finder.entry_states) {
+            opportunities.simplify();
             opportunities.apply();
         }
     }
@@ -796,6 +797,88 @@ impl<'a, 'tcx> OpportunitySet<'a, 'tcx> {
         Some(OpportunitySet { basic_blocks, entry_states, duplicates })
     }
 
+    /// Propagate fulfilled conditions forward in the CFG to reduce the amount of duplication.
+    #[instrument(level = "debug", skip(self))]
+    fn simplify(&mut self) {
+        let mut worklist = Vec::with_capacity(self.basic_blocks.len());
+        worklist.push(START_BLOCK);
+        let mut visited =
+            FxHashSet::with_capacity_and_hasher(self.basic_blocks.len(), Default::default());
+
+        // Start by computing the number of reachable predecessors for each block.
+        let mut predecessors = IndexVec::from_elem(0, &self.entry_states);
+        predecessors[START_BLOCK] = 1; // Account for the implicit entry edge.
+        while let Some(bb) = worklist.pop() {
+            if visited.insert(bb) {
+                let term = self.basic_blocks[bb].terminator();
+                for s in term.successors() {
+                    predecessors[s] += 1;
+                    worklist.push(s);
+                }
+            }
+        }
+
+        // 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, self.entry_states[bb].targets.len()),
+            self.entry_states.len(),
+        );
+
+        visited.clear();
+        worklist.push(START_BLOCK);
+        while let Some(bb) = worklist.pop() {
+            if !visited.insert(bb) {
+                continue;
+            }
+
+            trace!(?bb, preds = ?predecessors[bb]);
+            debug_assert_ne!(predecessors[bb], 0);
+
+            if let Some((condition, _)) =
+                fulfill_in_pred_count[bb].iter_enumerated().find(|&(_, &p)| p == predecessors[bb])
+            {
+                trace!(?condition);
+                self.entry_states[bb].fulfilled.push(condition);
+            }
+
+            let state = &self.entry_states[bb];
+            trace!(?state);
+
+            let mut targets: Vec<_> = state
+                .fulfilled
+                .iter()
+                .flat_map(|&index| state.targets[index].iter().copied())
+                .collect();
+            targets.sort();
+            targets.dedup();
+            trace!(?targets);
+
+            let mut successors =
+                self.basic_blocks[bb].terminator().successors().collect::<Vec<_>>();
+
+            targets.reverse();
+            while let Some(target) = targets.pop() {
+                match target {
+                    ConditionTarget::Goto(target) => {
+                        predecessors[target] += 1;
+                        for &s in successors.iter() {
+                            predecessors[s] -= 1;
+                        }
+                        targets.retain(|t| t.block() == target);
+                        successors.clear();
+                        successors.push(target);
+                    }
+                    ConditionTarget::Chain(target, condition) => {
+                        let count = successors.iter().filter(|&&s| s == target).count();
+                        fulfill_in_pred_count[target][condition] += count;
+                    }
+                }
+            }
+
+            worklist.extend(successors);
+        }
+    }
+
     /// Apply the opportunities on the graph.
     #[instrument(level = "debug", skip(self))]
     fn apply(mut self) {

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 -> bb37;
       }
 
       bb4: {
@@ -102,7 +102,7 @@
           _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: bb21, unwind unreachable];
       }
 
       bb6: {
@@ -144,7 +144,7 @@
           _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: bb27, unwind unreachable];
       }
 
       bb9: {
@@ -155,7 +155,7 @@
           StorageDead(_13);
           _1 = chained_conditions::BacktraceStyle::Short;
 -         goto -> bb18;
-+         goto -> bb36;
++         goto -> bb33;
       }
 
       bb10: {
@@ -218,118 +218,101 @@
 +     }
 + 
 +     bb21: {
-+         _1 = chained_conditions::BacktraceStyle::Off;
-+         goto -> bb40;
-+     }
-+ 
-+     bb22: {
 +         StorageDead(_5);
 +         StorageDead(_6);
-+         goto -> bb18;
++         goto -> bb23;
 +     }
 + 
-+     bb23: {
-+         StorageDead(_5);
-+         StorageDead(_6);
-+         goto -> bb25;
-+     }
-+ 
-+     bb24: {
++     bb22: {
 +         _24 = discriminant(_2);
 +         switchInt(move _24) -> [1: bb16, otherwise: bb15];
 +     }
 + 
-+     bb25: {
++     bb23: {
 +         _24 = discriminant(_2);
-+         switchInt(move _24) -> [1: bb28, otherwise: bb26];
++         switchInt(move _24) -> [1: bb26, otherwise: bb24];
 +     }
 + 
-+     bb26: {
++     bb24: {
 +         _22 = const false;
 +         _23 = const false;
 +         StorageDead(_2);
 +         _19 = discriminant(_1);
 +         goto -> bb11;
 +     }
 + 
-+     bb27: {
++     bb25: {
 +         switchInt(copy _23) -> [0: bb15, otherwise: bb17];
 +     }
 + 
-+     bb28: {
-+         goto -> bb26;
-+     }
-+ 
-+     bb29: {
-+         StorageDead(_12);
-+         StorageDead(_13);
-+         goto -> bb18;
++     bb26: {
++         goto -> bb24;
 +     }
 + 
-+     bb30: {
++     bb27: {
 +         StorageDead(_12);
 +         StorageDead(_13);
-+         goto -> bb32;
++         goto -> bb29;
 +     }
 + 
-+     bb31: {
++     bb28: {
 +         _24 = discriminant(_2);
 +         switchInt(move _24) -> [1: bb16, otherwise: bb15];
 +     }
 + 
-+     bb32: {
++     bb29: {
 +         _24 = discriminant(_2);
-+         switchInt(move _24) -> [1: bb35, otherwise: bb33];
++         switchInt(move _24) -> [1: bb32, otherwise: bb30];
 +     }
 + 
-+     bb33: {
++     bb30: {
 +         _22 = const false;
 +         _23 = const false;
 +         StorageDead(_2);
 +         _19 = discriminant(_1);
 +         goto -> bb13;
 +     }
 + 
-+     bb34: {
++     bb31: {
 +         switchInt(copy _23) -> [0: bb15, otherwise: bb17];
 +     }
 + 
-+     bb35: {
-+         goto -> bb33;
++     bb32: {
++         goto -> bb30;
 +     }
 + 
-+     bb36: {
++     bb33: {
 +         _24 = discriminant(_2);
-+         switchInt(move _24) -> [1: bb38, otherwise: bb37];
++         switchInt(move _24) -> [1: bb35, otherwise: bb34];
 +     }
 + 
-+     bb37: {
++     bb34: {
 +         _22 = const false;
 +         _23 = const false;
 +         StorageDead(_2);
 +         _19 = discriminant(_1);
 +         goto -> bb12;
 +     }
 + 
-+     bb38: {
-+         switchInt(copy _23) -> [0: bb37, otherwise: bb39];
++     bb35: {
++         switchInt(copy _23) -> [0: bb34, otherwise: bb36];
 +     }
 + 
-+     bb39: {
-+         drop(((_2 as Some).0: std::string::String)) -> [return: bb37, unwind unreachable];
++     bb36: {
++         drop(((_2 as Some).0: std::string::String)) -> [return: bb34, unwind unreachable];
 +     }
 + 
-+     bb40: {
++     bb37: {
 +         _24 = discriminant(_2);
-+         switchInt(move _24) -> [1: bb41, otherwise: bb33];
++         switchInt(move _24) -> [1: bb38, otherwise: bb30];
 +     }
 + 
-+     bb41: {
-+         goto -> bb42;
++     bb38: {
++         goto -> bb39;
 +     }
 + 
-+     bb42: {
-+         drop(((_2 as Some).0: std::string::String)) -> [return: bb33, unwind unreachable];
++     bb39: {
++         drop(((_2 as Some).0: std::string::String)) -> [return: bb30, unwind unreachable];
       }
   }