From 798512517dd87ddec7f8daee60089b32322a6473 Mon Sep 17 00:00:00 2001
From: Alexander Nozik <altavir@gmail.com>
Date: Sat, 5 Jan 2019 21:01:30 +0300
Subject: [PATCH] Added factories for Complex number buffers and nd-structures

---
 .../kmath/structures/ComplexBufferSpec.kt         | 15 +++++++++++++--
 .../scientifik/kmath/structures/ObjectBuffer.kt   | 10 ++++++++--
 .../kmath/structures/ComplexBufferSpecTest.kt     |  2 +-
 3 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/kmath-core/src/jvmMain/kotlin/scientifik/kmath/structures/ComplexBufferSpec.kt b/kmath-core/src/jvmMain/kotlin/scientifik/kmath/structures/ComplexBufferSpec.kt
index 35ee5e6b7..516653817 100644
--- a/kmath-core/src/jvmMain/kotlin/scientifik/kmath/structures/ComplexBufferSpec.kt
+++ b/kmath-core/src/jvmMain/kotlin/scientifik/kmath/structures/ComplexBufferSpec.kt
@@ -1,6 +1,7 @@
 package scientifik.kmath.structures
 
 import scientifik.kmath.operations.Complex
+import scientifik.kmath.operations.ComplexField
 import java.nio.ByteBuffer
 
 object ComplexBufferSpec : FixedSizeBufferSpec<Complex> {
@@ -19,8 +20,18 @@ object ComplexBufferSpec : FixedSizeBufferSpec<Complex> {
 }
 
 /**
- * Create a mutable buffer which ignores boxing
+ * Create a read-only/mutable buffer which ignores boxing
  */
-fun Complex.Companion.createBuffer(size: Int) = ObjectBuffer.create(ComplexBufferSpec, size)
+fun Buffer.Companion.complex(size: Int): Buffer<Complex> =
+    ObjectBuffer.create(ComplexBufferSpec, size)
+
+fun MutableBuffer.Companion.complex(size: Int) =
+    ObjectBuffer.create(ComplexBufferSpec, size)
+
+fun NDField.Companion.complex(shape: IntArray) =
+    BufferNDField(shape, ComplexField) { size, init -> ObjectBuffer.create(ComplexBufferSpec, size, init) }
+
+fun NDElement.Companion.complex(shape: IntArray, initializer: ComplexField.(IntArray) -> Complex) =
+    NDField.complex(shape).produce(initializer)
 
 
diff --git a/kmath-core/src/jvmMain/kotlin/scientifik/kmath/structures/ObjectBuffer.kt b/kmath-core/src/jvmMain/kotlin/scientifik/kmath/structures/ObjectBuffer.kt
index fceda1b25..8d7bece0c 100644
--- a/kmath-core/src/jvmMain/kotlin/scientifik/kmath/structures/ObjectBuffer.kt
+++ b/kmath-core/src/jvmMain/kotlin/scientifik/kmath/structures/ObjectBuffer.kt
@@ -23,7 +23,13 @@ class ObjectBuffer<T : Any>(private val buffer: ByteBuffer, private val spec: Fi
     }
 
     companion object {
-        fun <T : Any> create(spec: FixedSizeBufferSpec<T>, size: Int) =
-            ObjectBuffer<T>(ByteBuffer.allocate(size * spec.unitSize), spec)
+        fun <T : Any> create(spec: FixedSizeBufferSpec<T>, size: Int, initializer: ((Int) -> T)? = null) =
+            ObjectBuffer(ByteBuffer.allocate(size * spec.unitSize), spec).also { buffer ->
+                if (initializer != null) {
+                    (0 until size).forEach {
+                        buffer[it] = initializer(it)
+                    }
+                }
+            }
     }
 }
\ No newline at end of file
diff --git a/kmath-core/src/jvmTest/kotlin/scientifik/kmath/structures/ComplexBufferSpecTest.kt b/kmath-core/src/jvmTest/kotlin/scientifik/kmath/structures/ComplexBufferSpecTest.kt
index 350f06848..1a3aea10f 100644
--- a/kmath-core/src/jvmTest/kotlin/scientifik/kmath/structures/ComplexBufferSpecTest.kt
+++ b/kmath-core/src/jvmTest/kotlin/scientifik/kmath/structures/ComplexBufferSpecTest.kt
@@ -7,7 +7,7 @@ import kotlin.test.assertEquals
 class ComplexBufferSpecTest {
     @Test
     fun testComplexBuffer() {
-        val buffer = Complex.createBuffer(20)
+        val buffer = MutableBuffer.complex(20)
         (0 until 20).forEach {
             buffer[it] = Complex(it.toDouble(), -it.toDouble())
         }