env.rs

// env

use std::iter::FromIterator;
// use std::iter::{FromIterator, Map};

use crate::ast::*;
// use crate::grammar::*;
use std::collections::{HashMap, VecDeque};

pub type Error = String;

pub type FnEnv<'a> = HashMap<&'a str, &'a FnDecl>;

impl FnDecl {
    fn get_id(&self) -> &str {
        &self.id
    }
}

pub fn new_fn_env(v: &Vec<FnDecl>) -> FnEnv {
    let i = v.into_iter().map(|td| (td.get_id(), td));
    HashMap::from_iter(i)
}

// Option<Type>
// 
// Some<Type>, variable has a declared type
// None, variable has been introduced, but type not yet determined

pub type IdType = HashMap<Id, Option<Type>>;

#[derive(Debug)]
pub struct VarEnv(VecDeque<IdType>);

impl VarEnv {
    pub fn new() -> Self {
        VarEnv(VecDeque::new())
    }

    pub fn get(&self, id: String) -> Option<&(Option<Type>)> {
        self.0.iter().find_map(|hm| hm.get(id.as_str()))
    }

    pub fn get_mut(&mut self, id: String) -> Option<&mut (Option<Type>)> {
        self.0.iter_mut().find_map(|hm| hm.get_mut(id.as_str()))
    }

    pub fn new_id(&mut self, id: String, oty: Option<Type>) {
        let hm = self.0.front_mut().unwrap();
        println!("insert id {:?}, in scope {:?}", id, hm);

        hm.insert(id, (oty));
    }

    pub fn update(&mut self, id: String, ty: Type) -> Result<(), Error> {
        use Type::*;
        match self.get_mut(id.clone()) {
            Some(ot) => match ot {
                None => {
                    // no type assigned
                    *ot = Some(ty);
                    Ok(())
                }
                Some(t) => match *t == ty {
                    true => Ok(()),
                    false => Err(format!("types differ {} <> {}", t, ty)),
                },
            },
            None => Err("variable not found".to_string()),
        }
    }

    pub fn push_empty_scope(&mut self) {
        self.0.push_front(HashMap::new());
    }

    pub fn push_param_scope(&mut self, args: IdType) {
        self.0.push_front(args)
    }

    pub fn pop_scope(&mut self) {
        self.0.pop_front();
    }
}

#[test]
// trace and assert
fn type_env_test() {
    use Type::*;
    let mut var_env = VarEnv::new();
    var_env.push_empty_scope();
    var_env.new_id("a".to_owned(), None);
    println!("{:?}", var_env);

    println!("update a {:?}", var_env.update("a".to_owned(), I32));
    println!("get a {:?}", var_env.get("a".to_owned()));
    assert_eq!(var_env.get("a".to_owned()), Some(&Some(I32)));
    println!("update a {:?}", var_env.update("a".to_owned(), I32));
    assert!(var_env.update("a".to_owned(), I32).is_ok());

    println!("update a {:?}", var_env.update("a".to_owned(), Bool));
    assert!(var_env.update("a".to_owned(), Bool).is_err());

    println!("get b {:?}", var_env.get("b".to_owned()));
    assert_eq!(var_env.get("b".to_owned()), None);

    println!("update b {:?}", var_env.update("b".to_owned(), Bool));
    assert!(var_env.update("b".to_owned(), Bool).is_err());

    var_env.push_empty_scope();
    println!("push empty scope");
    println!("update a {:?}", var_env.update("a".to_owned(), Bool));
    assert!(var_env.update("a".to_owned(), Bool).is_err());
    println!("update a {:?}", var_env.update("a".to_owned(), I32));
    assert!(var_env.update("a".to_owned(), I32).is_ok());

    var_env.new_id("a".to_owned(), None);
    println!("get a {:?}", var_env.get("a".to_owned()));
    assert_eq!(var_env.get("a".to_owned()), Some(&None));
    println!("update a {:?}", var_env.update("a".to_owned(), Bool));
    println!("get a {:?}", var_env.get("a".to_owned()));
    assert_eq!(var_env.get("a".to_owned()), Some(&Some(Bool)));

    println!("update a {:?}", var_env.update("a".to_owned(), Bool));
    assert!(var_env.update("a".to_owned(), Bool).is_ok());
    println!("update a {:?}", var_env.update("a".to_owned(), Type::I32));
    assert!(var_env.update("a".to_owned(), I32).is_err());

    println!("pop_scope");
    var_env.pop_scope();
    println!("get a {:?}", var_env.get("a".to_owned()));
    assert_eq!(var_env.get("a".to_owned()), Some(&Some(I32)));

    println!("update a {:?}", var_env.update("a".to_owned(), Bool));
    assert!(var_env.update("a".to_owned(), Bool).is_err());
    println!("update a {:?}", var_env.update("a".to_owned(), I32));
    assert!(var_env.update("a".to_owned(), I32).is_ok());

    println!("get b {:?}", var_env.get("b".to_owned()));
    assert_eq!(var_env.get("b".to_owned()), None);
}

pub type TypeEnv<'a> = HashMap<&'a str, &'a TypeDecl>;

impl TypeDecl {
    pub fn get_id(&self) -> &str {
        match self {
            TypeDecl::Struct(s, _) => s,
            TypeDecl::Enum(s, _) => s,
        }
    }
}

pub fn new_type_env(v: &Vec<TypeDecl>) -> TypeEnv {
    let i = v.into_iter().map(|td| (td.get_id(), td));
    HashMap::from_iter(i)
}

#[test]
fn test_type_env() {
    use crate::grammar::ProgramParser;
    use crate::*;

    let p = ProgramParser::new().parse(mixed_types!()).unwrap();
    let type_env = new_type_env(&p.type_decls);

    assert_eq! {
        **type_env.get("A").unwrap(),
        TypeDecl::Struct("A".to_string(),vec![])
    };

    assert_eq! {
        **type_env.get("Point").unwrap(),
        TypeDecl::Struct(
            "Point".to_string(),
            vec![
                Field {
                    id: "x".to_string(),
                    ty: Type::I32,
                },
                Field {
                    id: "y".to_string(),
                    ty: Type::I32,
                },
                ],
        )
    };
}

#[test]
fn fn_decl_env() {
    use crate::*;
    let p = grammar::ProgramParser::new().parse(fn_mixed!()).unwrap();
    let fn_env = new_fn_env(&p.fn_decls);

    assert_eq! {
        **fn_env.get("main").unwrap(),
        FnDecl {
            id: "main".to_string(),
            params: Params(vec![]),
            result: Type::Unit,
            body: Stmts{ stmts: vec![], ret: true }
        }
    };

    println!("{}", p);
}