//! The owning container for a normalized Slonik IR body.

use cranelift_entity::{EntityRef, ListPool, PrimaryMap};

use crate::ir::{
    Block, BlockCall, BlockData, Stmt, StmtData, Type, Value, ValueData, ValueDef, ValueList,
};

/// A normalized SSA body.
#[derive(Default)]
pub struct Body {
    /// All blocks in the body.
    pub blocks: PrimaryMap<Block, BlockData>,

    /// All statements in the body, in creation order.
    pub stmts: PrimaryMap<Stmt, StmtData>,

    /// All SSA values in the body.
    pub values: PrimaryMap<Value, ValueData>,

    /// Pooled storage for compact value lists.
    pub value_lists: ListPool<Value>,

    /// Pooled storage for compact statement lists.
    pub stmt_lists: ListPool<Stmt>,
}

impl Body {
    /// Creates an empty body.
    pub fn new() -> Self {
        Self::default()
    }

    /// Returns whether the body contains no blocks, statements, or values.
    pub fn is_empty(&self) -> bool {
        self.blocks.is_empty() && self.stmts.is_empty() && self.values.is_empty()
    }

    /// Returns the number of blocks in the body.
    pub fn num_blocks(&self) -> usize {
        self.blocks.len()
    }

    /// Returns the number of statements in the body.
    pub fn num_stmts(&self) -> usize {
        self.stmts.len()
    }

    /// Returns the number of SSA values in the body.
    pub fn num_values(&self) -> usize {
        self.values.len()
    }

    /// Creates a new empty block.
    pub fn create_block(&mut self) -> Block {
        self.blocks.push(BlockData::default())
    }

    /// Returns the data for `block`.
    pub fn block_data(&self, block: Block) -> &BlockData {
        &self.blocks[block]
    }

    /// Returns the mutable data for `block`.
    pub fn block_data_mut(&mut self, block: Block) -> &mut BlockData {
        &mut self.blocks[block]
    }

    /// Returns an iterator over blocks in creation order.
    pub fn blocks(&self) -> Blocks {
        Blocks {
            next: 0,
            len: self.blocks.len(),
        }
    }

    /// Appends a block parameter of type `ty` to `block`.
    pub fn append_block_param(&mut self, block: Block, ty: Type) -> Value {
        let index = self.blocks[block].params.len(&self.value_lists) as u16;

        let value = self.values.push(ValueData {
            ty,
            def: ValueDef::Param(block, index),
        });

        self.blocks[block].params.push(value, &mut self.value_lists);
        value
    }

    /// Returns the parameter list of `block`.
    pub fn block_params(&self, block: Block) -> &[Value] {
        self.blocks[block].params.as_slice(&self.value_lists)
    }

    /// Returns the ordered statement list of `block`.
    pub fn block_stmts(&self, block: Block) -> &[Stmt] {
        self.blocks[block].stmts.as_slice(&self.stmt_lists)
    }

    /// Returns the last statement in `block`, if any.
    pub fn last_stmt(&self, block: Block) -> Option<Stmt> {
        self.block_stmts(block).last().copied()
    }

    /// Creates a pooled value list from `values`.
    pub fn make_value_list(&mut self, values: &[Value]) -> ValueList {
        ValueList::from_slice(values, &mut self.value_lists)
    }

    /// Appends a non-value-producing statement to `block` and returns its handle.
    pub fn append_stmt(&mut self, block: Block, data: StmtData) -> Stmt {
        let stmt = self.stmts.push(data);
        self.blocks[block].stmts.push(stmt, &mut self.stmt_lists);
        stmt
    }

    /// Appends a value-producing statement to `block`.
    ///
    /// `make` receives the pre-allocated result `Value` and must embed it in
    /// the returned `StmtData`.  The allocation uses `next_key` so no other
    /// statement may be pushed between `next_key` and `push` — call sites must
    /// be careful not to interleave insertions.
    pub fn append_value_inst(
        &mut self,
        block: Block,
        ty: Type,
        make: impl FnOnce(Value) -> StmtData,
    ) -> Value {
        let stmt_key = self.stmts.next_key();
        let value = self.values.push(ValueData {
            ty,
            def: ValueDef::Inst(stmt_key),
        });
        let stmt = self.stmts.push(make(value));
        debug_assert_eq!(stmt, stmt_key);
        self.blocks[block].stmts.push(stmt, &mut self.stmt_lists);
        value
    }

    /// Returns the data for `stmt`.
    pub fn stmt_data(&self, stmt: Stmt) -> &StmtData {
        &self.stmts[stmt]
    }

    /// Returns the mutable data for `stmt`.
    pub fn stmt_data_mut(&mut self, stmt: Stmt) -> &mut StmtData {
        &mut self.stmts[stmt]
    }

    /// Returns the result value of `stmt`, if it produces one.
    pub fn inst_result(&self, stmt: Stmt) -> Option<Value> {
        self.stmts[stmt].result()
    }

    /// Returns the full value record for `value`.
    pub fn value_data(&self, value: Value) -> &ValueData {
        &self.values[value]
    }

    /// Returns the semantic type of `value`.
    pub fn value_type(&self, value: Value) -> Type {
        self.values[value].ty
    }

    /// Returns the definition site of `value`.
    pub fn value_def(&self, value: Value) -> ValueDef {
        self.values[value].def
    }

    /// Returns the terminator statement of `block`, if the block ends in one.
    pub fn block_terminator(&self, block: Block) -> Option<Stmt> {
        let stmt = self.last_stmt(block)?;
        self.stmts[stmt].is_terminator().then_some(stmt)
    }

    /// Returns the terminator data for `block`, if present.
    pub fn block_terminator_data(&self, block: Block) -> Option<StmtData> {
        let stmt = self.block_terminator(block)?;
        Some(self.stmts[stmt])
    }

    /// Returns an iterator over the CFG successors of `block`.
    pub fn block_successors(&self, block: Block) -> BlockSuccessors {
        let Some(term) = self.block_terminator_data(block) else {
            return BlockSuccessors::Empty;
        };

        match term {
            StmtData::Jump { dst } => BlockSuccessors::One(Some(dst.block)),

            StmtData::BrIf {
                then_dst, else_dst, ..
            } => BlockSuccessors::Two {
                first: Some(then_dst.block),
                second: Some(else_dst.block),
            },

            StmtData::Return { .. } => BlockSuccessors::Empty,

            _ => BlockSuccessors::Empty,
        }
    }

    /// Returns the number of CFG successors of `block`.
    pub fn block_successor_count(&self, block: Block) -> usize {
        self.block_successors(block).count()
    }

    /// Creates a block call target with pooled SSA arguments.
    pub fn block_call(&mut self, block: Block, args: &[Value]) -> BlockCall {
        BlockCall {
            block,
            args: self.make_value_list(args),
        }
    }
}

/// Iterator over blocks in creation order.
#[derive(Clone, Debug)]
pub struct Blocks {
    next: usize,
    len: usize,
}

impl Iterator for Blocks {
    type Item = Block;

    fn next(&mut self) -> Option<Self::Item> {
        if self.next == self.len {
            return None;
        }

        let block = Block::new(self.next);
        self.next += 1;
        Some(block)
    }
}

/// Iterator over the CFG successors of a block.
#[derive(Clone, Debug)]
pub enum BlockSuccessors {
    Empty,
    One(Option<Block>),
    Two {
        first: Option<Block>,
        second: Option<Block>,
    },
}

impl Iterator for BlockSuccessors {
    type Item = Block;

    fn next(&mut self) -> Option<Self::Item> {
        match self {
            Self::Empty => None,
            Self::One(slot) => slot.take(),
            Self::Two { first, second } => {
                if let Some(block) = first.take() {
                    Some(block)
                } else {
                    second.take()
                }
            }
        }
    }
}
