Rewrite tests and remove some support functions

Author: Virtlink

.github/workflows/build.yaml

@@ -18,13 +18,10 @@ jobs:
           - 'ubuntu-latest'
         java:
           - version: '11'
+          - version: '17'
         gradle:
           - version: '7.6.4'      # Latest stable 7
             command: 'gradle'
-          - version: '8.8'        # Latest stable 8
-            command: 'gradle'
-          - version: 'current'    # Latest stable
-            command: 'gradle'
           - version: 'wrapper'
             command: './gradlew'  # Workaround for issue https://github.com/gradle/actions/issues/273
     runs-on: ${{ matrix.os }}

disjointmap/build.gradle.kts

@@ -7,8 +7,7 @@ val junitVersion = "5.8.1"
 
 dependencies {
     implementation      (libs.kotlinx.immutablecollections)
-
     compileOnly         (libs.jsr305)
-
     testImplementation  (libs.junit)
+    testImplementation  (libs.kotest)
 }

disjointmap/src/main/java/com/virtlink/collections/N.java

@@ -1,10 +1,10 @@
 package com.virtlink.collections
 
 /**
- * A map of key in disjoint sets with a value.
+ * A map of disjoint sets associated with a value.
  *
- * @param K the type of keys
- * @param V the type of values
+ * @param K The type of keys.
+ * @param V The type of values.
  */
 interface DisjointMap<K, out V> {
 
@@ -21,34 +21,34 @@ interface DisjointMap<K, out V> {
      * Gets the value associated with the set that contains the specified key;
      * or `null` when the key is not in the map.
      *
-     * @param key the key whose associated value to get
-     * @return the associated value; or `null` when the key is not in the map
+     * @param key The key whose associated value to get.
+     * @return The associated value; or `null` when the key is not in the map.
      */
     operator fun get(key: K): V?
 
     /**
      * Gets the value associated with the set that contains the specified key;
      * or the default value when the key is not in the map.
      *
-     * @param key the key whose associated value to get
-     * @return the associated value; or [defaultValue] when the key is not in the map
+     * @param key The key whose associated value to get.
+     * @return The associated value (which may be `null`); or [defaultValue] when the key is not in the map.
      */
     fun getOrDefault(key: K, defaultValue: @UnsafeVariance V): V
 
     /**
      * Finds the representative of the set that includes the given key.
      *
-     * @param key the key to look for
-     * @return the representative of the set that includes the given key;
-     * which may be the key itself; or `null` when the key is not in the map
+     * @param key The key to look for.
+     * @return The representative of the set that includes the given key;
+     * which may be the key itself; or `null` when the key is not in the map.
      */
     fun find(key: K): K?
 
     /**
      * Whether the collection contains the specified key.
      *
      * @return `true` when the collection contains the specified key;
-     * otherwise, `false`
+     * otherwise, `false`.
      */
     operator fun contains(key: K): Boolean {
         return find(key) != null
@@ -57,10 +57,10 @@ interface DisjointMap<K, out V> {
     /**
      * Determines whether the given keys are in the same set.
      *
-     * @param key1 one key
-     * @param key2 another key
+     * @param key1 One key.
+     * @param key2 Another key.
      * @return `true` when the keys are present and part of the same set;
-     * otherwise, `false`
+     * otherwise, `false`.
      */
     fun same(key1: K, key2: K): Boolean {
         val rep1 = find(key1)
@@ -71,7 +71,7 @@ interface DisjointMap<K, out V> {
     /**
      * Copies the sets from this disjoint map to a new map.
      *
-     * @return the new map
+     * @return The new map of disjoint sets to values.
      */
     fun toMap(): Map<Set<K>, V>
 

disjointmap/src/main/kotlin/com/virtlink/collections/DisjointMap.kt

@@ -29,8 +29,8 @@ internal fun <K, V> findMutable(
 
     // TODO: Can we do this iteratively, to avoid updating the parents and ranks maps so much?
     // Find the representative of the parent key
-    val representative = N.assumeNotNull(findMutable(parent, roots, parents, ranks))
-    if (parent != representative) {
+    val representative = findMutable(parent, roots, parents, ranks)
+    if (representative != null && parent != representative) {
         // Path compression:
         // Update the key to point directly to the representative key
         parents[key] = representative
@@ -64,6 +64,7 @@ internal fun <K, V> findMutable(
  * @param V the type of values
  * @return whether disjoint sets have been unified
  */
+@Suppress("UNCHECKED_CAST")
 internal fun <K, V> unionMutable(
     key1: K,
     key2: K,
@@ -92,14 +93,14 @@ internal fun <K, V> unionMutable(
         elem in roots.keys && repr in roots.keys -> {
             // NOTE: We know both `element` and `rep` are in the values map,
             //  so values.get() should only return null when their value happens to be null.
-            val repValue = N.of(roots[repr])
-            val elemValue = N.of(roots[elem])
+            val repValue = roots[repr] as V
+            val elemValue = roots[elem] as V
             unify(repValue, elemValue)
         }
         // - the value associated with the representative element;
-        repr in roots.keys -> N.of(roots[repr])
+        repr in roots.keys -> roots[repr] as V
         // - the value associated with the eliminated element; or
-        elem in roots.keys -> N.of(roots[elem])
+        elem in roots.keys -> roots[elem] as V
         // - the default value when neither element has an associated value.
         else -> default()
     }
@@ -241,7 +242,10 @@ internal fun <K, V> DisjointMap<K, V>.toMapImpl(
     roots.keys.forEach { k -> mapping[k] = mutableSetOf(k) }
     parents.keys.forEach { k -> mapping[find(k)]!!.add(k) }
 
-    return mapping.map { (rep, keys) -> keys.toPersistentSet() to N.of(roots[rep]) }
+    return mapping.map { (rep, keys) ->
+        @Suppress("UNCHECKED_CAST")
+        keys.toPersistentSet() to (roots[rep] as V)
+    }
         .toMap<Set<K>, V>()
 }
 

disjointmap/src/main/kotlin/com/virtlink/collections/DisjointMapImpl.kt

@@ -5,8 +5,8 @@ import kotlinx.collections.immutable.ImmutableMap
 /**
  * An immutable disjoint map.
  *
- * @param K the type of keys
- * @param V the type of values
+ * @param K The type of keys.
+ * @param V The type of values.
  */
 interface ImmutableDisjointMap<K, out V> : DisjointMap<K, V> {
 

disjointmap/src/main/kotlin/com/virtlink/collections/ImmutableDisjointMap.kt

@@ -29,7 +29,10 @@ class MutableUnionFindMap<K, V> internal constructor(
 
     override fun getOrDefault(key: K, defaultValue: @UnsafeVariance V): V {
         val rep = find(key) ?: return defaultValue
-        return this._roots[rep] ?: defaultValue
+        @Suppress("UNCHECKED_CAST")
+        // We know the key is in the map,
+        // so this only returns `null` if the value in the map is `null`.
+        return this._roots[rep] as V
     }
 
     override fun find(key: K): K? {

disjointmap/src/main/kotlin/com/virtlink/collections/MutableUnionFindMap.kt

@@ -39,7 +39,10 @@ class PersistentUnionFindMap<K, V> internal constructor(
 
     override fun getOrDefault(key: K, defaultValue: @UnsafeVariance V): V {
         val rep = find(key) ?: return defaultValue
-        return this._roots[rep] ?: defaultValue
+        @Suppress("UNCHECKED_CAST")
+        // We know the key is in the map,
+        // so this only returns `null` if the value in the map is `null`.
+        return this._roots[rep] as V
     }
 
     override fun find(key: K): K? {