examples/src/main/kotlin/scientifik/kmath/commons/prob/DistributionDemo.kt

@@ -2,17 +2,17 @@ package scientifik.kmath.commons.prob
 
 import kotlinx.coroutines.runBlocking
 import scientifik.kmath.chains.Chain
-import scientifik.kmath.chains.StatefulChain
+import scientifik.kmath.chains.mapWithState
 import scientifik.kmath.prob.Distribution
 import scientifik.kmath.prob.RandomGenerator
 
 data class AveragingChainState(var num: Int = 0, var value: Double = 0.0)
 
-fun Chain<Double>.mean(): Chain<Double> = StatefulChain(AveragingChainState(), 0.0) {
-    val next = this@mean.next()
+fun Chain<Double>.mean(): Chain<Double> = mapWithState(AveragingChainState(),{it.copy()}){chain->
+    val next = chain.next()
     num++
     value += next
-    return@StatefulChain value / num
+    return@mapWithState value / num
 }
 
 

kmath-coroutines/src/commonMain/kotlin/scientifik/kmath/chains/Chain.kt

@@ -17,6 +17,7 @@
 package scientifik.kmath.chains
 
 import kotlinx.atomicfu.atomic
+import kotlinx.atomicfu.updateAndGet
 import kotlinx.coroutines.FlowPreview
 import kotlinx.coroutines.flow.Flow
 
@@ -26,11 +27,6 @@ import kotlinx.coroutines.flow.Flow
  * @param R - the chain element type
  */
 interface Chain<out R> {
-    /**
-     * Last cached value of the chain. Returns null if [next] was not called
-     */
-    fun peek(): R?
-
     /**
      * Generate next value, changing state if needed
      */
@@ -53,84 +49,59 @@ val <R> Chain<R>.flow: Flow<R>
 fun <T> Iterator<T>.asChain(): Chain<T> = SimpleChain { next() }
 fun <T> Sequence<T>.asChain(): Chain<T> = iterator().asChain()
 
-/**
- * Map the chain result using suspended transformation. Initial chain result can no longer be safely consumed
- * since mapped chain consumes tokens. Accepts regular transformation function
- */
-fun <T, R> Chain<T>.map(func: (T) -> R): Chain<R> {
-    val parent = this;
-    return object : Chain<R> {
-        override fun peek(): R? = parent.peek()?.let(func)
-
-        override suspend fun next(): R {
-            return func(parent.next())
-        }
-
-        override fun fork(): Chain<R> {
-            return parent.fork().map(func)
-        }
-    }
-}
-
 /**
  * A simple chain of independent tokens
  */
 class SimpleChain<out R>(private val gen: suspend () -> R) : Chain<R> {
-    private val atomicValue = atomic<R?>(null)
-    override fun peek(): R? = atomicValue.value
-
-    override suspend fun next(): R = gen().also { atomicValue.lazySet(it) }
-
+    override suspend fun next(): R = gen()
     override fun fork(): Chain<R> = this
 }
 
-//TODO force forks on mapping operations?
-
 /**
  * A stateless Markov chain
  */
-class MarkovChain<out R : Any>(private val seed: () -> R, private val gen: suspend (R) -> R) :
-    Chain<R> {
+class MarkovChain<out R : Any>(private val seed: suspend () -> R, private val gen: suspend (R) -> R) : Chain<R> {
 
-    constructor(seed: R, gen: suspend (R) -> R) : this({ seed }, gen)
+    constructor(seedValue: R, gen: suspend (R) -> R) : this({ seedValue }, gen)
 
-    private val atomicValue = atomic<R?>(null)
-    override fun peek(): R = atomicValue.value ?: seed()
+    private val value = atomic<R?>(null)
 
     override suspend fun next(): R {
-        val newValue = gen(peek())
-        atomicValue.lazySet(newValue)
-        return peek()
+        return value.updateAndGet { prev -> gen(prev ?: seed()) }!!
     }
 
     override fun fork(): Chain<R> {
-        return MarkovChain(peek(), gen)
+        return MarkovChain(seed = { value.value ?: seed() }, gen = gen)
     }
 }
 
 /**
  * A chain with possibly mutable state. The state must not be changed outside the chain. Two chins should never share the state
  * @param S - the state of the chain
+ * @param forkState - the function to copy current state without modifying it
  */
 class StatefulChain<S, out R>(
     private val state: S,
     private val seed: S.() -> R,
+    private val forkState: ((S) -> S),
     private val gen: suspend S.(R) -> R
 ) : Chain<R> {
 
-    constructor(state: S, seed: R, gen: suspend S.(R) -> R) : this(state, { seed }, gen)
+    constructor(state: S, seedValue: R, forkState: ((S) -> S), gen: suspend S.(R) -> R) : this(
+        state,
+        { seedValue },
+        forkState,
+        gen
+    )
 
     private val atomicValue = atomic<R?>(null)
-    override fun peek(): R = atomicValue.value ?: seed(state)
 
     override suspend fun next(): R {
-        val newValue = gen(state, peek())
-        atomicValue.lazySet(newValue)
-        return peek()
+        return atomicValue.updateAndGet { prev -> state.gen(prev ?: state.seed()) }!!
     }
 
     override fun fork(): Chain<R> {
-        throw RuntimeException("Fork not supported for stateful chain")
+        return StatefulChain(forkState(state), seed, forkState, gen)
     }
 }
 
@@ -138,11 +109,32 @@ class StatefulChain<S, out R>(
  * A chain that repeats the same value
  */
 class ConstantChain<out T>(val value: T) : Chain<T> {
-    override fun peek(): T? = value
-
     override suspend fun next(): T = value
 
     override fun fork(): Chain<T> {
         return this
     }
 }
+
+/**
+ * Map the chain result using suspended transformation. Initial chain result can no longer be safely consumed
+ * since mapped chain consumes tokens. Accepts regular transformation function
+ */
+fun <T, R> Chain<T>.pipe(func: suspend (T) -> R): Chain<R> = object : Chain<R> {
+    override suspend fun next(): R = func(this@pipe.next())
+    override fun fork(): Chain<R> = this@pipe.fork().pipe(func)
+}
+
+/**
+ * Map the whole chain
+ */
+fun <T, R> Chain<T>.map(mapper: suspend (Chain<T>) -> R): Chain<R> = object : Chain<R> {
+    override suspend fun next(): R = mapper(this@map)
+    override fun fork(): Chain<R> = this@map.fork().map(mapper)
+}
+
+fun <T, S, R> Chain<T>.mapWithState(state: S, stateFork: (S) -> S, mapper: suspend S.(Chain<T>) -> R): Chain<R> =
+    object : Chain<R> {
+        override suspend fun next(): R = state.mapper(this@mapWithState)
+        override fun fork(): Chain<R> = this@mapWithState.fork().mapWithState(stateFork(state), stateFork, mapper)
+    }

kmath-coroutines/src/jvmMain/kotlin/scientifik/kmath/chains/ChainExt.kt

@@ -18,23 +18,3 @@ operator fun <R> Chain<R>.iterator() = object : Iterator<R> {
 fun <R> Chain<R>.asSequence(): Sequence<R> = object : Sequence<R> {
     override fun iterator(): Iterator<R> = this@asSequence.iterator()
 }
-
-
-/**
- * Map the chain result using suspended transformation. Initial chain result can no longer be safely consumed
- * since mapped chain consumes tokens. Accepts suspending transformation function.
- */
-fun <T, R> Chain<T>.map(func: suspend (T) -> R): Chain<R> {
-    val parent = this;
-    return object : Chain<R> {
-        override fun peek(): R? = runBlocking { parent.peek()?.let { func(it) } }
-
-        override suspend fun next(): R {
-            return func(parent.next())
-        }
-
-        override fun fork(): Chain<R> {
-            return parent.fork().map(func)
-        }
-    }
-}