refactor-ir-cake: removed Library trait (merged with Scalan)

sc/jvm/src/test/scala/sigmastate/InterpreterReflectionGeneratorTests.scala

@@ -37,7 +37,7 @@ class InterpreterReflectionGeneratorTests extends AnyPropSpec with Matchers {
   }
 
   property("inner class") {
-    val ctx = null.asInstanceOf[scalan.Library] // ok! type level only
+    val ctx = null.asInstanceOf[scalan.Scalan] // ok! type level only
     val clazz = classOf[ctx.Coll.CollElem[_, _]]
     registerClassEntry(clazz,
       constructors = Array(

sc/shared/src/main/scala/scalan/GraphIRReflection.scala

@@ -24,7 +24,7 @@ object GraphIRReflection {
   val reflection = ReflectionData
 
   { val clazz = classOf[wrappers.scala.WOptions#WOption[_]]
-    val ctx = null.asInstanceOf[scalan.Library] // ok! type level only
+    val ctx = null.asInstanceOf[scalan.Scalan] // ok! type level only
     registerClassEntry(clazz,
       methods = Map(
         mkMethod(clazz, "filter", Array[Class[_]](classOf[Base#Ref[_]])) { (obj, args) =>
@@ -123,7 +123,7 @@ object GraphIRReflection {
   }
 
   { val clazz = classOf[Colls#CollBuilder]
-    val ctx = null.asInstanceOf[Library] // ok! type level only
+    val ctx = null.asInstanceOf[Scalan] // ok! type level only
     registerClassEntry(clazz,
       methods = Map(
         mkMethod(clazz, "xor", Array[Class[_]](classOf[Base#Ref[_]], classOf[Base#Ref[_]])) { (obj, args) =>
@@ -145,7 +145,7 @@ object GraphIRReflection {
 
   {
     val clazz = classOf[Colls#Coll[_]]
-    val ctx = null.asInstanceOf[Library] // ok! type level only
+    val ctx = null.asInstanceOf[Scalan] // ok! type level only
     registerClassEntry(clazz,
       methods = Map(
         mkMethod(clazz, "append", Array[Class[_]](classOf[Base#Ref[_]])) { (obj, args) =>

sc/shared/src/main/scala/scalan/Scalan.scala

@@ -2,6 +2,10 @@ package scalan
 
 import scalan.primitives._
 import scalan.staged.Transforming
+import sigma.CollsModule
+import sigma.data.{Nullable, RType}
+import sigma.util.MemoizedFunc
+import special.wrappers.WrappersModule
 
 /** Aggregate cake with all inter-dependent modules assembled together.
   * Each instance of this class contains independent IR context, thus many
@@ -16,7 +20,7 @@ import scalan.staged.Transforming
   * in classed derived from `Scalan`, this is significant benefit over
   * *everything is global* design.
   */
-class Scalan
+ trait Scalan
   extends TypeDescs
   with MethodCalls
   with Tuples
@@ -33,4 +37,116 @@ class Scalan
   with Entities
   with Modules
   with DefRewriting
+  with WrappersModule
+  with CollsModule {
+
+  import WOption._
+  import WRType._
+  import Coll._; import CollBuilder._;
+  import WSpecialPredef._
+
+  type LazyRep[T] = MutableLazy[Ref[T]]
+
+  private val _liftElemMemo = new MemoizedFunc({
+    case eT: Elem[t] =>
+      val lT = Liftables.asLiftable[Any, t](eT.liftable)
+      liftableRType(lT).lift(eT.sourceType.asInstanceOf[RType[Any]])
+  })
+  implicit def liftElem[T](eT: Elem[T]): Ref[WRType[T]] = {
+    _liftElemMemo(eT).asInstanceOf[Ref[WRType[T]]]  // asRep cannot be used for AnyRef
+  }
+
+  private val _specialPredef: LazyRep[WSpecialPredefCompanionCtor] = MutableLazy(RWSpecialPredef.value)
+  def specialPredef: Ref[WSpecialPredefCompanionCtor] = _specialPredef.value
+
+  override protected def onReset(): Unit = {
+    _specialPredef.reset()
+    _liftElemMemo.reset()
+    super.onReset()
+  }
+
+  val CM = CollMethods
+  private val CBM = CollBuilderMethods
+  private val WOptionM = WOptionMethods
+  private val SPCM = WSpecialPredefCompanionMethods
+
+  def colBuilder: Ref[CollBuilder]
+
+  object IsNumericToInt {
+    def unapply(d: Def[_]): Nullable[Ref[A] forSome {type A}] = d match {
+      case ApplyUnOp(_: NumericToInt[_], x) => Nullable(x.asInstanceOf[Ref[A] forSome {type A}])
+      case _ => Nullable.None
+    }
+  }
+  object IsNumericToLong {
+    def unapply(d: Def[_]): Nullable[Ref[A] forSome {type A}] = d match {
+      case ApplyUnOp(_: NumericToLong[_], x) => Nullable(x.asInstanceOf[Ref[A] forSome {type A}])
+      case _ => Nullable.None
+    }
+  }
+
+  override def rewriteDef[T](d: Def[T]) = d match {
+    case CM.length(ys) => ys.node match {
+      // Rule: xs.map(f).length  ==> xs.length
+      case CM.map(xs, _) =>
+        xs.length
+      // Rule: replicate(len, v).length => len
+      case CBM.replicate(_, len, _) =>
+        len
+      // Rule: Const[Coll[T]](coll).length =>
+      case CollConst(coll, _) =>
+        coll.length
+      // Rule: Coll(items @ Seq(x1, x2, x3)).length => items.length
+      case CBM.fromItems(_, items, _) =>
+        items.length
+      case _ => super.rewriteDef(d)
+    }
+
+    // Rule: replicate(l, x).zip(replicate(l, y)) ==> replicate(l, (x,y))
+    case CM.zip(CBM.replicate(b1, l1, v1), CBM.replicate(b2, l2, v2)) if b1 == b2 && l1 == l2 =>
+      b1.replicate(l1, Pair(v1, v2))
+
+    case CM.map(xs, _f) => _f.node match {
+      case IdentityLambda() => xs
+      case _ => xs.node match {
+        // Rule: replicate(l, v).map(f) ==> replicate(l, f(v))
+        case CBM.replicate(b, l, v: Ref[a]) =>
+          val f = asRep[a => Any](_f)
+          b.replicate(l, Apply(f, v, false))
+
+        // Rule: xs.map(f).map(g) ==> xs.map(x => g(f(x)))
+        case CM.map(_xs, f: RFunc[a, b]) =>
+          implicit val ea = f.elem.eDom
+          val xs = asRep[Coll[a]](_xs)
+          val g  = asRep[b => Any](_f)
+          xs.map[Any](fun { x: Ref[a] => g(f(x)) })
+
+        case _ => super.rewriteDef(d)
+      }
+    }
+
+    case WOptionM.getOrElse(opt, _) => opt.node match {
+      // Rule: Some(x).getOrElse(_) ==> x
+      case SPCM.some(x) => x
+      case WOptionConst(Some(x), lA) => lA.lift(x)
+      case _ => super.rewriteDef(d)
+    }
+
+    case _ => super.rewriteDef(d)
+  }
+
+  override def invokeUnlifted(e: Elem[_], mc: MethodCall, dataEnv: DataEnv): Any = e match {
+    case _: CollElem[_,_] => mc match {
+      case CollMethods.map(_, f) =>
+        val newMC = mc.copy(args = mc.args :+ f.elem.eRange)(mc.resultType, mc.isAdapterCall)
+        super.invokeUnlifted(e, newMC, dataEnv)
+      case _ =>
+        super.invokeUnlifted(e, mc, dataEnv)
+    }
+    case _ =>
+      super.invokeUnlifted(e, mc, dataEnv)
+  }
+
+
+}
 

sc/shared/src/main/scala/sigmastate/eval/SigmaLibrary.scala

@@ -1,10 +1,10 @@
 package sigmastate.eval
 
-import scalan.Library
+import scalan.Scalan
 import sigma.SigmaDslModule
 
 /** Main trait which represents sigma operations in graph IR cake. */
-trait SigmaLibrary extends Library
+trait SigmaLibrary extends Scalan
     with sigma.wrappers.WrappersModule
     with SigmaDslModule
 {

sc/shared/src/main/scala/special/collection/CollsUnit.scala

@@ -8,7 +8,7 @@ package sigma {
     * The semantics of each method is the same as in the original class, please look there
     * for details.
     */
-  trait Colls extends Base { self: Library =>
+  trait Colls extends Base { self: Scalan =>
     trait Coll[A] extends Def[Coll[A]] {
       implicit def eA: Elem[A];
       def length: Ref[Int];

sc/shared/src/main/scala/special/collection/impl/CollsImpl.scala

@@ -9,8 +9,8 @@ package impl {
   import sigma.reflection.{RClass, RMethod}
 
   // Abs -----------------------------------
-trait CollsDefs extends scalan.Scalan with Colls {
-  self: Library =>
+trait CollsDefs extends Base with Colls {
+  self: Scalan =>
 
   registerModule(CollsModule)
 
@@ -639,4 +639,4 @@ object CollsModule extends scalan.ModuleInfo("sigma", "Colls") {
 }
 }
 
-trait CollsModule extends sigma.impl.CollsDefs {self: Library =>}
+trait CollsModule extends sigma.impl.CollsDefs {self: Scalan =>}

sc/shared/src/main/scala/special/sigma/wrappers/WrappersModule.scala

@@ -1,4 +1,6 @@
 package sigma.wrappers
 
+import scalan.Scalan
+
 trait WrappersModule
-  extends special.wrappers.WrappersModule
+  extends special.wrappers.WrappersModule { self: Scalan => }

sc/shared/src/main/scala/special/wrappers/WrappersModule.scala

@@ -1,10 +1,11 @@
 package special.wrappers
 
+import scalan.Scalan
 import wrappers.scala.WOptionsModule
 import wrappers.scalan.WRTypesModule
 import wrappers.special.WSpecialPredefsModule
 
 trait WrappersModule
   extends WSpecialPredefsModule
   with WOptionsModule
-  with WRTypesModule
+  with WRTypesModule { self: Scalan => }

sc/shared/src/main/scala/wrappers/scala/WOptions.scala

@@ -3,7 +3,7 @@ package wrappers.scala {
 
   import special.wrappers.WrappersModule
 
-  trait WOptions extends Base { self: WrappersModule =>
+  trait WOptions extends Base { self: Scalan =>
     trait WOption[A] extends Def[WOption[A]] {
       implicit def eA: Elem[A];
       def isDefined: Ref[Boolean];

sc/shared/src/main/scala/wrappers/scala/impl/WOptionsImpl.scala

@@ -11,8 +11,8 @@ package impl {
   import sigma.reflection.{RClass, RMethod}
 
   // Abs -----------------------------------
-trait WOptionsDefs extends scalan.Scalan with WOptions {
-  self: WrappersModule =>
+trait WOptionsDefs extends Base with WOptions {
+  self: Scalan =>
 
 class WOptionCls extends EntityObject("WOption") {
   // entityConst: single const for each entity
@@ -240,4 +240,4 @@ object WOptionsModule extends scalan.ModuleInfo("wrappers.scala", "WOptions") {
 }
 }
 
-trait WOptionsModule extends wrappers.scala.impl.WOptionsDefs {self: WrappersModule =>}
+trait WOptionsModule extends wrappers.scala.impl.WOptionsDefs {self: Scalan =>}

sc/shared/src/main/scala/wrappers/scalan/WRTypes.scala

@@ -1,9 +1,8 @@
 package wrappers.scalan {
   import scalan._
-  import special.wrappers.WrappersModule
 
 
-  trait WRTypes extends Base { self: WrappersModule =>
+  trait WRTypes extends Base { self: Scalan =>
     trait WRType[A] extends Def[WRType[A]] {
       implicit def eA: Elem[A];
       def name: Ref[String]