chore: Avoid lift SingleSource to InflightSource.

Author: He-Pin

akka-stream/src/main/scala/akka/stream/impl/fusing/StreamOfStreams.scala

@@ -159,19 +159,6 @@ private[akka] object FlattenConcat {
     def materialize(): Unit = ()
   }
 
-  private final class InflightSingleSource[T](elem: T) extends InflightSource[T] {
-    private var _hasNext = true
-    override def hasNext: Boolean = _hasNext
-    override def next(): T =
-      if (_hasNext) {
-        _hasNext = false
-        elem
-      } else throw new NoSuchElementException("next called after completion")
-    override def tryPull(): Unit = ()
-    override def cancel(cause: Throwable): Unit = ()
-    override def isClosed: Boolean = !hasNext
-  }
-
   private final class InflightIteratorSource[T](iterator: Iterator[T]) extends InflightSource[T] {
     override def hasNext: Boolean = iterator.hasNext
     override def next(): T = iterator.next()
@@ -235,7 +222,9 @@ private[akka] final class FlattenConcat[T, M](parallelism: Int)
   override def createLogic(enclosingAttributes: Attributes) =
     new GraphStageLogic(shape) with InHandler with OutHandler {
       import FlattenConcat._
-      private var queue: BufferImpl[InflightSource[T]] = _
+      // InflightSource[T] or SingleSource[T]
+      // AnyRef here to avoid lift the SingleSource[T] to InflightSource[T]
+      private var queue: BufferImpl[AnyRef] = _
       private val invokeCb: InflightSource[T] => Unit =
         getAsyncCallback[InflightSource[T]](futureSourceCompleted).invoke
 
@@ -274,17 +263,23 @@ private[akka] final class FlattenConcat[T, M](parallelism: Int)
 
       override def onPull(): Unit = {
         if (queue.nonEmpty) {
-          val currentSource = queue.peek()
           //purge if possible
-          if (currentSource.hasNext) {
-            push(out, currentSource.next())
-            if (currentSource.isClosed) {
-              handleCurrentSourceClosed(currentSource)
-            }
-          } else if (currentSource.isClosed) {
-            handleCurrentSourceClosed(currentSource)
-          } else {
-            currentSource.tryPull()
+          queue.peek() match {
+            case singleSource: SingleSource[T] @unchecked =>
+              push(out, singleSource.elem)
+              removeSource()
+            case currentSource: InflightSource[T] @unchecked =>
+              if (currentSource.hasNext) {
+                push(out, currentSource.next())
+                if (currentSource.isClosed) {
+                  handleCurrentSourceClosed(currentSource)
+                }
+              } else if (currentSource.isClosed) {
+                handleCurrentSourceClosed(currentSource)
+              } else {
+                currentSource.tryPull()
+              }
+            case _ => throw new IllegalStateException("Should not reach here.")
           }
         } else if (!hasBeenPulled(in)) {
           tryPull(in)
@@ -308,18 +303,18 @@ private[akka] final class FlattenConcat[T, M](parallelism: Int)
       private def cancelInflightSources(cause: Throwable): Unit = {
         if (queue.nonEmpty) {
           var source = queue.dequeue()
-          while (source ne null) {
-            source.cancel(cause)
+          while ((source ne null) && (source.isInstanceOf[InflightSource[T] @unchecked])) {
+            source.asInstanceOf[InflightSource[T]].cancel(cause)
             source = queue.dequeue()
           }
         }
       }
 
-      private def addSourceElem(elem: T): Unit = {
+      private def addSource(singleSource: SingleSource[T]): Unit = {
         if (isAvailable(out) && queue.isEmpty) {
-          push(out, elem)
+          push(out, singleSource.elem)
         } else {
-          queue.enqueue(new InflightSingleSource(elem))
+          queue.enqueue(singleSource)
         }
       }
 
@@ -397,7 +392,7 @@ private[akka] final class FlattenConcat[T, M](parallelism: Int)
         (TraversalBuilder.getValuePresentedSource(source): @nowarn("cat=lint-infer-any")) match {
           case OptionVal.Some(graph) =>
             graph match {
-              case single: SingleSource[T] @unchecked => addSourceElem(single.elem)
+              case single: SingleSource[T] @unchecked => addSource(single)
               case futureSource: FutureSource[T] @unchecked =>
                 val future = futureSource.future
                 future.value match {
@@ -415,32 +410,45 @@ private[akka] final class FlattenConcat[T, M](parallelism: Int)
 
       }
 
+      private def removeSource(): Unit = {
+        queue.dequeue()
+        pullIfNeeded()
+      }
+
       private def removeSource(source: InflightSource[T]): Unit = {
         if (queue.nonEmpty && (source eq queue.peek())) {
           //only dequeue if it's the current emitting source
           val s = queue.dequeue()
           if (s != source) {
             throw new IllegalStateException("Should not reach here.")
           }
-          if (isClosed(in)) {
-            if (queue.isEmpty) {
-              completeStage()
-            } else {
-              //pull the new emitting source
-              val nextSource = queue.peek()
-              nextSource.tryPull()
-            }
+          pullIfNeeded()
+        } //not the head source, just ignore
+      }
+
+      private def pullIfNeeded(): Unit = {
+        if (isClosed(in)) {
+          if (queue.isEmpty) {
+            completeStage()
           } else {
-            if (queue.nonEmpty) {
-              //pull the new emitting source
-              val nextSource = queue.peek()
-              nextSource.tryPull()
-            }
-            if (!hasBeenPulled(in)) {
-              tryPull(in)
-            }
+            tryPullNextSourceInQueue()
           }
-        } //not the head source, just ignore
+        } else {
+          if (queue.nonEmpty) {
+            tryPullNextSourceInQueue()
+          }
+          if (!hasBeenPulled(in)) {
+            tryPull(in)
+          }
+        }
+      }
+
+      private def tryPullNextSourceInQueue(): Unit = {
+        //pull the new emitting source
+        val nextSource = queue.peek()
+        if (nextSource.isInstanceOf[InflightSource[T] @unchecked]) {
+          nextSource.asInstanceOf[InflightSource[T]].tryPull()
+        }
       }
 
       setHandlers(in, out, this)