Streams: support flatten (issue #22)

nativelibs4java · Mar 22, 2015 · d36d231 · d36d231
1 parent 5155aff
commit d36d231
Show file tree

Hide file tree

Showing 5 changed files with 116 additions and 1 deletion.
diff --git a/Streams/src/main/scala/package.scala b/Streams/src/main/scala/package.scala
@@ -42,6 +42,16 @@ package streams
     /** For testing */
     private[streams] var quietWarnings = false
 
+    private[streams] def withQuietWarnings[A](a: => A): A = {
+      val old = quietWarnings
+      try {
+        quietWarnings = true
+        a
+      } finally {
+        quietWarnings = old
+      }
+    }
+
     def recursivelyOptimize[A : c.WeakTypeTag](c: Context)(a: c.Expr[A]): c.Expr[A] = {
       optimize[A](c)(a, recurse = true)
     }

diff --git a/Streams/src/main/scala/streams/matchers/Strippers.scala b/Streams/src/main/scala/streams/matchers/Strippers.scala
@@ -31,7 +31,7 @@ private[streams] trait Strippers
 
   private[this] lazy val OptionModule = rootMirror.staticModule("scala.Option")
 
-  private[this] object Option2Iterable {
+  private[streams] object Option2Iterable {
     def unapply(tree: Tree): Option[Tree] = Option(tree) collect {
       case q"$target.option2Iterable[${_}]($value)" if target.symbol == OptionModule =>
         value

diff --git a/Streams/src/main/scala/streams/ops/FlattenOps.scala b/Streams/src/main/scala/streams/ops/FlattenOps.scala
@@ -0,0 +1,85 @@
+package scalaxy.streams
+
+import impl.withQuietWarnings
+
+private[streams] trait FlattenOps
+    extends ClosureStreamOps
+    with CanBuildFromSinks
+    with Streams
+    with Strippers
+{
+  val global: scala.reflect.api.Universe
+  import global._
+
+  val SomeStream: Extractor[Tree, Stream]
+
+  object SomeFlattenOp extends StreamOpExtractor {
+    private[this] lazy val PredefSymbol = rootMirror.staticModule("scala.Predef")
+    override def unapply(tree: Tree) = Option(tree) collect {
+      case q"$target.flatten[$tpt]($asTraversable)" if {
+        asTraversable match {
+          case q"$predef.`$conforms`[$colTpt]" if predef.symbol == PredefSymbol =>
+            // println(s"colTpt = $colTpt, target.tpe = ${target.tpe}")
+            target.tpe match {
+              case TypeRef(_, _, List(internalColTpe)) =>
+                internalColTpe =:= colTpt.tpe
+              case _ =>
+                false
+            }
+
+          case Strip(Function(List(param), Option2Iterable(ref))) if param.symbol == ref.symbol =>
+            true
+        }
+      } =>
+        (target, FlattenOp(tpt.tpe))
+      // case _ if {
+      //   println("NOT A FLATTEN: " + tree)
+      //   false
+      // } =>
+      //   ???
+    }
+  }
+
+  case class FlattenOp(tpe: Type) extends StreamOp
+  {
+    override def describe = Some("flatten")
+
+    override def lambdaCount = 0
+
+    override def subTrees = Nil
+
+    override def canInterruptLoop = false
+
+    override def canAlterSize = true
+
+    override def transmitOutputNeedsBackwards(paths: Set[TuploidPath]) = paths
+
+    override val sinkOption = None
+
+    override def emit(input: StreamInput,
+                      outputNeeds: OutputNeeds,
+                      nextOps: OpsAndOutputNeeds): StreamOutput =
+    {
+      import input.{ vars, fresh, transform, typed, currentOwner }
+
+      val itemVal = fresh("item")
+      val Function(List(itemValDef @ ValDef(_, _, _, _)), itemValRef @ Ident(_)) = typed(q"""
+        ($itemVal: $tpe) => $itemVal
+      """)
+
+      val sub = emitSub(
+        input.copy(
+          vars = ScalarValue(tpe, alias = Some(itemValRef)),
+          outputSize = None,
+          index = None),
+        nextOps)
+
+      sub.copy(body = List(withQuietWarnings(transform(typed(q"""
+        // TODO: plug that lambda's symbol as the new owner of sub.body's decls.
+        ${vars.alias.get}.foreach(($itemValDef) => {
+          ..${sub.body};
+        })
+      """)))))
+    }
+  }
+}
diff --git a/Streams/src/main/scala/streams/ops/StreamOps.scala b/Streams/src/main/scala/streams/ops/StreamOps.scala
@@ -8,6 +8,7 @@ private[streams] trait StreamOps
     with FilterOps
     with FindOps
     with ExistsOps
+    with FlattenOps
     with FlatMapOps
     with ForeachOps
     with IsEmptyOps
@@ -33,6 +34,7 @@ private[streams] trait StreamOps
       SomeExistsOp,
       SomeFilterOp,
       SomeFindOp,
+      SomeFlattenOp,
       SomeFlatMapOp,
       SomeForeachOp,
       SomeIsEmptyOp,

diff --git a/Streams/src/test/scala/IntegrationTests.scala b/Streams/src/test/scala/IntegrationTests.scala
@@ -463,6 +463,24 @@ object IntegrationTests
       -> streamMsg("List.flatMap -> List", "Option.getOrElse"),
       // TODO: -> streamMsg("List.flatMap(Option.getOrElse) -> List"),
 
+    "Option(Some(2)).flatten"
+      -> streamMsg("Option.flatten -> Option", "Option.foreach"),
+
+    "Option(Option(1)).flatten"
+      -> streamMsg("Option.flatten -> Option", "Option.foreach"),
+
+    "Seq(Option(None).flatten)"
+      -> streamMsg(),
+
+    "List(List(1, 2), List(3, 4)).flatten"
+      -> streamMsg("List.flatten -> List", "List.foreach"),
+
+    "List(List(1, 2), Seq(3, 4), Set(5, 6)).flatten"
+      -> streamMsg("List.flatten -> List"),
+
+    "List(Option(1), None, Some(2)).flatten"
+      -> streamMsg("List.flatten -> List", "Option.foreach"),
+
     "var tot = 0; for (i <- 0 until 10; x = new AnyRef) { tot += i }; tot"
       -> streamMsg("Range.map.foreach"),