cubecl-opt 0.11.0-pre.1

use core::cell::RefCell;

use alloc::{rc::Rc, vec::Vec};
use cubecl_ir::{Instruction, Value};
use stable_vec::StableVec;

use crate::{ControlFlow, Function, GlobalState, version::PhiInstruction};

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BlockUse {
    ContinueTarget,
    Merge,
}

/// A basic block of instructions interrupted by control flow. Phi nodes are assumed to come before
/// any instructions. See <https://en.wikipedia.org/wiki/Basic_block>
#[derive(Default, Debug, Clone)]
pub struct BasicBlock {
    pub(crate) block_use: Vec<BlockUse>,
    /// The phi nodes that are required to be generated at the start of this block.
    pub phi_nodes: Rc<RefCell<Vec<PhiInstruction>>>,
    /// A stable list of operations performed in this block.
    pub ops: Rc<RefCell<StableVec<Instruction>>>,
    /// The control flow that terminates this block.
    pub control_flow: Rc<RefCell<ControlFlow>>,
}

impl Function {
    /// Visit all operations in the program with the specified read and write visitors.
    pub fn visit_all(
        &mut self,
        state: &GlobalState,
        mut visit_read: impl FnMut(&mut Self, &mut Value) + Clone,
        mut visit_write: impl FnMut(&mut Self, &mut Value) + Clone,
    ) {
        for node in self.node_indices().collect::<Vec<_>>() {
            let phi = self[node].phi_nodes.clone();
            let ops = self[node].ops.clone();
            let control_flow = self[node].control_flow.clone();

            for phi in phi.borrow_mut().iter_mut() {
                for elem in &mut phi.entries {
                    visit_read(self, &mut elem.value);
                }
                visit_write(self, &mut phi.out);
            }
            for op in ops.borrow_mut().values_mut() {
                self.visit_instruction(state, op, visit_read.clone(), visit_write.clone());
            }
            self.visit_control_flow(&mut control_flow.borrow_mut(), visit_read.clone());
        }
    }
}