nixjit/src/compile.rs

use crate::bytecode::*;
use crate::ir;

pub struct Compiler {
    opcodes: Vec<OpCode>,
}

pub fn compile(downgraded: ir::Downgraded) -> Program {
    Program {
        top_level: Compiler::new().compile(downgraded.top_level),
        thunks: downgraded
            .thunks
            .into_iter()
            .map(|thunk| Compiler::new().compile(thunk))
            .collect(),
        funcs: downgraded
            .funcs
            .into_iter()
            .map(|func| Func {
                param: func.param.into(),
                opcodes: Compiler::new().compile(*func.body),
            })
            .collect(),
        symbols: downgraded.symbols,
        symmap: downgraded.symmap,
        consts: downgraded.consts,
    }
}

impl Compiler {
    fn new() -> Self {
        Self {
            opcodes: Vec::new(),
        }
    }

    fn compile(mut self, ir: ir::Ir) -> OpCodes {
        ir.compile(&mut self);
        self.opcodes()
    }

    fn push(&mut self, code: OpCode) {
        self.opcodes.push(code);
    }

    fn idx(&self) -> usize {
        self.opcodes.len()
    }

    fn modify(&mut self, idx: usize, code: OpCode) {
        self.opcodes[idx] = code;
    }

    fn pop(&mut self) -> Option<OpCode> {
        self.opcodes.pop()
    }

    fn opcodes(self) -> OpCodes {
        self.opcodes.into()
    }
}

pub trait Compile {
    fn compile(self, comp: &mut Compiler);
}

pub trait CompileWithLength {
    fn compile_with_length(self, comp: &mut Compiler) -> usize;
}

impl<T: Compile> CompileWithLength for T {
    fn compile_with_length(self, comp: &mut Compiler) -> usize {
        let start = comp.idx();
        self.compile(comp);
        let end = comp.idx();
        end - start
    }
}

impl Compile for ir::Const {
    fn compile(self, comp: &mut Compiler) {
        comp.push(OpCode::Const { idx: self.idx });
    }
}

impl Compile for ir::Var {
    fn compile(self, comp: &mut Compiler) {
        comp.push(OpCode::LookUp { sym: self.sym });
    }
}

impl Compile for ir::Thunk {
    fn compile(self, comp: &mut Compiler) {
        comp.push(OpCode::LoadThunk { idx: self.idx });
    }
}

impl Compile for ir::Attrs {
    fn compile(self, comp: &mut Compiler) {
        comp.push(OpCode::AttrSet);
        for stc in self.stcs {
            stc.1.compile(comp);
            if !self.rec {
                comp.push(OpCode::CaptureEnv);
            }
            comp.push(OpCode::PushStaticAttr { name: stc.0 });
        }
        for dynamic in self.dyns {
            dynamic.0.compile(comp);
            dynamic.1.compile(comp);
            if !self.rec {
                comp.push(OpCode::CaptureEnv);
            }
            comp.push(OpCode::PushDynamicAttr)
        }
        if self.rec {
            comp.push(OpCode::FinalizeRec);
        }
    }
}

impl Compile for ir::List {
    fn compile(self, comp: &mut Compiler) {
        comp.push(OpCode::List);
        for item in self.items {
            item.compile(comp);
            comp.push(OpCode::PushElem);
        }
    }
}

impl Compile for ir::UnOp {
    fn compile(self, comp: &mut Compiler) {
        use ir::UnOpKind::*;
        match self.kind {
            Neg => {
                self.rhs.compile(comp);
                comp.push(OpCode::UnOp { op: UnOp::Neg });
            }
            Not => {
                self.rhs.compile(comp);
                comp.push(OpCode::UnOp { op: UnOp::Not });
            }
        }
    }
}

impl Compile for ir::BinOp {
    fn compile(self, comp: &mut Compiler) {
        use ir::BinOpKind::*;
        match self.kind {
            Add => {
                self.lhs.compile(comp);
                self.rhs.compile(comp);
                comp.push(OpCode::BinOp { op: BinOp::Add });
            }
            Mul => {
                self.lhs.compile(comp);
                self.rhs.compile(comp);
                comp.push(OpCode::BinOp { op: BinOp::Mul });
            }
            Div => {
                self.lhs.compile(comp);
                self.rhs.compile(comp);
                comp.push(OpCode::BinOp { op: BinOp::Div });
            }
            And => {
                self.lhs.compile(comp);
                self.rhs.compile(comp);
                comp.push(OpCode::BinOp { op: BinOp::And });
            }
            Or => {
                self.lhs.compile(comp);
                self.rhs.compile(comp);
                comp.push(OpCode::BinOp { op: BinOp::Or });
            }
            Eq => {
                self.lhs.compile(comp);
                self.rhs.compile(comp);
                comp.push(OpCode::BinOp { op: BinOp::Eq });
            }
            Lt => {
                self.lhs.compile(comp);
                self.rhs.compile(comp);
                comp.push(OpCode::BinOp { op: BinOp::Lt });
            }
            Con => {
                self.lhs.compile(comp);
                self.rhs.compile(comp);
                comp.push(OpCode::BinOp { op: BinOp::Con });
            }
            Upd => {
                self.lhs.compile(comp);
                self.rhs.compile(comp);
                comp.push(OpCode::BinOp { op: BinOp::Upd });
            }

            Sub => {
                self.lhs.compile(comp);
                self.rhs.compile(comp);
                comp.push(OpCode::BinOp { op: BinOp::Sub });
            }
            Impl => {
                self.lhs.compile(comp);
                comp.push(OpCode::UnOp { op: UnOp::Not });
                self.rhs.compile(comp);
                comp.push(OpCode::BinOp { op: BinOp::Or });
            }
            Neq => {
                self.lhs.compile(comp);
                self.rhs.compile(comp);
                comp.push(OpCode::BinOp { op: BinOp::Eq });
                comp.push(OpCode::UnOp { op: UnOp::Not });
            }
            Gt => {
                self.rhs.compile(comp);
                self.lhs.compile(comp);
                comp.push(OpCode::BinOp { op: BinOp::Lt });
            }
            Leq => {
                self.rhs.compile(comp);
                self.lhs.compile(comp);
                comp.push(OpCode::BinOp { op: BinOp::Lt });
                comp.push(OpCode::UnOp { op: UnOp::Not });
            }
            Geq => {
                self.lhs.compile(comp);
                self.rhs.compile(comp);
                comp.push(OpCode::BinOp { op: BinOp::Lt });
                comp.push(OpCode::UnOp { op: UnOp::Not });
            }

            PipeL => {
                self.lhs.compile(comp);
                self.rhs.compile(comp);
                comp.push(OpCode::Call { arity: 1 });
            }
            PipeR => {
                self.rhs.compile(comp);
                self.lhs.compile(comp);
                comp.push(OpCode::Call { arity: 1 });
            }
        }
    }
}

impl Compile for ir::HasAttr {
    fn compile(self, comp: &mut Compiler) {
        self.lhs.compile(comp);
        for attr in self.rhs {
            match attr {
                ir::Attr::Str(sym) => {
                    comp.push(OpCode::AttrSet);
                    comp.push(OpCode::SelectOrDefault { sym })
                }
                ir::Attr::Dynamic(dynamic) => {
                    dynamic.compile(comp);
                    comp.push(OpCode::AttrSet);
                    comp.push(OpCode::SelectDynamicOrDefault);
                }
                ir::Attr::Strs(string) => {
                    string.compile(comp);
                    comp.push(OpCode::AttrSet);
                    comp.push(OpCode::SelectDynamicOrDefault);
                }
            }
        }
        let last = comp.pop().unwrap();
        let _ = comp.pop();
        match last {
            OpCode::SelectOrDefault { sym } => comp.push(OpCode::HasAttr { sym }),
            OpCode::SelectDynamicOrDefault => comp.push(OpCode::HasDynamicAttr),
            _ => unreachable!(),
        }
    }
}

impl Compile for ir::Select {
    fn compile(self, comp: &mut Compiler) {
        self.expr.compile(comp);
        for attr in self.attrpath {
            match attr {
                ir::Attr::Str(sym) => {
                    comp.push(OpCode::AttrSet);
                    comp.push(OpCode::SelectOrDefault { sym })
                }
                ir::Attr::Dynamic(dynamic) => {
                    dynamic.compile(comp);
                    comp.push(OpCode::AttrSet);
                    comp.push(OpCode::SelectDynamicOrDefault);
                }
                ir::Attr::Strs(string) => {
                    string.compile(comp);
                    comp.push(OpCode::AttrSet);
                    comp.push(OpCode::SelectDynamicOrDefault);
                }
            }
        }
        match self.default {
            Some(default) => {
                let last = comp.pop().unwrap();
                let _ = comp.pop();
                default.compile(comp);
                match last {
                    OpCode::SelectOrDefault { sym } => comp.push(OpCode::SelectOrDefault { sym }),
                    OpCode::SelectDynamicOrDefault => comp.push(OpCode::SelectDynamicOrDefault),
                    _ => unreachable!(),
                }
            }
            None => {
                let last = comp.pop().unwrap();
                let _ = comp.pop();
                match last {
                    OpCode::SelectOrDefault { sym } => comp.push(OpCode::Select { sym }),
                    OpCode::SelectDynamicOrDefault => comp.push(OpCode::SelectDynamic),
                    _ => unreachable!(),
                }
            }
        }
    }
}

impl Compile for ir::ConcatStrings {
    fn compile(self, comp: &mut Compiler) {
        let mut iter = self.parts.into_iter();
        iter.next().unwrap().compile(comp);
        for item in iter {
            item.compile(comp);
            comp.push(OpCode::ConcatString);
        }
    }
}

impl Compile for ir::If {
    fn compile(self, comp: &mut Compiler) {
        self.cond.compile(comp);

        let idx_jmp_if_false = comp.idx();
        // place holder
        comp.push(OpCode::Illegal);

        let consq_length = self.consq.compile_with_length(comp);

        let idx_jmp = comp.idx();
        // place holder
        comp.push(OpCode::Illegal);

        let alter_length = self.alter.compile_with_length(comp);

        comp.modify(
            idx_jmp_if_false,
            OpCode::JmpIfFalse {
                step: consq_length + 1,
            },
        );
        comp.modify(idx_jmp, OpCode::Jmp { step: alter_length });
    }
}

impl Compile for ir::Let {
    fn compile(self, comp: &mut Compiler) {
        self.attrs.compile(comp);
        comp.push(OpCode::EnterEnv);
        self.expr.compile(comp);
        comp.push(OpCode::LeaveEnv);
    }
}

impl Compile for ir::With {
    fn compile(self, comp: &mut Compiler) {
        self.namespace.compile(comp);
        comp.push(OpCode::EnterEnv);
        self.expr.compile(comp);
        comp.push(OpCode::LeaveEnv);
    }
}

impl Compile for ir::Assert {
    fn compile(self, comp: &mut Compiler) {
        self.assertion.compile(comp);
        comp.push(OpCode::Assert);
        self.expr.compile(comp);
    }
}

impl Compile for ir::LoadFunc {
    fn compile(self, comp: &mut Compiler) {
        comp.push(OpCode::Func { idx: self.idx });
    }
}

impl Compile for ir::Call {
    fn compile(self, comp: &mut Compiler) {
        let arity = self.args.len();
        self.func.compile(comp);
        self.args.into_iter().for_each(|arg| {
            arg.compile(comp);
        });
        comp.push(OpCode::Call { arity });
    }
}

impl Compile for ir::Path {
    fn compile(self, comp: &mut Compiler) {
        self.expr.compile(comp);
        comp.push(OpCode::Path);
    }
}