not-python2/src/compile/compiler.rs

use std::collections::VecDeque;

use crate::compile::{scope::GlobalScope, thunk::Thunk};
use crate::obj::prelude::*;
use crate::obj::ObjPtr;
use crate::syn::ast::*;
use crate::vm::inst::*;

pub struct Compiler {
    scope: GlobalScope,
}

impl Compiler {
    pub fn new() -> Self {
        Self {
            scope: Default::default(),
        }
    }

    /// Gets the global scope.
    pub fn scope(&self) -> &GlobalScope {
        &self.scope
    }

    /// Gets the global scope, mutably.
    pub fn scope_mut(&mut self) -> &mut GlobalScope {
        &mut self.scope
    }

    /// Compile the given list of AST statements, returning the body.
    pub fn compile_stmts(&mut self, stmts: &Vec<SpStmt>) -> Vec<Inst> {
        self.gather_names(stmts);
        let mut thunks = VecDeque::new();
        for stmt in stmts {
            thunks.push_back(self.emit_stmt(stmt));
        }
        Thunk::List(thunks).flatten()
    }

    fn emit_body(&mut self, body: &SpBody) -> Thunk {
        let mut thunks = VecDeque::new();
        for stmt in body.inner() {
            thunks.push_back(self.emit_stmt(stmt));
        }
        Thunk::List(thunks)
    }

    fn emit_stmt(&mut self, stmt: &SpStmt) -> Thunk {
        match stmt.inner() {
            Stmt::Assign(lhs_expr, expr) => match lhs_expr.inner() {
                AssignLhs::Name(n) => {
                    let name = self.scope().lookup_scoped(n).unwrap();
                    let mut thunk = self.emit_expr(expr);
                    thunk.push_back(Inst::Store(name));
                    thunk
                }
                AssignLhs::Complex(_, _) => {
                    todo!("Complex LHS foo.bar assign")
                }
            },
            Stmt::Expr(expr) => {
                let mut thunk = self.emit_expr(expr);
                thunk.push_back(Inst::Pop);
                thunk
            }
            Stmt::If {
                if_true,
                elseif,
                else_body,
            } => {
                // compile the "else" thunk first - it doesn't need to jump
                let mut else_thunk = if let Some(else_body) = else_body {
                    self.emit_body(else_body)
                } else {
                    Thunk::empty()
                };
                else_thunk.push_front(Inst::Comment("else body".to_string()));

                // compile the "else if" thunks next. These need to jump forward N
                // instructions at the very end
                let mut jump_forward: isize = (else_thunk.len() + 1).try_into().unwrap();
                let mut thunks = VecDeque::with_capacity(
                    1 +             // else thunk
                    elseif.len() +  // elseif thunks
                    1, // if thunk
                );
                thunks.push_front(else_thunk);

                for cond_block in elseif.iter().rev() {
                    // do block, then condition
                    let mut body_thunk = self.emit_body(&cond_block.inner().body);
                    // jump forward by the previous jump forward distance
                    body_thunk.push_front(Inst::Comment("elseif body".to_string()));
                    body_thunk.push_back(Inst::JumpRelative(jump_forward, JumpCondition::Always));

                    let mut cond_thunk = self.emit_expr(&cond_block.inner().cond);
                    cond_thunk.push_front(Inst::Comment("elseif compare ".to_string()));
                    cond_thunk.push_back(Inst::Compare);
                    // jump forward by the length of the body
                    cond_thunk.push_back(Inst::JumpRelative(
                        isize::try_from(body_thunk.len()).unwrap() + 1,
                        JumpCondition::False,
                    ));

                    jump_forward +=
                        isize::try_from(body_thunk.len() + cond_thunk.len()).unwrap() + 1;

                    thunks.push_front(body_thunk);
                    thunks.push_front(cond_thunk);
                }

                // same as above: do block, then condition
                let mut body_thunk = self.emit_body(&if_true.inner().body);
                body_thunk.push_front(Inst::Comment("if body".to_string()));
                body_thunk.push_back(Inst::JumpRelative(jump_forward, JumpCondition::Always));

                let mut cond_thunk = self.emit_expr(&if_true.inner().cond);
                cond_thunk.push_front(Inst::Comment("if compare".to_string()));
                cond_thunk.push_back(Inst::Compare);
                cond_thunk.push_back(Inst::JumpRelative(
                    // + 1 because otherwise we won't end up on the next instruction
                    isize::try_from(body_thunk.len()).unwrap() + 1,
                    JumpCondition::False,
                ));

                // don't need to add this
                //jump_forward += isize::try_from(body_thunk.len() + cond_thunk.len()).unwrap();

                thunks.push_front(body_thunk);
                thunks.push_front(cond_thunk);
                thunks.push_back(Thunk::Block(
                    vec![Inst::Comment("end of if block".to_string())].into(),
                ));

                Thunk::List(thunks)
            }
            Stmt::Def { .. } => todo!(),
            Stmt::Import { .. } => todo!(),
        }
    }

    fn emit_expr(&mut self, expr: &SpExpr) -> Thunk {
        match expr.inner() {
            Expr::Call(expr, args) => {
                let mut thunk = VecDeque::with_capacity(args.len() + 1);
                thunk.push_back(self.emit_expr(expr));
                thunk.extend(args.iter().map(|arg| self.emit_expr(arg)));
                let mut thunk = Thunk::List(thunk);
                thunk.push_back(Inst::Call(args.len()));
                thunk
            }
            Expr::Get(expr, name) => {
                let mut thunk = self.emit_expr(expr);
                thunk.push_back(Inst::GetAttr(name.clone()));
                thunk
            }
            Expr::Atom(atom) => self.emit_atom(atom),
        }
    }

    fn emit_atom(&mut self, atom: &SpAtom) -> Thunk {
        match atom.inner() {
            Atom::Int(i) => Thunk::Block(vec![Inst::Push(ObjPtr::new(Int::new(*i)))].into()),
            Atom::Float(_) => todo!(), // Thunk::Block(vec![Inst::Push(Gc::new(Float::new(*f)))]),
            Atom::Str(s) => Thunk::Block(vec![Inst::Push(ObjPtr::new(Str::new(s.clone())))].into()),
            Atom::Sym(_) => todo!(),
            Atom::Name(n) => {
                // Look up the symbol first. It may already exist.
                // If it doesn't exist, create it. It may be created
                // Dynamically.
                let name = self
                    .scope()
                    .lookup_scoped(n)
                    .unwrap_or_else(|| self.scope_mut().insert_local(n));
                Thunk::Block(vec![Inst::Load(name)].into())
            }
        }
    }

    /// Gather all names from the list of statements and add them to the scope.
    fn gather_names(&mut self, stmts: &Vec<SpStmt>) {
        for stmt in stmts {
            match stmt.inner() {
                Stmt::Assign(lhs, _) => match lhs.inner() {
                    AssignLhs::Name(name) => {
                        self.scope_mut().insert_local(name);
                    }
                    AssignLhs::Complex(expr, _) => self.gather_expr_names(expr.inner()),
                },
                Stmt::Def { name, .. } => {
                    self.scope_mut().insert_local(name);
                }
                Stmt::Expr(expr) => {
                    // Get the names of the variables used in this expr
                    self.gather_expr_names(expr.inner());
                }
                Stmt::If {
                    if_true,
                    elseif,
                    else_body,
                } => {
                    self.gather_expr_names(if_true.inner().cond.inner());
                    self.gather_names(if_true.inner().body.inner());

                    for cond_body in elseif {
                        self.gather_expr_names(cond_body.inner().cond.inner());
                        self.gather_names(cond_body.inner().body.inner());
                    }

                    if let Some(else_body) = else_body {
                        self.gather_names(else_body.inner());
                    }
                }
                Stmt::Import { .. } => {
                    todo!()
                }
            }
        }
    }

    /// Gather the names used inside of an expression, recursively.
    fn gather_expr_names(&mut self, expr: &Expr) {
        match expr {
            Expr::Call(fun, args) => {
                self.gather_expr_names(fun.inner());
                for arg in args {
                    self.gather_expr_names(arg.inner());
                }
            }
            Expr::Get(head, _) => {
                self.gather_expr_names(head.inner());
            }
            Expr::Atom(atom) => match atom.inner() {
                Atom::Name(name) => {
                    self.scope_mut().insert_local(name);
                }
                // TODO symbols
                _ => { /* no-op on non-names */ }
            },
        }
    }
}