make fused graphs fusable

akka-stream-tests/src/test/scala/akka/stream/FusingSpec.scala

@@ -7,32 +7,82 @@ import akka.stream._
 import akka.stream.scaladsl._
 import akka.stream.testkit.AkkaSpec
 import org.scalactic.ConversionCheckedTripleEquals
+import akka.stream.Attributes._
+import akka.stream.Fusing.FusedGraph
+import scala.annotation.tailrec
+import akka.stream.impl.StreamLayout.Module
 
 class FusingSpec extends AkkaSpec with ConversionCheckedTripleEquals {
 
+  final val Debug = false
   implicit val materializer = ActorMaterializer()
 
+  def graph(async: Boolean) =
+    Source.unfoldInf(1)(x ⇒ (x, x)).filter(_ % 2 == 1)
+      .alsoTo(Flow[Int].fold(0)(_ + _).to(Sink.head.named("otherSink")).addAttributes(if (async) Attributes.asyncBoundary else Attributes.none))
+      .via(Flow[Int].fold(1)(_ + _).named("mainSink"))
+
+  def singlePath[S <: Shape, M](fg: FusedGraph[S, M], from: Attribute, to: Attribute): Unit = {
+    val starts = fg.module.info.allModules.filter(_.attributes.contains(from))
+    starts.size should ===(1)
+    val start = starts.head
+    val ups = fg.module.info.upstreams
+    val owner = fg.module.info.outOwners
+
+    @tailrec def rec(curr: Module): Unit = {
+      if (Debug) println(extractName(curr, "unknown"))
+      if (curr.attributes.contains(to)) () // done
+      else {
+        val outs = curr.inPorts.map(ups)
+        outs.size should ===(1)
+        val out = outs.head
+        val next = owner(out)
+        rec(next)
+      }
+    }
+
+    rec(start)
+  }
+
   "Fusing" must {
 
-    "fuse a moderately complex graph" in {
-      val g = Source.unfoldInf(1)(x ⇒ (x, x)).filter(_ % 2 == 1).alsoTo(Sink.fold(0)(_ + _)).to(Sink.fold(1)(_ + _))
-      val fused = Fusing.aggressive(g)
+    def verify[S <: Shape, M](fused: FusedGraph[S, M], modules: Int, downstreams: Int): Unit = {
       val module = fused.module
-      module.subModules.size should ===(1)
-      module.info.downstreams.size should be > 5
-      module.info.upstreams.size should be > 5
+      module.subModules.size should ===(modules)
+      module.downstreams.size should ===(modules - 1)
+      module.info.downstreams.size should be >= downstreams
+      module.info.upstreams.size should be >= downstreams
+      singlePath(fused, Attributes.Name("mainSink"), Attributes.Name("unfoldInf"))
+      singlePath(fused, Attributes.Name("otherSink"), Attributes.Name("unfoldInf"))
+    }
+
+    "fuse a moderately complex graph" in {
+      val g = graph(false)
+      val fused = Fusing.aggressive(g)
+      verify(fused, modules = 1, downstreams = 5)
     }
 
     "not fuse across AsyncBoundary" in {
-      val g =
-        Source.unfoldInf(1)(x ⇒ (x, x)).filter(_ % 2 == 1)
-          .alsoTo(Sink.fold(0)(_ + (_: Int)).addAttributes(Attributes.asyncBoundary))
-          .to(Sink.fold(1)(_ + _))
+      val g = graph(true)
       val fused = Fusing.aggressive(g)
-      val module = fused.module
-      module.subModules.size should ===(2)
-      module.info.downstreams.size should be > 5
-      module.info.upstreams.size should be > 5
+      verify(fused, modules = 2, downstreams = 5)
+    }
+
+    "not fuse a FusedGraph again" in {
+      val g = Fusing.aggressive(graph(false))
+      Fusing.aggressive(g) should be theSameInstanceAs g
+    }
+
+    "properly fuse a FusedGraph that has been extended (no AsyncBoundary)" in {
+      val src = Fusing.aggressive(graph(false))
+      val fused = Fusing.aggressive(Source.fromGraph(src).to(Sink.head))
+      verify(fused, modules = 1, downstreams = 6)
+    }
+
+    "properly fuse a FusedGraph that has been extended (with AsyncBoundary)" in {
+      val src = Fusing.aggressive(graph(true))
+      val fused = Fusing.aggressive(Source.fromGraph(src).to(Sink.head))
+      verify(fused, modules = 2, downstreams = 6)
     }
 
   }