alduin 0.0.1

use bitvec::bitvec;
use std::{
    fmt,
    marker::PhantomData,
    ops::{Index, IndexMut},
    sync::atomic::{AtomicUsize, Ordering},
};

use crate::backend::ISA;

use super::{node::NodeId, ty::Type, AnyTy, BaseOp, CtrlOp, InputKind, NodeData, OpCode, Ty, Use};

pub struct Graph {
    pub signature: Signature,
    pub(crate) nodes: Vec<NodeData>,
    pub mark: AtomicUsize,
}

impl<T: AnyTy> Index<NodeId<T>> for Graph {
    type Output = NodeData;

    fn index(&self, index: NodeId<T>) -> &Self::Output {
        &self.nodes[index.index()]
    }
}

impl<T: AnyTy> IndexMut<NodeId<T>> for Graph {
    fn index_mut(&mut self, index: NodeId<T>) -> &mut Self::Output {
        &mut self.nodes[index.index()]
    }
}

impl Graph {
    pub fn new(signature: Signature) -> Box<Self> {
        Box::new(Self {
            signature,
            nodes: vec![],
            mark: AtomicUsize::new(0),
        })
    }

    pub fn update_mark(&self) -> usize {
        let old = self.mark.fetch_add(1, Ordering::SeqCst);
        old + 1
    }

    pub fn new_untyped_node(
        &mut self,
        t: Type,
        op: impl OpCode,
        inputs: &[NodeId],
        controls: &[NodeId<()>],
        effects: &[NodeId],
    ) -> NodeId {
        let id = NodeId {
            index: self.nodes.len(),
            // ty: t,
            _p: PhantomData,
        };
        let n = NodeData::new(id, t, op);
        self.nodes.push(n);
        for i in inputs {
            self.add_input(id, InputKind::Data, *i);
        }
        for i in controls {
            self.add_input(id, InputKind::Control, i.cast());
        }
        for i in effects {
            self.add_input(id, InputKind::Effect, *i);
        }
        id
    }

    pub fn new_node<T: Ty>(&mut self, op: impl OpCode, inputs: &[NodeId]) -> NodeId {
        self.new_untyped_node(T::TYPE, op, inputs, &[], &[]).cast()
    }

    pub fn add_input(&mut self, node: NodeId, kind: InputKind, input: NodeId) {
        match kind {
            InputKind::Data => {
                let i = self[node].inputs.len();
                self[node].inputs.push(input);
                self[input].uses.push(Use {
                    user: node,
                    input: i,
                });
            }
            InputKind::Control => {
                let i = self[node].controls.len();
                self[node].controls.push(input.cast());
                self[input].control_uses.push(Use {
                    user: node,
                    input: i,
                });
            }
            InputKind::Effect => {
                let i = self[node].effects.len();
                self[node].effects.push(input);
                self[input].effect_uses.push(Use {
                    user: node,
                    input: i,
                });
            }
        }
    }

    pub fn update_input(
        &mut self,
        node: NodeId,
        kind: InputKind,
        index: usize,
        new_input: NodeId,
    ) -> NodeId {
        let old_input = match kind {
            InputKind::Data => self[node].inputs[index],
            InputKind::Control => self[node].controls[index].cast(),
            InputKind::Effect => self[node].effects[index],
        };
        // Update input
        match kind {
            InputKind::Data => self[node].inputs[index] = new_input,
            InputKind::Control => self[node].controls[index] = new_input.cast(),
            InputKind::Effect => self[node].effects[index] = new_input,
        }
        // Remove old use
        let use_edge = Use {
            user: node,
            input: index,
        };
        match kind {
            InputKind::Data => self[old_input].uses.retain(|u| u != &use_edge),
            InputKind::Control => self[old_input].control_uses.retain(|u| u != &use_edge),
            InputKind::Effect => self[old_input].effect_uses.retain(|u| u != &use_edge),
        };
        // Add new use
        match kind {
            InputKind::Data => self[new_input].uses.push(use_edge),
            InputKind::Control => self[new_input].control_uses.push(use_edge),
            InputKind::Effect => self[new_input].effect_uses.push(use_edge),
        };
        // Return old input
        old_input
    }

    fn transfer_data_uses(&mut self, src: NodeId, target: NodeId) {
        debug_assert_eq!(self[src].ty, self[target].ty);
        debug_assert!(self[target].uses.is_empty());
        for u in std::mem::take(&mut self[src].uses) {
            // self[u.user].inputs.set(u.input, target, false);
            self.update_input(u.user, InputKind::Data, u.input, target);
        }
        self[src].uses.clear();
    }

    /// Create a BaseOp::Move node to replace a node input . This only update use edges corresponds to the input edge
    pub(crate) fn new_move_input_node(&mut self, node: NodeId, input: usize) -> NodeId {
        assert!(input < self[node].inputs.len(), "{:?}", node);
        let src = self[node].inputs[input];
        let mov = self.new_untyped_node(self[src].ty, BaseOp::Move, &[src], &[], &[]);
        // replace the use edge corresponding to the input
        self.update_input(node, InputKind::Data, input, mov);
        self[node].move_node_before = Some(node);
        mov
    }

    /// Create a BaseOp::Move node act as the output of the src node. This transfers all the src's use edges to the move node.
    pub(crate) fn new_move_output_node(&mut self, src: NodeId) -> NodeId {
        let mov = self.new_untyped_node(self[src].ty, BaseOp::Move, &[], &[], &[]);
        // replace all use edges and pointing them to `mov`
        self.transfer_data_uses(src, mov);
        self.add_input(mov, InputKind::Data, src);
        self[mov].move_node_after = Some(src);
        mov
    }

    /// Insert move instruction for phi nodes. This is a necessary step for register allocation.
    pub(crate) fn insert_move_instructions<Isa: ISA>(
        &mut self,
        coalescable_values: &mut Vec<(NodeId, NodeId)>,
    ) {
        let all_labels = self
            .nodes
            .iter()
            .filter(|x| x.op::<BaseOp>().ctrl_op() == Some(CtrlOp::Region))
            .map(|n| n.node)
            .collect::<Vec<_>>();
        for n in all_labels {
            let controls = {
                let mut dedup = bitvec![0; self.nodes.len()];
                for c in &self[n].controls {
                    if !dedup[c.index()] {
                        dedup.set(c.index(), true);
                    }
                }
                dedup.iter_ones().collect::<Vec<_>>()
            };
            let mut phis = self[n]
                .control_uses
                .iter()
                .cloned()
                .filter(|x| self[x.user].op::<BaseOp>().is_phi())
                .map(|u| u.user)
                .collect::<Vec<_>>();
            for p in controls {
                let p = self.nodes[p].node.cast::<()>();
                let ctrl_index = self[n].controls.iter().position(|x| *x == p).unwrap();
                for phi in &mut phis {
                    let mov = self.new_move_input_node(*phi, ctrl_index);
                    self[mov].move_node_after = None;
                    self[mov].move_node_before = Some(p.cast());
                    // join values
                    coalescable_values.push((mov, phi.cast()));
                }
            }
        }
        let nodes = self.nodes.iter().map(|x| x.node).collect::<Vec<_>>();
        for n in nodes {
            Isa::pre_assign_registers(self, n);
        }
    }
}

impl fmt::Debug for Graph {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        writeln!(
            f,
            "Graph {:?} {{   # {} nodes",
            self.signature,
            self.nodes.iter().filter(|x| !x.dead).count()
        )?;
        for n in &self.nodes {
            if n.dead {
                continue;
            }
            writeln!(f, "  {:?}", n)?;
        }
        write!(f, "}}")
    }
}

#[derive(Clone, PartialEq)]
pub struct Signature(pub Vec<Type>, pub Type);

impl fmt::Debug for Signature {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let args = self
            .0
            .iter()
            .map(|t| format!("{:?}", t))
            .collect::<Vec<_>>();
        write!(f, "({}): {}", args.join(", "), format!("{:?}", self.1))
    }
}