From 380c76fe40d2cdaeca9731aeb0d5b54b225ea37a Mon Sep 17 00:00:00 2001
From: Alexander Nozik <altavir@gmail.com>
Date: Sun, 23 May 2021 20:01:07 +0300
Subject: [PATCH] Piecewise optimization

---
 build.gradle.kts                              |  2 +-
 .../space/kscience/kmath/misc/annotations.kt  |  4 +-
 .../kscience/kmath/functions/Piecewise.kt     | 49 ++++++++++++++-----
 3 files changed, 42 insertions(+), 13 deletions(-)

diff --git a/build.gradle.kts b/build.gradle.kts
index ef9c60c55..4de6d8bad 100644
--- a/build.gradle.kts
+++ b/build.gradle.kts
@@ -15,7 +15,7 @@ allprojects {
     }
 
     group = "space.kscience"
-    version = "0.3.0-dev-12"
+    version = "0.3.0-dev-13"
 }
 
 subprojects {
diff --git a/kmath-core/src/commonMain/kotlin/space/kscience/kmath/misc/annotations.kt b/kmath-core/src/commonMain/kotlin/space/kscience/kmath/misc/annotations.kt
index 18718de97..e521e6237 100644
--- a/kmath-core/src/commonMain/kotlin/space/kscience/kmath/misc/annotations.kt
+++ b/kmath-core/src/commonMain/kotlin/space/kscience/kmath/misc/annotations.kt
@@ -26,4 +26,6 @@ public annotation class UnstableKMathAPI
     "Refer to the documentation to use this API in performance-critical code",
     RequiresOptIn.Level.WARNING
 )
-public annotation class PerformancePitfall
+public annotation class PerformancePitfall(
+    val message: String = "Potential performance problem"
+)
diff --git a/kmath-functions/src/commonMain/kotlin/space/kscience/kmath/functions/Piecewise.kt b/kmath-functions/src/commonMain/kotlin/space/kscience/kmath/functions/Piecewise.kt
index dd9c571b0..73fa57c7b 100644
--- a/kmath-functions/src/commonMain/kotlin/space/kscience/kmath/functions/Piecewise.kt
+++ b/kmath-functions/src/commonMain/kotlin/space/kscience/kmath/functions/Piecewise.kt
@@ -5,6 +5,7 @@
 
 package space.kscience.kmath.functions
 
+import space.kscience.kmath.misc.PerformancePitfall
 import space.kscience.kmath.operations.Ring
 
 /**
@@ -24,17 +25,43 @@ public fun interface Piecewise<T, R> {
  * Represents piecewise-defined function where all the sub-functions are polynomials.
  * @param pieces An ordered list of range-polynomial pairs. The list does not in general guarantee that there are no "holes" in it.
  */
-public class PiecewisePolynomial<T : Comparable<T>>(
-    public val pieces: List<Pair<ClosedRange<T>, Polynomial<T>>>,
-) : Piecewise<T, Polynomial<T>> {
+public interface PiecewisePolynomial<T : Comparable<T>> : Piecewise<T, Polynomial<T>> {
+    public val pieces: Collection<Pair<ClosedRange<T>, Polynomial<T>>>
 
-    public override fun findPiece(arg: T): Polynomial<T>? {
-        return if (arg < pieces.first().first.start || arg >= pieces.last().first.endInclusive)
-            null
-        else {
-            pieces.firstOrNull { arg in it.first }?.second
+    public override fun findPiece(arg: T): Polynomial<T>?
+}
+
+/**
+ * A generic piecewise without constraints on how pieces are placed
+ */
+@PerformancePitfall("findPiece method of resulting piecewise is slow")
+public fun <T : Comparable<T>> PiecewisePolynomial(
+    pieces: Collection<Pair<ClosedRange<T>, Polynomial<T>>>,
+): PiecewisePolynomial<T> = object : PiecewisePolynomial<T> {
+    override val pieces: Collection<Pair<ClosedRange<T>, Polynomial<T>>> = pieces
+
+    override fun findPiece(arg: T): Polynomial<T>? = pieces.firstOrNull { arg in it.first }?.second
+}
+
+/**
+ * An optimized piecewise which uses not separate pieces, but a range separated by delimiters.
+ * The pices search is logarithmic
+ */
+private class OrderedPiecewisePolynomial<T : Comparable<T>>(
+    override val pieces: List<Pair<ClosedRange<T>, Polynomial<T>>>,
+) : PiecewisePolynomial<T> {
+
+    override fun findPiece(arg: T): Polynomial<T>? {
+        val index = pieces.binarySearch { (range, _) ->
+            when {
+                arg >= range.endInclusive -> -1
+                arg < range.start -> +1
+                else -> 0
+            }
         }
+        return if (index < 0) null else pieces[index].second
     }
+
 }
 
 /**
@@ -71,9 +98,9 @@ public class PiecewiseBuilder<T : Comparable<T>>(delimiter: T) {
         pieces.add(0, piece)
     }
 
-    public fun build(): PiecewisePolynomial<T> {
-        return PiecewisePolynomial(delimiters.zipWithNext { l, r -> l..r }.zip(pieces))
-    }
+    public fun build(): PiecewisePolynomial<T> = OrderedPiecewisePolynomial(delimiters.zipWithNext { l, r ->
+        l..r
+    }.zip(pieces))
 }
 
 /**