LMS Structs

TestStruct.scala

import scala.virtualization.lms.common._
import scala.reflect.SourceContext
import scala.virtualization.lms.internal.GenericCodegen
import lms._

import java.io.PrintWriter
import java.io.StringWriter
import java.io.FileOutputStream

trait TestStruct extends Base with IfThenElse with Equal {

  /**
   * a ghost type, for demo
   * purposes
   */
  abstract class Box[T]

  /**
   * constructor for a box struct
   */
  def mkBox[T: Manifest](t: Rep[T]): Rep[Box[T]]

  /**
   * Simple conditional
   * Should generate precisely 3 branches
   * and not collapse the redundant condition
   */
  def conditionalTest(in: Rep[Int]): Rep[Int] = {
    if (in == unit(3)) unit(2)
    else if (in == unit(3)) unit(2)
    else unit(0)
  }

  /**
   * extended ifThenElse
   */
  def conditionalBoxTest(in: Rep[Int]): Rep[Box[Int]] = {
    if (in == unit(3)) mkBox(in)
    else if (in == unit(3)) mkBox(in)
    else mkBox(unit(2))
  }
}

trait TestStructExp
    extends TestStruct
    with BaseExp
    with IfThenElseExp
    with StructExp {

  def mkBox[T: Manifest](t: Rep[T]) =
    struct(classTag[Box[T]], "t" -> t)

}

trait TestStructGenBase extends GenericCodegen with BaseGenStruct {
  val IR: TestStructExp

  override def remap[A](m: Manifest[A]) = m.erasure.getSimpleName match {
    case "Box" => IR.structName(m)
    case _ => super.remap(m)
  }
}

class StructSuite extends FileDiffSuite {

  val prefix = "test-out/"

  def testStruct = {
    withOutFile(prefix + "box") {
      new TestStruct
          with EqualExpBridge
          with BooleanOpsExpOpt
          with IfThenElseFatExp
          with TestStructExp
          with MyScalaCompile { self =>

        val codegen = new ScalaGenBase
            //with ScalaGenFatStruct
            with TestStructGenBase
            with ScalaGenStruct
            with ScalaGenIfThenElseFat
            with ScalaGenEqual
            with ScalaGenBooleanOps { val IR: self.type = self }

        codegen.emitSource(conditionalTest _, "conditionalTest", new java.io.PrintWriter(System.out))
        codegen.reset

        /**
         * ======== GENERATES THE FOLLOWING CODE =====
         */
        /*****************************************
          Emitting Generated Code
        *******************************************/
        //  class conditionalTest extends ((Int)=>(Int)) {
        //  def apply(x0:Int): Int = {
        //  val x1 = x0 == 3
        //  val x3 = if (x1) {
        //  2
        //  } else {
        //  val x2 = if (x1) {
        //  2
        //  } else {
        //  0
        //  }
        //  x2
        //  }
        //  x3
        //  }
        //  }
        /*****************************************
          End of Generated Code
        *******************************************/


        codegen.emitSource(conditionalBoxTest _, "conditionalBoxTest", new java.io.PrintWriter(System.out))
        codegen.reset

        /**
         * ======== GENERATES THE FOLLOWING CODE =====
         */
        /*****************************************
          Emitting Generated Code
        *******************************************/
        // class conditionalBoxTest extends ((Int)=>(BoxInt)) {
        // def apply(x0:Int): BoxInt = {
        // val x1 = x0 == 3
        // val x5 = if (x1) {
        // val x2 = new BoxInt(x0)
        // x2
        // } else {
        // val x4 = if (x1) {
        // val x2 = new BoxInt(x0)
        // x2
        // } else {
        // val x3 = new BoxInt(2)
        // x3
        // }
        // x4
        // }
        // x5
        // }
        // }
        /*****************************************
          End of Generated Code
        *******************************************/

      }

      /**
       * A different, optimized codegen
       */
      new TestStruct
          with EqualExpBridge
          with BooleanOpsExpOpt
          with IfThenElseFatExp
          with TestStructExp
          with StructFatExpOptCommon
          with MyScalaCompile { self =>

        val codegen = new ScalaGenBase
            with ScalaGenFatStruct
            with TestStructGenBase
            with ScalaGenStruct
            with ScalaGenIfThenElseFat
            with ScalaGenEqual
            with ScalaGenBooleanOps { val IR: self.type = self }

        codegen.emitSource(conditionalTest _, "conditionalTestBis", new java.io.PrintWriter(System.out))
        codegen.reset

        /**
         * ======== GENERATES THE FOLLOWING CODE =====
         */
        /*****************************************
          Emitting Generated Code
        *******************************************/
        // class conditionalTestBis extends ((Int)=>(Int)) {
        // def apply(x0:Int): Int = {
        // val x1 = x0 == 3
        // val x3 = if (x1) {
        // 2
        // } else {
        // val x2 = if (x1) {
        // 2
        // } else {
        // 0
        // }
        // x2
        // }
        // x3
        // }
        // }
        /*****************************************
          End of Generated Code
        *******************************************/

        codegen.emitSource(conditionalBoxTest _, "conditionalBoxTestBis", new java.io.PrintWriter(System.out))
        codegen.reset

        /**
         * ======== GENERATES THE FOLLOWING CODE =====
         *
         * This code does not compile, note x5 is used
         * recursively inside the conditional expression
         */

        /*****************************************
          Emitting Generated Code
        *******************************************/
        // class conditionalBoxTestBis extends ((Int)=>(BoxInt)) {
        // def apply(x0:Int): BoxInt = {
        // val x1 = x0 == 3
        // // TODO: use vars instead of tuples to return multiple values
        // val (x5,x7) = if (x1) {
        // (x0,x0)
        // } else {
        // (2,x5)
        // }
        // val x8 = new BoxInt(x7)
        // x8
        // }
        // }
        /*****************************************
          End of Generated Code
        *******************************************/

      }
    }

    assertFileEqualsCheck(prefix + "box")
  }
}