not-python/src/compile/mod.rs

pub mod basic_block;
pub mod error;
mod locals;
pub mod thunk;

use crate::{syn::ast::Body, obj::prelude::*, vm::consts::*};
use std::collections::{BTreeMap, HashMap};

pub struct Compile<'t> {
    text: &'t str,
    const_data: ConstData,
    globals: BTreeMap<Sym, Local>,
    locals: Vec<BTreeMap<Sym, Local>>,
    next_local: Local,
}

impl<'t> Compile<'t> {
    pub fn new(text: &'t str) -> Self {
        Compile {
            text,
            const_data: Default::default(),
            globals: Default::default(),
            locals: Default::default(),
            next_local: Default::default(),
        }
    }

    pub fn text(&self) -> &'t str {
        self.text
    }

    pub fn const_data(&self) -> &ConstData {
        &self.const_data
    }

    pub fn const_data_mut(&mut self) -> &mut ConstData {
        &mut self.const_data
    }

    /// Gets or inserts a static int reference.
    pub fn const_int(&mut self, int: i64) -> (ConstHandle, IntRef) {
        self.const_data_mut().const_int(int)
    }

    /// Gets or inserts a static string reference.
    pub fn const_str(&mut self, s: impl AsRef<str>) -> (ConstHandle, StrRef) {
        self.const_data_mut().const_str(s)
    }

    /// Looks up a variable name.
    ///
    /// This will search up the locals stack until the given name is found, ultimately ending with
    /// a global name lookup.
    pub fn lookup_scope(&mut self, sym: Sym) -> Option<Local> {
        self.locals
            .iter()
            .rev()
            .filter_map(|locals| locals.get(&sym))
            .next()
            .or_else(|| self.globals.get(&sym))
            .copied()
    }

    /// Creates a new local variable if it does not exist in the current local scope.
    pub fn create_local(&mut self, sym: Sym) -> Local {
        let locals = self.locals.last_mut().expect("scope");
        if let Some(local) = locals.get(&sym) {
            *local
        } else {
            // wish I could use mem::replace here, oh well
            let local = self.next_local;
            self.next_local = self.next_local.next();
            locals.insert(sym, local);
            local
        }
    }

    /// Creates a new global variable if it does not exist in the current global scope.
    pub fn create_global(&mut self, sym: Sym) -> Local {
        if let Some(global) = self.globals.get(&sym) {
            *global
        } else {
            let global = self.next_local;
            self.next_local = self.next_local.next();
            self.globals.insert(sym, global);
            global
        }
    }

    /// Pushes an empty scope layer of local variables.
    pub fn push_scope_layer(&mut self) {
        self.locals.push(Default::default());
    }

    /// Pops a scope layer of local variables, if any are available.
    pub fn pop_scope_layer(&mut self) -> Option<BTreeMap<Sym, Local>> {
        self.locals.pop()
    }

    /// Collects local variables for the given AST body non-recursively.
    pub fn collect_locals(&mut self, body: &Body) {
        locals::CollectLocals::new(self).collect(body);
    }
}

#[derive(Debug, Default)]
pub struct ConstData {
    ints: HashMap<i64, (ConstHandle, IntRef)>,
    strs: HashMap<String, (ConstHandle, StrRef)>,
    const_pool: ConstPool,
}

impl ConstData {
    /// Gets or inserts a static int reference.
    pub fn const_int(&mut self, int: i64) -> (ConstHandle, IntRef) {
        if let Some(pair) = self.ints.get(&int) {
            pair.clone()
        } else {
            let obj = Int::new_obj(int);
            let hdl = self.const_pool.push(obj.clone());
            self.ints.insert(int, (hdl, obj.clone()));
            (hdl, obj)
        }
    }

    /// Gets or inserts a static string reference.
    pub fn const_str(&mut self, s: impl AsRef<str>) -> (ConstHandle, StrRef) {
        let s = s.as_ref();
        if let Some(pair) = self.strs.get(s) {
            pair.clone()
        } else {
            let obj = Str::new_obj(s.to_string());
            let hdl = self.const_pool.push(obj.clone());
            self.strs.insert(s.to_string(), (hdl, obj.clone()));
            (hdl, obj)
        }
    }
}