crsn/crsn/src/asm/instr/flatten.rs

use std::collections::HashMap;
use std::fmt::Debug;
use std::sync::atomic::AtomicU32;

use sexp::SourcePosition;

use crate::asm::data::{Rd, RdData};
use crate::asm::data::literal::{Label, Value};
use crate::asm::error::{AsmError, CrsnError};
use crate::asm::instr::{Cond, InstrWithBranches, Op, Routine};
use crate::asm::instr::op::OpKind;
use crate::builtin::defs::Barrier;
use crate::builtin::defs::BuiltinOp;

/// A trait for something that can turn into multiple instructions
pub trait Flatten: Debug {
    fn flatten(self: Box<Self>, label_num: &AtomicU32) -> Result<Vec<Op>, CrsnError>;

    fn pos(&self) -> SourcePosition;
}

impl Flatten for () {
    fn flatten(self: Box<Self>, _label_num: &AtomicU32) -> Result<Vec<Op>, CrsnError> {
        Ok(vec![])
    }

    fn pos(&self) -> SourcePosition {
        SourcePosition::default()
    }
}

impl Flatten for Op {
    fn flatten(self: Box<Self>, _label_num: &AtomicU32) -> Result<Vec<Op>, CrsnError> {
        Ok(vec![*self])
    }

    fn pos(&self) -> SourcePosition {
        self.pos.clone()
    }
}

impl Flatten for InstrWithBranches {
    fn pos(&self) -> SourcePosition {
        self.pos.clone()
    }

    fn flatten(self: Box<Self>, label_num: &AtomicU32) -> Result<Vec<Op>, CrsnError> {
        let mut ops = vec![self.op];
        let parent_pos = self.pos;

        if let Some(branches) = self.branches {
            // trace!("Branches {:?}", branches);

            let labels = HashMap::<Cond, u32>::new();
            let branch_count = branches.len();
            let end_lbl = Label::unique(label_num);
            for (cnt, (cond, branch)) in branches.into_iter().enumerate() {
                if labels.contains_key(&cond) {
                    return Err(CrsnError::Asm(AsmError::ConditionalAlreadyUsed(cond), branch.pos()));
                }

                if cnt != branch_count - 1 && cond == Cond::True {
                    warn!("\"Else\" conditional used in non-final branch at {}", branch.pos());
                }

                let next_lbl = if cnt == branch_count - 1 {
                    end_lbl.clone()
                } else {
                    Label::unique(label_num)
                };

                let pos = branch.pos().clone();
                let mut flattened = branch.flatten(label_num)?;

                if flattened.len() == 0 {
                    ops.push(Op { cond: Some(cond), pos: pos.clone(), kind: BuiltinOp::Jump(end_lbl.clone()).into() });
                } else if flattened.len() == 1 && flattened[0].cond.is_none() && branch_count == 1 {
                    // optimization for single-branch conditionals with a single instruction
                    ops.push(Op { cond: Some(cond), pos: pos.clone(), kind: flattened.remove(0).kind });
                } else {
                    if cond != Cond::True { // evoid emiting `op.never`
                        ops.push(Op {
                            kind: OpKind::BuiltIn(BuiltinOp::Jump(next_lbl.clone())),
                            pos: pos.clone(),
                            cond: Some(!cond),
                        });
                    }
                    ops.extend(flattened);
                }

                if cnt != branch_count - 1 {
                    ops.push(BuiltinOp::Jump(end_lbl.clone()).into_op(pos.clone()));
                    ops.push(BuiltinOp::Label(next_lbl).into_op(pos.clone()));
                }
            }
            ops.push(BuiltinOp::Label(end_lbl).into_op(parent_pos));
        }

        Ok(ops)
    }
}

impl Flatten for Vec<Box<dyn Flatten>> {
    fn pos(&self) -> SourcePosition {
        match self.first() {
            None => {
                Default::default()
            }
            Some(f) => {
                f.pos()
            }
        }
    }

    fn flatten(self: Box<Self>, label_num: &AtomicU32) -> Result<Vec<Op>, CrsnError> {
        let mut ops = vec![];
        for item in self.into_iter() {
            ops.extend(item.flatten(label_num)?);
        }
        Ok(ops)
    }
}

impl Flatten for Vec<Op> {
    fn pos(&self) -> SourcePosition {
        match self.first() {
            None => {
                Default::default()
            }
            Some(f) => {
                f.pos()
            }
        }
    }

    fn flatten(self: Box<Self>, _label_num: &AtomicU32) -> Result<Vec<Op>, CrsnError> {
        Ok(*self)
    }
}

impl Flatten for Routine {
    fn pos(&self) -> SourcePosition {
        self.pos.clone()
    }

    fn flatten(self: Box<Self>, label_num: &AtomicU32) -> Result<Vec<Op>, CrsnError> {
        let skip_label = Label::unique(label_num);

        let self_pos = self.pos();

        let mut ops: Vec<Op> = vec![
            BuiltinOp::Barrier {
                kind: Barrier::Open(skip_label.clone()),
                msg: Some(format!("proc {} start", self.name).into()),
            }.into_op(self.pos()),
            BuiltinOp::Routine(self.name.clone()).into_op(self.pos()),
        ];

        ops.extend(self.body.flatten(label_num)?);

        ops.push(
            BuiltinOp::Barrier {
                kind: Barrier::Close(skip_label.clone()),
                msg: Some(format!("proc {} end", self.name).into()),
            }.into_op(self_pos)
        );

        jumps_to_skips(ops)
    }
}

/// Convert jumps to relative skips
pub fn jumps_to_skips(ops: Vec<Op>) -> Result<Vec<Op>, CrsnError> {
    let mut label_positions = HashMap::<Label, usize>::new();
    for (n, op) in ops.iter().enumerate() {
        if let OpKind::BuiltIn(BuiltinOp::Label(name)) = &op.kind {
            if label_positions.contains_key(name) {
                return Err(CrsnError::Asm(AsmError::LabelDuplicate(name.clone()), op.pos.clone()));
            }

            label_positions.insert(name.clone(), n - label_positions.len());
        }
    }

    let mut cleaned = vec![];
    let mut skipped = 0;
    for (n, op) in ops.into_iter().enumerate() {
        match op.kind {
            OpKind::BuiltIn(BuiltinOp::Label(_)) => {
                skipped += 1;
            }
            OpKind::BuiltIn(BuiltinOp::Jump(target)) => {
                if let Some(dest) = label_positions.get(&target) {
                    let skip = *dest as isize - n as isize + skipped;
                    cleaned.push(Op {
                        cond: op.cond,
                        pos: op.pos.clone(),
                        kind: OpKind::BuiltIn(BuiltinOp::Skip(Rd::new(RdData::Immediate(skip as Value)))),
                    });
                } else {
                    return Err(CrsnError::Asm(AsmError::LabelNotDefined(target), op.pos));
                }
            }
            _ => {
                cleaned.push(op);
            }
        }
    }

    Ok(cleaned)
}