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)
     }
 
   }

akka-stream-tests/src/test/scala/akka/stream/impl/GraphStageLogicSpec.scala

@@ -5,6 +5,7 @@ package akka.stream.impl
 
 import akka.stream.testkit.AkkaSpec
 import akka.stream._
+import akka.stream.Fusing.aggressive
 import akka.stream.scaladsl._
 import akka.stream.stage._
 import akka.stream.testkit.Utils.assertAllStagesStopped
@@ -66,12 +67,6 @@ class GraphStageLogicSpec extends AkkaSpec with GraphInterpreterSpecKit with Con
     }
   }
 
-  class FusedGraph[S <: Shape](ga: GraphAssembly, s: S, a: Attributes = Attributes.none) extends Graph[S, Unit] {
-    override def shape = s
-    override val module = GraphModule(ga, s, a, ga.stages.map(_.module))
-    override def withAttributes(attr: Attributes) = new FusedGraph(ga, s, attr)
-  }
-
   "A GraphStageLogic" must {
 
     "emit all things before completing" in assertAllStagesStopped {
@@ -84,38 +79,21 @@ class GraphStageLogicSpec extends AkkaSpec with GraphInterpreterSpecKit with Con
     }
 
     "emit all things before completing with two fused stages" in assertAllStagesStopped {
-      new Builder {
-        val g = new FusedGraph(
-          builder(emit1234, emit5678)
-            .connect(Upstream, emit1234.in)
-            .connect(emit1234.out, emit5678.in)
-            .connect(emit5678.out, Downstream)
-            .buildAssembly(),
-          FlowShape(emit1234.in, emit5678.out))
+      val g = aggressive(Flow[Int].via(emit1234).via(emit5678))
 
-        Source.empty.via(g).runWith(TestSink.probe)
-          .request(9)
-          .expectNextN(1 to 8)
-          .expectComplete()
-      }
+      Source.empty.via(g).runWith(TestSink.probe)
+        .request(9)
+        .expectNextN(1 to 8)
+        .expectComplete()
     }
 
     "emit all things before completing with three fused stages" in assertAllStagesStopped {
-      new Builder {
-        val g = new FusedGraph(
-          builder(emit1234, passThrough, emit5678)
-            .connect(Upstream, emit1234.in)
-            .connect(emit1234.out, passThrough.in)
-            .connect(passThrough.out, emit5678.in)
-            .connect(emit5678.out, Downstream)
-            .buildAssembly(),
-          FlowShape(emit1234.in, emit5678.out))
+      val g = aggressive(Flow[Int].via(emit1234).via(passThrough).via(emit5678))
 
-        Source.empty.via(g).runWith(TestSink.probe)
-          .request(9)
-          .expectNextN(1 to 8)
-          .expectComplete()
-      }
+      Source.empty.via(g).runWith(TestSink.probe)
+        .request(9)
+        .expectNextN(1 to 8)
+        .expectComplete()
     }
 
     "invoke lifecycle hooks in the right order" in assertAllStagesStopped {

akka-stream-tests/src/test/scala/akka/stream/scaladsl/GraphMatValueSpec.scala

@@ -3,7 +3,7 @@
  */
 package akka.stream.scaladsl
 
-import akka.stream.{ ClosedShape, SourceShape, ActorMaterializer, ActorMaterializerSettings }
+import akka.stream._
 import akka.stream.testkit._
 
 import scala.concurrent.Await
@@ -106,5 +106,15 @@ class GraphMatValueSpec extends AkkaSpec {
 
     }
 
+    "work also when the source’s module is copied" in {
+      val foldFlow: Flow[Int, Int, Future[Int]] = Flow.fromGraph(GraphDSL.create(Sink.fold[Int, Int](0)(_ + _)) {
+        implicit builder ⇒
+          fold ⇒
+            FlowShape(fold.inlet, builder.materializedValue.mapAsync(4)(identity).outlet)
+      })
+
+      Await.result(Source(1 to 10).via(foldFlow).runWith(Sink.head), 3.seconds) should ===(55)
+    }
+
   }
 }