From 9c21f69ec02d94ba4b3740f5ed205e000356453f Mon Sep 17 00:00:00 2001
From: Alexander Nozik <altavir@gmail.com>
Date: Thu, 30 Apr 2020 09:38:42 +0300
Subject: [PATCH] Fix for #87

---
 .../scientifik/kmath/prob/Distribution.kt       |  8 ++++++--
 .../kotlin/scientifik/kmath/prob/RandomChain.kt |  3 +--
 .../kotlin/scientifik/kmath/prob/SamplerTest.kt | 17 +++++++++++++++++
 3 files changed, 24 insertions(+), 4 deletions(-)
 create mode 100644 kmath-prob/src/jvmTest/kotlin/scientifik/kmath/prob/SamplerTest.kt

diff --git a/kmath-prob/src/commonMain/kotlin/scientifik/kmath/prob/Distribution.kt b/kmath-prob/src/commonMain/kotlin/scientifik/kmath/prob/Distribution.kt
index 44af84d5d..3b874adaa 100644
--- a/kmath-prob/src/commonMain/kotlin/scientifik/kmath/prob/Distribution.kt
+++ b/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) {
-            tmp[i] = chain.next()
+        //clear list from previous run
+        tmp.clear()
+        //Fill list
+        repeat(size){
+            tmp.add(chain.next())
         }
+        //return new buffer with elements from tmp
         bufferFactory(size) { tmp[it] }
     }
 }
diff --git a/kmath-prob/src/commonMain/kotlin/scientifik/kmath/prob/RandomChain.kt b/kmath-prob/src/commonMain/kotlin/scientifik/kmath/prob/RandomChain.kt
index 9b581afd7..47fc6e4c5 100644
--- a/kmath-prob/src/commonMain/kotlin/scientifik/kmath/prob/RandomChain.kt
+++ b/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.flow(gen: suspend RandomGenerator.() -> R) = chain(gen).fork()
\ No newline at end of file
+fun <R> RandomGenerator.chain(gen: suspend RandomGenerator.() -> R): RandomChain<R> = RandomChain(this, gen)
\ No newline at end of file
diff --git a/kmath-prob/src/jvmTest/kotlin/scientifik/kmath/prob/SamplerTest.kt b/kmath-prob/src/jvmTest/kotlin/scientifik/kmath/prob/SamplerTest.kt
new file mode 100644
index 000000000..1152f3057
--- /dev/null
+++ b/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())
+        }
+    }
+}
\ No newline at end of file