Concurrent map for flows

kmath-coroutines/src/commonMain/kotlin/scientifik/kmath/CoroutinesExtra.kt

@@ -1,12 +1,8 @@
 package scientifik.kmath
 
 import kotlinx.coroutines.*
-import kotlinx.coroutines.channels.ReceiveChannel
 import kotlinx.coroutines.channels.produce
-import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.*
-import kotlinx.coroutines.flow.FlowCollector
-import kotlinx.coroutines.flow.collect
-import kotlinx.coroutines.flow.map
 
 val Dispatchers.Math: CoroutineDispatcher get() = Dispatchers.Default
 
@@ -17,7 +13,7 @@ internal class LazyDeferred<T>(val dispatcher: CoroutineDispatcher, val block: s
     private var deferred: Deferred<T>? = null
 
     internal fun start(scope: CoroutineScope) {
-        if(deferred==null) {
+        if (deferred == null) {
             deferred = scope.async(dispatcher, block = block)
         }
     }
@@ -60,7 +56,7 @@ suspend fun <T> AsyncFlow<T>.collect(concurrency: Int, collector: FlowCollector<
     require(concurrency >= 1) { "Buffer size should be more than 1, but was $concurrency" }
     coroutineScope {
         //Starting up to N deferred coroutines ahead of time
-        val channel = produce(capacity = concurrency-1) {
+        val channel = produce(capacity = concurrency - 1) {
             deferredFlow.collect { value ->
                 value.start(this@coroutineScope)
                 send(value)
@@ -91,31 +87,16 @@ suspend fun <T> AsyncFlow<T>.collect(concurrency: Int, action: suspend (value: T
     })
 }
 
-//suspend fun <T> Flow<T>.collect(concurrency: Int, dispatcher: CoroutineDispatcher, collector: FlowCollector<T>){
+@FlowPreview
-//    require(concurrency >= 1) { "Buffer size should be more than 1, but was $concurrency" }
+fun <T, R> Flow<T>.map(
-//    coroutineScope {
+    concurrencyLevel: Int,
-//        //Starting up to N deferred coroutines ahead of time
+    dispatcher: CoroutineDispatcher = Dispatchers.Default,
-//        val channel = produce(capacity = concurrency-1) {
+    bufferSize: Int = concurrencyLevel,
-//            this@collect.
+    transform: suspend (T) -> R
-//            deferredFlow.collect { value ->
+): Flow<R> {
-//                value.start(this@coroutineScope)
+    return flatMapMerge(concurrencyLevel, bufferSize) { value ->
-//                send(value)
+        flow { emit(transform(value)) }
-//            }
+    }.flowOn(dispatcher)
-//        }
+}
-//
+
-//        (channel as Job).invokeOnCompletion {
-//            if (it is CancellationException && it.cause == null) cancel()
-//        }
-//
-//        for (element in channel) {
-//            collector.emit(element.await())
-//        }
-//
-//        val producer = channel as Job
-//        if (producer.isCancelled) {
-//            producer.join()
-//            //throw producer.getCancellationException()
-//        }
-//    }
-//}
 

kmath-coroutines/src/jvmTest/kotlin/scientifik/kmath/streaming/BufferFlowTest.kt

@@ -1,20 +1,36 @@
 package scientifik.kmath.streaming
 
 import kotlinx.coroutines.*
-import kotlinx.coroutines.flow.asFlow
+import kotlinx.coroutines.flow.*
 import org.junit.Test
 import scientifik.kmath.async
 import scientifik.kmath.collect
+import scientifik.kmath.map
+
 
 @ExperimentalCoroutinesApi
 @InternalCoroutinesApi
 @FlowPreview
 class BufferFlowTest {
 
+    @Test(timeout = 2000)
+    fun concurrentMap() {
+        runBlocking {
+            (1..20).asFlow().map(4) {
+                println("Started $it")
+                @Suppress("BlockingMethodInNonBlockingContext")
+                Thread.sleep(200)
+                it
+            }.collect {
+                println("Completed $it")
+            }
+        }
+    }
+
     @Test(timeout = 2000)
     fun mapParallel() {
         runBlocking {
-            (1..20).asFlow().async(Dispatchers.Default) {
+            (1..20).asFlow().async {
                 println("Started $it")
                 @Suppress("BlockingMethodInNonBlockingContext")
                 Thread.sleep(200)