Playing around with "Finally Tagless, Partially Evaluated" in Scala

gistfile1.scala

// Finally tagless lambda calculus
object Final {
  def varZ[A,B](env1: A, env2: B): A = env1
  def varS[A,B,T](vp: B => T)(env1: A, env2: B): T = vp(env2)
  def b[E](bv: Boolean)(env: E): Boolean = bv
  def lam[A,E,T](e: (A, E) => T)(env: E): A => T = x => e(x, env)
  def app[A,E,T](e1: E => A => T, e2: E => A)(env: E): T = e1(env)(e2(env))

  def testf1[E,A](env: E): Boolean =
    app(lam[Boolean,E,Boolean](varZ) _, b(true))(env)

  def testf3[A,B](env: (A, B)): A =
    app(lam(varS((varZ(_:A,_:B)).tupled)), b(true))(env)
}

// Lambda calculus represented by `R`, using HOAS
abstract class Symantics[R[_]] {
  def int(n: Int): R[Int]
  def bool(b: Boolean): R[Boolean]

  def lam[A,B](f: R[A] => R[B]): R[A => B]
  def app[A,B](f: R[A => B], a: R[A]): R[B]
  def fix[A](f: R[A] => R[A]): R[A]

  def add(a: R[Int], b: R[Int]): R[Int]
  def mul(a: R[Int], b: R[Int]): R[Int]
  def leq(a: R[Int], b: R[Int]): R[Boolean]
  def if_[A](b: R[Boolean], t: () => R[A], f: () => R[A]): R[A]
}

// Lambda calculus represented by `R`, using de Bruijn indexed variables
// represented by `VR`.
abstract class SymanticsB[R[_, _], VR[_]] {
  def vz[D,H]: R[(VR[D], H), D]
  def vs[A,D,H](r: R[H, D]): R[(A, H), D]

  def int[H](n: Int): R[H, Int]
  def bool[H](b: Boolean): R[H, Boolean]

  def lam[H,DA,DB](s: R[(VR[DA], H), DB]): R[H, DA => DB]
  def app[H,DA,DB](f: R[H, DA => DB], a: R[H, DA]): R[H, DB]
  def fix[H,DA,DB](f: R[(VR[DA => DB], H), DA => DB]): R[H, DA => DB]

  def add[H](a: R[H,Int], b: R[H,Int]): R[H,Int]
  def mul[H](a: R[H,Int], b: R[H,Int]): R[H,Int]
  def leq[H](a: R[H,Int], b: R[H,Int]): R[H,Boolean]
  def if_[H,DA](b: R[H,Boolean],
                  t: () => R[H,DA],
                  f: () => R[H,DA]): R[H,DA]
}

import language.higherKinds

// Test the HOAS encoding, given a representation
class TestHOAS[R[_]](S: Symantics[R]) {
  import S._

  def test1: R[Boolean] = app(lam((x: R[Boolean]) => x), bool(true))

  def testpowfix: R[Int => Int => Int] =
    lam((x: R[Int]) => fix((self: R[Int => Int]) => lam((n: R[Int]) =>
        if_(leq(n, int(0)), () => int(1),
            () => mul(x, app(self, add(n, (int(-1)))))))))

  def testpowfix7: R[Int => Int] =
    lam((x: R[Int]) => app(app(testpowfix, x), int(7)))
}

// Representing HOAS programs as Scala values
object R extends Symantics[Function0] {
  def int(n: Int) = () => n
  def bool(b: Boolean) = () => b

  def lam[A,B](f: (() => A) => (() => B)) = () => a => f(() => a)()
  def app[A,B](f: () => A => B, a: () => A) = () => f()(a())
  def fix[A](f: (() => A) => (() => A)): () => A = {
    lazy val x: () => A = () => f(x)()
    x
  }
  def add(e1: () => Int, e2: () => Int) = () => e1() + e2()
  def mul(e1: () => Int, e2: () => Int) = () => e1() * e2()
  def leq(e1: () => Int, e2: () => Int) = () => e1() <= e2()
  def if_[A](eb: () => Boolean, et: () => () => A, ee: () => () => A) = if (eb()) et() else ee()
}

// Some definitions
object Defs {
  type IntF[A] = Int
  type Repr[H,DV] = H => DV
  type Vr[D] = D
}

import Defs._

// Representing HOAS programs as their length
object L extends Symantics[IntF] {
  def int(n: Int) = 1
  def bool(b: Boolean) = 1
  def lam[A,B](f: Int => Int) = f(0) + 1
  def app[A,B](f: Int, n: Int) = f + n + 1
  def fix[A](f: Int => Int) = f(0) + 1
  def add(e1: Int, e2: Int) = e1 + e2 + 1
  def mul(e1: Int, e2: Int) = e1 + e2 + 1
  def leq(e1: Int, e2: Int) = e1 + e2 + 1
  def if_[A](eb: Int, et: () => Int, ef: () => Int) = eb + et() + ef() + 1
}

// Representing de Bruijn programs as Scala functions from an environment
object RB extends SymanticsB[Repr, Vr] {
  type R[A,B] = Repr[A,B]
  type VR[A] = Vr[A]

  def vz[D,H]: R[(VR[D], H), D] = {
    case (x, _) => x
  }
  def vs[A,D,H](v: R[H, D]): R[(A, H), D] = {
    case (_, h) => v(h)
  }

  def int[H](n: Int): R[H, Int] = _ => n
  def bool[H](b: Boolean): R[H, Boolean] = _ => b

  def lam[H,DA,DB](f: R[(VR[DA], H), DB]): R[H, DA => DB] =
    h => x => f((x,h))
  def app[H,DA,DB](f: R[H, DA => DB], a: R[H, DA]): R[H, DB] =
    h => f(h)(a(h))
  def fix[H,DA,DB](f: R[(VR[DA => DB], H), DA => DB]): R[H, DA => DB] =
    h => {
      def self(n: DA): DB = f((self, h))(n)
      self(_)
    }

  def add[H](a: R[H,Int], b: R[H,Int]): R[H,Int] = h => a(h) + b(h)
  def mul[H](a: R[H,Int], b: R[H,Int]): R[H,Int] = h => a(h) * b(h)
  def leq[H](a: R[H,Int], b: R[H,Int]): R[H,Boolean] = h => a(h) <= b(h)
  def if_[H,DA](b: R[H,Boolean],
                  t: () => R[H,DA],
                  f: () => R[H,DA]): R[H,DA] =
    h => if (b(h)) t()(h) else f()(h)
}

// Testing a de Bruijn representation
class TestB[R[_,_], VR[_]](S: SymanticsB[R, VR]) {
  import S._

  def test1[H]: R[H, Boolean] =
    app(lam(vz), bool(true))

  def testpowfix[H]: R[H, Int => Int => Int] =
    lam(fix(lam(
      if_(leq(vz, int(0)), () => int(1),
        () => mul(vs(vs(vz)),
            app(vs(vz), add(vz, (int(-1)))))))))

  def testpowfix7[H]: R[H, Int => Int] =
    lam(app(app(testpowfix, vz), int(7)))
}

import scala.language.experimental.macros
import scala.reflect.macros.whitebox.Context


// *** THIS TIME WITH MACROS! *** //


// A tagless HOASed lambda calculus that evaluates to Scala at compile time
class Compiled[C <: Context](val c: C) {
  import c.universe._
  def int(x: Int) = c.Expr[Int](q"$x")
  def bool(x: Boolean) = c.Expr[Boolean](q"$x")
  def lam[A:c.WeakTypeTag,B](f: c.Expr[A] => c.Expr[B]) = {
    val x = newTermName(c.fresh)
    c.Expr[A => B](q"($x: ${c.weakTypeOf[A]}) => ${f(c.Expr(q"$x"))}")
  }
  def app[A,B](f: c.Expr[A => B], a: c.Expr[A]) =
    c.Expr[B](q"$f($a)")
  // Doesn't work. Suspect bug in macro expansion.
  //def fix[A:c.WeakTypeTag](f: c.Expr[A] => c.Expr[A]) =
  //  c.Expr[A](q"{lazy val x: ${c.weakTypeOf[A]} = ${f(c.Expr(q"x"))}; x}")
  def add(e1: c.Expr[Int], e2: c.Expr[Int]) =
    c.Expr[Int](q"$e1 + $e2")
  def mul(e1: c.Expr[Int], e2: c.Expr[Int]) =
    c.Expr[Int](q"$e1 * $e2")
  def leq(e1: c.Expr[Int], e2: c.Expr[Int]) =
    c.Expr[Boolean](q"$e1 <= $e2")
  def if_[A](eb: c.Expr[Boolean], et: () => c.Expr[A], ef: () => c.Expr[A]) =
    c.Expr[A](q"if ($eb) ${et()} else ${ef()}")
}

object TestC {
  def test1: Boolean = macro test1Impl
  def test1Impl(ctx: Context): ctx.Expr[Boolean] = {
    val C = new Compiled[ctx.type](ctx)
    import C._

    app(lam((x: c.Expr[Boolean]) => x), bool(true))
  }

  def test2: Int => Int = macro test2Impl
  def test2Impl(ctx: Context): ctx.Expr[Int => Int] = {
    val C = new Compiled[ctx.type](ctx)
    import C._

    lam((n: c.Expr[Int]) =>
      if_(leq(n, int(0)), () => int(1), () => mul(n, add(n, (int(-1))))))
  }
}

// Tagless de Bruijn indexed calculus compiled to Scala at compile time
class CompiledDB[C <: Context](val c: C) {
  import c.universe._
  type R[A,B] = A => c.Expr[B]
  type V[A] = c.Expr[A]

  def vz[D,H]: R[(V[D], H), D] = {
    case (x, _) => x
  }

  def vs[A,D,H](v: R[H, D]): R[(A, H), D] = {
    case (_, h) => v(h)
  }

  def int[H](n: Int): R[H, Int] = _ => c.Expr[Int](q"$n")
  def bool[H](b: Boolean): R[H, Boolean] = _ => c.Expr[Boolean](q"$b")

  def lam[H,DA:c.WeakTypeTag,DB](f: R[(V[DA], H), DB]): R[H, DA => DB] =
    h => {
      val x = newTermName(c.fresh)
      c.Expr[DA => DB](q"($x: ${c.weakTypeOf[DA]}) => ${f((c.Expr(q"$x"), h))}")
    }

  def app[H,DA,DB](f: R[H, DA => DB], a: R[H, DA]): R[H, DB] =
    h => c.Expr[DB](q"${f(h)}(${a(h)})")

  def fix[H,DA:c.WeakTypeTag,DB:c.WeakTypeTag](
    f: R[(V[DA => DB], H), DA => DB]): R[H, DA => DB] =
      h => c.Expr[DA => DB](
        q"{ def self(n: ${c.weakTypeOf[DA]}): ${c.weakTypeOf[DB]} = ${f((c.Expr(q"self(_)"), h))}(n); self _}")

  def add[H](a: R[H,Int], b: R[H,Int]): R[H,Int] =
    h => c.Expr[Int](q"${a(h)} + ${b(h)}")
  def mul[H](a: R[H,Int], b: R[H,Int]): R[H,Int] =
    h => c.Expr[Int](q"${a(h)} * ${b(h)}")
  def leq[H](a: R[H,Int], b: R[H,Int]): R[H,Boolean] =
    h => c.Expr[Boolean](q"${a(h)} <= ${b(h)}")
  def if_[H,DA](b: R[H,Boolean],
                  t: () => R[H,DA],
                  f: () => R[H,DA]): R[H,DA] =
    h => c.Expr[DA](q"if (${b(h)}) ${t()(h)} else ${f()(h)}")
}

object TestCDB {
  def test1: Boolean = macro test1Impl
  def test1Impl(ctx: Context): ctx.Expr[Boolean] = {
    val C = new CompiledDB[ctx.type](ctx)
    import C._

    app[Unit, Boolean, Boolean](lam(vz), bool(true))(())
  }

  def test2: Int => Int = macro test2Impl
  def test2Impl(ctx: Context): ctx.Expr[Int => Int] = {
    val C = new CompiledDB[ctx.type](ctx)
    import C._

    lam[Unit, Int, Int](if_(leq(vz, int(0)),
            () => int(1),
            () => mul(vz, add(vz, (int(-1)))))).apply(())
  }

  def testpowfix: Int => Int => Int = macro testpowfixImpl

  def testpowfixImpl(ctx: Context): ctx.Expr[Int => Int => Int] = {
    val C = new CompiledDB[ctx.type](ctx)
    import C._

    lam[Unit, Int, Int => Int](fix(lam(
      if_(leq(vz, int(0)), () => int(1),
        () => mul(vs(vs(vz)),
            app(vs(vz), add(vz, (int(-1))))))))).apply(())
  }
}