From 4b91481e96669b8ea1654cf535c467a38ede2651 Mon Sep 17 00:00:00 2001
From: Gabriel Tofvesson <contact@w1zzrd.dev>
Date: Fri, 24 Sep 2021 02:14:17 +0200
Subject: [PATCH] Add useful functions to MultiSortedList

---
 src/main/kotlin/MultiSortedList.kt | 34 ++++++++++++++++++++++++++++++
 1 file changed, 34 insertions(+)

diff --git a/src/main/kotlin/MultiSortedList.kt b/src/main/kotlin/MultiSortedList.kt
index c075236..8eaa4f1 100644
--- a/src/main/kotlin/MultiSortedList.kt
+++ b/src/main/kotlin/MultiSortedList.kt
@@ -44,6 +44,10 @@ class MultiSortedList<E> constructor(
         return true
     }
 
+    fun contains(element: E, comparator: Comparator<E>) =
+        if (comparator == this.comparator) contains(element)
+        else extraLists[comparator]!!.contains(element)
+
     override fun remove(element: E): Boolean {
         if (super.remove(element)) {
             extraLists.values.forEach {
@@ -90,6 +94,36 @@ class MultiSortedList<E> constructor(
         if (comparator == this.comparator) this[index]
         else extraLists[comparator]!![index]
 
+    /**
+     * Get all contiguous elements that match a comparison
+     * @return Iterable object of all matching elements, or null if none exist
+     */
+    fun getAll(comparator: Comparator<in E>, comparison: (E) -> Int): Iterable<E>? {
+        var index = binSearch(comparator, comparison)
+        if (index < 0) return null
+
+        val result = LinkedList<E>()
+
+        var element: E
+        do {
+            element = get(index++, comparator)
+            if (comparison(element) != 0)
+                break
+
+            result += element
+        } while (index < size)
+
+        return result
+    }
+
+    fun findValueOrNull(comparator: Comparator<in E>, comparison: (E) -> Int): E? {
+        val index = binSearch(comparator, comparison)
+
+        if (index < 0) return null
+
+        return get(index, comparator)
+    }
+
     fun binSearch(element: E, comparator: Comparator<in E>) =
         if (comparator == this.comparator) binarySearch(element, comparator)
         else extraLists[comparator]!!.binarySearch(element, comparator)