schala/schala-lang/language/src/typechecking.rs

use std::rc::Rc;

use ast::*;
use util::ScopeStack;

pub type TypeName = Rc<String>;

pub struct TypeContext<'a> {
  variable_map: ScopeStack<'a, Rc<String>, Type<()>>
}

type InferResult<T> = Result<T, TypeError>;

#[derive(Debug, Clone)]
struct TypeError { msg: String }

impl TypeError {
  fn new<A>(msg: &str) -> InferResult<A> {
    Err(TypeError { msg: msg.to_string() })
  }
}

#[derive(Debug, Clone)]
enum Type<a> {
  Var(a),
  Const(TConst),
  Arrow(Box<Type<a>>, Box<Type<a>>),
  ExistentialVar(u32)
}

impl TypeIdentifier {
  fn to_monotype(&self) -> Type<()> {
    match self {
      TypeIdentifier::Tuple(items) => unimplemented!(),
      TypeIdentifier::Singleton(TypeSingletonName { name, .. }) => {
        match &name[..] {
          "Nat" => Type::Const(TConst::Nat),
          "Int" => Type::Const(TConst::Int),
          "Float" => Type::Const(TConst::Float),
          "Bool" => Type::Const(TConst::Bool),
          "String" => Type::Const(TConst::StringT),
          _ => unimplemented!()
        }
      }
    }
  }
}

#[derive(Debug, Clone)]
enum TConst {
  User(Rc<String>),
  Unit,
  Nat,
  Int,
  Float,
  StringT,
  Bool,
}

impl TConst {
  fn user(name: &str) -> TConst {
    TConst::User(Rc::new(name.to_string()))
  }
}

#[derive(Debug, Clone)]
struct PolyType {
  vars: Vec<Rc<String>>,
  ty: Type<()>
}

impl<'a> TypeContext<'a> {
  pub fn new() -> TypeContext<'a> {
    TypeContext { 
      variable_map: ScopeStack::new(None),
    }
  }

  pub fn typecheck(&mut self, ast: &AST) -> Result<String, String> {
    match self.infer_ast(ast) {
      Ok(t) => Ok(format!("{:?}", t)),
      Err(err) => Err(format!("Type error: {:?}", err))
    }
  }
}

impl<'a> TypeContext<'a> {
  fn infer_ast(&mut self, ast: &AST) -> InferResult<Type<()>> {
    let mut output = Type::Const(TConst::Unit);
    for statement in ast.0.iter() {
      output = self.infer_statement(statement)?;
    }
    Ok(output)
  }

  fn infer_statement(&mut self, stmt: &Statement) -> InferResult<Type<()>> {
    match stmt {
      Statement::ExpressionStatement(ref expr) => self.infer_expr(expr),
      Statement::Declaration(ref decl) => self.infer_decl(decl),
    }
  }

  fn infer_expr(&mut self, expr: &Expression) -> InferResult<Type<()>> {
    match expr {
      Expression(expr_type, Some(type_anno)) => {
        let tx = self.infer_expr_type(expr_type)?;
        let ty = type_anno.to_monotype();
        self.unify(&ty, &tx)
      },
      Expression(expr_type, None) => self.infer_expr_type(expr_type)
    }
  }

  fn infer_decl(&mut self, expr: &Declaration) -> InferResult<Type<()>> {
    Ok(Type::Const(TConst::user("unimplemented")))
  }

  fn infer_expr_type(&mut self, expr_type: &ExpressionType) -> InferResult<Type<()>> {
    use self::ExpressionType::*;
    Ok(match expr_type {
      NatLiteral(_) => Type::Const(TConst::Nat),
      FloatLiteral(_) => Type::Const(TConst::Float),
      StringLiteral(_) => Type::Const(TConst::StringT),
      BoolLiteral(_) => Type::Const(TConst::Bool),
      Value(name) => {
        //TODO handle the distinction between 0-arg constructors and variables at some point
        // need symbol table for that
        match self.variable_map.lookup(name) {
          Some(ty) => ty.clone(),
          None => return TypeError::new(&format!("Variable {} not found", name))
        }
      },
      IfExpression { discriminator, body } => self.infer_if_expr(discriminator, body)?,
      Call { f, arguments } => {
        let tf: Type<()> = self.infer_expr(f)?; //has to be an Arrow Type
        let targ = self.infer_expr(&arguments[0])?; // TODO make this work with functions with more than one arg
        match tf {
          Type::Arrow(t1, t2) => {
            self.unify(&t1, &targ)?;
            *t2.clone()
          },
          _ => return TypeError::new("not a function")
        }
      },

      Lambda { params, type_anno, body } => {

        let arg_type = unimplemented!();
        let result_type = unimplemented!();

        Type::Arrow(Box::new(arg_type), Box::new(result_type))
      }
      _ => Type::Const(TConst::user("unimplemented"))
    })
  }

  fn infer_if_expr(&mut self, discriminator: &Discriminator, body: &IfExpressionBody) -> InferResult<Type<()>> {
    let test = match discriminator {
      Discriminator::Simple(expr) => expr,
      _ => return TypeError::new("Dame desu")
    };

    let (then_clause, maybe_else_clause) = match body {
      IfExpressionBody::SimpleConditional(a, b) => (a, b),
      _ => return TypeError::new("Dont work")
    };

    unimplemented!()
  }

  fn infer_block(&mut self, block: &Block) -> InferResult<Type<()>> {
    let mut output = Type::Const(TConst::Unit);
    for statement in block.iter() {
      output = self.infer_statement(statement)?;
    }
    Ok(output)
  }

  fn unify(&mut self, t1: &Type<()>, t2: &Type<()>) -> InferResult<Type<()>> {
    unimplemented!()
  }

  fn allocate_existential(&mut self) -> Type<()> {
    Type::ExistentialVar(0)
  }
}

#[cfg(test)]
mod tests {
  use super::*;

  fn parse(input: &str) -> AST {
    let tokens: Vec<::tokenizing::Token> = ::tokenizing::tokenize(input);
    let mut parser = ::parsing::Parser::new(tokens);
    parser.parse().unwrap()
  }

  macro_rules! type_test {
    ($input:expr, $correct:expr) => {
      {
      let mut tc = TypeContext::new();
      let ast = parse($input);
      tc.add_symbols(&ast);
      assert_eq!($correct, tc.type_check(&ast).unwrap())
      }
    }
  }



  #[test]
  fn basic_inference() {

  }
}