Type checker for the simply typed lambda calculus using cofree

case class Cofree[F[_], A](head: A, tail: F[Cofree[F, A]])

sealed abstract class Exp[A]
case class Var[A](name: String) extends Exp[A]
case class Abs[A](param: String, paramT: Type, body: A) extends Exp[A]
case class App[A](func: A, arg: A) extends Exp[A]

sealed abstract class Type
case object BoolT extends Type
case object StringT extends Type
case class FuncT(from: Type, to: Type) extends Type

object TypeCheck {
    type Env = Map[String, Type]

    def typeCheck[A](env: Env)(exp: Cofree[Exp, A]): Option[Cofree[Exp, Type]] =
        exp.tail match {
            case Var(name) =>
                env.get(name).map(t => Cofree(t, Var(name)))
            case Abs(param, paramT, body) =>
                val env2 = env + (param -> paramT)
                typeCheck(env2)(body).map(body2 =>
                    Cofree(FuncT(paramT, body2.head), Abs(param, paramT, body2)))
            case App(func, arg) =>
                for {
                    func2 <- typeCheck(env)(func)
                    arg2 <- typeCheck(env)(arg)
                    r <- func2.head match {
                        case FuncT(from, to) if from == arg2.head =>
                            Some(Cofree(to, App(func2, arg2)))
                        case _ => None
                    }
                } yield r
        }
}

object Main {
    import TypeCheck._

    def mkVar(name: String) = Cofree((), Var[Cofree[Exp, Unit]](name))
    def mkAbs(param: String, paramT: Type, body: Cofree[Exp, Unit]) =
        Cofree((), Abs(param, paramT, body))
    def mkApp(func: Cofree[Exp, Unit], arg: Cofree[Exp, Unit]) =
        Cofree((), App(func, arg))

    def test(env: Env, exp: Cofree[Exp, Unit]) =
        println(exp.toString + " : " + typeCheck(env)(exp).toString)

    def main(args: Array[String]): Unit = {
        test(
            Map(),
            mkVar("x")
        )
        test(
            Map("x" -> BoolT),
            mkVar("x")
        )
        test(
            Map("x" -> BoolT, "not" -> FuncT(BoolT, BoolT)),
            mkApp(mkVar("not"), mkVar("x"))
        )
        test(
            Map("x" -> StringT, "not" -> FuncT(BoolT, BoolT)),
            mkAbs("x", BoolT, mkApp(mkVar("not"), mkVar("x")))
        )
    }
}