Merge pull request #603 from egraphs-good/oflatt-simpler-jump-optimizer

Simplify jump optimizer and make slightly more general

Makefile

@@ -5,16 +5,19 @@ DIRS = . dag_in_context
 all: nits test
 
 test:
-	$(foreach dir,$(DIRS),(cd $(dir) && cargo insta test --release --unreferenced=reject) &&) :
+	cargo insta test --release --unreferenced=reject
+	cd dag_in_context && cargo insta test --release --unreferenced=reject
 
 test-clean:
-	$(foreach dir,$(DIRS),(cd $(dir) && cargo insta test --release --unreferenced=delete) &&) :
+	cargo insta test --release --unreferenced=delete
+	cd dag_in_context && cargo insta test --release --unreferenced=delete
 
 nits:
 	npx prettier infra/nightly-resources/*.js --check
 	@rustup component add clippy
 	@rustup component add rustfmt
-	$(foreach dir,$(DIRS),(cd $(dir) && cargo clippy --tests -- -D warnings && cargo fmt --check) &&) :
+	cargo clippy --tests -- -D warnings && cargo fmt --check
+	cd dag_in_context && cargo clippy --tests -- -D warnings && cargo fmt --check
 
 
 

src/rvsdg/optimize_direct_jumps.rs

@@ -8,7 +8,7 @@ use hashbrown::HashMap;
 use petgraph::{
     graph::EdgeIndex,
     stable_graph::{NodeIndex, StableDiGraph, StableGraph},
-    visit::{Bfs, EdgeRef},
+    visit::{DfsPostOrder, EdgeRef},
     Direction,
 };
 
@@ -19,10 +19,13 @@ use crate::Optimizer;
 
 impl SimpleCfgFunction {
     pub fn optimize_jumps(&self) -> Self {
-        self.single_in_single_out().collapse_empty_blocks()
+        self.fuse_down().collapse_empty_blocks()
     }
 
-    fn single_in_single_out(&self) -> SimpleCfgFunction {
+    /// Find cases where a block jumps directly to another block A -> B where B
+    /// A has only one outgoing edge and B has one incoming edge
+    /// Turn it into one block AB
+    fn fuse_down(&self) -> SimpleCfgFunction {
         let mut resulting_graph: StableGraph<BasicBlock, Branch> = StableDiGraph::new();
 
         // a map from nodes in the old graph to nodes in the
@@ -31,55 +34,56 @@ impl SimpleCfgFunction {
         // it points to the new, fused node
         let mut node_mapping: HashMap<NodeIndex, NodeIndex> = HashMap::new();
 
-        // we use a bfs so that previous nodes are mapped to new nodes
-        // before their children.
-        // This ensures that `node_mapping[&previous]` succeeds.
-        let mut bfs = Bfs::new(&self.graph, self.entry);
+        // we use a dfs post order
+        // so dependencies are visited before parents
+        // This ensures that `node_mapping[&next]` succeeds.
+        let mut dfs = DfsPostOrder::new(&self.graph, self.entry);
 
         let mut edges_to_add = vec![];
 
         // copy the graph without the edges
         // also choose which nodes get fused to which
         // by re-assigning in the node map
-        while let Some(node) = bfs.next(&self.graph) {
-            let mut collapse_node = false;
-            let edges = self
+        while let Some(node) = dfs.next(&self.graph) {
+            let outgoing_from_node = self
                 .graph
-                .edges_directed(node, Direction::Incoming)
+                .edges_directed(node, Direction::Outgoing)
                 .collect::<Vec<_>>();
-            // single incoming edge to node
-            if let &[single_edge] = edges.as_slice() {
-                let previous = single_edge.source();
-                let previous_outgoing = self
+            let target = if let &[single_edge] = outgoing_from_node.as_slice() {
+                let target = single_edge.target();
+                let incoming_to_next = self
                     .graph
-                    .edges_directed(previous, Direction::Outgoing)
-                    .collect::<Vec<_>>();
-                // single outgoing edge from previous
-                // and two distinct nodes
-                if previous_outgoing.len() == 1 && previous != node {
-                    let previous_new = node_mapping[&previous];
-
-                    // this node will be mapped to the previous
-                    node_mapping.insert(node, previous_new);
-
-                    // add instructions to the end of the previous node
-                    resulting_graph[previous_new]
-                        .instrs
-                        .extend(self.graph[node].instrs.to_vec());
-                    resulting_graph[previous_new]
-                        .footer
-                        .extend(self.graph[node].footer.to_vec());
-
-                    collapse_node = true;
+                    .edges_directed(target, Direction::Incoming)
+                    .count();
+                if incoming_to_next == 1 && target != node {
+                    Some(target)
+                } else {
+                    None
                 }
-            }
-
-            if !collapse_node {
+            } else {
+                None
+            };
+            // single outgoing edge
+            if let Some(next) = target {
+                let new_target = node_mapping[&next];
+
+                // this node will be mapped to the previous
+                node_mapping.insert(node, new_target);
+
+                // add instructions to the beginning of the next node
+                let mut new_instrs = self.graph[node].instrs.to_vec();
+                let mut new_footer = self.graph[node].footer.to_vec();
+                new_instrs.extend(resulting_graph[new_target].instrs.to_vec());
+                new_footer.extend(resulting_graph[new_target].footer.to_vec());
+
+                resulting_graph[new_target].instrs = new_instrs;
+                resulting_graph[new_target].footer = new_footer;
+            } else {
                 // add the node
                 let new_node = resulting_graph.add_node(self.graph[node].clone());
                 node_mapping.insert(node, new_node);
 
-                edges_to_add.extend(self.graph.edges_directed(node, Direction::Incoming));
+                edges_to_add.extend(self.graph.edges_directed(node, Direction::Outgoing));
             }
         }
 
@@ -89,15 +93,6 @@ impl SimpleCfgFunction {
             resulting_graph.add_edge(source, target, edge.weight().clone());
         }
 
-        let mut label_mapping = HashMap::<String, String>::new();
-        for (old, new) in node_mapping.iter() {
-            let old_entry = label_mapping.insert(
-                self.graph[*old].name.to_string(),
-                resulting_graph[*new].name.to_string(),
-            );
-            assert!(old_entry.is_none(), "Duplicate labels in graph");
-        }
-
         SimpleCfgFunction {
             name: self.name.clone(),
             args: self.args.clone(),

src/rvsdg/snapshots/eggcc__rvsdg__optimize_direct_jumps__add_block_ind_test.snap

@@ -3,7 +3,7 @@ source: src/rvsdg/optimize_direct_jumps.rs
 expression: cfg.optimize_jumps().to_bril().to_string()
 ---
 @main {
-.entry___:
+.exit___:
   v0: int = const 1;
   v1: int = const 2;
   v2: int = add v0 v1;

src/rvsdg/snapshots/eggcc__rvsdg__tests__rvsdg_state_mem_to_cfg_more_blocks_snapshot.snap

@@ -3,24 +3,24 @@ source: src/rvsdg/tests.rs
 expression: prog.to_string()
 ---
 @main {
-.__0__:
+.__21__:
   v1: int = const 1;
   v4: ptr<int> = alloc v1;
   v7: int = const 10;
   store v4 v7;
   v11: int = load v4;
   v14: bool = lt v11 v7;
-  br v14 .__33__ .__22__;
-.__33__:
+  br v14 .__39__ .__32__;
+.__39__:
   print v11;
   free v4;
   v31: int = id v7;
-  jmp .__40__;
-.__22__:
+  jmp .__42__;
+.__32__:
   v24: int = add v11 v1;
   free v4;
   print v24;
   v31: int = id v7;
-.__40__:
+.__42__:
   print v31;
 }

src/rvsdg/snapshots/eggcc__rvsdg__tests__rvsdg_state_mem_to_cfg_snapshot.snap

@@ -3,7 +3,7 @@ source: src/rvsdg/tests.rs
 expression: prog.to_string()
 ---
 @main {
-.__0__:
+.__16__:
   v1: int = const 1;
   v4: ptr<int> = alloc v1;
   v8: int = const 10;