//! Program order and containment for Slonik IR.

use cranelift_entity::{SecondaryMap, packed_option::PackedOption};

use crate::ir::{Block, Inst};

/// Structural layout information for one block.
#[derive(Copy, Clone, Debug, Default)]
pub struct BlockNode {
    /// The previous block in layout order.
    pub prev: PackedOption<Block>,

    /// The next block in layout order.
    pub next: PackedOption<Block>,

    /// The first instruction in this block, if any.
    pub first_inst: PackedOption<Inst>,

    /// The last instruction in this block, if any.
    pub last_inst: PackedOption<Inst>,
}

/// Structural layout information for one instruction.
#[derive(Copy, Clone, Debug, Default)]
pub struct InstNode {
    /// The block containing this instruction.
    pub block: PackedOption<Block>,

    /// The previous instruction in the containing block.
    pub prev: PackedOption<Inst>,

    /// The next instruction in the containing block.
    pub next: PackedOption<Inst>,
}

/// Program order and block membership for a body.
#[derive(Default)]
pub struct Layout {
    /// The first block in layout order.
    pub first_block: PackedOption<Block>,

    /// The last block in layout order.
    pub last_block: PackedOption<Block>,

    /// Per-block layout links.
    pub blocks: SecondaryMap<Block, BlockNode>,

    /// Per-instruction layout links.
    pub insts: SecondaryMap<Inst, InstNode>,
}

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

    /// Returns whether the layout contains no blocks.
    pub fn is_empty(&self) -> bool {
        self.first_block.is_none()
    }

    /// Appends `block` to the end of the block list.
    pub fn append_block(&mut self, block: Block) {
        let prev = self.last_block;

        self.blocks[block].prev = prev;
        self.blocks[block].next = PackedOption::default();

        if let Some(prev_block) = prev.expand() {
            self.blocks[prev_block].next = block.into();
        } else {
            self.first_block = block.into();
        }

        self.last_block = block.into();
    }

    /// Returns the first block in layout order.
    pub fn first_block(&self) -> Option<Block> {
        self.first_block.expand()
    }

    /// Returns the last block in layout order.
    pub fn last_block(&self) -> Option<Block> {
        self.last_block.expand()
    }

    /// Returns the previous block before `block`, if any.
    pub fn prev_block(&self, block: Block) -> Option<Block> {
        self.blocks[block].prev.expand()
    }

    /// Returns the next block after `block`, if any.
    pub fn next_block(&self, block: Block) -> Option<Block> {
        self.blocks[block].next.expand()
    }

    /// Returns the first instruction in `block`, if any.
    pub fn first_inst(&self, block: Block) -> Option<Inst> {
        self.blocks[block].first_inst.expand()
    }

    /// Returns the last instruction in `block`, if any.
    pub fn last_inst(&self, block: Block) -> Option<Inst> {
        self.blocks[block].last_inst.expand()
    }

    /// Returns the block containing `inst`, if the instruction is laid out.
    pub fn inst_block(&self, inst: Inst) -> Option<Block> {
        self.insts[inst].block.expand()
    }

    /// Returns the previous instruction before `inst` in its containing block.
    pub fn prev_inst(&self, inst: Inst) -> Option<Inst> {
        self.insts[inst].prev.expand()
    }

    /// Returns the next instruction after `inst` in its containing block.
    pub fn next_inst(&self, inst: Inst) -> Option<Inst> {
        self.insts[inst].next.expand()
    }

    /// Returns whether `inst` is currently in the layout.
    pub fn is_inst_inserted(&self, inst: Inst) -> bool {
        self.insts[inst].block.is_some()
    }

    /// Appends `inst` to the end of `block`.
    pub fn append_inst(&mut self, block: Block, inst: Inst) {
        let prev = self.blocks[block].last_inst;

        self.insts[inst].block = block.into();
        self.insts[inst].prev = prev;
        self.insts[inst].next = PackedOption::default();

        if let Some(prev_inst) = prev.expand() {
            self.insts[prev_inst].next = inst.into();
        } else {
            self.blocks[block].first_inst = inst.into();
        }

        self.blocks[block].last_inst = inst.into();
    }
}

/// Iterator over blocks in layout order.
#[derive(Clone)]
pub struct Blocks<'a> {
    layout: &'a Layout,
    cur: Option<Block>,
}

impl<'a> Blocks<'a> {
    /// Creates a block iterator starting at the first block.
    pub fn new(layout: &'a Layout) -> Self {
        Self {
            layout,
            cur: layout.first_block(),
        }
    }
}

impl Iterator for Blocks<'_> {
    type Item = Block;

    fn next(&mut self) -> Option<Self::Item> {
        let block = self.cur?;
        self.cur = self.layout.next_block(block);
        Some(block)
    }
}

/// Iterator over instructions in one block.
#[derive(Clone)]
pub struct Insts<'a> {
    layout: &'a Layout,
    cur: Option<Inst>,
}

impl<'a> Insts<'a> {
    /// Creates an instruction iterator starting at the first instruction in `block`.
    pub fn new(layout: &'a Layout, block: Block) -> Self {
        Self {
            layout,
            cur: layout.first_inst(block),
        }
    }
}

impl Iterator for Insts<'_> {
    type Item = Inst;

    fn next(&mut self) -> Option<Self::Item> {
        let inst = self.cur?;
        self.cur = self.layout.next_inst(inst);
        Some(inst)
    }
}

impl Layout {
    /// Returns an iterator over all blocks in layout order.
    pub fn blocks(&self) -> Blocks<'_> {
        Blocks::new(self)
    }

    /// Returns an iterator over the instructions in `block`.
    pub fn insts(&self, block: Block) -> Insts<'_> {
        Insts::new(self, block)
    }
}
