Full-gist.scala

import com.scalakata._

import scala.reflect._
@instrument class Playground {
  
  Prog.main()

  
}

/*
* Internal
*********************************************
*/
trait Expr {
  type Exp[+T]
  def const[T](x:T): Exp[T]
}

trait StagedExpr extends Expr {
  sealed trait Exp[+T]
  case class Const[T](x:T) extends Exp[T]
  def const[T](x: T) = Const(x)
}

trait UnstagedExpr extends Expr {
  type Exp[+T] = T
  def const[T](x:T) = x
}


trait Lift[A, B] {
  def lift(x:A):B
}

trait Top {
  
  this: Expr => 
  
  trait StagedTop  {
    type Internal
    def exp: Exp[Internal]
  }
  
  trait Stagable[S <: StagedTop] {
    def ct: ClassTag[S#Internal]
    def fromExp(exp: Exp[S#Internal]): S
    val lift: Lift[S#Internal, S] = new Lift[S#Internal, S] {
      def lift(x: S#Internal):S = fromExp(const[S#Internal](x))
    }
  }
}


trait Staged extends Top with StagedExpr
trait Unstaged extends Top with UnstagedExpr


/*
* Booleans
*********************************************
*/
trait Booleans {
  this: Top => 
  type Bool <: BooleanOps with StagedTop { type Internal = Boolean}

  trait BooleanOps {
    def &&(b: Bool): Bool
    def ||(b: Bool): Bool
    def unary_!(): Bool
  }
  
  def not(b: Bool): Bool
  
  def stagableBool: Stagable[Bool] 
  def liftBool: Lift[Boolean, Bool] 

}

trait BooleansExp extends Booleans {
  this: Staged => 
  case class BoolAnd(x: Exp[Boolean], y: Exp[Boolean]) extends Exp[Boolean]
  case class BoolOr(x: Exp[Boolean], y: Exp[scala.Boolean]) extends Exp[scala.Boolean]
  case class BoolNot(x: Exp[scala.Boolean]) extends Exp[scala.Boolean]


  case class Bool(exp: Exp[scala.Boolean]) extends BooleanOps with StagedTop {
    type Internal = Boolean
    def &&(b: Bool) = Bool(bool_and applyOrElse((this.exp, b.exp), (BoolAnd.apply _).tupled)) 
    def ||(b: Bool) = Bool(bool_or applyOrElse((this.exp, b.exp), (BoolOr.apply _).tupled)) 
    def unary_! = Bool(bool_not applyOrElse(this.exp , BoolNot.apply)) 
  }
  
  val stagableBool = new Stagable[Bool] {
    def ct = classTag[Boolean]
    def fromExp(exp: Exp[scala.Boolean]): Bool = Bool(exp)
  }

  def liftBool = stagableBool.lift

  def bool_and = PartialFunction.empty[(Exp[scala.Boolean], Exp[scala.Boolean]), Exp[scala.Boolean]]
  def bool_or = PartialFunction.empty[(Exp[scala.Boolean], Exp[scala.Boolean]), Exp[scala.Boolean]]
  def bool_not = PartialFunction.empty[Exp[scala.Boolean], Exp[scala.Boolean]]

  def not(b: Bool) = !b
}

trait BooleansOptImpl extends BooleansExp {
  this: Staged =>
  override def bool_and = {
    case (Const(cx), Const(cy)) => Const(cx && cy)
  }
}

trait BooleansLib extends Booleans {
  this: Unstaged =>
  
  case class Bool(exp: Boolean) extends BooleanOps with StagedTop {
    type Internal = Boolean
    def &&(b: Bool) = Bool(exp && b.exp)
    def ||(b: Bool) = Bool(exp || b.exp)
    def unary_! = Bool(!exp)
  }
  
  val stagableBool = new Stagable[Bool] {
    def ct = classTag[Boolean]
    def fromExp(exp: Exp[Boolean]): Bool = Bool(exp)
  }

  def liftBool = stagableBool.lift
  
  def not(b: Bool) = !b

}

trait BooleansAPI {
  val x: Booleans with Top
  type Boolean = x.Bool
  def not(b: Boolean): Boolean = x.not(b)
  implicit val stagableBool = x.stagableBool
  implicit val liftBool: Lift[scala.Boolean, Boolean]  = x.liftBool
  //If not present, plain typecheck fail at lifting array, not sure why ? 
  //TODO: inspect why
  implicit val boolCT: ClassTag[x.Bool] = classTag[Boolean]
}


/*
* Arrays
*********************************************
* The array indices are Boolean instead of Int for sake of simplicity
* Else I would have to implement Int as well
*/

trait Arrays {
  this: Booleans with Top => 
  type AArray[T  <: StagedTop] <: AArrayOps[T] with StagedTop { type Internal = scala.Array[T#Internal]}

  trait AArrayOps[T] {
    def update(x: Bool, y: T): AArrayOps[T]
    def get(x: Bool): T
  }
  
  def stagableArray[T <: StagedTop : Stagable]: Stagable[AArray[T]]
  def liftArray[T <: StagedTop : Stagable]: Lift[scala.Array[T#Internal], AArray[T]] 
  def liftStagedArray[T <: StagedTop : Stagable]: Lift[scala.Array[T], AArray[T]] 

}

trait ArraysExp extends Arrays {
  this: BooleansExp with Staged => 

  case class ArrayUpdate[T <: StagedTop](x: Exp[scala.Array[T#Internal]], y: Exp[scala.Boolean], z: Exp[T#Internal]) extends Exp[scala.Array[T#Internal]]
  case class ArrayGet[T <: StagedTop](x: Exp[scala.Array[T#Internal]], y: Exp[scala.Boolean]) extends Exp[T#Internal]
  case class ToArray[T <: StagedTop](x: scala.Array[Exp[T#Internal]]) extends Exp[scala.Array[T#Internal]]
  
  case class AArray[T <: StagedTop : Stagable](exp: Exp[scala.Array[T#Internal]]) extends AArrayOps[T] with StagedTop {
    type Internal = scala.Array[T#Internal]
    val s = implicitly[Stagable[T]]
    def get(b: Bool) = s.fromExp(array_get[T] applyOrElse((this.exp, b.exp), (ArrayGet.apply[T] _).tupled)) 
    def update(b: Bool, y: T) = AArray(array_update[T] applyOrElse((this.exp , b.exp, y.exp), (ArrayUpdate.apply[T] _).tupled))
    
  }
  
  def stagableArray[T <: StagedTop : Stagable] = new Stagable[AArray[T]] {
    val s = implicitly[Stagable[T]]
    implicit val tct = s.ct
    def ct = classTag[AArray[T]#Internal]
    def fromExp(exp: Exp[scala.Array[T#Internal]]): AArray[T] = AArray(exp)
  }

  def liftArray[T <: StagedTop : Stagable]: Lift[scala.Array[T#Internal], AArray[T]] = stagableArray[T].lift
  
  def liftStagedArray[T <: StagedTop : Stagable] = new Lift[scala.Array[T], AArray[T]] {
    val s = implicitly[Stagable[T]]
    implicit val tct = s.ct
    def lift(x:scala.Array[T]) = AArray(ToArray(x.map(_.exp)))
  }

  def array_get[T <: StagedTop] = PartialFunction.empty[(Exp[scala.Array[T#Internal]], Exp[scala.Boolean]), Exp[T#Internal]]
  def array_update[T <: StagedTop]  = PartialFunction.empty[(Exp[scala.Array[T#Internal]], Exp[scala.Boolean], Exp[T#Internal]), Exp[scala.Array[T#Internal]]]
  

}

trait ArraysLib extends Arrays {
  this: BooleansLib with Unstaged =>
  
 case class AArray[T <: StagedTop: Stagable](exp: Array[T#Internal])  extends AArrayOps[T] with StagedTop {
    type Internal = scala.Array[T#Internal]
    val s = implicitly[Stagable[T]]
    def get(b: Bool) = {
      val e = 
        if (b.exp == true)
          exp(0)
        else
          exp(1)
      s.fromExp(e)
    }
   def update(b, y) = {
      if (b.exp == true) 
        exp(0) = y.exp
      else
        exp(1) = y.exp
      this
   }

 }
  
  def stagableArray[T <: StagedTop : Stagable] = new Stagable[AArray[T]] {
    val s = implicitly[Stagable[T]]
    implicit val tct = s.ct
    def ct = classTag[AArray[T]#Internal]
    def fromExp(exp: Exp[scala.Array[T#Internal]]): AArray[T] = AArray(exp)
  }

  def liftArray[T <: StagedTop : Stagable]: Lift[scala.Array[T#Internal], AArray[T]] = 
    stagableArray[T].lift
  
  def liftStagedArray[T <: StagedTop : Stagable] = new Lift[scala.Array[T], AArray[T]] {
    val s = implicitly[Stagable[T]]
    implicit val tct = s.ct
    def lift(x:scala.Array[T]) = AArray(x.map(_.exp).toArray)
  }
  
}


trait ArraysAPI {
  val x: Arrays with Top
  type Array[T <: x.StagedTop] = x.AArray[T]  
  implicit def stagableArray[T <: x.StagedTop : x.Stagable] = x.stagableArray[T]
  implicit def liftArray[T <: x.StagedTop : x.Stagable]  = x.liftArray[T]
  implicit def liftStagedArray[T <: x.StagedTop : x.Stagable]  = x.liftStagedArray[T]
}

/*
* If-then-else
*********************************************
*/
trait IfThenElse {
  this: Booleans with Top =>
  def __ifThenElse[T <: StagedTop : Stagable](cond: Bool, thenp: =>T, elsep: =>T): T
}

trait IfThenElseExp extends IfThenElse {
  this: Booleans with Staged =>
  
  case class IfThenElse[T <: StagedTop](cond: Exp[Boolean], thenp: Exp[T#Internal], elsep: Exp[T#Internal]) extends Exp[T#Internal]
  
  def __ifThenElse[T <: StagedTop : Stagable](cond: Bool, thenp: =>T, elsep: =>T): T = {
    val s = implicitly[Stagable[T]]
    s.fromExp(ifthenelse[T] applyOrElse((cond.exp, thenp.exp, elsep.exp), (IfThenElse.apply[T] _).tupled))
  }
  
  def ifthenelse[T <: StagedTop : Stagable] = PartialFunction.empty[(Exp[Boolean], Exp[T#Internal], Exp[T#Internal]), Exp[T#Internal]]
    
}

trait IfThenElseOptImpl extends IfThenElseExp {
  this: Booleans with Staged => 
  override def ifthenelse[T <: StagedTop : Stagable] = {
    case (Const(true), x, _) => x
    case (Const(false), _, x) => x 
  }
}

trait IfThenElseLib extends IfThenElse {
  this: BooleansLib with Unstaged => 
  
  def __ifThenElse[T <: StagedTop : Stagable](cond: Bool, thenp: =>T, elsep: =>T): T = {
    if (cond.exp)
      thenp
    else
      elsep
  }
}

trait IfThenElseAPI {
  val x: IfThenElse with Booleans with Top
  
  //The lifting doesn't rely on an implicit lift but on the knowledge that a StagedTop can Lift.
  //This is nice
  def __ifThenElse[T <: x.StagedTop : x.Stagable](cond: x.Bool, thenp: T, elsep: T): T = 
    x.__ifThenElse(cond, thenp, elsep)
  
  def __ifThenElse[T <: x.StagedTop : x.Stagable](cond: x.Bool, thenp: T#Internal, elsep: T): T = {
    val s = implicitly[x.Stagable[T]]    
    x.__ifThenElse(cond, s.lift.lift(thenp), elsep)
  }
  
  def __ifThenElse[T <: x.StagedTop : x.Stagable](cond: x.Bool, thenp: T, elsep: T#Internal): T = {
    val s = implicitly[x.Stagable[T]] 
    x.__ifThenElse(cond, thenp, s.lift.lift(elsep))
  }
  
  def __ifThenElse[T <: x.StagedTop : x.Stagable](cond: x.Bool, thenp: T#Internal, elsep: T#Internal): T = {
    val s = implicitly[x.Stagable[T]] 
    x.__ifThenElse(cond, s.lift.lift(thenp), s.lift.lift(elsep))
  }
}

/*
* Equal
*********************************************
*/

trait Equals {
  this: Booleans with Top =>
  def __equals[T <: StagedTop](a: T, b: T): Bool 
}

trait EqualsExp extends Equals {
  this: BooleansExp with Staged =>
  
  case class Equals[T <: StagedTop](x: Exp[T#Internal], y: Exp[T#Internal]) extends Exp[Boolean]
  
  def __equals[T <: StagedTop](x: T, y:T): Bool = {
    Bool(equals[T] applyOrElse((x.exp, y.exp), (Equals.apply[T] _).tupled))
  }
  
  def equals[T <: StagedTop] = PartialFunction.empty[(Exp[T#Internal], Exp[T#Internal]), Exp[Boolean]]
    
}

trait EqualsOptImpl extends EqualsExp {
  this: BooleansExp with Staged => 
  override def equals[T <: StagedTop] = {
    case (x, y) if x == y => const(true)
  }
}

trait EqualsLib extends Equals{
  this: BooleansLib with Unstaged => 
  
  def __equals[T <: StagedTop](a: T, b: T): Bool = 
    Bool(a.exp == b.exp)
  
}

trait EqualsAPI {
  val x: Equals with Booleans with Top
  
  def __equals[T <: x.StagedTop](e1: T, e2: T): x.Bool = 
    x.__equals(e1, e2)
  
  def __equals[T <: x.StagedTop : x.Stagable](e1: T, e2: T#Internal): x.Bool = {
    val s = implicitly[x.Stagable[T]]    
    x.__equals(e1, s.lift.lift(e2))
  }
  
  def __equals[T <: x.StagedTop : x.Stagable](e1: T#Internal, e2: T): x.Bool = {
    val s = implicitly[x.Stagable[T]] 
    x.__equals(s.lift.lift(e1), e2)
  }
  
  //No need to handle T#Internal T#Internal as this should be the default scala == and handled by the macro virtualization
  
}

/*
* Frontend
*********************************************
*/


//Common DSL
trait DSL extends BooleansAPI with ArraysAPI with IfThenElseAPI with EqualsAPI {
  val x: Arrays with Booleans with IfThenElse with Equals with Top
  
  implicit def liftF[A,B](x:A)(implicit l: Lift[A, B]):B =
    l.lift(x)
  
  //That almost work except for array and then the combination of the two above confuse the compiler :<
  //But it shows how flexible is the type system
  //implicit def liftS[A, B <: x.StagedTop{ type Internal = A }](y:A)(implicit l: x.Stagable[B]):B =
  //  l.lift.lift(y)
}


//Stage backend
trait Stage extends DSL {
  override lazy val x = new BooleansOptImpl with ArraysExp with IfThenElseOptImpl with EqualsOptImpl with Staged {}//with ...
}

//Lib backend
trait Lib extends DSL {
  override lazy val x = new BooleansLib with ArraysLib with IfThenElseLib with EqualsLib with Unstaged {}
}


object Prog {
  val ir = new Lib {}
  //To switch to Staging, uncomment next line
  //val ir = new Stage {}
  import ir._
  def main() = {
    val a: Boolean = true
    val b: Boolean = false
    val c: Array[Boolean] = Array(a, b)
    val d = a && c.get(b)
    __ifThenElse(a, __equals(d, d), not(d))

  }
  //compiler(ir)(main)
}


//For plain typecheck we don't need even need any implementation
trait TypeCheck extends DSL {
  def main() = {
    val a: Boolean = true
    val b: Boolean = false
    val c: Array[Boolean] = Array(a, b)
    val d = a && c.get(b)
    __ifThenElse(a, __equals(d, d), not(d))
  }
}