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
133
134
135
136
137
138
139
140
141
142
143
144
145
146

use std::collections::HashMap;

use itertools::Itertools;

use crate::blocks::FlatBlocksBuilder;
use crate::borrow_check::analysis::{Analyzer, BackAnalysis, StatementLocation};
use crate::utils::{Rebuilder, RebuilderEx};
use crate::{
    BlockId, FlatBlock, FlatBlockEnd, FlatLowered, MatchInfo, VarRemapping, VarUsage, VariableId,
};

/// Reorganizes the blocks in lowered function.
///
/// Removes unreachable blocks.
/// Blocks that are reachable only through goto are combined with the block that does the goto.
/// The order of the blocks is changed to be a topologically sorted.
/// Assumes that there are not remappings on gotos to a block with 1 incoming edge.
pub fn reorganize_blocks(lowered: &mut FlatLowered) {
    if lowered.blocks.is_empty() {
        return;
    }
    let ctx = TopSortContext {
        old_block_rev_order: Default::default(),
        incoming_gotos: vec![0; lowered.blocks.len()],
        can_be_merged: vec![true; lowered.blocks.len()],
    };
    let mut analysis =
        BackAnalysis { lowered: &*lowered, block_info: Default::default(), analyzer: ctx };
    analysis.get_root_info();
    let ctx = analysis.analyzer;

    // Rebuild the blocks in the correct order.
    let mut new_blocks = FlatBlocksBuilder::default();

    // Keep only blocks that can't be merged or have more than 1 incoming
    // goto.
    // Note that unreachable block were not added to `ctx.old_block_rev_order` during
    // the analysis above.
    let mut old_block_rev_order = ctx
        .old_block_rev_order
        .into_iter()
        .filter(|block_id| !ctx.can_be_merged[block_id.0] || ctx.incoming_gotos[block_id.0] > 1)
        .collect_vec();

    // Add the root block as it was filtered above.
    old_block_rev_order.push(BlockId::root());

    let n_visited_blocks = old_block_rev_order.len();

    let mut rebuilder = RebuildContext {
        block_remapping: HashMap::from_iter(
            old_block_rev_order
                .iter()
                .enumerate()
                .map(|(idx, block_id)| (*block_id, BlockId(n_visited_blocks - idx - 1))),
        ),
    };
    for block_id in old_block_rev_order.into_iter().rev() {
        let mut statements = vec![];

        let mut block = &lowered.blocks[block_id];
        loop {
            for stmt in &block.statements {
                statements.push(rebuilder.rebuild_statement(stmt));
            }
            if let FlatBlockEnd::Goto(target_block_id, remappings) = &block.end {
                if rebuilder.block_remapping.get(target_block_id).is_none() {
                    assert!(remappings.is_empty(), "Remapping should be empty.");
                    block = &lowered.blocks[*target_block_id];
                    continue;
                }
            }
            break;
        }

        let end = rebuilder.rebuild_end(&block.end);
        new_blocks.alloc(FlatBlock { statements, end });
    }

    lowered.blocks = new_blocks.build().unwrap();
}

pub struct TopSortContext {
    old_block_rev_order: Vec<BlockId>,
    // The number of incoming gotos, indexed by block_id.
    incoming_gotos: Vec<usize>,

    // True if the block can be merged with the block that goes to it.
    can_be_merged: Vec<bool>,
}

impl Analyzer<'_> for TopSortContext {
    type Info = ();

    fn visit_block_start(&mut self, _info: &mut Self::Info, block_id: BlockId, _block: &FlatBlock) {
        self.old_block_rev_order.push(block_id);
    }

    fn visit_goto(
        &mut self,
        _info: &mut Self::Info,
        _statement_location: StatementLocation,
        target_block_id: BlockId,
        _remapping: &VarRemapping,
    ) {
        self.incoming_gotos[target_block_id.0] += 1;
    }

    fn merge_match(
        &mut self,
        _statement_location: StatementLocation,
        match_info: &MatchInfo,
        _infos: &[Self::Info],
    ) -> Self::Info {
        for arm in match_info.arms().iter() {
            self.can_be_merged[arm.block_id.0] = false;
        }
    }

    fn info_from_return(
        &mut self,
        _statement_location: StatementLocation,
        _vars: &[VarUsage],
    ) -> Self::Info {
    }

    fn info_from_panic(
        &mut self,
        _statement_location: StatementLocation,
        _data: &VarUsage,
    ) -> Self::Info {
    }
}

pub struct RebuildContext {
    block_remapping: HashMap<BlockId, BlockId>,
}
impl Rebuilder for RebuildContext {
    fn map_var_id(&mut self, var: VariableId) -> VariableId {
        var
    }

    fn map_block_id(&mut self, block: BlockId) -> BlockId {
        self.block_remapping[&block]
    }
}