Clean up todos in traversal builder

stream-tests/src/test/scala/org/apache/pekko/stream/impl/TraversalBuilderSpec.scala

@@ -14,11 +14,11 @@
 package org.apache.pekko.stream.impl
 
 import scala.concurrent.Promise
-
 import org.apache.pekko
 import pekko.NotUsed
 import pekko.stream._
 import pekko.stream.impl.Stages.DefaultAttributes
+import pekko.stream.impl.StreamLayout.AtomicModule
 import pekko.stream.impl.TraversalTestUtils._
 import pekko.stream.impl.fusing.{ IterableSource, IteratorSource, RangeSource }
 import pekko.stream.impl.fusing.GraphStages.{ FutureSource, RepeatSource, SingleSource }
@@ -451,6 +451,45 @@ class TraversalBuilderSpec extends PekkoSpec {
           (flow1, Attributes.name("test") and Attributes.name("flow"), TestIsland1),
           (sink, Attributes.none, TestDefaultIsland)))
     }
+
+    "mark only the assigned output as wired when a submodule still has other unwired outputs" in {
+      class CompositeTestFanOut2 extends AtomicModule[FanOutShape2[Any, Any, Any], Any] {
+        val in = Inlet[Any]("testFanOut2.in")
+        val out0 = Outlet[Any]("testFanOut2.out0")
+        val out1 = Outlet[Any]("testFanOut2.out1")
+
+        override val shape: FanOutShape2[Any, Any, Any] = new FanOutShape2(in, out0, out1)
+        override val traversalBuilder = TraversalBuilder.atomic(this, Attributes.name("testFanOut2"))
+
+        override def withAttributes(attributes: Attributes): AtomicModule[FanOutShape2[Any, Any, Any], Any] = ???
+        override def toString = "TestFanOut2"
+      }
+
+      val fanOut = new CompositeTestFanOut2
+      val sink1 = new CompositeTestSink
+      val sink2 = new CompositeTestSink
+
+      val partial = source.traversalBuilder
+        .add(fanOut.traversalBuilder, fanOut.shape, Keep.left)
+        .add(sink1.traversalBuilder, sink1.shape, Keep.left)
+        .add(sink2.traversalBuilder, sink2.shape, Keep.left)
+        .wire(source.out, fanOut.in)
+        .wire(fanOut.out0, sink1.in)
+
+      partial.isUnwired(fanOut.out0) should ===(false)
+      partial.isUnwired(fanOut.out1) should ===(true)
+
+      val builder = partial.wire(fanOut.out1, sink2.in)
+
+      val mat = testMaterialize(builder)
+      mat.connections should ===(3)
+      mat.outlets(0) should ===(source.out)
+      mat.inlets(0) should ===(fanOut.in)
+      mat.outlets(1) should ===(fanOut.out0)
+      mat.inlets(1) should ===(sink1.in)
+      mat.outlets(2) should ===(fanOut.out1)
+      mat.inlets(2) should ===(sink2.in)
+    }
   }
 
   "find Source.empty via TraversalBuilder with isEmptySource" in {

stream/src/main/scala/org/apache/pekko/stream/impl/TraversalBuilder.scala

@@ -633,7 +633,6 @@ import pekko.util.OptionVal
     extends TraversalBuilder {
 
   override def add(submodule: TraversalBuilder, shape: Shape, combineMat: AnyFunction2): TraversalBuilder = {
-    // TODO: Use automatically a linear builder if applicable
     // Create a composite, add ourselves, then the other.
     CompositeTraversalBuilder(attributes = attributes)
       .add(this, module.shape, Keep.right)
@@ -689,7 +688,6 @@ import pekko.util.OptionVal
  */
 @InternalApi private[pekko] object LinearTraversalBuilder {
 
-  // TODO: Remove
   private val cachedEmptyLinear =
     LinearTraversalBuilder(OptionVal.None, OptionVal.None, 0, 0, PushNotUsed, OptionVal.None, Attributes.none)
 
@@ -1318,26 +1316,13 @@ import pekko.util.OptionVal
 
     // Do the assignment in the submodule
     val result = submodule.assign(out.mappedTo, relativeSlot)
-    val wired = if (result.isTraversalComplete) {
-      // Remove the builder (and associated data).
-      // We can't simply append its Traversal as there might be uncompleted builders that come earlier in the
-      // final traversal (remember, input ports are assigned in traversal order of modules, and the inOffsets
-      // and inBaseOffseForOut Maps are updated when adding a module; we must respect addition order).
-
+    val wired =
       copy(
         inBaseOffsetForOut = inBaseOffsetForOut - out,
         outOwners = outOwners - out,
-        // TODO Optimize Map access
         pendingBuilders = pendingBuilders.updated(builderKey, result),
         // pendingBuilders = pendingBuilders - builderKey,
         unwiredOuts = unwiredOuts - 1)
-    } else {
-      // Update structures with result
-      copy(
-        inBaseOffsetForOut = inBaseOffsetForOut - out,
-        unwiredOuts = unwiredOuts - 1,
-        pendingBuilders = pendingBuilders.updated(builderKey, result))
-    }
 
     // If we have no more unconnected outputs, we can finally build the Traversal and shed most of the auxiliary data.
     wired.completeIfPossible
@@ -1378,7 +1363,6 @@ import pekko.util.OptionVal
         val in = inIterator.next()
         // Calculate offset in the current scope. This is the our first unused input slot plus
         // the relative offset of the input port in the submodule.
-        // TODO Optimize Map access
         newInOffsets = newInOffsets.updated(in, inSlots + submodule.offsetOf(in.mappedTo))
       }
 
@@ -1399,7 +1383,6 @@ import pekko.util.OptionVal
       while (inIterator.hasNext) {
         val in = inIterator.next()
         // Calculate offset in the current scope
-        // TODO Optimize Map access
         newInOffsets = newInOffsets.updated(in, inSlots + submodule.offsetOf(in.mappedTo))
       }
 
@@ -1409,7 +1392,6 @@ import pekko.util.OptionVal
         // Record the base offsets of all the modules we included and which have unwired output ports. We need
         // to adjust their offset by inSlots as that would be their new position in this module.
         newBaseOffsetsForOut = newBaseOffsetsForOut.updated(out, inSlots + submodule.offsetOfModule(out.mappedTo))
-        // TODO Optimize Map access
         newOutOwners = newOutOwners.updated(out, builderKey)
       }