Kotlin Contracts with generics!

build.gradle.kts

@@ -194,6 +194,7 @@ allprojects {
             // enables support for kotlin.time.Instant as kotlinx.datetime.Instant was deprecated; Issue #1350
             // Can be removed once kotlin.time.Instant is marked "stable".
             optIn.add("kotlin.time.ExperimentalTime")
+            freeCompilerArgs.addAll("-Xallow-contracts-on-more-functions")
         }
     }
 

core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/DataColumnType.kt

@@ -1,8 +1,9 @@
-@file:OptIn(ExperimentalContracts::class)
+@file:OptIn(ExperimentalContracts::class, ExperimentalExtendedContracts::class)
 
 package org.jetbrains.kotlinx.dataframe.api
 
 import org.jetbrains.kotlinx.dataframe.AnyCol
+import org.jetbrains.kotlinx.dataframe.DataColumn
 import org.jetbrains.kotlinx.dataframe.columns.ColumnGroup
 import org.jetbrains.kotlinx.dataframe.columns.ColumnKind
 import org.jetbrains.kotlinx.dataframe.columns.FrameColumn
@@ -17,6 +18,7 @@ import org.jetbrains.kotlinx.dataframe.util.IS_COMPARABLE
 import org.jetbrains.kotlinx.dataframe.util.IS_COMPARABLE_REPLACE
 import org.jetbrains.kotlinx.dataframe.util.IS_INTER_COMPARABLE_IMPORT
 import kotlin.contracts.ExperimentalContracts
+import kotlin.contracts.ExperimentalExtendedContracts
 import kotlin.contracts.contract
 import kotlin.reflect.KType
 import kotlin.reflect.full.isSubtypeOf
@@ -40,21 +42,36 @@ public fun AnyCol.isValueColumn(): Boolean {
 public fun AnyCol.isSubtypeOf(type: KType): Boolean =
     this.type.isSubtypeOf(type) && (!this.type.isMarkedNullable || type.isMarkedNullable)
 
-public inline fun <reified T> AnyCol.isSubtypeOf(): Boolean = isSubtypeOf(typeOf<T>())
+public inline fun <reified T> AnyCol.isSubtypeOf(): Boolean {
+    contract { returns(true) implies (this@isSubtypeOf is DataColumn<T>) }
+    return isSubtypeOf(typeOf<T>())
+}
 
-public inline fun <reified T> AnyCol.isType(): Boolean = type() == typeOf<T>()
+public inline fun <reified T> AnyCol.isType(): Boolean {
+    contract { returns(true) implies (this@isType is DataColumn<T>) }
+    return type() == typeOf<T>()
+}
 
 /** Returns `true` when this column's type is a subtype of `Number?` */
-public fun AnyCol.isNumber(): Boolean = isSubtypeOf<Number?>()
+public fun AnyCol.isNumber(): Boolean {
+    contract { returns(true) implies (this@isNumber is ValueColumn<Number?>) }
+    return isSubtypeOf<Number?>()
+}
 
 /** Returns `true` only when this column's type is exactly `Number` or `Number?`. */
-public fun AnyCol.isMixedNumber(): Boolean = type().isMixedNumber()
+public fun AnyCol.isMixedNumber(): Boolean {
+    contract { returns(true) implies (this@isMixedNumber is ValueColumn<Number?>) }
+    return type().isMixedNumber()
+}
 
 /**
  * Returns `true` when this column has the (nullable) type of either:
  * [Byte], [Short], [Int], [Long], [Float], or [Double].
  */
-public fun AnyCol.isPrimitiveNumber(): Boolean = type().isPrimitiveNumber()
+public fun AnyCol.isPrimitiveNumber(): Boolean {
+    contract { returns(true) implies (this@isPrimitiveNumber is ValueColumn<Number?>) }
+    return type().isPrimitiveNumber()
+}
 
 /**
  * Returns `true` when this column has the (nullable) type of either:
@@ -63,9 +80,15 @@ public fun AnyCol.isPrimitiveNumber(): Boolean = type().isPrimitiveNumber()
  * Careful: Will return `true` if the column contains multiple number types that
  * might NOT be primitive.
  */
-public fun AnyCol.isPrimitiveOrMixedNumber(): Boolean = type().isPrimitiveOrMixedNumber()
+public fun AnyCol.isPrimitiveOrMixedNumber(): Boolean {
+    contract { returns(true) implies (this@isPrimitiveOrMixedNumber is ValueColumn<Number?>) }
+    return type().isPrimitiveOrMixedNumber()
+}
 
-public fun AnyCol.isList(): Boolean = typeClass == List::class
+public fun AnyCol.isList(): Boolean {
+    contract { returns(true) implies (this@isList is ValueColumn<List<*>>) }
+    return typeClass == List::class
+}
 
 /** @include [valuesAreComparable] */
 @Deprecated(
@@ -84,4 +107,7 @@ public fun AnyCol.isComparable(): Boolean = valuesAreComparable()
  *
  * Technically, this means the values' common type `T(?)` is a subtype of [Comparable]`<in T>(?)`
  */
-public fun AnyCol.valuesAreComparable(): Boolean = isValueColumn() && type().isIntraComparable()
+public fun <T> DataColumn<T>.valuesAreComparable(): Boolean {
+    contract { returns(true) implies (this@valuesAreComparable is ValueColumn<Comparable<T>>) }
+    return isValueColumn() && type().isIntraComparable()
+}

core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/typeConversions.kt

@@ -89,13 +89,13 @@ public fun DataColumn<Any>.asNumbers(): ValueColumn<Number> {
 
 public fun <T : Any> DataColumn<T>.asComparable(): DataColumn<Comparable<T>> {
     require(valuesAreComparable())
-    return this as DataColumn<Comparable<T>>
+    return this
 }
 
 @JvmName("asComparableNullable")
 public fun <T : Any?> DataColumn<T?>.asComparable(): DataColumn<Comparable<T>?> {
     require(valuesAreComparable())
-    return this as DataColumn<Comparable<T>?>
+    return this
 }
 
 public fun <T> ColumnReference<T?>.castToNotNullable(): ColumnReference<T> = cast()

core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/describe.kt

@@ -55,8 +55,8 @@ internal fun describeImpl(cols: List<AnyCol>): DataFrame<ColumnDescription> {
                 ?.key
         }
         if (hasNumericCols) {
-            ColumnDescription::mean from { if (it.isNumber()) it.asNumbers().mean() else null }
-            ColumnDescription::std from { if (it.isNumber()) it.asNumbers().std() else null }
+            ColumnDescription::mean from { if (it.isNumber()) it.mean() else null }
+            ColumnDescription::std from { if (it.isNumber()) it.std() else null }
         }
         if (hasComparableCols || hasNumericCols) {
             ColumnDescription::min from inferType {
@@ -113,10 +113,10 @@ private fun List<AnyCol>.collectAll(atAnyDepth: Boolean): List<AnyCol> =
 @Suppress("UNCHECKED_CAST")
 private fun DataColumn<Any?>.convertToComparableOrNull(): DataColumn<Comparable<Any>?>? {
     return when {
-        valuesAreComparable() -> asComparable()
+        valuesAreComparable() -> this
 
         // Found incomparable number types, convert all to Double first
-        isNumber() -> cast<Number?>().map {
+        isNumber() -> map {
             if (it?.isPrimitiveNumber() == false) {
                 // Cannot calculate statistics of a non-primitive number type
                 return@convertToComparableOrNull null

core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/parse.kt

@@ -736,7 +736,7 @@ internal fun <T> DataFrame<T>.parseImpl(options: ParserOptions?, columns: Column
 
             // Base case, parse the column if it's a `String?` column
             col.isSubtypeOf<String?>() ->
-                col.cast<String?>().tryParseImpl(options)
+                col.tryParseImpl(options)
 
             else -> col
         }

core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/html.kt

@@ -170,7 +170,7 @@ internal fun AnyFrame.toHtmlData(
         configuration: DisplayConfiguration,
     ): ColumnDataForJs {
         val values = if (rowsLimit != null) rows().take(rowsLimit) else rows()
-        val scale = if (col.isNumber()) col.asNumbers().scale() else 1
+        val scale = if (col.isNumber()) col.scale() else 1
         val format = if (scale > 0) {
             RendererDecimalFormat.fromPrecision(scale)
         } else {

core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/string.kt

@@ -48,7 +48,7 @@ public fun AnyFrame.renderToString(
     }
     val values = cols.map {
         val top = it.take(rowsLimit)
-        val precision = if (top.isNumber()) top.asNumbers().scale() else 0
+        val precision = if (top.isNumber()) top.scale() else 0
         val decimalFormat =
             if (precision >= 0) RendererDecimalFormat.fromPrecision(precision) else RendererDecimalFormat.of("%e")
         top.values().map {

dataframe-json/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/io/writeJson.kt

@@ -176,7 +176,7 @@ internal fun encodeValue(col: AnyCol, index: Int, customEncoders: List<CustomEnc
         matchingEncoder != null -> matchingEncoder.encode(col[index])
 
         col.isList() -> col[index]?.let { list ->
-            val values = (list as List<*>).map { convert(it) }
+            val values = list.map { convert(it) }
             JsonArray(values)
         } ?: JsonArray(emptyList())