Tagless.scala

import org.parboiled2._

/*
http://okmij.org/ftp/tagless-final/course/lecture.pdf
*/

object Tagless extends scala.App {

  type Id[A] = A
  trait ~>[F[_], G[_]] {
    def apply[A](fa: F[A]): G[A]
  }

  trait AParser {
    self: Parser =>
    import CharPredicate._
    def nl: Rule0 = rule('\n')
    def sp: Rule0 = rule(quiet(zeroOrMore(' ' | nl)))
    def Name: Rule1[String] = rule(atomic(
      capture(Alpha ~ (AlphaNum).*)))
  }

  trait Term[T] {
    def ref(n: String): T
    def lam(n: String, t: T): T
    def app(f: T, a: T): T
  }

  trait ShowTerm extends Term[String] {
    def ref(n: String) = n
    def lam(n: String, t: String) = s"$n -> $t"
    def app(f: String, a: String) = s"$f $a"
  }

  trait TermAst[F[_]] extends Term[F ~> Id] {
    def alg[A](f: F[A]): Term[A]
    object TermAst {
      type Ast = F ~> Id
      case class Ref(n: String) extends Ast {
        def apply[A](fa: F[A]): Id[A] = alg(fa).ref(n)
      }
      case class Lam(n: String, t: Ast) extends Ast {
        def apply[A](fa: F[A]): Id[A] = alg(fa).lam(n, t(fa))
      }
      case class App(f: Ast, a: Ast) extends Ast {
        def apply[A](fa: F[A]): Id[A] = alg(fa).app(f(fa), a(fa))
      }
    }
    import TermAst._
    def ref(n: String): Ast = Ref(n)
    def lam(n: String, t: Ast): Ast = Lam(n, t)
    def app(f: Ast, a: Ast): Ast = App(f, a)
  }

  trait TermParser[F[_]] extends TermAst[F] with AParser {
    self: Parser =>
    import TermAst._
    def Term: Rule1[Ast] = rule(App | NonApp)
    def App: Rule1[Ast] = rule(NonApp ~ oneOrMore(' ' ~ sp ~ NonApp ~> (app _)))
    def NonApp: Rule1[Ast] = rule(Lam | Ref)
    def Lam: Rule1[Ast] = rule(Name ~ sp ~ "->" ~ sp ~ Term ~> (lam _))
    def Ref: Rule1[Ast] = rule(Name ~> (ref _))
  }

  trait Expr[T] {
    def oper(a: T, f: String, b: T): T
    def parens(a: T): T
  }

  trait ShowExpr extends Expr[String] {
    def oper(a: String, f: String, b: String) = s"$a $f $b"
    def parens(a: String) = s"( $a )"
  }

  trait ExprAst[F[_]] extends Expr[F ~> Id] {
    def alg[A](f: F[A]): Expr[A]
    object ExprAst {
      type Ast = F ~> Id
      case class Oper(a: Ast, f: String, b: Ast) extends Ast {
        def apply[A](fa: F[A]): Id[A] = alg(fa).oper(a(fa), f, b(fa))
      }
      case class Parens(a: Ast) extends Ast {
        def apply[A](fa: F[A]): Id[A] = alg(fa).parens(a(fa))
      }
    }
    import ExprAst._
    def parens(a: Ast): Ast = Parens(a)
    def oper(a: Ast, f: String, b: Ast): Ast = (a, f) match {
      case (Oper(a1, op1@("+" | "-"), b1), op2@("*" | "/")) =>
        oper(a1, op1, oper(b1, op2, b))
      case _ => Oper(a, f, b)
    }

  }

  trait ExprParser[F[_]] extends ExprAst[F] with AParser {
    self: Parser =>
    import ExprAst._
    def operChar = CharPredicate('~', '!', '@', '#', '$', '%',
      'ˆ', '&', '*', '-', '+', '/', '?', '<', '>', ':', '.', ',')
    def Parens(p: => Rule1[Ast]): Rule1[Ast] =
      rule('(' ~ sp ~ p ~ sp ~ ')' ~> (parens _))
    def Oper: Rule1[String] = rule(atomic(capture(operChar.+)))
    def Expr: Rule1[Ast] = rule(Operand ~
      oneOrMore(sp ~ Oper ~ sp ~ Operand ~> (oper _)))
    def Operand: Rule1[Ast]
  }

  trait Literal[T] {
    def int(i: Long): T
    def bool(b: Boolean): T
  }

  trait ShowLiteral extends Literal[String] {
    def int(i: Long) = s"$i"
    def bool(b: Boolean): String = s"$b"
  }

  trait LiteralAst[F[_]] extends Literal[F ~> Id] {
    def alg[A](f: F[A]): Literal[A]
    def int(i: Long): F ~> Id = IntV(i)
    def bool(b: Boolean): F ~> Id = BoolV(b)
    case class IntV(i: Long) extends (F ~> Id) {
      def apply[A](fa: F[A]): Id[A] = alg(fa).int(i)
    }
    case class BoolV(b: Boolean) extends (F ~> Id) {
      def apply[A](fa: F[A]): Id[A] = alg(fa).bool(b)
    }
  }

  trait LiteralParser[F[_]] extends LiteralAst[F] {
    self: Parser =>
    import CharPredicate._
    def Number: Rule1[F ~> Id] = rule(capture(Digit.+) ~>
      ((s: String) => int(s.toInt)))
    def True: Rule1[F ~> Id] = rule(capture("true") ~>
      ((_: String) => bool(true)))
    def False: Rule1[F ~> Id] = rule(capture("false") ~>
      ((_: String) => bool(false)))
    def Bool: Rule1[F ~> Id] = rule(True | False)
    def Literal: Rule1[F ~> Id] = rule(Number | Bool)
  }

  trait LetIn[T] {
    def let(n: String, v: T): T
    def in(s: Seq[T], e: T): T
  }

  trait LetAst[F[_]] extends LetIn[F ~> Id] {
    def alg[A](f: F[A]): LetIn[A]
    object LetAst {
      type Ast = F ~> Id
      case class Let(n: String, v: Ast) extends Ast {
        def apply[A](fa: F[A]): Id[A] =
          alg(fa).let(n, v(fa))
      }
      case class In(s: Seq[Ast], e: Ast) extends Ast {
        def apply[A](fa: F[A]): Id[A] =
          alg(fa).in(s.map(_ (fa)), e(fa))
      }
    }
    import LetAst._
    def let(n: String, v: Ast): Ast = Let(n, v)
    def in(s: Seq[Ast], e: Ast): Ast = In(s, e)
  }

  trait LetParser[F[_]] extends LetAst[F] with AParser {
    self: Parser =>
    import LetAst._
    def Term: Rule1[Ast]
    def Let: Rule1[Ast] = rule(Name ~
      sp ~ "=" ~ sp ~ Term ~> (let _))
    def LetIn: Rule1[Ast] = rule("let" ~ sp ~
      oneOrMore(Let) ~ sp ~ "in" ~ sp ~ Term ~> (in _))
  }

  trait ShowLet extends LetIn[String] {
    override def let(n: String, v: String): String = s"$n = $v"
    override def in(s: Seq[String], e: String): String =
      s"let ${s.mkString("\n")} in $e"
  }

  trait IfSym[T] {
    def _if(boo: T, tru: T, fals: T): T
  }

  trait IfAst[F[_]] extends IfSym[F ~> Id] {
    def alg[A](f: F[A]): IfSym[A]
    object IfAst {
      type Ast = F ~> Id
      case class If(boo: Ast, tru: Ast, fals: Ast) extends Ast {
        def apply[A](fa: F[A]): Id[A] =
          alg(fa)._if(boo(fa), tru(fa), fals(fa))
      }
    }
    import IfAst._
    def _if(boo: Ast, tru: Ast, fals: Ast): Ast =
      If(boo, tru, fals)
  }

  trait IfParser[F[_]] extends IfAst[F] with AParser {
    self: Parser =>
    import IfAst._
    def Term: Rule1[Ast]
    def If: Rule1[Ast] = rule("if" ~ sp ~ "(" ~ sp ~ Term ~ sp ~ ")"
      ~ sp ~ Term ~ sp ~ "else" ~ sp ~ Term ~> (_if _))
  }

  trait ShowIf extends IfSym[String] {
    def _if(boo: String, tru: String,
            fals: String): String =
      s"if($boo) ${tru} else ${fals}"
  }

  trait Sym[T] extends Term[T]
    with Expr[T] with Literal[T]
    with LetIn[T] with IfSym[T]

  type SymAst = Sym ~> Id

  case class SymParser(input: ParserInput)
    extends Parser
      with TermParser[Sym]
      with ExprParser[Sym]
      with LiteralParser[Sym]
      with LetParser[Sym]
      with IfParser[Sym] {
    self: Parser =>
    def alg[A](f: Sym[A]): Sym[A] = f
    def parser = rule(sp ~ Term ~ sp ~ quiet(EOI))
    override def Term = rule(LetIn | If | App | Expr | NonApp)
    override def NonApp = rule(Parens(Term) | super.NonApp | Literal)
    override def Operand = NonApp
  }

  def symParser(input: ParserInput) =
    SymParser(input).parser.run()

  object ShowSym extends Sym[String]
    with ShowTerm with ShowExpr
    with ShowLiteral with ShowLet with ShowIf {
    def apply(ast: SymAst): String = ast(this)
  }

  case class VM(env: Map[String, VMVal] = Map(),
                st: List[VMVal] = List())

  type EvalVM = VM => VM
  type VMVal = (EvalVM, Long)

  trait TermVM extends Term[EvalVM] {
    def ref(n: String): EvalVM = m => m.copy(st = m.env(n) :: m.st)
    def lam(n: String, t: EvalVM): EvalVM = m =>
      t(m.copy(env = m.env + (n -> m.st.head), st = m.st.tail))
    def app(f: EvalVM, a: EvalVM): EvalVM = m => f(m).st.head._1(a(m))
  }

  trait ExprVM extends Expr[EvalVM] {
    def oper(a: EvalVM, f: String, b: EvalVM): EvalVM =
      m => m.copy(st = (identity[VM] _,
        intOp(f)(a(m).st.head._2, b(m).st.head._2)) :: m.st)
    def parens(a: EvalVM): EvalVM = m => a(m)
    def intOp(f: String)(a: Long, b: Long): Long = f match {
      case "+" => a + b
      case "-" => a - b
      case "*" => a * b
      case "/" => a / b
    }
  }

  trait LiteralVM extends Literal[EvalVM] {
    def int(i: Long): EvalVM = m => m.copy(st =
      (identity[VM] _, i) :: m.st)
    def bool(b: Boolean): EvalVM = int(if (b) 1 else 0)
  }

  trait LetVM extends LetIn[EvalVM] {
    def let(n: String, v: EvalVM): EvalVM = m =>
      m.copy(env = m.env + (n -> (v, 0)))
    def in(s: Seq[EvalVM], e: EvalVM): EvalVM = m =>
      e(s.foldRight(m)(_ (_)))
  }

  trait IfVM extends IfSym[EvalVM] {
    def _if(boo: EvalVM, tru: EvalVM, fals: EvalVM): EvalVM =
      m => if (boo(m).st.head._2 != 0) tru(m) else fals(m)
  }

  object SymVM extends Sym[EvalVM]
    with TermVM with ExprVM
    with LiteralVM with LetVM with IfVM {
    def apply(ast: SymAst, args: Long*) =
      ast(this)(VM(st = args.toList.map((x: Long) =>
        (identity[VM] _, x)))).st.head._2
  }

  def showAndEval(s: String, args: Long*) =
    symParser(s).map(x => ShowSym(x) +
      " " + args.mkString(" ") +
      " = " + SymVM(x, args: _*))


  val main = "(x -> y -> z -> y * ( x + y * z))"
  println(showAndEval(main, 1, 2, 10))

  val fact =
    """
      | let
      |   fact = x ->
      |     if(x)
      |       x * (fact (x - 1))
      |     else 1
      | in fact 20
    """.stripMargin

  println(showAndEval(fact))
}

//Success(( x -> y -> z -> y * ( x + y * z ) ) 1 2 10 = 42)
//Success(let fact = x -> if(x) x * ( fact ( x - 1 ) ) else 1 in fact 20  = 2432902008176640000)