not-python/src/compile/mod.rs

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

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

#[derive(Default)]
pub struct Compile {
    const_data: ConstData,
    scope: Scope,
}

impl Compile {
    /// Creates a new compiler using the given text.
    pub fn new() -> Self {
        Default::default()
    }

    /// Compiles the given AST body.
    pub fn compile(mut self, body: &Body) -> error::Result<(ConstPool, UserFunRef)> {
        self.push_scope_layer();
        let main = thunk::CompileBody::new(&mut self)
            .compile(body)?
            .flatten()
            .to_vec();
        let globals_syms: std::collections::BTreeMap<_, _> = self.pop_scope_layer().unwrap()
            .into_iter()
            .map(|(sym, name)| (name, sym))
            .collect();
        let globals = globals_syms.into_iter()
            .enumerate()
            .map(|(index, (name, sym))| {
                assert_eq!(index, name.index());
                sym
            })
            .collect();
        Ok((self.const_data.const_pool, UserFun::new_obj(main, globals)))
    }

    /// Gets the constant data that is interned in this compile session.
    pub fn const_data(&self) -> &ConstData {
        &self.const_data
    }

    /// Mutably gets the constant data that is interned in this compile session.
    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)
    }

    /// Inserts a constant object without deduplication.
    pub fn push_const(&mut self, obj: ObjRef) -> (ConstHandle, ObjRef) {
        self.const_data_mut().push(obj)
    }

    /// Looks up a local variable name.
    pub fn lookup_local(&mut self, sym: Sym) -> Option<Name> {
        self.scope.lookup_local(sym)
    }

    /// Looks up a global variable name.
    pub fn lookup_global(&mut self, sym: Sym) -> Option<Name> {
        self.scope.lookup_global(sym)
    }

    /// Creates a new local variable if it does not exist in the current local scope.
    pub(crate) fn create_local(&mut self, sym: Sym) -> Name {
        self.scope.create_local(sym)
    }

    /// Creates a new global variable if it does not exist in the current global scope.
    pub fn create_global(&mut self, sym: Sym) -> Name {
        self.scope.create_global(sym)
    }

    /// Pushes an empty scope layer of local variables.
    pub fn push_scope_layer(&mut self) {
        self.scope.push_layer()
    }

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

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

/// Constant data pool used for building the const pool.
///
/// This distinguishes between the different types while building a const pool to avoid duplicates.
#[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.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.push(obj.clone());
            self.strs.insert(s.to_string(), (hdl, obj.clone()));
            (hdl, obj)
        }
    }

    /// Inserts a user-defined function constant object.
    ///
    /// This function provides no de-duplication. Use `const_str` or `const_int` for objects that
    /// may be repeated.
    pub fn push(&mut self, obj: ObjRef) -> (ConstHandle, ObjRef) {
        let hdl = self.const_pool.push(obj.clone());
        (hdl, obj)
    }
}