fix a stupid bug in toposort (#53)

src/ir/editor/analyzer/control_flow/control_flow_loop.rs

@@ -8,6 +8,16 @@ pub enum LoopContent {
     SubLoop(Box<Loop>),
     Node(usize),
 }
+
+impl LoopContent {
+    pub fn is_node_in(&self, node: NodeIndex<usize>) -> bool {
+        match self {
+            LoopContent::SubLoop(it) => it.is_node_in(node),
+            LoopContent::Node(it) => node.index() == *it,
+        }
+    }
+}
+
 #[derive(Debug, PartialEq)]
 pub struct Loop {
     pub entries: Vec<usize>,
@@ -54,7 +64,7 @@ impl Loop {
                     .collect(),
             };
         }
-        let content = sccs
+        let mut content: Vec<_> = sccs
             .into_iter()
             .map(|mut scc| {
                 if scc.len() == 1 {
@@ -65,6 +75,11 @@ impl Loop {
                 }
             })
             .collect();
+        for entry in &entries {
+            if !content.iter().any(|it| it.is_node_in(*entry)) {
+                content.push(LoopContent::Node(entry.index()));
+            }
+        }
         Self {
             entries: entries.into_iter().map(|it| it.index()).collect(),
             content,
@@ -136,12 +151,22 @@ impl Loop {
             .find_map(|it| it.smallest_loop_node_in(node))
     }
 
-    pub fn is_in_loop(&self, node: NodeIndex<usize>) -> bool {
+    pub fn is_node_in(&self, node: NodeIndex<usize>) -> bool {
         self.content.iter().any(|it| match it {
-            LoopContent::SubLoop(sub_loop) => sub_loop.is_in_loop(node),
+            LoopContent::SubLoop(sub_loop) => sub_loop.is_node_in(node),
             LoopContent::Node(n) => node.index() == *n,
         })
     }
+
+    pub fn node_count(&self) -> usize {
+        self.content
+            .iter()
+            .map(|it| match it {
+                LoopContent::SubLoop(sub_loop) => sub_loop.node_count(),
+                LoopContent::Node(_) => 1,
+            })
+            .sum()
+    }
 }
 
 #[cfg(test)]

src/ir/editor/analyzer/control_flow/mod.rs

@@ -7,7 +7,10 @@ use std::{
 use bimap::BiMap;
 use itertools::Itertools;
 use petgraph::{
-    algo::{self, dominators::simple_fast},
+    algo::{
+        self,
+        dominators::{simple_fast, Dominators},
+    },
     prelude::*,
 };
 
@@ -25,6 +28,7 @@ pub struct ControlFlowGraphContent {
     graph: DiGraph<(), (), usize>,
     frontiers: HashMap<usize, Vec<usize>>,
     bb_name_index_map: BiMap<usize, String>,
+    dominators: Dominators<NodeIndex<usize>>,
     // fixme: remove this refcell!
     from_to_may_pass_blocks: RefCell<HashMap<(usize, usize), Vec<usize>>>,
 }
@@ -73,6 +77,7 @@ impl ControlFlowGraphContent {
         Self {
             graph,
             frontiers,
+            dominators: dorminators,
             bb_name_index_map,
             from_to_may_pass_blocks: RefCell::new(HashMap::new()),
         }
@@ -83,6 +88,31 @@ impl ControlFlowGraphContent {
         self.frontiers.get(&bb_index).unwrap()
     }
 
+    fn dominates_calculate(&self, visiting: usize, visited: &mut Vec<usize>) {
+        if visited.contains(&visiting) {
+            return;
+        }
+        visited.push(visiting);
+        let mut imm_dominates: Vec<usize> = self.immediately_dominates(visiting);
+        imm_dominates.retain(|it| !visited.contains(it));
+        for it in imm_dominates {
+            self.dominates_calculate(it, visited);
+        }
+    }
+
+    fn immediately_dominates(&self, node: usize) -> Vec<usize> {
+        self.dominators
+            .immediately_dominated_by(node.into())
+            .map(|it| it.index())
+            .collect()
+    }
+
+    pub fn dominates(&self, node: usize) -> Vec<usize> {
+        let mut visited = Vec::new();
+        self.dominates_calculate(node, &mut visited);
+        visited
+    }
+
     /// Get the index of basic block named `name`.
     pub fn basic_block_index_by_name(&self, name: &str) -> usize {
         *self.bb_name_index_map.get_by_right(name).unwrap()
@@ -139,6 +169,9 @@ impl ControlFlowGraph {
     ) -> Ref<Vec<usize>> {
         self.content(content).may_pass_blocks(from, to)
     }
+    fn dominate(&self, content: &ir::FunctionDefinition, bb_index: usize) -> Vec<usize> {
+        self.content(content).dominates(bb_index)
+    }
     // todo: cache it
     fn loops(&self, content: &FunctionDefinition) -> Loop {
         let graph = &self.content(content).graph;
@@ -171,6 +204,9 @@ impl<'item, 'bind: 'item> BindedControlFlowGraph<'item, 'bind> {
     pub fn graph(&self) -> &DiGraph<(), (), usize> {
         &self.item.content(self.bind_on).graph
     }
+    pub fn dominates(&self, bb_index: usize) -> Vec<usize> {
+        self.item.dominate(self.bind_on, bb_index)
+    }
 }
 
 impl<'item, 'bind: 'item> IsAnalyzer<'item, 'bind> for ControlFlowGraph {

src/ir/editor/mod.rs

@@ -59,8 +59,7 @@ impl Editor {
     ) {
         let mut indexes = indexes.into_iter().collect::<Vec<_>>();
         indexes.sort();
-        while !indexes.is_empty() {
-            let index = indexes.pop().unwrap();
+        while let Some(index) = indexes.pop() {
             self.remove_statement(index);
         }
     }

src/ir/optimize/pass/memory_to_register.rs

@@ -71,8 +71,7 @@ fn insert_phi_positions(
         let mut pending_bb_indexes = memory_access_info.store.iter().map(|it| it.0).collect_vec();
         pending_bb_indexes.dedup();
         let mut done_bb_index = Vec::new();
-        while !pending_bb_indexes.is_empty() {
-            let considering_bb_index = pending_bb_indexes.pop().unwrap();
+        while let Some(considering_bb_index) = pending_bb_indexes.pop() {
             done_bb_index.push(considering_bb_index);
             let dominator_frontier_bb_indexes =
                 control_flow_graph.dominance_frontier(considering_bb_index);

src/ir/optimize/pass/topological_sort.rs

@@ -4,7 +4,7 @@ use itertools::Itertools;
 use petgraph::prelude::*;
 
 use crate::ir::{
-    analyzer::{IsAnalyzer, Loop},
+    analyzer::{BindedControlFlowGraph, IsAnalyzer, Loop},
     optimize::pass::fix_irreducible::FixIrreducible,
 };
 
@@ -17,8 +17,7 @@ impl IsPass for TopologicalSort {
         let analyzer = editor.analyzer.bind(&editor.content);
         let graph = analyzer.control_flow_graph();
         let loops = graph.loops();
-        let graph = graph.graph();
-        let content: Vec<_> = topological_order(graph, &loops)
+        let content: Vec<_> = topological_order(&graph, &loops)
             .into_iter()
             .map(|it| mem::take(&mut editor.content.content[it]))
             .collect();
@@ -35,54 +34,58 @@ impl IsPass for TopologicalSort {
 }
 
 fn topological_order_dfs(
-    graph: &DiGraph<(), (), usize>,
+    graph: &BindedControlFlowGraph,
     top_level: &Loop,
     current_at: NodeIndex<usize>,
+    visited: &mut Vec<NodeIndex<usize>>,
     result: &mut Vec<NodeIndex<usize>>,
 ) {
-    if result.contains(&current_at) {
+    if visited.contains(&current_at) {
         return;
     }
-    result.push(current_at);
+    visited.push(current_at);
     let in_loop = top_level.smallest_loop_node_in(current_at);
-    assert!({
-        if let Some(in_loop) = in_loop {
-            if let Some(&root) = in_loop.entries.first() {
-                let root: NodeIndex<usize> = root.into();
-                result.contains(&root)
-            } else {
-                true
-            }
-        } else {
-            true
-        }
-    });
     let mut to_visit = graph
+        .graph()
         .neighbors_directed(current_at, Direction::Outgoing)
         .filter(|it| !result.contains(it))
         .collect::<Vec<_>>();
-    to_visit.sort_unstable();
     to_visit.sort_by_cached_key(|to_visit_node| {
-        let forward_of_to_visit_node =
-            graph.neighbors_directed(*to_visit_node, Direction::Incoming);
+        let forward_of_to_visit_node = graph
+            .graph()
+            .neighbors_directed(*to_visit_node, Direction::Incoming)
+            .filter(|forward_node| {
+                !graph
+                    .dominates(to_visit_node.index())
+                    .contains(&forward_node.index())
+            });
         // first visit those nodes which current_at is the only parent of to_visit_node
-        if forward_of_to_visit_node.count() == 1 {
-            return 0;
+        if forward_of_to_visit_node.clone().count() == 1 {
+            let forward_node = forward_of_to_visit_node.into_iter().next().unwrap();
+            let at_index = graph
+                .graph()
+                .neighbors_directed(forward_node, Direction::Outgoing)
+                .position(|it| it == *to_visit_node)
+                .unwrap();
+            return 2 + at_index;
         }
         // we should visit all nodes in this loop before the others
-        if let Some(in_loop) = in_loop && in_loop.is_in_loop(*to_visit_node) {
+        if let Some(in_loop) = in_loop && in_loop.is_node_in(*to_visit_node) {
             return 1;
         }
-        2
+        0
     });
     for to_visit_node in to_visit {
-        topological_order_dfs(graph, top_level, to_visit_node, result);
+        topological_order_dfs(graph, top_level, to_visit_node, visited, result);
     }
+    result.push(current_at);
 }
 
-pub fn topological_order(graph: &DiGraph<(), (), usize>, top_level: &Loop) -> Vec<usize> {
+pub fn topological_order(graph: &BindedControlFlowGraph, top_level: &Loop) -> Vec<usize> {
     let mut order = vec![];
-    topological_order_dfs(graph, top_level, 0.into(), &mut order);
+    let mut visited = vec![];
+    topological_order_dfs(graph, top_level, 0.into(), &mut visited, &mut order);
+    order.reverse();
     let mut order: Vec<usize> = order.into_iter().map(NodeIndex::index).collect();
     let exit_block_position = order.iter().position_max().unwrap();
     order.remove(exit_block_position);
@@ -210,5 +213,19 @@ mod tests {
             .position(|it| it.name == Some("bb3".to_string()))
             .unwrap();
         assert_eq!(bb2_pos + 1, bb3_pos);
+        let bb14_pos = editor
+            .content
+            .content
+            .iter()
+            .position(|it| it.name == Some("bb14".to_string()))
+            .unwrap();
+        let bb13_pos = editor
+            .content
+            .content
+            .iter()
+            .position(|it| it.name == Some("bb13".to_string()))
+            .unwrap();
+        assert!(bb1_pos < bb14_pos);
+        assert!(bb13_pos < bb14_pos);
     }
 }