alduin 0.0.1

WIP: A toy compiler backend
Documentation
use std::marker::PhantomData;

use super::CFG;
use crate::{
    backend::ISA,
    compiler::graph::{cfg::Block, CtrlOp, Graph, NodeId, OpCode, Top},
};

struct Scheduler<Isa: ISA> {
    cfg: CFG<Isa::Op>,
    mark: usize,
    _p: PhantomData<Isa>,
}

impl<Isa: ISA> Scheduler<Isa> {
    fn new(g: Box<Graph>) -> Self {
        Self {
            cfg: CFG::new(g),
            mark: 0,
            _p: PhantomData,
        }
    }

    fn end_block(&mut self, end: NodeId<()>) {
        if self.cfg.g[end].block.is_some() {
            return;
        }
        assert_eq!(self.cfg.g[end].controls.len(), 1);
        let start = self.cfg.g[end].controls[0];
        assert!(self.cfg.g[start]
            .op::<Isa::Op>()
            .ctrl_op()
            .unwrap()
            .is_label());
        let index = self.cfg.blocks.len();
        let mut b = Box::new(Block::new(start, end));
        b.succs
            .resize(self.cfg.g[end].control_uses.len(), usize::MAX);
        self.cfg.blocks.push(b);
        self.cfg.g[start].block = Some(index);
        self.cfg.g[end].block = Some(index);
        if self.cfg.g[start].op::<Isa::Op>().ctrl_op() == Some(CtrlOp::Start) {
            self.cfg.entry_block = index;
        }
        let max = self.cfg.g[start].controls.len();
        for i in 0..max {
            let pred_terminal = self.cfg.g[start].controls[i];
            assert!(self.cfg.g[pred_terminal]
                .op::<Isa::Op>()
                .ctrl_op()
                .unwrap()
                .is_terminal());
            self.end_block(pred_terminal);
            assert_eq!(self.cfg.g[pred_terminal].controls.len(), 1);
            let pred_label = self.cfg.g[pred_terminal].controls[0];
            assert!(self.cfg.g[pred_label]
                .op::<Isa::Op>()
                .ctrl_op()
                .unwrap()
                .is_label());
            let pred_block = self.cfg.g[pred_label].block.unwrap();
            self.cfg.blocks[index].preds.push(pred_block);
        }
    }

    fn construct_cfg(&mut self) {
        self.mark = self.cfg.g.update_mark();
        let end_node = self
            .cfg
            .g
            .nodes
            .iter()
            .find(|x| x.op::<Isa::Op>().ctrl_op() == Some(CtrlOp::Return))
            .unwrap();
        self.end_block(end_node.node.cast());
        for i in 0..self.cfg.blocks.len() {
            let terminal = self.cfg.blocks[i].terminal;
            assert_eq!(
                self.cfg.g[terminal].control_uses.len(),
                self.cfg.blocks[i].succs.len()
            );
            for (j, u) in self.cfg.g[terminal].control_uses.iter().enumerate() {
                self.cfg.blocks[i].succs[j] = self.cfg.g[u.user].block.unwrap();
            }
        }
    }

    fn pin_node(&mut self, n: NodeId) {
        debug_assert_eq!(self.cfg.g[n].controls.len(), 1);
        debug_assert!(!self.cfg.g[n].op::<Isa::Op>().is_label_or_terminal());
        let ctrl = self.cfg.g[n].controls[0];
        debug_assert!(self.cfg.g[ctrl]
            .op::<Isa::Op>()
            .ctrl_op()
            .unwrap()
            .is_label());
        let block = self.cfg.g[ctrl].block.unwrap();
        self.cfg.g[n].block = Some(block);
    }

    fn find_lca(&self, a: Option<usize>, mut b: usize) -> Option<usize> {
        if a.is_none() {
            return Some(b);
        }
        let mut a = a.unwrap();
        while self.cfg.blocks[a].dom_depth.unwrap() > self.cfg.blocks[b].dom_depth.unwrap() {
            a = self.cfg.blocks[a].idom.unwrap();
        }
        while self.cfg.blocks[b].dom_depth.unwrap() > self.cfg.blocks[a].dom_depth.unwrap() {
            b = self.cfg.blocks[b].idom.unwrap();
        }
        let mut a = Some(a);
        let mut b = Some(b);
        while a != b {
            a = self.cfg.blocks[a.unwrap()].idom;
            b = self.cfg.blocks[b.unwrap()].idom;
        }
        a
    }

    fn schedule_early(&mut self, n: NodeId) {
        let max = self.cfg.g[n].inputs.len();
        for i in 0..max {
            let input = self.cfg.g[n].inputs[i];
            if self.cfg.g[input].block.is_none() {
                self.schedule_early(input);
            }
        }
        if self.cfg.g[n].block.is_some() {
            return;
        }
        if self.cfg.g[n].inputs.len() == 0 {
            self.cfg.g[n].block = Some(self.cfg.entry_block);
            return;
        }
        let mut b = self.cfg.g[self.cfg.g[n].inputs[0]].block.unwrap();
        for i in 1..self.cfg.g[n].inputs.len() {
            let input = self.cfg.g[n].inputs[i];
            let inb = self.cfg.g[input].block.unwrap();
            if self.cfg.blocks[b].dom_depth.unwrap() < self.cfg.blocks[inb].dom_depth.unwrap() {
                b = inb;
            }
        }
        for i in 0..self.cfg.g[n].effects.len() {
            let effect = self.cfg.g[n].effects[i];
            let effectb = self.cfg.g[effect].block.unwrap();
            if self.cfg.blocks[b].dom_depth.unwrap() < self.cfg.blocks[effectb].dom_depth.unwrap() {
                b = effectb;
            }
        }
        self.cfg.g[n].block = Some(b);
    }

    fn schedule_late(&mut self, n: NodeId, mark: usize) {
        if self.cfg.g[n].mark == mark {
            return;
        }
        if self.cfg.g[n].op::<Isa::Op>().is_move()
            && (self.cfg.g[n].move_node_after.is_some() || self.cfg.g[n].move_node_before.is_some())
        {
            debug_assert!(self.cfg.g[n].block.is_some());
            return;
        }
        self.cfg.g[n].mark = mark;
        let max = self.cfg.g[n].uses.len();
        for i in 0..max {
            let u = self.cfg.g[n].uses[i].user;
            if self.cfg.g[n].controls.is_empty() {
                self.schedule_late(u, mark);
            }
        }
        let effect_max = self.cfg.g[n].effect_uses.len();
        for i in 0..effect_max {
            let u = self.cfg.g[n].effect_uses[i].user;
            if self.cfg.g[n].controls.is_empty() {
                self.schedule_late(u, mark);
            }
        }
        if !self.cfg.g[n].controls.is_empty() {
            return;
        }
        let mut lca = None;
        for i in 0..max {
            let u = self.cfg.g[n].uses[i].user;
            let use_block = match self.cfg.g[u].block {
                Some(x) => x,
                _ => continue,
            };
            if self.cfg.g[u].op::<Isa::Op>().is_phi() {
                let control_block = self.cfg.g[self.cfg.g[u].controls[0]].block.unwrap();
                for j in 0..self.cfg.g[u].inputs.len() {
                    if self.cfg.g[u].inputs[j] == n {
                        let pred_node =
                            self.cfg.g[self.cfg.blocks[control_block].label].controls[j];
                        let pred_block = self.cfg.g[pred_node].block.unwrap();
                        lca = self.find_lca(lca, pred_block);
                        break;
                    }
                }
            } else {
                lca = self.find_lca(lca, use_block);
            }
        }
        for i in 0..effect_max {
            let u = self.cfg.g[n].effect_uses[i].user;
            let use_block = match self.cfg.g[u].block {
                Some(x) => x,
                _ => continue,
            };
            if self.cfg.g[u].op::<Isa::Op>().is_effect_phi() {
                let control_block = self.cfg.g[self.cfg.g[u].controls[0]].block.unwrap();
                for j in 0..self.cfg.g[u].inputs.len() {
                    if self.cfg.g[u].inputs[j] == n {
                        let pred_node =
                            self.cfg.g[self.cfg.blocks[control_block].label].controls[j];
                        let pred_block = self.cfg.g[pred_node].block.unwrap();
                        lca = self.find_lca(lca, pred_block);
                        break;
                    }
                }
            } else {
                lca = self.find_lca(lca, use_block);
            }
        }

        if lca.is_none() {
            return;
        }

        let mut best = lca.unwrap();
        while lca != self.cfg.g[n].block && lca.is_some() {
            let lca_value = lca.unwrap();
            if self.cfg.blocks[lca_value].loop_depth < self.cfg.blocks[best].loop_depth {
                best = lca_value;
            }
            lca = self.cfg.blocks[lca_value].idom;
        }
        self.cfg.g[n].block = Some(best);
    }

    fn schedule_nodes(&mut self) {
        let mark = self.cfg.g.update_mark();
        let num_nodes = self.cfg.g.nodes.len();
        // schedule nodes with control deps
        for n in (0..num_nodes).map(|i| NodeId::<Top>::from(i)) {
            if self.cfg.g[n].controls.len() > 0 && self.cfg.g[n].block.is_none() {
                self.pin_node(n)
            }
        }
        // put all nodes in a correct block
        for n in (0..num_nodes).map(|i| NodeId::<Top>::from(i)) {
            if self.cfg.g[n].block.is_none() {
                self.schedule_early(n)
            }
        }
        // move nodes to uses' common lca
        for n in (0..num_nodes).map(|i| NodeId::<Top>::from(i)) {
            self.schedule_late(n, mark)
        }
        // put move nodes close to the src/tgt
        for n in (0..num_nodes).map(|i| NodeId::<Top>::from(i)) {
            let n_data = &self.cfg.g[n];
            if n_data.move_node_before.is_some() || n_data.move_node_after.is_some() {
                debug_assert!(
                    !(n_data.move_node_before.is_some() && n_data.move_node_after.is_some())
                );
                let block = self.cfg.g[n_data.move_node_before.or(n_data.move_node_after).unwrap()]
                    .block
                    .unwrap();
                self.cfg.g[n].block = Some(block);
            }
        }
        for n in (0..num_nodes).map(|i| NodeId::<Top>::from(i)) {
            let b = self.cfg.g[n].block.unwrap();
            if self.cfg.g[n].op::<Isa::Op>().is_label_or_terminal() {
                continue;
            }
            self.cfg.blocks[b].nodes.push(n);
        }
        let mark = self.cfg.g.update_mark();
        for block in &mut self.cfg.blocks {
            block.sort_nodes_rpo::<Isa::Op>(&mut self.cfg.g, mark);
        }
    }

    fn run(&mut self) {
        let mut coalescable_values = vec![];
        self.cfg
            .g
            .insert_move_instructions::<Isa>(&mut coalescable_values);
        self.construct_cfg();
        self.cfg.sort_blocks_rpo();
        self.cfg.build_dom_tree();
        self.cfg.compute_loops();
        self.schedule_nodes();
        self.cfg.renumber_nodes();
        self.cfg.compute_liveness::<Isa>(coalescable_values);
        trace!(target: "cfg", "\n{:?}", self.cfg);
    }
}

pub fn schedule<Isa: ISA>(graph: Box<Graph>) -> CFG<Isa::Op> {
    let mut scheduler = Scheduler::<Isa>::new(graph);
    scheduler.run();
    scheduler.cfg
}