RIR dominator graph pass

compiler/qsc_rir/src/passes.rs

@@ -1,11 +1,13 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
+mod build_dominator_graph;
 mod defer_meas;
 mod reindex_qubits;
 mod remap_block_ids;
 mod unreachable_code_check;
 
+pub use build_dominator_graph::build_dominator_graph;
 pub use defer_meas::defer_measurements;
 pub use reindex_qubits::reindex_qubits;
 pub use remap_block_ids::remap_block_ids;

compiler/qsc_rir/src/passes/build_dominator_graph.rs

@@ -0,0 +1,91 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+use qsc_data_structures::index_map::IndexMap;
+
+use crate::{
+    rir::{BlockId, Program},
+    utils::build_predecessors_map,
+};
+
+#[cfg(test)]
+mod tests;
+
+/// Given a program, return a map from block IDs to the block ID of its immediate dominator. From this,
+/// the dominator tree can be constructed by treating the map as a directed graph where the keys are the
+/// children and the values are the parents.
+/// This algorithm is from [A Simple, Fast Dominance Algorithm](http://www.hipersoft.rice.edu/grads/publications/dom14.pdf)
+/// by Cooper, Harvey, and Kennedy, with two notable differences:
+/// - Blocks are assumed to be sequentially numbered starting from 0 in reverse postorder rather than depth first order.
+/// - Given that reversal, intersection between nodes uses the lesser of the two nodes rather than the greater.
+#[must_use]
+pub fn build_dominator_graph(program: &Program) -> IndexMap<BlockId, BlockId> {
+    let mut doms = IndexMap::default();
+    let entry_block_id = program
+        .get_callable(program.entry)
+        .body
+        .expect("entry point should have a body");
+
+    // Predecessors are needed to compute the dominator tree.
+    let preds = build_predecessors_map(program);
+
+    // The entry block dominates itself.
+    doms.insert(entry_block_id, entry_block_id);
+
+    // The algorithm needs to run until the dominance map stabilizes, ie: no block's immediate dominator changes.
+    let mut changed = true;
+    while changed {
+        changed = false;
+        // Always skip the entry block, as it is the only block that by definition dominates itself.
+        for (block_id, _) in program.blocks.iter().skip(1) {
+            // The immediate dominator of a block is the intersection of the dominators of its predecessors.
+            // Start from an assumption that the first predecessor is the dominator, and intersect with the rest.
+            let (first_pred, rest_preds) = preds
+                .get(block_id)
+                .expect("block should be present")
+                .split_first()
+                .expect("every block should have at least one predecessor");
+            let mut new_dom = *first_pred;
+
+            // If there are no other predecessors, the immediate dominator is the first predecessor.
+            for pred in rest_preds {
+                // For each predecessor whose dominator is known, intersect with the current best guess.
+                // Note that the dominator of the predecessor may be a best guess that gets updated in
+                // a later iteration.
+                if doms.contains_key(*pred) {
+                    new_dom = intersect(&doms, new_dom, *pred);
+                }
+            }
+
+            // If the immediate dominator has changed, update the map and mark that the map has changed
+            // so that the algorithm will run again.
+            if doms.get(block_id) != Some(&new_dom) {
+                doms.insert(block_id, new_dom);
+                changed = true;
+            }
+        }
+    }
+
+    doms
+}
+
+/// Calculates the closest intersection of two blocks in the current dominator tree.
+/// This is the block that dominates both block1 and block2, and is the closest to both.
+/// This is done by walking up the dominator tree from both blocks until they meet, and
+/// can take advantage of the ordering in the the block ids to walk only as far as necessary
+/// and avoid membership checks in favor of simple comparisons.
+fn intersect(
+    doms: &IndexMap<BlockId, BlockId>,
+    mut block1: BlockId,
+    mut block2: BlockId,
+) -> BlockId {
+    while block1 != block2 {
+        while block1 > block2 {
+            block1 = *doms.get(block1).expect("block should be present");
+        }
+        while block2 > block1 {
+            block2 = *doms.get(block2).expect("block should be present");
+        }
+    }
+    block1
+}

compiler/qsc_rir/src/passes/build_dominator_graph/tests.rs

@@ -0,0 +1,292 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+#![allow(clippy::too_many_lines, clippy::needless_raw_string_hashes)]
+
+use expect_test::expect;
+use qsc_data_structures::index_map::IndexMap;
+
+use crate::{
+    passes::remap_block_ids,
+    rir::{
+        Block, BlockId, Callable, CallableId, CallableType, Instruction, Literal, Program, Value,
+    },
+};
+
+use super::build_dominator_graph;
+
+/// Creates a new program with a single, entry callable that has block 0 as its body.
+fn new_program() -> Program {
+    let mut program = Program::new();
+    program.entry = CallableId(0);
+    program.callables.insert(
+        CallableId(0),
+        Callable {
+            name: "main".to_string(),
+            input_type: Vec::new(),
+            output_type: None,
+            body: Some(BlockId(0)),
+            call_type: CallableType::Regular,
+        },
+    );
+    program
+}
+
+fn display_dominator_graph(doms: &IndexMap<BlockId, BlockId>) -> String {
+    let mut result = String::new();
+    for (block_id, dom) in doms.iter() {
+        result.push_str(&format!(
+            "Block {} dominated by block {},\n",
+            block_id.0, dom.0
+        ));
+    }
+    result
+}
+
+#[test]
+fn test_dominator_graph_single_block_dominates_itself() {
+    let mut program = new_program();
+    program
+        .blocks
+        .insert(BlockId(0), Block(vec![Instruction::Return]));
+
+    remap_block_ids(&mut program);
+    let doms = build_dominator_graph(&program);
+
+    expect![[r#"
+        Block 0 dominated by block 0,
+    "#]]
+    .assert_eq(&display_dominator_graph(&doms));
+}
+
+#[test]
+fn test_dominator_graph_sequential_blocks_dominated_by_predecessor() {
+    let mut program = new_program();
+    program
+        .blocks
+        .insert(BlockId(0), Block(vec![Instruction::Jump(BlockId(1))]));
+    program
+        .blocks
+        .insert(BlockId(1), Block(vec![Instruction::Jump(BlockId(2))]));
+    program
+        .blocks
+        .insert(BlockId(2), Block(vec![Instruction::Return]));
+
+    remap_block_ids(&mut program);
+    let doms = build_dominator_graph(&program);
+
+    expect![[r#"
+        Block 0 dominated by block 0,
+        Block 1 dominated by block 0,
+        Block 2 dominated by block 1,
+    "#]]
+    .assert_eq(&display_dominator_graph(&doms));
+}
+
+#[test]
+fn test_dominator_graph_branching_blocks_dominated_by_common_predecessor() {
+    let mut program = new_program();
+    program
+        .blocks
+        .insert(BlockId(0), Block(vec![Instruction::Jump(BlockId(1))]));
+    program.blocks.insert(
+        BlockId(1),
+        Block(vec![Instruction::Branch(
+            Value::Literal(Literal::Bool(true)),
+            BlockId(2),
+            BlockId(3),
+        )]),
+    );
+    program
+        .blocks
+        .insert(BlockId(2), Block(vec![Instruction::Return]));
+    program
+        .blocks
+        .insert(BlockId(3), Block(vec![Instruction::Return]));
+
+    remap_block_ids(&mut program);
+    let doms = build_dominator_graph(&program);
+
+    expect![[r#"
+        Block 0 dominated by block 0,
+        Block 1 dominated by block 0,
+        Block 2 dominated by block 1,
+        Block 3 dominated by block 1,
+    "#]]
+    .assert_eq(&display_dominator_graph(&doms));
+}
+
+#[test]
+fn test_dominator_graph_infinite_loop() {
+    let mut program = new_program();
+    program
+        .blocks
+        .insert(BlockId(0), Block(vec![Instruction::Jump(BlockId(1))]));
+    program
+        .blocks
+        .insert(BlockId(1), Block(vec![Instruction::Jump(BlockId(1))]));
+
+    remap_block_ids(&mut program);
+    let doms = build_dominator_graph(&program);
+
+    expect![[r#"
+        Block 0 dominated by block 0,
+        Block 1 dominated by block 0,
+    "#]]
+    .assert_eq(&display_dominator_graph(&doms));
+}
+
+#[test]
+fn test_dominator_graph_branch_and_loop() {
+    let mut program = new_program();
+    program
+        .blocks
+        .insert(BlockId(0), Block(vec![Instruction::Jump(BlockId(1))]));
+    program.blocks.insert(
+        BlockId(1),
+        Block(vec![Instruction::Branch(
+            Value::Literal(Literal::Bool(true)),
+            BlockId(2),
+            BlockId(3),
+        )]),
+    );
+    program
+        .blocks
+        .insert(BlockId(2), Block(vec![Instruction::Jump(BlockId(4))]));
+    program
+        .blocks
+        .insert(BlockId(3), Block(vec![Instruction::Jump(BlockId(1))]));
+    program
+        .blocks
+        .insert(BlockId(4), Block(vec![Instruction::Return]));
+
+    remap_block_ids(&mut program);
+    let doms = build_dominator_graph(&program);
+
+    expect![[r#"
+        Block 0 dominated by block 0,
+        Block 1 dominated by block 0,
+        Block 2 dominated by block 1,
+        Block 3 dominated by block 1,
+        Block 4 dominated by block 2,
+    "#]]
+    .assert_eq(&display_dominator_graph(&doms));
+}
+
+#[test]
+fn test_dominator_graph_complex_structure_only_dominated_by_entry() {
+    // This example comes from the paper from [A Simple, Fast Dominance Algorithm](http://www.hipersoft.rice.edu/grads/publications/dom14.pdf)
+    // by Cooper, Harvey, and Kennedy and uses the node numbering from the paper. However, the resulting dominator graph
+    // is different due to the numbering of the blocks, such that each block is numbered in reverse postorder.
+    let mut program = new_program();
+    program
+        .callables
+        .get_mut(CallableId(0))
+        .expect("callable should be present")
+        .body = Some(BlockId(6));
+    program.blocks.insert(
+        BlockId(6),
+        Block(vec![Instruction::Branch(
+            Value::Literal(Literal::Bool(true)),
+            BlockId(5),
+            BlockId(4),
+        )]),
+    );
+    program
+        .blocks
+        .insert(BlockId(5), Block(vec![Instruction::Jump(BlockId(1))]));
+    program.blocks.insert(
+        BlockId(4),
+        Block(vec![Instruction::Branch(
+            Value::Literal(Literal::Bool(true)),
+            BlockId(2),
+            BlockId(3),
+        )]),
+    );
+    program
+        .blocks
+        .insert(BlockId(1), Block(vec![Instruction::Jump(BlockId(2))]));
+    program.blocks.insert(
+        BlockId(2),
+        Block(vec![Instruction::Branch(
+            Value::Literal(Literal::Bool(true)),
+            BlockId(3),
+            BlockId(1),
+        )]),
+    );
+    program
+        .blocks
+        .insert(BlockId(3), Block(vec![Instruction::Jump(BlockId(2))]));
+
+    remap_block_ids(&mut program);
+    let doms = build_dominator_graph(&program);
+
+    expect![[r#"
+        Block 0 dominated by block 0,
+        Block 1 dominated by block 0,
+        Block 2 dominated by block 0,
+        Block 3 dominated by block 0,
+        Block 4 dominated by block 0,
+        Block 5 dominated by block 0,
+    "#]]
+    .assert_eq(&display_dominator_graph(&doms));
+}
+
+#[test]
+fn test_dominator_graph_with_node_having_many_predicates() {
+    let mut program = new_program();
+    program.blocks.insert(
+        BlockId(0),
+        Block(vec![Instruction::Branch(
+            Value::Literal(Literal::Bool(true)),
+            BlockId(1),
+            BlockId(2),
+        )]),
+    );
+    program.blocks.insert(
+        BlockId(1),
+        Block(vec![Instruction::Branch(
+            Value::Literal(Literal::Bool(true)),
+            BlockId(3),
+            BlockId(4),
+        )]),
+    );
+    program.blocks.insert(
+        BlockId(2),
+        Block(vec![Instruction::Branch(
+            Value::Literal(Literal::Bool(true)),
+            BlockId(5),
+            BlockId(6),
+        )]),
+    );
+    program
+        .blocks
+        .insert(BlockId(3), Block(vec![Instruction::Jump(BlockId(7))]));
+    program
+        .blocks
+        .insert(BlockId(4), Block(vec![Instruction::Jump(BlockId(7))]));
+    program
+        .blocks
+        .insert(BlockId(5), Block(vec![Instruction::Jump(BlockId(7))]));
+    program
+        .blocks
+        .insert(BlockId(6), Block(vec![Instruction::Jump(BlockId(7))]));
+    program
+        .blocks
+        .insert(BlockId(7), Block(vec![Instruction::Return]));
+
+    remap_block_ids(&mut program);
+    let doms = build_dominator_graph(&program);
+
+    expect![[r#"
+        Block 0 dominated by block 0,
+        Block 1 dominated by block 0,
+        Block 2 dominated by block 0,
+        Block 3 dominated by block 1,
+        Block 4 dominated by block 1,
+        Block 5 dominated by block 2,
+        Block 6 dominated by block 2,
+        Block 7 dominated by block 0,
+    "#]]
+    .assert_eq(&display_dominator_graph(&doms));
+}