not-python/src/compile/thunk.rs

use crate::{
    compile::{error::*, Compile},
    syn::{ast::*, visit::*},
    obj::prelude::*,
    vm::inst::*
};
use std::mem;

pub enum Thunk {
    Body(Vec<Inst>),
    List(Vec<Thunk>),
    Branch {
        thunk_true: Box<Thunk>,
        thunk_false: Box<Thunk>,
    },
    Nop,
}

impl Thunk {
    pub fn push(&mut self, append: impl Into<Inst>) {
        self.push_thunk(Thunk::Body(vec![append.into()]))
    }

    pub fn push_thunk(&mut self, append: impl Into<Thunk>) {
        let append = append.into();
        let inner = mem::replace(self, Thunk::Nop);

        *self = match (inner, append) {
            // X + Nop = X
            (lhs, Thunk::Nop) => lhs,
            (Thunk::Nop, rhs) => rhs,

            // Body + Body = Body
            (Thunk::Body(mut lhs), Thunk::Body(rhs)) => {
                lhs.extend(rhs);
                Thunk::Body(lhs)
            }
            // List + List = List
            (Thunk::List(mut lhs), Thunk::List(rhs)) => {
                lhs.extend(rhs);
                Thunk::List(lhs)
            }
            // List + X = List
            (Thunk::List(mut lhs), rhs) => {
                lhs.push(rhs);
                Thunk::List(lhs)
            }
            // X + List = List
            (lhs, Thunk::List(mut rhs)) => {
                rhs.insert(0, lhs);
                Thunk::List(rhs)
            }
            // X + X = List
            (lhs, rhs) => Thunk::List(vec![lhs, rhs]),
        };
    }
}

impl From<Inst> for Thunk {
    fn from(other: Inst) -> Self {
        Self::from(vec![other])
    }
}

impl From<Vec<Inst>> for Thunk {
    fn from(other: Vec<Inst>) -> Self {
        Thunk::Body(other)
    }
}

impl From<Vec<Thunk>> for Thunk {
    fn from(other: Vec<Thunk>) -> Self {
        Thunk::List(other)
    }
}

pub struct CompileBody<'c, 't> {
    compile: &'c mut Compile<'t>,
}

impl<'c, 't> CompileBody<'c, 't> {
    pub fn new(compile: &'c mut Compile<'t>) -> Self {
        CompileBody { compile }
    }

    pub fn compile(&mut self, body: &'c Body) -> Result<Thunk> {
        self.visit_body(body)
    }
}

impl Visit for CompileBody<'_, '_> {
    type Out = Result<Thunk>;

    fn visit_body(&mut self, body: &Body) -> Self::Out {
        let mut thunk = Thunk::Nop;

        for stmt in body.body.iter() {
            thunk.push_thunk(stmt.accept(self)?);
        }

        Ok(thunk)
    }

    fn visit_stmt(&mut self, stmt: &Stmt) -> Self::Out {
        DefaultAccept::default_accept(stmt, self)
    }

    fn visit_assign_stmt(&mut self, assign: &AssignStmt) -> Self::Out {
        // - push rhs
        // - push lhs (which handles the assignment)
        let mut thunk = self.visit_expr(&assign.rhs)?;
        thunk.push_thunk(self.visit_lhs_expr(&assign.lhs)?);
        Ok(thunk)
    }

    fn visit_lhs_expr(&mut self, lhs_expr: &LhsExpr) -> Self::Out {
        // Do different things depending on the LHS
        let mut thunk;
        match &lhs_expr {
            LhsExpr::SetAttr(expr) => {
                // - push lhs expression (without accessor)
                // - setattr (access)  NOTE : rhs should already be on stack
                thunk = self.visit_expr(&expr.expr)?;
                let attr = global_sym(expr.access.to_string());
                thunk.push(Inst::SetAttr(attr));
            }
            LhsExpr::Local(local) => {
                let local = global_sym(local.to_string());
                thunk = Inst::Pop(Some(local)).into();
            }
        }
        Ok(thunk)
    }

    fn visit_expr(&mut self, expr: &Expr) -> Self::Out {
        DefaultAccept::default_accept(expr, self)
    }

    fn visit_bin_expr(&mut self, expr: &BinExpr) -> Self::Out {
        // - push lhs
        // - push rhs
        // - call operator's function
        let mut thunk = self.visit_expr(&expr.lhs)?;
        thunk.push_thunk(self.visit_expr(&expr.rhs)?);
        let inst = match expr.op {
            BinOp::Plus => Inst::BinPlus,
            BinOp::Minus => Inst::BinMinus,
            BinOp::Times => Inst::BinMul,
            BinOp::Div => Inst::BinDiv,
            BinOp::Eq => Inst::BinEq,
            BinOp::Neq => Inst::BinNeq,
            BinOp::Lt => Inst::BinLt,
            BinOp::Le => Inst::BinLe,
            BinOp::Gt => Inst::BinGt,
            BinOp::Ge => Inst::BinGe,
            BinOp::And => Inst::BinAnd,
            BinOp::Or => Inst::BinOr,
        };

        thunk.push(inst);
        Ok(thunk)
    }

    fn visit_un_expr(&mut self, expr: &UnExpr) -> Self::Out {
        // - push expr
        // - call operator's function
        let mut thunk = self.visit_expr(&expr.expr)?;
        match expr.op {
            UnOp::Plus => thunk.push(Inst::UnPos),
            UnOp::Minus => thunk.push(Inst::UnNeg),
        }
        Ok(thunk)
    }

    fn visit_call_expr(&mut self, expr: &CallExpr) -> Self::Out {
        // - push expr
        // - push args in order
        // - call function
        let mut thunk = self.visit_expr(&expr.expr)?;
        for arg in expr.args.iter() {
            thunk.push_thunk(self.visit_expr(&arg)?);
        }
        Ok(thunk)
    }

    fn visit_index_expr(&mut self, expr: &IndexExpr) -> Self::Out {
        // - eval expr
        // - eval index
        // - index
        let mut thunk = self.visit_expr(&expr.expr)?;
        thunk.push_thunk(self.visit_expr(&expr.index)?);
        thunk.push(Inst::Index);
        Ok(thunk)
    }

    fn visit_access_expr(&mut self, expr: &AccessExpr) -> Self::Out {
        // - eval expr
        // - getattr (expr.access)
        let mut thunk = self.visit_expr(&expr.expr)?;
        thunk.push_thunk(Thunk::Body(vec![
                Inst::GetAttr(global_sym(expr.access.to_string())),
        ]));
        Ok(thunk)
    }

    fn visit_atom(&mut self, atom: &Atom) -> Self::Out {
        let thunk = match atom {
            Atom::Ident(ident) => {
                // get local
                Inst::PushLocal(global_sym(ident.to_string())).into()
            }
            Atom::Sym(sym) => {
                // push symbol
                Inst::PushSym(global_sym(sym.clone())).into()
            }
            Atom::Num(num) => {
                // push const
                let (hdl, _) = self.compile.const_int(*num);
                Inst::PushConst(hdl).into()
            }
            Atom::String(s) => {
                // push const
                let (hdl, _) = self.compile.const_str(s);
                Inst::PushConst(hdl).into()
            }
        };
        Ok(thunk)
    }
}