1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132

// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

use super::cfg::{BlockCFG, CFG};
use crate::{
    cfgir::ast::remap_labels,
    hlir::ast::{BasicBlocks, Command_, Label},
};
use std::collections::{BTreeMap, BTreeSet};

/// returns true if anything changed
pub fn optimize(cfg: &mut BlockCFG) -> bool {
    let changed = optimize_(cfg.start_block(), cfg.blocks_mut());
    if changed {
        let dead_blocks = cfg.recompute();
        assert!(dead_blocks.is_empty())
    }
    changed
}

fn optimize_(start: Label, blocks: &mut BasicBlocks) -> bool {
    let single_target_labels = find_single_target_labels(start, blocks);
    inline_single_target_blocks(&single_target_labels, start, blocks)
}

fn find_single_target_labels(start: Label, blocks: &BasicBlocks) -> BTreeSet<Label> {
    use Command_ as C;
    let mut counts = BTreeMap::new();
    // 'start' block starts as one as it is the entry point of the function. In some sense,
    // there is an implicit "jump" to this label to begin executing the function
    counts.insert(start, 1);
    for block in blocks.values() {
        match &block.back().unwrap().value {
            C::JumpIf {
                cond: _cond,
                if_true,
                if_false,
            } => {
                *counts.entry(*if_true).or_insert(0) += 1;
                *counts.entry(*if_false).or_insert(0) += 1
            }
            C::Jump { target, .. } => *counts.entry(*target).or_insert(0) += 1,
            _ => (),
        }
    }
    counts
        .into_iter()
        .filter(|(_, count)| *count == 1)
        .map(|(lbl, _)| lbl)
        .collect()
}

#[allow(clippy::needless_collect)]
fn inline_single_target_blocks(
    single_jump_targets: &BTreeSet<Label>,
    start: Label,
    blocks: &mut BasicBlocks,
) -> bool {
    //cleanup of needless_collect would result in mut and non mut borrows, and compilation warning.
    let labels_vec = blocks.keys().cloned().collect::<Vec<_>>();
    let mut labels = labels_vec.into_iter();
    let mut next = labels.next();

    let mut working_blocks = std::mem::take(blocks);
    let finished_blocks = blocks;

    let mut remapping = BTreeMap::new();
    while let Some(cur) = next {
        let mut block = match working_blocks.remove(&cur) {
            None => {
                next = labels.next();
                continue;
            }
            Some(b) => b,
        };

        match block.back().unwrap() {
            // Do not need to worry about infinitely unwrapping loops as loop heads will always
            // be the target of at least 2 jumps: the jump to the loop and the "continue" jump
            // This is always true as long as we start the count for the start label at 1
            sp!(_, Command_::Jump { target, .. }) if single_jump_targets.contains(target) => {
                remapping.insert(cur, *target);
                let target_block = working_blocks.remove(target).unwrap();
                block.pop_back();
                block.extend(target_block);
                working_blocks.insert(cur, block);
            }
            _ => {
                finished_blocks.insert(cur, block);
                next = labels.next();
            }
        }
    }

    let changed = !remapping.is_empty();
    remap_to_last_target(remapping, start, finished_blocks);
    changed
}

/// In order to preserve loop invariants at the bytecode level, when a block is "inlined", that
/// block needs to be relabelled as the "inlined" block
/// Without this, blocks that were outside of loops could be inlined into the loop-body, breaking
/// invariants needed at the bytecode level.
/// For example:
/// Before:
///   A: block_a; jump B
///   B: block_b
///
///   s.t. A is the only one jumping to B
///
/// After:
///   B: block_a; block_b
fn remap_to_last_target(
    mut remapping: BTreeMap<Label, Label>,
    start: Label,
    blocks: &mut BasicBlocks,
) {
    // The start block can't be relabelled in the current CFG API.
    // But it does not need to be since it will always be the first block, thus it will not run
    // into issues in the bytecode verifier
    remapping.remove(&start);
    if remapping.is_empty() {
        return;
    }
    // populate remapping for non changed blocks
    for label in blocks.keys() {
        remapping.entry(*label).or_insert(*label);
    }
    let owned_blocks = std::mem::take(blocks);
    let (_start, remapped_blocks) = remap_labels(&remapping, start, owned_blocks);
    *blocks = remapped_blocks;
}