Type-Changing-Example.scala

package ddbtoaster

import squid.utils._
import stagerwr2.compiler.Embedding
import Embedding.Predef._
import Embedding.Quasicodes._
import squid.ir.BottomUpTransformer
import squid.ir.FixPointRuleBasedTransformer
import squid.ir.FixPointTransformer
import squid.ir.RewriteAbort
import squid.ir.SimpleRuleBasedTransformer
import squid.lib.transparencyPropagating

import scala.language.existentials
import scala.collection.mutable.ArrayBuffer

class TypedSlice[+T] {
  def foreach(f: T => Unit): Unit = ???
}
import TypeChanging.Slice

class Store {
  def get(k: Any*): Slice = ???
  def getTyped[K,R](k: K): TypedSlice[R] = ???
}
object Store {
  @transparencyPropagating
  def apply(name:String): Store = ???
}

object TypeChanging extends App {
  
  type Slice = TypedSlice[Any]
  @transparencyPropagating def asEntryOf[T](e: Any, name: String, idx: Int): T = e.asInstanceOf[T]
  
  @transparencyPropagating def tupleElem[T](tup: Product, idx: Int): T = tup.productElement(idx).asInstanceOf[T]

{
  
  val pgrm = ir{
    val s = Store("store0")
    val ab = ArrayBuffer[Any]()
    val ab2 = ArrayBuffer[Any]()
    s.get(1,"xyz").foreach { e =>
      ab += asEntryOf[Double](e, "store0", 2)
    }
    //ab foreach println
    ab foreach { e =>
      ab2 += e
      ()
    }
    ab2 foreach println
  }
  
  println(pgrm)
  
  val schema = Map[String, List[IRType[_]] -> IRType[_]] {
    "store0" -> (List(irTypeOf[Int],irTypeOf[String]) -> irTypeOf[(Int,String,Double)])
  }
  
  def typeGet[C](name: String, args: Seq[IR[Any,C]]) = {
    val schKey -> (schVal:IRType[schVal]) = schema(name)
    assert(schKey.size == args.size)
    schKey -> args match {
      case List(t0:IRType[t0],t1:IRType[t1]) -> Seq(a0,a1) =>
        ir"Store(${Const(name)}).getTyped[($t0,$t1),$schVal](($a0.asInstanceOf[$t0],$a1.asInstanceOf[$t1]))"
    }
  }
  
  def typeArrayBuffer[T:IRType,C](arr: IR[ArrayBuffer[Any],C{val arr: ArrayBuffer[Any]}], body: IR[T,C{val arr: ArrayBuffer[Any]}]): IR[T,C] = {
    var arrTyp = Option.empty[IRType[_]]
    body analyse {
      case ir"$$arr += (($v:$t):Any)" =>
        // TODO check if arrTyp already defined
        arrTyp = Some(t)
    }
    arrTyp match {
      case Some(arrTyp:IRType[arrTyp]) =>
        val newArr = ir"newArr? : ArrayBuffer[arrTyp]"
        val body0 = body rewrite {
          case ir"$$arr += $v"           =>  ir"$newArr += $v.asInstanceOf[arrTyp]; $arr"
          case ir"$$arr.foreach[$t](x => $fbody)" =>  ir"$newArr foreach {e => val x = e; $fbody}"
          // ^ we have a special case above because the problem of the following is that it leaves the type of `f` as (Any=>Unit), which prevents type-specialization of consumers of the value
          case ir"$$arr.foreach[$t]($f)" =>  ir"$newArr foreach $f"
        }
        val body1 = body0 subs 'arr -> {System.err.println(s"Oops $body0"); throw RewriteAbort()}
        ir"val newArr = ArrayBuffer[arrTyp](); $body1"
      case None => ???
    }
  }
  
  object Transf extends Embedding.SelfTransformer with SimpleRuleBasedTransformer with BottomUpTransformer with FixPointTransformer {
    
    rewrite {
      
      case ir"Store(${Const(name)}).get($args*)" if schema isDefinedAt name =>
        typeGet(name, args)

      case ir"val $slice = ($s:Store).getTyped[$t,$r]($arg) : Slice; $body:$bt" if slice.typ =:= irTypeOf[Slice] =>
        println(s"Slice $slice")
        val newSlice = ir"slice? : TypedSlice[$r]" // Note: using the same 'slice' name here not a problem! it refines the hole type...
        val body0 = body rewrite {
          case ir"$$slice.foreach(e => $febody)" =>
            val febody0 = febody rewrite {
              case ir"asEntryOf[$tt]($e, ${Const(name)}, ${Const(idx)})" =>
                ir"tupleElem[$tt]((g? : $r).asInstanceOf[Product], ${Const(idx)})"
            }
            val febody1 = febody0 subs 'e -> {
              System.err.println(s"Oops2 $febody0"); Abort()
            }
            ir"$newSlice.foreach((g:$r) => $febody1)"
        }
        ir"val slice = ($s:Store).getTyped[$t,$r]($arg); $body0"

      case ir"val $arr: ArrayBuffer[Any] = ArrayBuffer(); $body:$bt" if arr.typ =:= irTypeOf[ArrayBuffer[Any]] =>
        typeArrayBuffer(arr, body)
      
    }
  }
  val pgrm0 = pgrm transformWith Transf
  
  def tupleElemRemove[T:IRType,C](tup:IR[Any,C],idx:Int) = {
    val res = tup match {
      case ir"$tup: ($ta,$tb)" => idx match {
        case 0 => ir"$tup._1"
        case 1 => ir"$tup._2"
      }
      case ir"$tup: ($ta,$tb,$tc)" => idx match {
        case 0 => ir"$tup._1"
        case 1 => ir"$tup._2"
        case 2 => ir"$tup._3"
      }
      case ir"$tup: Product" => assert(false, s"$tup $idx"); ??? // TODO
      case _ => tup
    }
    ir"$res.asInstanceOf[T]"
  }
  val pgrm1 = pgrm0 rewrite {
    case ir"tupleElem[$t]($tup,${Const(idx)})" => tupleElemRemove(tup,idx)(t)
  }
  
  println(pgrm1)
  
}}