Make SortedMap implement Map interface. Not changing the name yet. (#534

)

Some work on https://github.com/korlibs/kds/issues/20

kds/src/commonMain/kotlin/com/soywiz/kds/SortedMap.kt

@@ -1,31 +1,67 @@
 package com.soywiz.kds
 
 import com.soywiz.kds.comparator.*
+import com.soywiz.kds.map.*
 import kotlin.math.*
 
-open class SortedMap<K, V>(val comparator: Comparator<K>) {
+fun <K, V> sortedMapOf(comparator: Comparator<K>, vararg values: Pair<K, V>): SortedMap<K, V> =
+    SortedMap<K, V>(comparator).also { it.putAll(values) }
+
+fun <K : Comparable<K>, V> sortedMapOf(vararg values: Pair<K, V>): SortedMap<K, V> =
+    SortedMap<K, V>().also { it.putAll(values) }
+
+open class SortedMap<K, V>(val comparator: Comparator<K>) : MutableMapExt<K, V> {
     private val keysToIndex = hashMapOf<K, Int>() // @TODO: Maybe we could just try to perform a binary search?
-    private val keys = FastArrayList<K>()
-    private val values = FastArrayList<V>()
+    private val _keys = FastArrayList<K>()
+    private val _values = FastArrayList<V>()
     private var isSorted = true
-    val size get() = keys.size
+    override val size get() = _keys.size
+
+    override fun containsKey(key: K): Boolean = keysToIndex.contains(key) // O(1)
+    override fun containsValue(value: V): Boolean = _values.contains(value) // O(n)
+
+    override val keys: MutableSet<K> get() {
+        ensureSorted()
+        return _keys.toMutableSet()
+    }
+    override val values: MutableCollection<V> get() {
+        ensureSorted()
+        return _values.toMutableList()
+    }
+    override val entries: MutableSet<MutableMap.MutableEntry<K, V>> get() {
+        ensureSorted()
+        return MutableEntryExt.fromMap(this, _keys)
+    }
+
+    override fun clear() {
+        keysToIndex.clear()
+        _keys.clear()
+        _values.clear()
+        isSorted = true
+    }
+
+    override fun put(key: K, value: V): V? {
+        val oldValue = get(key)
+        set(key, value)
+        return oldValue
+    }
 
     companion object {
         inline operator fun <K : Comparable<K>, V> invoke() = SortedMap<K, V>(ComparatorComparable())
     }
 
     fun swap(indexL: Int, indexR: Int) {
-        val keyL = keys[indexL]
-        val keyR = keys[indexR]
-        keys.swap(indexL, indexR)
-        values.swap(indexL, indexR)
+        val keyL = _keys[indexL]
+        val keyR = _keys[indexR]
+        _keys.swap(indexL, indexR)
+        _values.swap(indexL, indexR)
         keysToIndex[keyL] = indexR
         keysToIndex[keyR] = indexL
     }
 
     object Sorting : SortOps<SortedMap<Any, Any>>() {
         override fun compare(subject: SortedMap<Any, Any>, l: Int, r: Int): Int {
-            return subject.comparator.compare(subject.keys[l], subject.keys[r])
+            return subject.comparator.compare(subject._keys[l], subject._keys[r])
         }
 
         override fun swap(subject: SortedMap<Any, Any>, indexL: Int, indexR: Int) {
@@ -35,15 +71,18 @@ open class SortedMap<K, V>(val comparator: Comparator<K>) {
 
     fun removeAt(index: Int) {
         isSorted = false
-        val key = keys[index]
+        val key = _keys[index]
         swap(index, size - 1)
-        keys.removeAt(size - 1)
-        values.removeAt(size - 1)
+        _keys.removeAt(size - 1)
+        _values.removeAt(size - 1)
         keysToIndex.remove(key)
     }
 
-    fun remove(key: K) {
-        removeAt(keysToIndex[key] ?: return)
+    override fun remove(key: K): V? {
+        val index = keysToIndex[key] ?: return null
+        val value = _values[index]
+        removeAt(index)
+        return value
     }
 
     @PublishedApi
@@ -56,27 +95,27 @@ open class SortedMap<K, V>(val comparator: Comparator<K>) {
     operator fun set(key: K, value: V) {
         val index = keysToIndex[key]
         if (index != null) {
-            values[index] = value
+            _values[index] = value
         } else {
             isSorted = false
-            keysToIndex[key] = keys.size
-            keys.add(key)
-            values.add(value)
+            keysToIndex[key] = _keys.size
+            _keys.add(key)
+            _values.add(value)
         }
     }
 
-    operator fun get(key: K): V? {
-        return values[keysToIndex[key] ?: return null]
+    override operator fun get(key: K): V? {
+        return _values[keysToIndex[key] ?: return null]
     }
 
     fun getKeyAt(index: Int): K {
         ensureSorted()
-        return keys[index]
+        return _keys[index]
     }
 
     fun getValueAt(index: Int): V {
         ensureSorted()
-        return values[index]
+        return _values[index]
     }
 
     inline fun fastForEach(block: (index: Int, key: K, value: V) -> Unit) {
@@ -87,43 +126,43 @@ open class SortedMap<K, V>(val comparator: Comparator<K>) {
 
     fun nearestLowHighIndex(key: K, doHigh: Boolean): Int {
         ensureSorted()
-        return genericBinarySearch(0, keys.size, { from, to, low, high -> if (doHigh) max(low, high) else min(low, high) }) { comparator.compare(keys[it], key) }
+        return genericBinarySearch(0, _keys.size, { from, to, low, high -> if (doHigh) max(low, high) else min(low, high) }) { comparator.compare(_keys[it], key) }
     }
 
     fun nearestLowIndex(key: K): Int = nearestLowHighIndex(key, doHigh = false)
     fun nearestHighIndex(key: K): Int = nearestLowHighIndex(key, doHigh = true)
 
-    fun nearestLow(key: K): K? = keys.getOrNull(nearestLowIndex(key))
-    fun nearestHigh(key: K): K? = keys.getOrNull(nearestHighIndex(key))
+    fun nearestLow(key: K): K? = _keys.getOrNull(nearestLowIndex(key))
+    fun nearestHigh(key: K): K? = _keys.getOrNull(nearestHighIndex(key))
 
     fun nearestLowExcludingExact(key: K): K? {
         val bindex = nearestLowIndex(key)
         val index = if (key in keysToIndex) bindex - 1 else bindex
-        return keys.getOrNull(index)
+        return _keys.getOrNull(index)
     }
     fun nearestHighExcludingExact(key: K): K? {
         val bindex = nearestHighIndex(key)
         val index = if (key in keysToIndex) bindex + 1 else bindex
-        return keys.getOrNull(index)
+        return _keys.getOrNull(index)
     }
 
     fun keysToList(): List<K> {
         ensureSorted()
-        return keys.toFastList()
+        return _keys.toFastList()
     }
 
     fun valuesToList(): List<V> {
         ensureSorted()
-        return values.toFastList()
+        return _values.toFastList()
     }
 
     fun toList(): List<Pair<K, V>> {
         ensureSorted()
-        return (0 until size).map { keys[it] to values[it] }
+        return (0 until size).map { _keys[it] to _values[it] }
     }
 
     fun toMap(): Map<K, V> {
         ensureSorted()
-        return (0 until size).map { keys[it] to values[it] }.toLinkedMap()
+        return (0 until size).map { _keys[it] to _values[it] }.toLinkedMap()
     }
 }

kds/src/commonMain/kotlin/com/soywiz/kds/map/MutableMapExt.kt

@@ -0,0 +1,31 @@
+package com.soywiz.kds.map
+
+import com.soywiz.kds.*
+
+interface MutableMapExt<K, V> : MutableMap<K, V> {
+    override fun isEmpty(): Boolean = size == 0
+
+    override fun putAll(from: Map<out K, V>) {
+        from.forEach { set(it.key, it.value) }
+    }
+}
+
+class MutableEntryExt<K, V>(
+    val map: MutableMap<K, V>,
+    override val key: K,
+) : MutableMap.MutableEntry<K, V> {
+    override val value: V get() = map[key]!!
+
+    override fun setValue(newValue: V): V {
+        val oldValue = value
+        map[key] = newValue
+        return oldValue
+    }
+
+    companion object {
+        fun <K, V> fromMap(map: MutableMap<K, V>, keys: Collection<K>): MutableSet<MutableMap.MutableEntry<K, V>> {
+            return keys.map { MutableEntryExt(map, it) }.toMutableSet()
+        }
+    }
+}
+

kds/src/commonTest/kotlin/com/soywiz/kds/SortedMapTest.kt

@@ -99,4 +99,10 @@ class SortedMapTest {
             assertEquals(null, map.nearestHigh(30))
         }
     }
+
+    @Test
+    fun testSortedMapOf() {
+        val sortedValues = sortedMapOf(2 to "two", 3 to "three", 1 to "one").values.toList()
+        assertEquals(listOf("one", "two", "three"), sortedValues)
+    }
 }