duskphantom_middle/transform/
block_fuse.rs

1// Copyright 2024 Duskphantom Authors
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14//
15// SPDX-License-Identifier: Apache-2.0
16
17use anyhow::{Ok, Result};
18
19use crate::{
20    ir::{BBPtr, FunPtr},
21    Program,
22};
23
24use super::Transform;
25
26pub fn optimize_program(program: &mut Program) -> Result<bool> {
27    BlockFuse::new(program).run_and_log()
28}
29
30pub struct BlockFuse<'a> {
31    program: &'a mut Program,
32}
33
34impl<'a> Transform for BlockFuse<'a> {
35    fn get_program_mut(&mut self) -> &mut Program {
36        self.program
37    }
38
39    fn name() -> String {
40        "block_fuse".to_string()
41    }
42
43    fn run(&mut self) -> Result<bool> {
44        let mut changed = false;
45        for func in self.program.module.functions.clone() {
46            if func.is_lib() {
47                continue;
48            }
49            for bb in func.rpo_iter() {
50                changed |= self.fuse_block(bb, func)?;
51            }
52        }
53        Ok(changed)
54    }
55}
56
57impl<'a> BlockFuse<'a> {
58    pub fn new(program: &'a mut Program) -> Self {
59        Self { program }
60    }
61
62    /// If block has only one predecessor, and that predecessor has only one successor,
63    /// these two blocks can be fused as one.
64    fn fuse_block(&mut self, mut bb: BBPtr, func: FunPtr) -> Result<bool> {
65        let Some(mut pred) = bb.get_pred_bb().first().cloned() else {
66            return Ok(false);
67        };
68        if pred.get_succ_bb().len() == 1 && bb.get_pred_bb().len() == 1 {
69            // Last instruction is "br", move the rest to successor block
70            for inst in pred.iter_rev().skip(1) {
71                bb.push_front(inst);
72            }
73
74            // Replace `pred -> bb` with `bb`
75            // TODO remove requirement of func in replace_entry
76            pred.replace_entry(bb, func);
77
78            // Remove `pred`
79            pred.remove_self();
80            return Ok(true);
81        }
82        Ok(false)
83    }
84}