schala/schala-lang/src/typechecking.rs

use std::rc::Rc;
use std::collections::HashMap;
//use std::char;
use std::fmt;
use std::fmt::Write;

use itertools::Itertools;

use parsing;

pub struct TypeContext {
  //type_var_count: u64,
  bindings: HashMap<Rc<String>, Type>,
  pub symbol_table: SymbolTable
}

//cf. p. 150 or so of Language Implementation Patterns
pub struct SymbolTable {
  pub values: HashMap<Rc<String>, Symbol> //TODO this will eventually have real type information
}

impl SymbolTable {
  fn new() -> SymbolTable {
    SymbolTable { values: HashMap::new() }
  }
}

#[derive(Debug)]
pub struct Symbol {
  pub name: Rc<String>,
  pub spec: SymbolSpec,
}

#[derive(Debug)]
pub enum SymbolSpec {
  Func, Custom(String)
}


/* real meat of type stuff here */


#[derive(Debug, PartialEq, Clone)]
pub enum Type {
  Const(TConst),
  Var(TVar),
  Func(Box<Type>, Box<Type>),
  //UVar(String),
  //EVar(u64),
  Sum(Vec<Type>),
  Void
}

#[derive(Debug, PartialEq, Clone)]
pub struct TVar(String);

#[derive(Debug, PartialEq, Clone)]
pub enum TConst {
  Unit,
  Nat,
  Int,
  Float,
  StringT,
  Bool,
  Custom(String),
}

impl fmt::Display for Type {
  fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
    write!(f, "{:?}", self)
  }
}

/* TODO this should just check the name against a map, and that map should be pre-populated with
 * types */
impl parsing::TypeName {
  fn to_type(&self) -> TypeResult<Type> {
	use self::parsing::TypeSingletonName;
	use self::parsing::TypeName::*;
	use self::Type::*; use self::TConst::*;
    Ok(match self {
      Tuple(_) => return Err(format!("Tuples not yet implemented")),
      Singleton(name) => match name {
        TypeSingletonName { name, .. } => match &name[..] {
          "Nat" => Const(Nat),
          "Int" => Const(Int),
          "Float" => Const(Float),
          "Bool" => Const(Bool),
          "String" => Const(StringT),
          n => Const(Custom(n.to_string()))
        }
      }
    })
  }
}

pub type TypeResult<T> = Result<T, String>;

impl TypeContext {
  pub fn new() -> TypeContext {
    TypeContext { bindings: HashMap::new(), /*type_var_count: 0*/ symbol_table: SymbolTable::new() }
  }
}

impl TypeContext {
  /* note: this adds names for *forward reference* but doesn't actually create any types. solve that problem
   * later */
  pub fn add_top_level_types(&mut self, ast: &parsing::AST) -> TypeResult<()> {
    use self::parsing::{Statement, TypeName, Variant, TypeSingletonName, TypeBody};
    use self::parsing::Declaration::*;
    use self::Type::*;
    for statement in ast.0.iter() {
      if let Statement::Declaration(decl) = statement {
        match decl {
          FuncSig(signature) | FuncDecl(signature, _) => {
            self.symbol_table.values.insert(
              signature.name.clone(),
              Symbol { name: signature.name.clone(), spec: SymbolSpec::Func }
            );
          },
          TypeDecl(TypeSingletonName { name, ..}, TypeBody(variants)) => {
            for var in variants {
              match var {
                Variant::UnitStruct(variant_name) => {
                  //TODO will have to make this a function to this type eventually
                  let spec = SymbolSpec::Custom(format!("{}", name));
                  self.symbol_table.values.insert(variant_name.clone(), Symbol { name: variant_name.clone(), spec });
                },
                e => return Err(format!("{:?} not supported in typing yet", e)),
              }
            }
          },
          _ => ()
        }
      }
    }
    Ok(())
  }
  pub fn debug_symbol_table(&self) -> String {
    let mut output = format!("Symbol table\n");
    for (sym, ty) in &self.symbol_table.values {
      write!(output, "{} -> {:?}\n", sym, ty).unwrap();
    }
    write!(output, "\nBindings\n").unwrap();
    for (sym, ty) in &self.bindings {
      write!(output, "{} : {:?}\n", sym, ty).unwrap();
    }
    output
  }

  pub fn type_check_ast(&mut self, ast: &parsing::AST) -> TypeResult<Type> {
    let ref block = ast.0;
    Ok(self.infer_block(block)?)
  }
}

impl TypeContext {
  fn infer_block(&mut self, statements: &Vec<parsing::Statement>) -> TypeResult<Type> {
    let mut ret_type = Type::Const(TConst::Unit);
    for statement in statements {
      ret_type = self.infer_statement(statement)?;
    }
    Ok(ret_type)
  }

  fn infer_statement(&mut self, statement: &parsing::Statement) -> TypeResult<Type> {
    use self::parsing::Statement::*;
    match statement {
      ExpressionStatement(expr) => self.infer(expr),
      Declaration(decl) => self.add_declaration(decl),
    }
  }
  fn add_declaration(&mut self, decl: &parsing::Declaration) -> TypeResult<Type> {
    use self::parsing::Declaration::*;
    use self::Type::*;
    match decl {
      Binding { name, expr, .. } => {
        let ty = self.infer(expr)?;
        self.bindings.insert(name.clone(), ty);
      },
      _ => return Err(format!("other formats not done"))
    }
    Ok(Void)
  }
  fn infer(&mut self, expr: &parsing::Expression) -> TypeResult<Type> {
    use self::parsing::Expression;
    match expr {
      Expression(e, Some(anno)) => {
        let anno_ty = anno.to_type()?;
        let ty = self.infer_exprtype(&e)?;
        self.unify(ty, anno_ty)
      },
      Expression(e, None) => self.infer_exprtype(e)
    }
  }
  fn infer_exprtype(&mut self, expr: &parsing::ExpressionType) -> TypeResult<Type> {
    use self::parsing::ExpressionType::*;
    use self::Type::*; use self::TConst::*;
    match expr {
      NatLiteral(_) => Ok(Const(Nat)),
      FloatLiteral(_) => Ok(Const(Float)),
      StringLiteral(_) => Ok(Const(StringT)),
      BoolLiteral(_) => Ok(Const(Bool)),
      BinExp(op, lhs, rhs) => { /* remember there are both the haskell convention talk and the write you a haskell ways to do this! */
        match op.get_type()? {
          Func(box t1, box Func(box t2, box t3)) => {
            let lhs_ty = self.infer(lhs)?;
            let rhs_ty = self.infer(rhs)?;
            self.unify(t1, lhs_ty)?;
            self.unify(t2, rhs_ty)?;
            Ok(t3)
          },
          other => Err(format!("{:?} is not a binary function type", other))
        }
      },
      PrefixExp(op, expr) => match op.get_type()? {
        Func(box t1, box t2) => {
          let expr_ty = self.infer(expr)?;
          self.unify(t1, expr_ty)?;
          Ok(t2)
        },
        other => Err(format!("{:?} is not a prefix op function type", other))
      },
      Value(name) => {
        match self.bindings.get(name) {
          Some(ty) => Ok(ty.clone()),
          None => Err(format!("No binding found for variable: {}", name)),
        }
      },
      Call { f, arguments } => {
        let mut tf = self.infer(f)?;
        for arg in arguments.iter() {
          match tf {
            Func(box t, box rest) => {
              let t_arg = self.infer(arg)?;
              self.unify(t, t_arg)?;
              tf = rest;
            },
            other => return Err(format!("Function call failed to unify; last type: {:?}", other)),
          }
        }
        Ok(tf)
      },
      TupleLiteral(expressions) => {
        let mut types = vec![];
        for expr in expressions {
          types.push(self.infer(expr)?);
        }
        Ok(Sum(types))
      },
      _ => Err(format!("Type not yet implemented"))
    }
  }
  fn unify(&mut self, t1: Type, t2: Type) -> TypeResult<Type> {
    use self::Type::*;// use self::TConst::*;
    match (t1, t2) {
      (Const(ref a), Const(ref b)) if a == b => Ok(Const(a.clone())),
      (a, b) => Err(format!("Types {:?} and {:?} don't unify", a, b))
    }
  }
}