macroInterpreter.scala

package eu.shiftforward.macroInterpreter

import scala.annotation.tailrec
import scala.reflect.runtime.universe._
import scala.collection._

sealed trait Expression[-ST, +A]

abstract class Atom[-ST, A](implicit val ev: Liftable[A]) extends Expression[ST, A] { val value: A }

case class IntegerAtom[-ST](value: Int) extends Atom[ST, Int]
case class StringAtom[-ST](value: String) extends Atom[ST, String]
case class BooleanAtom[-ST](value: Boolean) extends Atom[ST, Boolean]

sealed trait Operator[-ST, +A] extends Expression[ST, A]

case class Accessor[+B, A](k: Expression[B => A, B]) extends Operator[B => A, A]

sealed abstract class BinaryOp[-ST, +A](lhs: Expression[ST, A], rhs: Expression[ST, A]) extends Operator[ST, A]

case class And[-ST](lhs: Expression[ST, Boolean], rhs: Expression[ST, Boolean]) extends BinaryOp[ST, Boolean](lhs, rhs)
case class Sum[-ST](lhs: Expression[ST, Int], rhs: Expression[ST, Int]) extends BinaryOp[ST, Int](lhs, rhs)
case class Minus[-ST](lhs: Expression[ST, Int], rhs: Expression[ST, Int]) extends BinaryOp[ST, Int](lhs, rhs)
case class Concat[-ST](lhs: Expression[ST, String], rhs: Expression[ST, String]) extends BinaryOp[ST, String](lhs, rhs)
case class In[-ST, A](elem: Expression[ST, A], set: Set[A])(implicit val ev: Liftable[A]) extends Operator[ST, Boolean]

object DSL {
  import scala.reflect.runtime.currentMirror
  import scala.tools.reflect.ToolBox

  private[this] val toolbox = currentMirror.mkToolBox()
  private[this] val symbolTableTerm = TermName(s"r")

  @tailrec def optimize[ST, A](exp: Expression[ST, A]): Expression[ST, A] = {
    val optimized = optimizeStep(exp)
    if (optimized == exp) exp
    else optimize(optimized)
  }

  private def optimizeStep[ST, A](x: Expression[ST, A]): Expression[ST, A] = x match {
    case Sum(IntegerAtom(a), IntegerAtom(b)) => IntegerAtom(a + b)
    case Minus(IntegerAtom(a), IntegerAtom(b)) => IntegerAtom(a - b)
    case Sum(a, b) => Sum(optimizeStep[ST, Int](a), optimizeStep[ST, Int](b))
    case Minus(a, b) => Minus(optimizeStep[ST, Int](a), optimizeStep[ST, Int](b))
    case op @ In(a: Atom[_, _], b) => BooleanAtom(b.contains(a.value))
    case op @ In(a: Expression[ST @unchecked, A @unchecked], b) => implicit val ev = op.ev; In(optimizeStep(a), b)
    case _ => x
  }

  def execute[ST, A](x: Expression[ST, A])(implicit scope: ST): A = x match {
    case node: Atom[ST, A] => node.value
    case Sum(a, b) => execute[ST, Int](a) + execute[ST, Int](b)
    case Minus(a, b) => execute[ST, Int](a) - execute[ST, Int](b)
    case Concat(a, b) => execute[ST, String](a) + execute[ST, String](b)
    case Accessor(k) => scope.asInstanceOf[Any => A](execute(k))
    case In(k, x) => x.contains(k)
  }

  private def genCode[ST, A](x: Expression[ST, A])(implicit preamble: mutable.ListBuffer[Tree]): Tree = x match {
    case node: Atom[ST, A] => implicit val ev = node.ev; q"${node.value}"
    case Sum(a, b) => q"${genCode(a)} + ${genCode(b)}"
    case Minus(a, b) => q"${genCode(a)} - ${genCode(b)}"
    case Concat(a, b) => q"${genCode(a)} + ${genCode(b)}"
    case Accessor(k) => q"$symbolTableTerm(${genCode(k)})"
    case op @ In(elem, set) =>
      implicit val ev = op.ev
      val term = internal.reificationSupport.freshTermName("s$")
      preamble += q"val $term = Set(..$set)"
      q"$term.contains(${genCode(elem)})"
  }

  def compile[ST: TypeTag, A](x: Expression[ST, A]): (ST => A) = {
    implicit val preamble = mutable.ListBuffer[Tree]()
    val generatedCode = genCode(optimize(x))
    val stType = typeOf[ST]
    // Optimization step may be redundant due to Scala's compiler optimizations
    val code = q"""
      ..$preamble
      ($symbolTableTerm: $stType) => { ..$generatedCode }
    """
    println(showCode(code))
    toolbox.eval(code).asInstanceOf[ST => A]
  }
}

object Main extends App {
  type ST = (String) => Int

  val exp = Sum(Minus(Sum(IntegerAtom(8), IntegerAtom(7)), Sum(IntegerAtom(6), Accessor[String, Int](StringAtom("height")))), Accessor[String, Int](StringAtom("age")))
  val exp2 = In(Sum(Accessor[String, Int](StringAtom("cenas")), Sum(IntegerAtom(1), IntegerAtom(2))), Set(5, 4, 3, 2, 1))
  val fc = DSL.compile(exp)
  val fc2 = DSL.compile(exp2)
  val row = Map("age" -> 25, "height" -> 180, "cenas" -> 0)

  println("-----")
  (1 to 5).foreach { run =>
    val t = System.currentTimeMillis
    var i = 0
    while (i < 10000000) { i += 1; DSL.execute(exp2)(row) }
    println(s"Elapsed millis run {$run}: ${System.currentTimeMillis - t}")
  }

  println("-----")
  (1 to 5).foreach { run =>
    val t = System.currentTimeMillis
    var i = 0
    while (i < 10000000) { i += 1; fc2(row) }
    println(s"Elapsed millis run {$run}: ${System.currentTimeMillis - t}")
  }
}