relooper/
lib.rs

1/*
2
3Relooper library
4================
5
6Copyright (c) 2021 Dannii Willis
7MIT licenced
8https://github.com/curiousdannii/if-decompiler
9
10Inspired by the Cheerp Stackifier algorithm
11https://medium.com/leaningtech/solving-the-structured-control-flow-problem-once-and-for-all-5123117b1ee2
12
13And the Relooper algorithm paper by Alon Zakai
14https://github.com/emscripten-core/emscripten/blob/master/docs/paper.pdf
15
16*/
17
18#![forbid(unsafe_code)]
19
20use core::hash::Hash;
21use std::fmt::Debug;
22use std::iter::FromIterator;
23
24use fnv::{FnvHashMap, FnvHashSet};
25use petgraph::prelude::*;
26use petgraph::algo;
27use petgraph::visit::{EdgeFiltered, Visitable, VisitMap};
28
29#[cfg(test)]
30mod tests;
31
32// Common traits for labels
33pub trait RelooperLabel: Copy + Debug + Eq + Hash + Ord {}
34impl<T> RelooperLabel for T
35where T: Copy + Debug + Eq + Hash + Ord {}
36
37type LoopId = u16;
38
39// The Relooper accepts a map of block labels to the labels each block can branch to
40pub fn reloop<L: RelooperLabel>(blocks: Vec<(L, Vec<L>)>, first_label: L) -> Box<ShapedBlock<L>> {
41    let mut relooper = Relooper::new(blocks, first_label);
42    relooper.process_loops();
43    relooper.process_rejoined_branches();
44    relooper.output(vec![relooper.graph_root], false).unwrap()
45}
46
47// And returns a ShapedBlock tree
48#[derive(Debug, PartialEq)]
49pub enum ShapedBlock<L: RelooperLabel> {
50    Simple(SimpleBlock<L>),
51    Loop(LoopBlock<L>),
52    Multiple(MultipleBlock<L>),
53}
54
55#[derive(Debug, PartialEq)]
56pub struct SimpleBlock<L: RelooperLabel> {
57    pub label: L,
58    pub immediate: Option<Box<ShapedBlock<L>>>,
59    pub branches: FnvHashMap<L, BranchMode>,
60    pub next: Option<Box<ShapedBlock<L>>>,
61}
62
63// Branch modes
64#[derive(Clone, Copy, Debug, PartialEq)]
65pub enum BranchMode {
66    LoopBreak(LoopId),
67    LoopBreakIntoMulti(LoopId),
68    LoopContinue(LoopId),
69    LoopContinueIntoMulti(LoopId),
70    MergedBranch,
71    MergedBranchIntoMulti,
72    SetLabelAndBreak,
73}
74
75#[derive(Debug, PartialEq)]
76pub struct LoopBlock<L: RelooperLabel> {
77    pub loop_id: LoopId,
78    pub inner: Box<ShapedBlock<L>>,
79    pub next: Option<Box<ShapedBlock<L>>>,
80}
81
82#[derive(Debug, PartialEq)]
83pub struct MultipleBlock<L: RelooperLabel> {
84    // It would be nicer to use a Hashmap here, but if the graph ever has triple branches it's possible you'd have a Multiple going into a LoopMulti, so we need a Vec of handled labels
85    pub handled: Vec<HandledBlock<L>>,
86}
87
88#[derive(Debug, PartialEq)]
89pub struct HandledBlock<L: RelooperLabel> {
90    pub labels: Vec<L>,
91    pub inner: ShapedBlock<L>,
92    pub break_after: bool,
93}
94
95/* =======================
96   Internal implementation
97   ======================= */
98
99#[derive(Clone, Copy, Debug, PartialEq)]
100enum Node<L> {
101    Root,
102    Basic(L),
103    Multiple(L),
104    Loop(LoopId),
105    LoopMulti(LoopId),
106}
107
108#[derive(Clone, Copy, Debug, PartialEq)]
109enum Edge<L> {
110    // Edges that form the acyclic CFG
111    Forward,
112    ForwardMulti(L),
113    Next(bool), // Whether or not the next block must be a Multiple
114    // Edges that will be filtered out when traversing the graph
115    ForwardMultiViaNext(L),
116    LoopBreak(LoopId),
117    LoopBreakIntoMulti(LoopId),
118    LoopContinue(LoopId),
119    LoopContinueIntoMulti(LoopId),
120    MergedBranch,
121    MergedBranchIntoMulti,
122    SetLabelAndBreak,
123    SwitchFallThrough,
124    Removed,
125}
126
127fn filter_edges<L>(edge: petgraph::graph::EdgeReference<Edge<L>>) -> bool {
128    use Edge::*;
129    match edge.weight() {
130        Forward | ForwardMulti(_) | Next(_) => true,
131        _ => false,
132    }
133}
134
135fn filter_edges_including_processed<L>(edge: petgraph::graph::EdgeReference<Edge<L>>) -> bool {
136    use Edge::*;
137    match edge.weight() {
138        Forward | ForwardMulti(_) | Next(_) | ForwardMultiViaNext(_)
139            | LoopBreak(_) | LoopBreakIntoMulti(_) | MergedBranch
140            | MergedBranchIntoMulti | SetLabelAndBreak | SwitchFallThrough => true,
141        _ => false,
142    }
143}
144
145// The Relooper algorithm
146struct Relooper<L: RelooperLabel> {
147    counter: LoopId,
148    graph: Graph<Node<L>, Edge<L>>,
149    graph_root: NodeIndex,
150    root: NodeIndex,
151}
152
153impl<L: RelooperLabel> Relooper<L> {
154    fn new(blocks: Vec<(L, Vec<L>)>, root_label: L) -> Relooper<L> {
155        let mut graph = Graph::new();
156        let mut nodes = FnvHashMap::default();
157
158        // Check the blocks are sorted
159        // Replace with ._is_sorted() when stable (https://github.com/rust-lang/rust/issues/53485)
160        let mut label = blocks[0].0;
161        for block in &blocks[1..] {
162            assert!(block.0 > label, "Blocks were not provided in sorted order");
163            label = block.0;
164        }
165
166        // Add a root node to the graph, in order to handle when the first label is a loop
167        // Do this now so that it won't be invalidated by deleted orphan nodes
168        let graph_root = graph.add_node(Node::Root);
169
170        // Add nodes for each block
171        for (label, branches) in &blocks {
172            nodes.insert(*label, graph.add_node(if branches.len() > 1 { Node::Multiple(*label) } else { Node::Basic(*label) }));
173        }
174
175        // Add the edges
176        for (label, branches) in &blocks {
177            for branch in branches {
178                graph.add_edge(nodes[&label], nodes[&branch], Edge::Forward);
179            }
180        }
181
182        // Connect the root node to the first label
183        graph.add_edge(graph_root, nodes[&root_label], Edge::Forward);
184
185        // Remove orphan nodes
186        let dominators = algo::dominators::simple_fast(&graph, graph_root);
187        for node in graph.node_indices() {
188            if node != graph_root && dominators.immediate_dominator(node).is_none() {
189                graph.remove_node(node);
190            }
191        }
192
193        Relooper {
194            counter: 0,
195            graph,
196            graph_root,
197            root: nodes[&root_label],
198        }
199    }
200
201    // Process loops by adding loop nodes and converting back edges
202    fn process_loops(&mut self) {
203        // Loop until we have no more SCCs
204        loop {
205            let mut found_loop = false;
206
207            // Get the strongly connected components
208            let sccs = self.graph_sccs();
209            for scc in sccs {
210                if scc.len() == 1 {
211                    // Test for self-loops
212                    let node = scc[0];
213                    let mut found_self_loop = false;
214                    for edge in self.graph.edges(node) {
215                        if let Edge::Forward = edge.weight() {
216                            if edge.target() == node {
217                                found_self_loop = true;
218                                break;
219                            }
220                        }
221                    }
222                    if !found_self_loop {
223                        continue;
224                    }
225                }
226                found_loop = true;
227
228                // Go through all the incoming edges and find the loop headers
229                let mut loop_headers = FnvHashSet::default();
230                let mut loop_parents = FnvHashSet::default();
231                for &node in &scc {
232                    for edge in self.graph.edges_directed(node, Incoming) {
233                        if !scc.contains(&edge.source()) {
234                            loop_headers.insert(node);
235                            loop_parents.insert(edge.source());
236                        }
237                    }
238                }
239                let loop_headers = Vec::from_iter(loop_headers);
240                let loop_parents = Vec::from_iter(loop_parents);
241
242                let scc_test = |i| !&scc.contains(&i);
243                let (loop_node, loop_id) = self.make_loop(&loop_headers, &loop_parents, scc_test);
244
245                // Fix edges which are branching outside the loop
246                let filtered_graph = EdgeFiltered::from_fn(&self.graph, filter_edges);
247                let dominators = algo::dominators::simple_fast(&filtered_graph, self.graph_root);
248
249                // Walk through the graph manually
250                let loop_at_root = loop_headers.contains(&self.root);
251                let mut stack = vec![loop_node];
252                let mut discovered = self.graph.visit_map();
253
254                while let Some(node) = stack.pop() {
255                    if discovered.visit(node) {
256                        let mut edges = self.graph.neighbors(node).detach();
257                        'edge_loop: while let Some((edge, target)) = edges.next(&self.graph) {
258                            // Look for not just Forward edges, but also LoopBreaks from a previous loop
259                            if let Edge::Forward | Edge::ForwardMulti(_) | Edge::LoopBreak(_) | Edge::LoopBreakIntoMulti(_) = self.graph[edge] {
260                                // When the root node is a loop there can't be any un-dominated nodes, so just push to the stack
261                                if loop_at_root {
262                                    if !discovered.is_visited(&target) {
263                                        stack.push(target);
264                                    }
265                                    continue 'edge_loop;
266                                }
267
268                                let target_dominators = dominators.strict_dominators(target).unwrap();
269                                for dom in target_dominators {
270                                    if dom == loop_node {
271                                        // This node is dominated by the structural dominator, so add it to the stack
272                                        if !discovered.is_visited(&target) {
273                                            stack.push(target);
274                                        }
275                                        continue 'edge_loop;
276                                    }
277                                }
278
279                                // Not dominated, so convert the edges
280                                // Add a next edge to the dominator if there isn't one already
281                                let dominator = dominators.immediate_dominator(target).unwrap();
282                                if !self.graph.contains_edge(dominator, target) {
283                                    self.graph.add_edge(dominator, target, Edge::Next(false));
284                                }
285                                // Check if the loop already has a next to something else
286                                let mut into_multi = false;
287                                for edge in self.graph.edges(dominator) {
288                                    if let Edge::Next(_) = edge.weight() {
289                                        if edge.target() != target {
290                                            into_multi = true;
291                                            break;
292                                        }
293                                    }
294                                }
295                                self.graph[edge] = if into_multi { Edge::LoopBreakIntoMulti(loop_id) } else { Edge::LoopBreak(loop_id) };
296                            }
297                        }
298                    }
299                }
300            }
301
302            if found_loop == false {
303                break;
304            }
305        }
306
307        let filtered_graph = EdgeFiltered::from_fn(&self.graph, filter_edges);
308        assert!(!algo::is_cyclic_directed(&filtered_graph), "Graph should not contain any cycles");
309    }
310
311    // Handle branches that merge back together
312    fn process_rejoined_branches(&mut self) {
313        struct MergingBranches {
314            dominator: NodeIndex,
315            dominator_order: usize,
316            dominated_nodes: Vec<NodeIndex>,
317            parent_nodes: FnvHashSet<NodeIndex>,
318        }
319
320        // Get the list of nodes in topological order and the dominators list
321        let filtered_graph = EdgeFiltered::from_fn(&self.graph, filter_edges_including_processed);
322        let mut space = algo::DfsSpace::new(&filtered_graph);
323        let dominators = algo::dominators::simple_fast(&filtered_graph, self.graph_root);
324        let sorted_nodes = algo::toposort(&filtered_graph, None).unwrap();
325        let mut nodes_to_process = FnvHashMap::default();
326
327        // Now go through the nodes in order, looking for those that have multiple parents
328        for &node in sorted_nodes.iter() {
329            let mut parent_nodes = FnvHashSet::default();
330            for edge in self.graph.edges_directed(node, Incoming) {
331                match edge.weight() {
332                    Edge::Forward | Edge::ForwardMulti(_) | Edge::LoopBreak(_) | Edge::LoopBreakIntoMulti(_) | Edge::Next(_) => {
333                        parent_nodes.insert(edge.source());
334                    },
335                    _ => {},
336                }
337            }
338            if parent_nodes.len() > 1 {
339                let dominator_id = dominators.immediate_dominator(node).unwrap();
340                if !nodes_to_process.contains_key(&dominator_id) {
341                    // Insert into the map the dominator's topological position and an empty vec for the dominated nodes
342                    nodes_to_process.insert(dominator_id, MergingBranches {
343                        dominator: dominator_id,
344                        dominator_order: sorted_nodes.iter().position(|&n| n == dominator_id).unwrap(),
345                        dominated_nodes: Vec::default(),
346                        parent_nodes: FnvHashSet::default(),
347                    });
348                }
349                let branch = nodes_to_process.get_mut(&dominator_id).unwrap();
350                branch.dominated_nodes.push(node);
351                branch.dominated_nodes.sort();
352                for parent in parent_nodes {
353                    branch.parent_nodes.insert(parent);
354                }
355            }
356        }
357        // Sort the dominator nodes in topological order
358        let mut nodes_to_process = Vec::from_iter(nodes_to_process.values());
359        nodes_to_process.sort_by(|a, b| a.dominator_order.cmp(&b.dominator_order));
360
361        // Now go through the dominator nodes, processing each merging node it dominates
362        for merged_branch in nodes_to_process {
363            let dominator = merged_branch.dominator;
364            let dominated_nodes = &merged_branch.dominated_nodes;
365            let into_multi = dominated_nodes.len() > 1;
366
367            // Handle pre-existing next nodes from loop breaks
368            let mut dominator_edges = self.graph.neighbors(dominator).detach();
369            while let Some((edge_id, target)) = dominator_edges.next(&self.graph) {
370                if let Edge::Next(_) = self.graph[edge_id] {
371                    if !dominated_nodes.contains(&target) {
372                        panic!("Loop break to a node that is not one of this dominator's branch merges");
373                    }
374                    if into_multi {
375                        // Convert the LoopBreak nodes into LoopBreakIntoMulti if this is a multi node
376                        let mut target_edges = self.graph.neighbors_directed(target, Incoming).detach();
377                        while let Some((edge_id, _)) = target_edges.next(&self.graph) {
378                            match self.graph[edge_id] {
379                                Edge::LoopBreak(loop_id) => { self.graph[edge_id] = Edge::LoopBreakIntoMulti(loop_id); },
380                                _ => {},
381                            };
382                        }
383                    }
384                }
385            }
386
387            // Simple case - only one merged branch, branches that can't reach each other, or branches that only reach the next branch in order
388            if !into_multi || self.can_merged_nodes_use_multiple(&dominated_nodes) {
389                let dominated_nodes_count = dominated_nodes.len();
390                for index in 0..dominated_nodes_count {
391                    let node = dominated_nodes[index];
392                    let next_node = dominated_nodes.get(index + 1);
393                    // Add the next edge
394                    self.graph.add_edge(dominator, node, Edge::Next(false));
395                    // Patch the existing incoming edges
396                    let mut incoming_edges = self.graph.neighbors_directed(node, Incoming).detach();
397                    while let Some((edge, _)) = incoming_edges.next(&self.graph) {
398                        match self.graph[edge] {
399                            Edge::Forward => self.graph[edge] = if into_multi { Edge::MergedBranchIntoMulti } else { Edge::MergedBranch },
400                            Edge::LoopBreak(loop_id) => self.graph[edge] = if into_multi { Edge::LoopBreakIntoMulti(loop_id) } else { Edge::LoopBreak(loop_id) },
401                            _ => {},
402                        };
403                    }
404
405                    // Check if this node can reach the next
406                    if let Some(&next_node) = next_node {
407                        let filtered_graph = EdgeFiltered::from_fn(&self.graph, filter_edges_including_processed);
408                        if algo::has_path_connecting(&filtered_graph, node, next_node, Some(&mut space)) {
409                            // Turn any MergedBranch that go outside the dominator into SetLabelAndBreak edges
410                            // That means another manual search
411                            let mut stack = vec![node];
412                            let mut discovered = self.graph.visit_map();
413                            while let Some(node) = stack.pop() {
414                                if discovered.visit(node) {
415                                    let mut edges = self.graph.neighbors(node).detach();
416                                    'edges_loop: while let Some((edge, target)) = edges.next(&self.graph) {
417                                        // If the target is the next node, convert the edge to a SwitchFallThrough
418                                        if target == next_node {
419                                            self.graph[edge] = Edge::SwitchFallThrough;
420                                            continue;
421                                        }
422                                        match self.graph[edge] {
423                                            Edge::Forward | Edge::ForwardMulti(_) | Edge::Next(_) => {
424                                                if !discovered.is_visited(&target) {
425                                                    stack.push(target);
426                                                }
427                                            },
428                                            Edge::MergedBranch | Edge::MergedBranchIntoMulti => {
429                                                // If the target node is dominated by the top node then add it to the stack
430                                                let target_dominators = dominators.strict_dominators(target).unwrap();
431                                                for dom in target_dominators {
432                                                    if dom == node && !discovered.is_visited(&target) {
433                                                        stack.push(target);
434                                                        continue 'edges_loop;
435                                                    }
436                                                }
437                                                // This edge branches outside the top node, so convert to a SetLabelAndBreak
438                                                self.graph[edge] = Edge::SetLabelAndBreak;
439                                            },
440                                            _ => {},
441                                        };
442                                    }
443                                }
444                            }
445                        }
446                    }
447                }
448                continue;
449            }
450
451            // Multiple nodes get turned into a LoopMulti
452            let loop_headers = dominated_nodes.iter().map(|&i| i).collect();
453            let loop_parents: Vec<NodeIndex> = merged_branch.parent_nodes.difference(&FnvHashSet::from_iter(dominated_nodes.iter().map(|&i| i))).map(|&i| i).collect();
454
455            let loop_parents_test = |i| loop_parents.contains(&i);
456            self.make_loop(&loop_headers, &loop_parents, loop_parents_test);
457        }
458
459        // Look for MergedBranch|LoopBreak edges which don't go to the right place
460        for &node in sorted_nodes.iter() {
461            let mut edges = self.graph.neighbors(node).detach();
462            while let Some((edge, target)) = edges.next(&self.graph) {
463                if let Edge::MergedBranch | Edge::MergedBranchIntoMulti | Edge::LoopBreak(_) | Edge::LoopBreakIntoMulti(_) = self.graph[edge] {
464                    // Go through the node's dominators (including itself)
465                    let mut change_to_multi = false;
466                    let mut processed_nodes = Vec::new();
467                    let node_dominators = dominators.dominators(node).unwrap();
468                    'dominator_loop: for dominator in node_dominators {
469                        // If we reach the root node then something has gone wrong
470                        if dominator == self.graph_root {
471                            panic!("No dominator of {:?} found with Next", target);
472                        }
473                        processed_nodes.push(dominator);
474
475                        // Get all the next nodes of this dominator
476                        let mut next_nodes = Vec::new();
477                        let mut dominator_edges = self.graph.neighbors(dominator).detach();
478                        while let Some((edge, target)) = dominator_edges.next(&self.graph) {
479                            if let Edge::Next(_) = self.graph[edge] {
480                                next_nodes.push(target);
481                            }
482                        }
483                        next_nodes.sort();
484
485                        // No next nodes is fine
486                        if next_nodes.len() == 0 {
487                            continue 'dominator_loop;
488                        }
489                        // If there is one next node...
490                        if next_nodes.len() == 1 {
491                            let next_target = next_nodes[0];
492                            // And it's the target, great!
493                            if next_target == target {
494                                break 'dominator_loop;
495                            }
496                            // If it's a node we've already processed, that's okay too
497                            else if processed_nodes.contains(&next_target) {
498                                continue 'dominator_loop;
499                            }
500                            // If it isn't, then force that next to be a Multiple
501                            else {
502                                change_to_multi = true;
503                                // Patch the incoming edges to the node
504                                let mut incoming_edges = self.graph.neighbors_directed(next_target, Incoming).detach();
505                                while let Some((edge, _)) = incoming_edges.next(&self.graph) {
506                                    match self.graph[edge] {
507                                        Edge::MergedBranch => { self.graph[edge] = Edge::MergedBranchIntoMulti; },
508                                        Edge::LoopBreak(loop_id) => { self.graph[edge] = Edge::LoopBreakIntoMulti(loop_id); },
509                                        Edge::Next(_) => { self.graph[edge] = Edge::Next(true); },
510                                        _ => {},
511                                    };
512                                }
513                            }
514                        }
515                        // If there are 2+ next nodes
516                        if next_nodes.len() > 1 {
517                            // And the target is one of them, great!
518                            for &edge_target in &next_nodes {
519                                if edge_target == target {
520                                    break 'dominator_loop;
521                                }
522                            }
523                            change_to_multi = true;
524                        }
525                    }
526
527                    // Patch the edges to Target if necessary
528                    if change_to_multi {
529                        let mut incoming_edges = self.graph.neighbors_directed(target, Incoming).detach();
530                        while let Some((edge, _)) = incoming_edges.next(&self.graph) {
531                            match self.graph[edge] {
532                                Edge::MergedBranch => { self.graph[edge] = Edge::MergedBranchIntoMulti; },
533                                Edge::LoopBreak(loop_id) => { self.graph[edge] = Edge::LoopBreakIntoMulti(loop_id); },
534                                Edge::Next(_) => { self.graph[edge] = Edge::Next(true); },
535                                _ => {},
536                            };
537                        }
538                    }
539                }
540            }
541        }
542    }
543
544    // Output the graph as blocks
545    fn output(&self, entries: Vec<NodeIndex>, force_multi: bool) -> Option<Box<ShapedBlock<L>>> {
546        if entries.len() == 0 {
547            return None
548        }
549
550        // If we have one entry, then return the appropriate block
551        if !force_multi && entries.len() == 1 {
552            let node_id = entries[0];
553            let node = &self.graph[node_id];
554            let mut immediate_entries = FnvHashSet::default();
555            let mut next_entries = FnvHashSet::default();
556            let mut outgoing_branches = FnvHashMap::default();
557            let mut next_multi = false;
558            for edge in self.graph.edges(node_id) {
559                if let Edge::Removed = edge.weight() {
560                    continue;
561                }
562                let target = edge.target();
563                let mut add_branch = |target, branch| outgoing_branches.insert(self.get_basic_node_label(target), branch);
564                match edge.weight() {
565                    Edge::Forward | Edge::ForwardMulti(_) => { immediate_entries.insert(target); },
566                    Edge::Next(is_multi) => { next_entries.insert(target); next_multi |= is_multi; },
567                    Edge::ForwardMultiViaNext(label) => { outgoing_branches.insert(*label, BranchMode::MergedBranchIntoMulti); },
568                    Edge::LoopBreak(loop_id) => { add_branch(target, BranchMode::LoopBreak(*loop_id)); },
569                    Edge::LoopBreakIntoMulti(loop_id) => { add_branch(target, BranchMode::LoopBreakIntoMulti(*loop_id)); },
570                    Edge::LoopContinue(loop_id) => { add_branch(target, BranchMode::LoopContinue(*loop_id)); },
571                    Edge::LoopContinueIntoMulti(loop_id) => { add_branch(target, BranchMode::LoopContinueIntoMulti(*loop_id)); },
572                    Edge::MergedBranch | Edge::SwitchFallThrough => { add_branch(target, BranchMode::MergedBranch); },
573                    Edge::MergedBranchIntoMulti => { add_branch(target, BranchMode::MergedBranchIntoMulti); },
574                    Edge::SetLabelAndBreak => { add_branch(target, BranchMode::SetLabelAndBreak); },
575                    Edge::Removed => {},
576                };
577            }
578            // Dedup the ForwardMulti edges
579            let immediate_entries = Vec::from_iter(immediate_entries);
580            let next_entries = Vec::from_iter(next_entries);
581            let is_multi = match node {
582                Node::Multiple(_) | Node::LoopMulti(_) => true,
583                _ => false,
584            };
585
586            return match node {
587                Node::Root => {
588                    assert_eq!(next_entries.len(), 0, "Root node should have no next entries");
589                    self.output(immediate_entries, false)
590                },
591                Node::Basic(label) | Node::Multiple(label) => {
592                    Some(Box::new(ShapedBlock::Simple(SimpleBlock {
593                        label: *label,
594                        immediate: self.output(immediate_entries, is_multi),
595                        next: self.output(next_entries, next_multi),
596                        branches: outgoing_branches,
597                    })))
598                },
599                Node::Loop(loop_id) | Node::LoopMulti(loop_id) => {
600                    Some(Box::new(ShapedBlock::Loop(LoopBlock {
601                        loop_id: *loop_id,
602                        inner: self.output(immediate_entries, is_multi).unwrap(),
603                        next: self.output(next_entries, next_multi),
604                    })))
605                },
606            }
607        }
608
609        // Multiples
610        let handled = self.output_multiple_handled(entries);
611        Some(Box::new(ShapedBlock::Multiple(MultipleBlock {
612            handled,
613        })))
614    }
615
616    // Can any of these nodes reach any other?
617    fn can_merged_nodes_use_multiple(&self, nodes: &Vec<NodeIndex>) -> bool {
618        // Filter the graph to ignore processed back edges
619        let filtered_graph = EdgeFiltered::from_fn(&self.graph, filter_edges);
620        let mut space = algo::DfsSpace::new(&filtered_graph);
621        for (x_index, &x) in nodes.iter().enumerate() {
622            for (y_index, &y) in nodes.iter().enumerate() {
623                // Skip nodes that are the same, or when y is immediately after x.
624                if x_index == y_index || y_index == x_index + 1 {
625                    continue;
626                }
627                if algo::has_path_connecting(&filtered_graph, x, y, Some(&mut space)) {
628                    return false;
629                }
630            }
631        }
632        return true;
633    }
634
635    fn get_basic_node_label(&self, id: NodeIndex) -> L {
636        match self.graph[id] {
637            Node::Basic(label) | Node::Multiple(label) => label,
638            Node::Loop(_) => {
639                let mut edges = self.graph.neighbors(id).detach();
640                while let Some((edge, target)) = edges.next(&self.graph) {
641                    if let Edge::Forward = self.graph[edge] {
642                        match self.graph[target] {
643                            Node::Basic(label) | Node::Multiple(label) => return label,
644                            other => panic!("Cannot get label of node inside loop: {:?}", other),
645                        }
646                    }
647                }
648                panic!("No inner node within loop node: {:?}", self.graph[id]);
649            },
650            other => panic!("Cannot get label of node: {:?}", other),
651        }
652    }
653
654    // Return the SCCs over a filtered graph
655    fn graph_sccs(&self) -> Vec<Vec<NodeIndex>> {
656        // Filter the graph to ignore processed back edges
657        let filtered_graph = EdgeFiltered::from_fn(&self.graph, filter_edges);
658        algo::kosaraju_scc(&filtered_graph)
659    }
660
661    // Make a loop
662    fn make_loop<F: Fn(NodeIndex) -> bool>(&mut self, loop_headers: &Vec<NodeIndex>, loop_parents: &Vec<NodeIndex>, loop_parent_filter: F) -> (NodeIndex, LoopId) {
663        let multi_loop = loop_headers.len() > 1;
664
665        // Add the new node
666        let loop_id = self.counter;
667        self.counter += 1;
668        let loop_node = self.graph.add_node(if multi_loop { Node::LoopMulti(loop_id) } else { Node::Loop(loop_id) });
669
670        // Process the incoming edges
671        for &node in loop_headers {
672            let mut edges = self.graph.neighbors_directed(node, Incoming).detach();
673            while let Some((edge_id, parent)) = edges.next(&self.graph) {
674                if loop_parent_filter(parent) {
675                    match self.graph[edge_id] {
676                        // Forward edges get replaced
677                        Edge::Forward => {
678                            self.graph.add_edge(parent, loop_node, if multi_loop { Edge::ForwardMulti(self.get_basic_node_label(node)) } else { Edge::Forward });
679                            self.graph[edge_id] = Edge::Removed;
680                        },
681                        Edge::ForwardMulti(_) => unreachable!("A loop node should never be a loop header of a second loop"),
682                        // Next edges get removed
683                        Edge::Next(_) => {
684                            self.graph[edge_id] = Edge::Removed;
685                        },
686                        // Other edges are left as they are
687                        _ => {},
688                    };
689                }
690            }
691        }
692
693        // Now add edges from the new node to the loop header(s)
694        for &header in loop_headers {
695            self.graph.add_edge(loop_node, header, Edge::Forward);
696        }
697
698        // If the loop parent of a LoopMulti was a Multiple, check if we can turn it into a Simple now.
699        if multi_loop {
700            for &node in loop_parents {
701                if let Node::Multiple(_) = self.graph[node] {
702                    let mut neighbors = FnvHashSet::default();
703                    for edge in self.graph.edges(node) {
704                        match edge.weight() {
705                            // Ignore Next and Removed edges, but include everything else
706                            Edge::Next(_) | Edge::Removed => {},
707                            _ => { neighbors.insert(edge.target()); },
708                        };
709                    }
710                    if neighbors.len() == 1 {
711                        self.graph[node] = Node::Basic(self.get_basic_node_label(node));
712                    }
713                }
714            }
715        }
716
717        let filtered_graph = EdgeFiltered::from_fn(&self.graph, filter_edges);
718        let dominators = algo::dominators::simple_fast(&filtered_graph, self.graph_root);
719
720        // If branching to a loop header, convert to a back edge
721        // Walk through the graph manually
722        let mut stack = loop_headers.clone();
723        let mut discovered = self.graph.visit_map();
724        while let Some(node) = stack.pop() {
725            if discovered.visit(node) {
726                let mut edges = self.graph.neighbors(node).detach();
727                while let Some((edge, target)) = edges.next(&self.graph) {
728                    if loop_headers.contains(&target) {
729                        self.graph[edge] = if multi_loop { Edge::LoopContinueIntoMulti(loop_id) } else { Edge::LoopContinue(loop_id) };
730                    }
731                    else {
732                        match self.graph[edge] {
733                            Edge::Forward | Edge::ForwardMulti(_) | Edge::Next(_) => {
734                                // If the target node is dominated by the loop node then add it to the stack
735                                let target_dominators = dominators.strict_dominators(target).unwrap();
736                                for dom in target_dominators {
737                                    if dom == loop_node && !discovered.is_visited(&target) {
738                                        stack.push(target);
739                                    }
740                                }
741                            },
742                            _ => {},
743                        };
744                    }
745                }
746            }
747        }
748
749        // If we have multiple parents, fix the ForwardMulti edges so that the loop won't be outputted multiple times
750        if loop_parents.len() > 1 {
751            let dominator = dominators.immediate_dominator(loop_node).unwrap();
752            for &node in loop_parents {
753                let mut edges = self.graph.neighbors(node).detach();
754                while let Some((edge, _)) = edges.next(&self.graph) {
755                    if let Edge::ForwardMulti(label) = self.graph[edge] {
756                        self.graph[edge] = Edge::ForwardMultiViaNext(label);
757                    }
758                };
759            }
760            self.graph.add_edge(dominator, loop_node, Edge::Next(false));
761        }
762
763        (loop_node, loop_id)
764    }
765
766    fn output_multiple_handled(&self, mut entries: Vec<NodeIndex>) -> Vec<HandledBlock<L>> {
767        let filtered_graph = EdgeFiltered::from_fn(&self.graph, filter_edges_including_processed);
768        let mut space = algo::DfsSpace::new(&filtered_graph);
769        let mut handled = Vec::default();
770        let entries_count = entries.len();
771        // TODO: require input to be sorted and just check instead?
772        entries.sort();
773        for index in 0..entries_count {
774            let entry = entries[index];
775            let next_entry = entries.get(index + 1);
776            handled.push(HandledBlock {
777                labels: match self.graph[entry] {
778                    Node::Basic(label) | Node::Multiple(label) => vec![label],
779                    Node::Loop(_) | Node::LoopMulti(_) => {
780                        let mut labels = Vec::default();
781                        let mut edges = self.graph.neighbors(entry).detach();
782                        while let Some((edge, target)) = edges.next(&self.graph) {
783                            if let Edge::Forward = self.graph[edge] {
784                                labels.push(self.get_basic_node_label(target));
785                            }
786                        }
787                        labels.sort();
788                        labels
789                    },
790                    _ => unimplemented!(),
791                },
792                inner: *self.output(vec![entry], false).unwrap(),
793                // false if this entry can reach the next, otherwise true
794                break_after: next_entry.map_or(true, |next| !algo::has_path_connecting(&filtered_graph, entry, *next, Some(&mut space))),
795            });
796        }
797        // Sort so that the tests will work
798        handled.sort_by(|a, b| a.labels.cmp(&b.labels));
799        handled
800    }
801}