examples/src/main/kotlin/scientifik/kmath/structures/typeSafeDimensions.kt

@@ -14,6 +14,7 @@ fun DMatrixContext<Double, RealField>.simple() {
     m1.transpose() + m2
 }
 
+
 object D5 : Dimension {
     override val dim: UInt = 5u
 }

kmath-dimensions/build.gradle.kts

@@ -7,8 +7,13 @@ description = "A proof of concept module for adding typ-safe dimensions to struc
 kotlin.sourceSets {
     commonMain {
         dependencies {
-            implementation(kotlin("reflect"))
             api(project(":kmath-core"))
         }
     }
+
+    jvmMain{
+        dependencies{
+            api(kotlin("reflect"))
+        }
+    }
 }

kmath-dimensions/src/commonMain/kotlin/scientifik/kmath/dimensions/Dimensions.kt

@@ -1,37 +1,35 @@
 package scientifik.kmath.dimensions
 
+import kotlin.reflect.KClass
+
 /**
- * An interface which is not used in runtime. Designates a size of some structure.
- * All descendants should be singleton objects.
+ * An abstract class which is not used in runtime. Designates a size of some structure.
+ * Could be replaced later by fully inline constructs
  */
 interface Dimension {
-    val dim: UInt
 
+    val dim: UInt
     companion object {
-        @Suppress("NOTHING_TO_INLINE")
-        inline fun of(dim: UInt): Dimension {
-            return when (dim) {
-                1u -> D1
-                2u -> D2
-                3u -> D3
-                else -> object : Dimension {
-                    override val dim: UInt = dim
-                }
-            }
-        }
+
     }
 }
 
-expect inline fun <reified D : Dimension>  Dimension.Companion.dim(): UInt
+fun <D : Dimension> KClass<D>.dim(): UInt = Dimension.resolve(this).dim
+
+expect fun <D : Dimension> Dimension.Companion.resolve(type: KClass<D>): D
+
+expect fun Dimension.Companion.of(dim: UInt): Dimension
+
+inline fun <reified D : Dimension> Dimension.Companion.dim(): UInt = D::class.dim()
 
 object D1 : Dimension {
-    override val dim: UInt = 1u
+    override val dim: UInt get() = 1U
 }
 
 object D2 : Dimension {
-    override val dim: UInt = 2u
+    override val dim: UInt get() = 2U
 }
 
 object D3 : Dimension {
-    override val dim: UInt = 3u
+    override val dim: UInt get() = 31U
 }

kmath-dimensions/src/commonMain/kotlin/scientifik/kmath/dimensions/Wrappers.kt

@@ -29,7 +29,7 @@ interface DMatrix<T, R : Dimension, C : Dimension> : Structure2D<T> {
         /**
          * The same as [coerce] but without dimension checks. Use with caution
          */
-        inline fun <T, reified R : Dimension, reified C : Dimension> coerceUnsafe(structure: Structure2D<T>): DMatrix<T, R, C> {
+        fun <T, R : Dimension, C : Dimension> coerceUnsafe(structure: Structure2D<T>): DMatrix<T, R, C> {
             return DMatrixWrapper(structure)
         }
     }
@@ -57,7 +57,7 @@ interface DPoint<T, D : Dimension> : Point<T> {
             return DPointWrapper(point)
         }
 
-        inline fun <T, reified D : Dimension> coerceUnsafe(point: Point<T>): DPoint<T, D> {
+        fun <T, D : Dimension> coerceUnsafe(point: Point<T>): DPoint<T, D> {
             return DPointWrapper(point)
         }
     }
@@ -81,8 +81,15 @@ inline class DPointWrapper<T, D : Dimension>(val point: Point<T>) :
  */
 inline class DMatrixContext<T : Any, Ri : Ring<T>>(val context: GenericMatrixContext<T, Ri>) {
 
-    inline fun <reified R : Dimension, reified C : Dimension> Matrix<T>.coerce(): DMatrix<T, C, R> =
-        DMatrix.coerceUnsafe(this)
+    inline fun <reified R : Dimension, reified C : Dimension> Matrix<T>.coerce(): DMatrix<T, R, C> {
+        if (rowNum != Dimension.dim<R>().toInt()) {
+            error("Row number mismatch: expected ${Dimension.dim<R>()} but found $rowNum")
+        }
+        if (colNum != Dimension.dim<C>().toInt()) {
+            error("Column number mismatch: expected ${Dimension.dim<C>()} but found $colNum")
+        }
+        return DMatrix.coerceUnsafe(this)
+    }
 
     /**
      * Produce a matrix with this context and given dimensions
@@ -90,7 +97,7 @@ inline class DMatrixContext<T : Any, Ri : Ring<T>>(val context: GenericMatrixCon
     inline fun <reified R : Dimension, reified C : Dimension> produce(noinline initializer: (i: Int, j: Int) -> T): DMatrix<T, R, C> {
         val rows = Dimension.dim<R>()
         val cols = Dimension.dim<C>()
-        return context.produce(rows.toInt(), cols.toInt(), initializer).coerce()
+        return context.produce(rows.toInt(), cols.toInt(), initializer).coerce<R,C>()
     }
 
     inline fun <reified D : Dimension> point(noinline initializer: (Int) -> T): DPoint<T, D> {

kmath-dimensions/src/commonTest/kotlin/scientifik/dimensions/DMatrixContextTest.kt

@@ -1,5 +1,8 @@
-package scientifik.kmath.dimensions
+package scientifik.dimensions
 
+import scientifik.kmath.dimensions.D2
+import scientifik.kmath.dimensions.D3
+import scientifik.kmath.dimensions.DMatrixContext
 import kotlin.test.Test
 
 

kmath-dimensions/src/jsMain/kotlin/scientifik/kmath/dimensions/dim.kt

@@ -1,5 +1,22 @@
 package scientifik.kmath.dimensions
 
-actual inline fun <reified D : Dimension> Dimension.Companion.dim(): UInt {
-    TODO("KClass::objectInstance does not work")
+import kotlin.reflect.KClass
+
+private val dimensionMap = hashMapOf<UInt, Dimension>(
+    1u to D1,
+    2u to D2,
+    3u to D3
+)
+
+@Suppress("UNCHECKED_CAST")
+actual fun <D : Dimension> Dimension.Companion.resolve(type: KClass<D>): D {
+    return dimensionMap.entries.find { it.value::class == type }?.value as? D ?: error("Can't resolve dimension $type")
+}
+
+actual fun Dimension.Companion.of(dim: UInt): Dimension {
+    return dimensionMap.getOrPut(dim) {
+        object : Dimension {
+            override val dim: UInt get() = dim
+        }
+    }
 }

kmath-dimensions/src/jvmMain/kotlin/scientifik/kmath/dimensions/dim.kt

@@ -1,4 +1,18 @@
 package scientifik.kmath.dimensions
 
-actual inline fun <reified D : Dimension> Dimension.Companion.dim(): UInt =
-    D::class.objectInstance?.dim ?: error("Dimension object must be a singleton")
+import kotlin.reflect.KClass
+
+actual fun <D:Dimension> Dimension.Companion.resolve(type: KClass<D>): D{
+    return type.objectInstance ?: error("No object instance for dimension class")
+}
+
+actual fun Dimension.Companion.of(dim: UInt): Dimension{
+    return when(dim){
+        1u -> D1
+        2u -> D2
+        3u -> D3
+        else -> object : Dimension {
+            override val dim: UInt get() = dim
+        }
+    }
+}