rudus/src/compiler.rs

use crate::parser::Ast;
use crate::spans::Spanned;
use crate::value::*;
use chumsky::prelude::SimpleSpan;
use num_derive::{FromPrimitive, ToPrimitive};
use num_traits::FromPrimitive;

#[derive(Copy, Clone, Debug, PartialEq, Eq, FromPrimitive, ToPrimitive)]
pub enum Op {
    Constant,
    Jump,
    JumpIfFalse,
    Pop,
    PushBinding,
    Store,
    Load,
}

impl std::fmt::Display for Op {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        use Op::*;
        match self {
            Constant => write!(f, "constant"),
            Jump => write!(f, "jump"),
            JumpIfFalse => write!(f, "jump_if_false"),
            Pop => write!(f, "pop"),
            PushBinding => write!(f, "push_binding"),
            Store => write!(f, "store"),
            Load => write!(f, "load"),
        }
    }
}

#[derive(Clone, Debug, PartialEq)]
pub struct Binding {
    name: &'static str,
    depth: isize,
}

#[derive(Clone, Debug, PartialEq)]
pub struct Chunk<'a> {
    pub bindings: Vec<Binding>,
    scope_depth: isize,
    num_bindings: usize,
    pub constants: Vec<Value>,
    pub bytecode: Vec<u8>,
    pub spans: Vec<SimpleSpan>,
    pub strings: Vec<&'static str>,
    pub keywords: Vec<&'static str>,
    pub nodes: Vec<&'a Ast>,
    pub ast: &'a Ast,
    pub span: SimpleSpan,
    pub src: &'static str,
    pub name: &'static str,
}

impl<'a> Chunk<'a> {
    pub fn new(ast: &'a Spanned<Ast>, name: &'static str, src: &'static str) -> Chunk<'a> {
        Chunk {
            bindings: vec![],
            scope_depth: -1,
            num_bindings: 0,
            constants: vec![],
            bytecode: vec![],
            spans: vec![],
            strings: vec![],
            keywords: vec![
                "nil", "bool", "number", "keyword", "string", "tuple", "list", "dict", "box", "fn",
            ],
            nodes: vec![],
            ast: &ast.0,
            span: ast.1,
            src,
            name,
        }
    }

    pub fn kw_from(&self, kw: &str) -> Option<Value> {
        self.keywords
            .iter()
            .position(|s| *s == kw)
            .map(Value::Keyword)
    }

    pub fn visit(&mut self, node: &'a Spanned<Ast>) {
        let root_node = self.ast;
        let root_span = self.span;
        let (ast, span) = node;
        self.ast = ast;
        self.span = *span;
        self.compile();
        self.ast = root_node;
        self.span = root_span;
    }

    fn emit_constant(&mut self, val: Value) {
        let constant_index = self.constants.len();
        if constant_index > u8::MAX as usize {
            panic!(
                "internal Ludus compiler error: too many constants in chunk:{}:: {}",
                self.span, self.ast
            )
        }
        self.constants.push(val);
        self.bytecode.push(Op::Constant as u8);
        self.spans.push(self.span);
        self.bytecode.push(constant_index as u8);
        self.spans.push(self.span);
    }

    fn emit_op(&mut self, op: Op) {
        self.bytecode.push(op as u8);
        self.spans.push(self.span);
    }

    fn bind(&mut self, name: &'static str) {
        println!("binding {name} at depth {}", self.scope_depth);
        self.bindings.push(Binding {
            name,
            depth: self.scope_depth,
        });
        println!("{:?}", self.bindings)
    }

    pub fn compile(&mut self) {
        use Ast::*;
        match self.ast {
            Nil => self.emit_constant(Value::Nil),
            Number(n) => self.emit_constant(Value::Number(*n)),
            Boolean(b) => self.emit_constant(if *b { Value::True } else { Value::False }),
            String(s) => {
                let str_index = self.strings.len();
                self.strings.push(s);
                self.emit_constant(Value::Interned(str_index));
            }
            Keyword(s) => {
                let existing_kw = self.keywords.iter().position(|kw| kw == s);
                let kw_index = match existing_kw {
                    Some(index) => index,
                    None => self.keywords.len(),
                };
                self.keywords.push(s);
                self.emit_constant(Value::Keyword(kw_index));
            }
            Block(lines) => {
                self.scope_depth += 1;
                println!("now entering scope level {}", self.scope_depth);
                for expr in lines {
                    self.visit(expr);
                    self.emit_op(Op::Pop);
                }
                self.emit_op(Op::Store);
                self.scope_depth -= 1;
                while let Some(binding) = self.bindings.last() {
                    if binding.depth > self.scope_depth {
                        self.emit_op(Op::Pop);
                        self.bindings.pop();
                    } else {
                        break;
                    }
                }
                self.emit_op(Op::Load);
            }
            If(cond, then, r#else) => {
                self.visit(cond);
                let jif_idx = self.bytecode.len();
                self.emit_op(Op::JumpIfFalse);
                self.bytecode.push(0xff);
                self.visit(then);
                let jump_idx = self.bytecode.len();
                self.emit_op(Op::Jump);
                self.bytecode.push(0xff);
                self.visit(r#else);
                let end_idx = self.bytecode.len();
                let jif_offset = jump_idx - jif_idx;
                let jump_offset = end_idx - jump_idx;
                self.bytecode[jif_idx + 1] = jif_offset as u8;
                self.bytecode[jump_idx + 1] = jump_offset as u8;
            }
            Let(patt, expr) => {
                println!("let binding!");
                self.visit(expr);
                self.visit(patt);
            }
            WordPattern(name) => {
                self.bind(name);
            }
            Word(name) => {
                println!("resolving binding {name}");
                println!("current bindings {:?}", self.bindings);
                self.emit_op(Op::PushBinding);
                let biter = self.bindings.iter().enumerate().rev();
                for (i, binding) in biter {
                    println!("at index {i}");
                    if binding.name == *name {
                        self.bytecode.push(i as u8);
                        break;
                    }
                }
            }
            PlaceholderPattern => {
                self.bind("_");
            }
            _ => todo!(),
        }
    }

    pub fn disassemble(&self) {
        println!("=== chunk: {} ===", self.name);
        println!("IDX | CODE           | INFO");
        let mut codes = self.bytecode.iter().enumerate();
        while let Some((i, byte)) = codes.next() {
            let op = Op::from_u8(*byte).unwrap();
            use Op::*;
            match op {
                Pop | Store | Load => println!("{i:04}: {op}"),
                Constant => {
                    let (_, next) = codes.next().unwrap();
                    let value = &self.constants[*next as usize].show(self);
                    println!("{i:04}: {:16} {next:04}: {value}", op.to_string());
                }
                PushBinding => {
                    let (_, next) = codes.next().unwrap();
                    println!("{i:04}: {:16} {next:04}", op.to_string());
                }
                Jump | JumpIfFalse => {
                    let (_, next) = codes.next().unwrap();
                    println!("{i:04}: {:16} {next:04}", op.to_string())
                }
            }
        }
    }
}