kasl-ir 0.3.5

//
//  Copyright 2026 Shuntaro Kasatani
//
//  Licensed under the Apache License, Version 2.0 (the "License");
//  you may not use this file except in compliance with the License.
//  You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
//  Unless required by applicable law or agreed to in writing, software
//  distributed under the License is distributed on an "AS IS" BASIS,
//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//  See the License for the specific language governing permissions and
//  limitations under the License.
//

mod inst_builder;

pub use inst_builder::InstBuilder;

use crate::{Block, BlockData, IRType, Inst, Value, Variable, function::Function};

/// A builder for constructing IR functions.
#[derive(Default)]
pub struct IRBuilder {
    /// The function which is being built.
    function: Option<Function>,
    /// The current block being built.
    current_block: Option<Block>,

    /// The next value ID.
    next_val_id: u32,
    /// The next variable ID.
    next_var_id: u32,
    /// The next block ID.
    next_block_id: u32,
}

impl IRBuilder {
    // --- NEW ---

    /// Creates a new IRBuilder with a function set.
    pub fn new() -> Self {
        Self {
            function: Some(Function::default()),
            ..Default::default()
        }
    }

    // --- FUNCTION CREATION ---

    /// Start translating a new function, overwriting the currently built function if exists.
    pub fn start_new_func(&mut self) {
        self.function = Some(Function::default())
    }

    // --- ID GENERATION ---

    /// Generates a new block ID and returns it.
    fn gen_block_id(&mut self) -> Block {
        let id = Block(self.next_block_id);
        self.next_block_id += 1;
        id
    }

    /// Checks if the given value has the specified type.
    /// Calling this function before starting a function can cause a panic.
    fn is_val_type(&self, val: Value, ty: IRType) -> bool {
        self.function.as_ref().unwrap().is_val_type(val, ty)
    }

    /// Returns the type of the given value.
    /// The caller must call this method with the value generated by this builder.
    /// Calling this function before starting a function can cause a panic.
    pub fn get_val_type(&self, val: Value) -> IRType {
        self.function.as_ref().unwrap().get_val_type(val)
    }

    // --- VALUE MANAGEMENT ---

    /// Creates a new value with the given type.
    /// Calling this function before starting a function can cause a panic.
    fn create_val(&mut self, ty: IRType) -> Value {
        let val = Value(self.next_val_id);
        self.next_val_id += 1;

        self.function.as_mut().unwrap().values.insert(val, ty);
        val
    }

    /// Returns the type of the given variable.
    /// The caller must call this method with the variable generated by this builder.
    /// Calling this function before starting a function can cause a panic.
    pub fn get_var_type(&self, var: Variable) -> IRType {
        self.function.as_ref().unwrap().get_var_type(var)
    }

    // --- VARIABLE MANAGEMENT ---

    /// Creates a new variable with the given type.
    /// Calling this function before starting a function can cause a panic.
    pub fn create_var(&mut self, ty: IRType) -> Variable {
        let var = Variable(self.next_var_id);
        self.next_var_id += 1;

        self.function.as_mut().unwrap().variables.insert(var, ty);
        var
    }

    // --- BLOCK OPERATION ---

    /// Creates a new block and returns the newly created block.
    /// Calling this function before starting a function can cause a panic.
    pub fn create_block(&mut self, param_types: &[IRType]) -> Block {
        let block = self.gen_block_id();

        // Create values for block arguments
        let params = param_types
            .iter()
            .map(|ty| self.create_val(*ty))
            .collect::<Vec<_>>();
        // Create block data with the generated arguments and an empty instruction list
        let block_data = BlockData {
            params,
            insts: Vec::new(),
        };

        // Insert the new block into the function's block map
        self.function
            .as_mut()
            .unwrap()
            .blocks
            .insert(block, block_data);
        block
    }

    /// Build instruction in the specified block from now.
    pub fn switch_to_block(&mut self, block: Block) {
        self.current_block = Some(block);
    }

    /// Returns the values for the block parameters of the specified block.
    /// The caller must call this method with the block generated by this builder.
    /// Calling this function before starting a function can cause a panic.
    pub fn get_block_params(&self, block: Block) -> &[Value] {
        self.function
            .as_ref()
            .unwrap()
            .get_block(&block)
            .unwrap()
            .params
            .as_slice()
    }

    /// Sets the entry block of the function.
    /// Calling this function before starting a function can cause a panic.
    pub fn set_entry_block(&mut self, block: Block) {
        self.function.as_mut().unwrap().entry_block = Some(block);
    }

    // --- INSTRUCTION UTILITY ---

    /// Adds an instruction to the current block.
    /// Calling this function before starting a function can cause a panic.
    fn push_inst(&mut self, inst: Inst) {
        if let Some(target_block) = self
            .current_block
            .as_ref()
            .and_then(|block| self.function.as_mut().unwrap().get_block_mut(block))
        {
            target_block.insts.push(inst);
        }
    }

    // --- FINALIZATION ---

    /// Finalizes the function and return the translated function.
    /// Calling this function before starting a function can cause a panic.
    pub fn finalize_func(&mut self) -> Function {
        self.function.take().unwrap()
    }
}