path: root/sources/examples/typeinf.scala

package examples;

trait Term {}

case class Var(x: String)                   extends Term {
  override def toString() = x
}
case class Lam(x: String, e: Term)          extends Term {
  override def toString() = "(\\" + x + "." + e + ")"
}
case class App(f: Term, e: Term)            extends Term {
  override def toString() = "(" + f + " " + e + ")"
}
case class Let(x: String, e: Term, f: Term) extends Term {
  override def toString() = "let " + x + " = " + e + " in " + f;
}

sealed trait Type {}
case class Tyvar(a: String) extends Type {
  override def toString() = a
}
case class Arrow(t1: Type, t2: Type) extends Type {
  override def toString() = "(" + t1 + "->" + t2 + ")"
}
case class Tycon(k: String, ts: List[Type]) extends Type {
  override def toString() =
    k + (if (ts.isEmpty) "" else ts.mkString("[", ",", "]"))
}

object typeInfer {

  private var n: Int = 0;
  def newTyvar(): Type = { n = n + 1 ; Tyvar("a" + n) }

  trait Subst with Function1[Type,Type] {
    def lookup(x: Tyvar): Type;
    def apply(t: Type): Type = t match {
      case tv @ Tyvar(a) => val u = lookup(tv); if (t == u) t else apply(u);
      case Arrow(t1, t2) => Arrow(apply(t1), apply(t2))
      case Tycon(k, ts) => Tycon(k, ts map apply)
    }
    def extend(x: Tyvar, t: Type) = new Subst {
      def lookup(y: Tyvar): Type = if (x == y) t else Subst.this.lookup(y);
    }
  }

  val emptySubst = new Subst { def lookup(t: Tyvar): Type = t }

  case class TypeScheme(tyvars: List[Tyvar], tpe: Type) {
    def newInstance: Type =
      (emptySubst /: tyvars) ((s, tv) => s.extend(tv, newTyvar())) (tpe);
  }

  type Env = List[Pair[String, TypeScheme]];

  def lookup(env: Env, x: String): TypeScheme = env match {
    case List() => null
    case Pair(y, t) :: env1 => if (x == y) t else lookup(env1, x)
  }

  def gen(env: Env, t: Type): TypeScheme =
    TypeScheme(tyvars(t) diff tyvars(env), t);

  def tyvars(t: Type): List[Tyvar] = t match {
    case tv @ Tyvar(a) => List(tv)
    case Arrow(t1, t2) => tyvars(t1) union tyvars(t2)
    case Tycon(k, ts) => (List[Tyvar]() /: ts) ((tvs, t) => tvs union tyvars(t));
  }

  def tyvars(ts: TypeScheme): List[Tyvar] =
    tyvars(ts.tpe) diff ts.tyvars;

  def tyvars(env: Env): List[Tyvar] =
    (List[Tyvar]() /: env) ((tvs, nt) => tvs union tyvars(nt._2));

  def mgu(t: Type, u: Type, s: Subst): Subst = Pair(s(t), s(u)) match {
    case Pair(Tyvar(a), Tyvar(b)) if (a == b) =>
      s
    case Pair(Tyvar(a), _) if !(tyvars(u) contains a) =>
      s.extend(Tyvar(a), u)
    case Pair(_, Tyvar(a)) =>
      mgu(u, t, s)
    case Pair(Arrow(t1, t2), Arrow(u1, u2)) =>
      mgu(t1, u1, mgu(t2, u2, s))
    case Pair(Tycon(k1, ts), Tycon(k2, us)) if (k1 == k2) =>
      (s /: (ts zip us)) ((s, tu) => mgu(tu._1, tu._2, s))
    case _ => throw new TypeError("cannot unify " + s(t) + " with " + s(u))
  }

  case class TypeError(s: String) extends Exception(s) {}

  def tp(env: Env, e: Term, t: Type, s: Subst): Subst = {
    current = e;
    e match {
      case Var(x) =>
	val u = lookup(env, x);
	if (u == null) throw new TypeError("undefined: " + x);
	else mgu(u.newInstance, t, s)

      case Lam(x, e1) =>
	val a, b = newTyvar();
	val s1 = mgu(t, Arrow(a, b), s);
	val env1 = Pair(x, TypeScheme(List(), a)) :: env;
	tp(env1, e1, b, s1)

      case App(e1, e2) =>
	val a = newTyvar();
	val s1 = tp(env, e1, Arrow(a, t), s);
	tp(env, e2, a, s1)

      case Let(x, e1, e2) =>
	val a = newTyvar();
	val s1 = tp(env, e1, a, s);
	tp(Pair(x, gen(env, s1(a))) :: env, e2, t, s1)
    }
  }
  var current: Term = null;

  def typeOf(env: Env, e: Term): Type = {
    val a = newTyvar();
    tp(env, e, a, emptySubst)(a)
  }
}

object predefined {
  val booleanType = Tycon("Boolean", List());
  val intType = Tycon("Int", List());
  def listType(t: Type) = Tycon("List", List(t));

  private def gen(t: Type): typeInfer.TypeScheme = typeInfer.gen(List(), t);
  private val a = typeInfer.newTyvar();
  val env = List(
/*
    Pair("true", gen(booleanType)),
    Pair("false", gen(booleanType)),
    Pair("if", gen(Arrow(booleanType, Arrow(a, Arrow(a, a))))),
    Pair("zero", gen(intType)),
    Pair("succ", gen(Arrow(intType, intType))),
    Pair("nil", gen(listType(a))),
    Pair("cons", gen(Arrow(a, Arrow(listType(a), listType(a))))),
    Pair("isEmpty", gen(Arrow(listType(a), booleanType))),
    Pair("head", gen(Arrow(listType(a), a))),
    Pair("tail", gen(Arrow(listType(a), listType(a)))),
*/
    Pair("fix", gen(Arrow(Arrow(a, a), a)))
  )
}

abstract class MiniMLParsers[intype] extends CharParsers[intype] {

  /** whitespace */
  def whitespace = rep{chr(' ') ||| chr('\t') ||| chr('\n')};

  /** A given character, possible preceded by whitespace */
  def wschr(ch: char) = whitespace &&& chr(ch);

  /** identifiers or keywords */
  def id: Parser[String] =
    for (
      val c: char <- rep(chr(' ')) &&& chr(Character.isLetter);
      val cs: List[char] <- rep(chr(Character.isLetterOrDigit))
    ) yield (c :: cs).mkString("", "", "");

  /** Non-keyword identifiers */
  def ident: Parser[String] =
    for (val s <- id; s != "let" && s != "in") yield s;

  /** term = '\' ident '.' term | term1 {term1} | let ident "=" term in term */
  def term: Parser[Term] =
    ( for (
	val _ <- wschr('\\');
	val x <- ident;
	val _ <- wschr('.');
	val t <- term)
      yield Lam(x, t): Term )
    |||
    ( for (
        val letid <- id; letid == "let";
        val x <- ident;
	val _ <- wschr('=');
        val t <- term;
	val inid <- id; inid == "in";
	val c <- term)
      yield Let(x, t, c) )
    |||
    ( for (
	val t <- term1;
	val ts <- rep(term1))
      yield (t /: ts)((f, arg) => App(f, arg)) );

  /** term1 = ident | '(' term ')' */
  def term1: Parser[Term] =
    ( for (val s <- ident)
      yield Var(s): Term )
    |||
    ( for (
	val _ <- wschr('(');
	val t <- term;
	val _ <- wschr(')'))
      yield t );

  /** all = term ';' */
  def all: Parser[Term] =
    for (
      val t <- term;
      val _ <- wschr(';'))
    yield t;
}

object testInfer {

  def showType(e: Term): String =
    try {
      typeInfer.typeOf(predefined.env, e).toString();
    } catch {
      case typeInfer.TypeError(msg) =>
	"\n cannot type: " + typeInfer.current +
	"\n reason: " + msg;
    }

  def main(args: Array[String]): unit = {
    val ps = new MiniMLParsers[int] with ParseString(args(0));
    ps.all(ps.input) match {
      case Some(Pair(term, _)) =>
	System.out.println("" + term + ": " + showType(term));
      case None =>
	System.out.println("syntax error");
    }
  }
}