chore: Fix zipWithIndex when use in graphdsl. (#32626)

Author: He-Pin

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

@@ -4,8 +4,10 @@
 
 package akka.stream.scaladsl
 
-import akka.stream.{ ActorMaterializer, ActorMaterializerSettings, Materializer }
+import akka.stream.testkit.scaladsl.TestSink
+import akka.stream.{ ActorMaterializer, ActorMaterializerSettings, Materializer, UniformFanInShape }
 import akka.stream.testkit.{ StreamSpec, TestSubscriber }
+
 import scala.annotation.nowarn
 
 @nowarn // keep unused imports
@@ -39,6 +41,38 @@ class FlowZipWithIndexSpec extends StreamSpec {
       probe.expectComplete()
     }
 
+    "will works in GraphDSL" in {
+      import akka.stream.ClosedShape
+      val pickMaxOfThree = GraphDSL.create() { implicit b =>
+        import GraphDSL.Implicits._
+
+        val zip1 = b.add(ZipWith[Int, Int, Int](math.max))
+        val zip2 = b.add(ZipWith[Int, Int, Int](math.max))
+        zip1.out ~> zip2.in0
+
+        UniformFanInShape(zip2.out, zip1.in0, zip1.in1, zip2.in1)
+      }
+
+      val resultSink = TestSink[(Int, Long)]()
+
+      val g = RunnableGraph.fromGraph(GraphDSL.createGraph(resultSink) { implicit b => sink =>
+        import GraphDSL.Implicits._
+
+        // importing the partial graph will return its shape (inlets & outlets)
+        val pm3 = b.add(pickMaxOfThree)
+
+        Source.single(1) ~> pm3.in(0)
+        Source.single(2) ~> pm3.in(1)
+        Source.single(3) ~> pm3.in(2)
+        pm3.out.zipWithIndex ~> sink.in
+        ClosedShape
+      })
+      val p = g.run()
+      p.request(1)
+      p.expectNext((3, 0L))
+      p.expectComplete()
+    }
+
     "work in fruit example" in {
       //#zip-with-index
       Source(List("apple", "orange", "banana")).zipWithIndex.runWith(Sink.foreach(println))

akka-stream/src/main/scala/akka/stream/scaladsl/Flow.scala

@@ -1080,8 +1080,9 @@ trait FlowOps[+Out, +Mat] {
    * @param onComplete a function that transforms the ongoing state into an optional output element
    */
   def statefulMap[S, T](create: () => S)(f: (S, Out) => (S, T), onComplete: S => Option[T]): Repr[T] =
-    via(new StatefulMap[S, Out, T](create, f, onComplete))
-      .withAttributes(DefaultAttributes.statefulMap and SourceLocation.forLambda(f))
+    via(
+      new StatefulMap[S, Out, T](create, f, onComplete)
+        .withAttributes(DefaultAttributes.statefulMap and SourceLocation.forLambda(f)))
 
   /**
    * Transform each stream element with the help of a resource.
@@ -1116,8 +1117,8 @@ trait FlowOps[+Out, +Mat] {
    */
   def mapWithResource[R, T](create: () => R)(f: (R, Out) => T, close: R => Option[T]): Repr[T] =
     via(
-      new StatefulMap[R, Out, T](create, (resource, out) => (resource, f(resource, out)), resource => close(resource)))
-      .withAttributes(DefaultAttributes.mapWithResource and SourceLocation.forLambda(f))
+      new StatefulMap[R, Out, T](create, (resource, out) => (resource, f(resource, out)), resource => close(resource))
+        .withAttributes(DefaultAttributes.mapWithResource and SourceLocation.forLambda(f)))
 
   /**
    * Transform each input element into an `Iterable` of output elements that is
@@ -2981,8 +2982,11 @@ trait FlowOps[+Out, +Mat] {
    * '''Cancels when''' downstream cancels
    */
   def zipWithIndex: Repr[(Out, Long)] =
-    statefulMap[Long, (Out, Long)](() => 0L)((index, out) => (index + 1L, (out, index)), _ => None)
-      .withAttributes(DefaultAttributes.zipWithIndex)
+    via(
+      new StatefulMap[Long, Out, (Out, Long)](
+        () => 0L,
+        (index, out) => (index + 1L, (out, index)),
+        ConstantFun.scalaAnyToNone).withAttributes(DefaultAttributes.zipWithIndex))
 
   /**
    * Interleave is a deterministic merge of the given [[Source]] with elements of this [[Flow]].