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

use std::collections::HashMap;
use std::sync::atomic::{AtomicU32, Ordering};

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

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

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

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

        if let Some(branches) = self.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)));
                }

                let next_lbl = if cnt == branch_count - 1 {
                    end_lbl.clone()
                } else {
                    Label::unique(label_num)
                };
                ops.push(BuiltinOp::JumpIf(!cond, next_lbl.clone()).into());
                ops.extend(branch.flatten(label_num)?);

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

        Ok(ops)
    }
}

impl Flatten for Vec<Box<dyn Flatten>> {
    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)?);
        }
        labels_to_skips(ops)
    }
}

impl Flatten for Routine {
    fn flatten(self: Box<Self>, label_num: &AtomicU32) -> Result<Vec<Op>, CrsnError> {
        let skip_label = Label::Numbered(label_num.fetch_add(1, Ordering::Relaxed));

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

            BuiltinOp::Routine(self.name.clone()).into(),
        ];

        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()
        );

        labels_to_skips(ops)
    }
}

/// Convert jumps to relative skips
pub fn labels_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 Op::BuiltIn(BuiltinOp::Label(name)) = op {
            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 {
            Op::BuiltIn(BuiltinOp::Label(_)) => {
                skipped += 1;
            }
            Op::BuiltIn(BuiltinOp::Jump(target)) => {
                if let Some(dest) = label_positions.get(&target) {
                    let skip = *dest as isize - n as isize + skipped;
                    cleaned.push(Op::BuiltIn(BuiltinOp::Skip(Rd::new(RdData::Immediate(skip as Value)))));
                } else {
                    return Err(CrsnError::Asm(AsmError::LabelNotDefined(target)));
                }
            }
            Op::BuiltIn(BuiltinOp::JumpIf(cond, target)) => {
                if let Some(dest) = label_positions.get(&target) {
                    let skip = *dest as isize - n as isize + skipped;
                    cleaned.push(Op::BuiltIn(BuiltinOp::SkipIf(cond, Rd::new(RdData::Immediate(skip as Value)))));
                } else {
                    return Err(CrsnError::Asm(AsmError::LabelNotDefined(target)));
                }
            }
            other => {
                cleaned.push(other);
            }
        }
    }

    Ok(cleaned)
}