kmath-prob/src/commonMain/kotlin/scientifik/kmath/prob/Distribution.kt

@@ -58,9 +58,13 @@ fun <T : Any> Sampler<T>.sampleBuffer(
     //creating temporary storage once
     val tmp = ArrayList<T>(size)
     return sample(generator).collect { chain ->
-        for (i in tmp.indices) {
+        //clear list from previous run
-            tmp[i] = chain.next()
+        tmp.clear()
+        //Fill list
+        repeat(size){
+            tmp.add(chain.next())
         }
+        //return new buffer with elements from tmp
         bufferFactory(size) { tmp[it] }
     }
 }

kmath-prob/src/commonMain/kotlin/scientifik/kmath/prob/RandomChain.kt

@@ -11,5 +11,4 @@ class RandomChain<out R>(val generator: RandomGenerator, private val gen: suspen
     override fun fork(): Chain<R> = RandomChain(generator.fork(), gen)
 }
 
-fun <R> RandomGenerator.chain(gen: suspend RandomGenerator.() -> R): RandomChain<R> = RandomChain(this, gen)
+fun <R> RandomGenerator.chain(gen: suspend RandomGenerator.() -> R): RandomChain<R> = RandomChain(this, gen)
-fun <R> RandomGenerator.flow(gen: suspend RandomGenerator.() -> R) = chain(gen).fork()

kmath-prob/src/jvmTest/kotlin/scientifik/kmath/prob/SamplerTest.kt

@@ -0,0 +1,17 @@
+package scientifik.kmath.prob
+
+import kotlinx.coroutines.runBlocking
+import kotlin.test.Test
+
+class SamplerTest {
+
+    @Test
+    fun bufferSamplerTest(){
+        val sampler: Sampler<Double> =
+            BasicSampler { it.chain { nextDouble() } }
+        val data = sampler.sampleBuffer(RandomGenerator.default, 100)
+        runBlocking {
+            println(data.next())
+        }
+    }
+}