ir.rs

use std::cell::RefCell;
use std::collections::{HashMap, HashSet};
use std::fmt;

use crate::ast::*;

use super::{symbol_table::Def, Context};

#[derive(PartialEq, Eq, Copy, Clone, Debug, Hash)]
pub struct BlockId(pub usize);

impl fmt::Display for BlockId {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "#{}", self.0)
    }
}

#[derive(PartialEq, Eq, Copy, Clone, Debug, Hash)]
pub struct RegisterId(pub usize);

impl fmt::Display for RegisterId {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "reg(#{})", self.0)
    }
}

#[derive(PartialEq, Eq, Copy, Clone, Debug, Hash)]
pub struct ParamIdx(pub usize);

impl fmt::Display for ParamIdx {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "param(#{})", self.0)
    }
}

#[derive(PartialEq, Eq, Copy, Clone, Debug, Hash)]
pub struct BlockFunkId(pub usize);

impl fmt::Display for BlockFunkId {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "#{}", self.0)
    }
}

#[derive(Debug)]
struct BlockFunk {
    id: BlockFunkId,
    node_id: NodeId,

    blocks: HashMap<BlockId, Block>,
    entry: Option<BlockId>,

    params: HashMap<NodeId, ParamIdx>,
    locals: HashMap<NodeId, RegisterId>,
}

impl fmt::Display for BlockFunk {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "block_funk {} (node_id: {})\n", self.id, self.node_id)?;
        for block in self.blocks.values() {
            write!(f, "{}", block)?;
        }
        Ok(())
    }
}

#[derive(Debug)]
struct Block {
    id: BlockId,
    block_funk_id: BlockFunkId,
    instructions: Vec<Instruction>,
    terminator: Option<Terminator>,
    // TODO: This is unused
    successors: HashSet<BlockId>,
    // TODO: This is unused
    predecessors: HashSet<BlockId>,
}

impl fmt::Display for Block {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "block {}\n", self.id)?;
        for ins in self.instructions.iter() {
            write!(f, "\t{}\n", ins)?;
        }
        Ok(())
    }
}

#[derive(Debug)]
struct Instruction {
    target: Place,
    kind: InstructionKind,
}

impl fmt::Display for Instruction {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{} <- {}", self.target, self.kind)
    }
}

#[derive(Debug, PartialEq)]
enum InstructionKind {
    Move(Operand),
    Binary(BinOp, Operand, Operand),
    Call(Operand, Vec<Operand>),
}

impl fmt::Display for InstructionKind {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            InstructionKind::Move(op) => write!(f, "{}", op),
            InstructionKind::Binary(op, lhs, rhs) => write!(f, "{} {} {}", lhs, op, rhs),
            InstructionKind::Call(target, args) => write!(f, "{} {:?}", target, args),
        }
    }
}

#[derive(Clone, Debug, PartialEq)]
enum Terminator {
    Halt,
    Return {
        value: Operand,
    },

    If {
        cond: Operand,
        targets: (BlockId, BlockId),
    },

    Goto {
        target: BlockId,
    },
}

impl Terminator {
    #[cfg(test)]
    fn is_return(&self) -> bool {
        match self {
            Terminator::Return { .. } => true,
            _ => false,
        }
    }
}

#[derive(Clone, Debug, PartialEq)]
enum Operand {
    None,

    Constant(Constant),

    Register(RegisterId),
    Param(ParamIdx),

    Funk(BlockFunkId),
}

impl fmt::Display for Operand {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Operand::Funk(block_funk_id) => write!(f, "{}", block_funk_id),
            Operand::Register(reg) => write!(f, "{}", reg),
            Operand::Param(p) => write!(f, "{}", p),
            Operand::Constant(cons) => write!(f, "{}", cons),
            Operand::None => write!(f, "none"),
        }
    }
}

#[derive(Clone, Debug, PartialEq)]
enum Place {
    Register(RegisterId),
    Param(ParamIdx),
    Funk(BlockFunkId),
}

impl Place {
    #[cfg(test)]
    fn reg_id(&self) -> RegisterId {
        match self {
            Place::Register(id) => id.clone(),
            _ => panic!("unreachable!"),
        }
    }

    fn to_op(&self) -> Operand {
        match self {
            Place::Register(reg) => Operand::Register(reg.clone()),
            Place::Param(p) => Operand::Param(p.clone()),
            Place::Funk(p) => Operand::Funk(p.clone()),
        }
    }
}

impl fmt::Display for Place {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Place::Register(reg) => write!(f, "{}", reg),
            Place::Param(p) => write!(f, "{}", p),
            Place::Funk(p) => write!(f, "{}", p),
        }
    }
}

#[derive(Clone, Debug, PartialEq)]
pub enum Constant {
    Int(u64),
    String(String),
    Bool(bool),
}

impl fmt::Display for Constant {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Constant::Int(i) => write!(f, "{}", i),
            Constant::String(s) => write!(f, "{}", s),
            Constant::Bool(b) => write!(f, "{}", b),
        }
    }
}