chore(arrow-atomic): Cleanup atomic code (#3318)

Co-authored-by: Alejandro Serrano <trupill@gmail.com>
arrow-kt · Dec 20, 2023 · a4c4053 · a4c4053
1 parent 46ad6fa
commit a4c4053
Show file tree

Hide file tree

Showing 8 changed files with 87 additions and 140 deletions.
diff --git a/arrow-libs/core/arrow-atomic/api/arrow-atomic.api b/arrow-libs/core/arrow-atomic/api/arrow-atomic.api
@@ -12,7 +12,9 @@ public final class arrow/atomic/AtomicBooleanKt {
 	public static final fun getAndUpdate (Larrow/atomic/AtomicBoolean;Lkotlin/jvm/functions/Function1;)Z
 	public static final fun loop (Larrow/atomic/AtomicBoolean;Lkotlin/jvm/functions/Function1;)Ljava/lang/Void;
 	public static final fun tryUpdate (Larrow/atomic/AtomicBoolean;Lkotlin/jvm/functions/Function1;)Z
+	public static final fun tryUpdate (Larrow/atomic/AtomicBoolean;Lkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function2;)Z
 	public static final fun update (Larrow/atomic/AtomicBoolean;Lkotlin/jvm/functions/Function1;)V
+	public static final fun update (Larrow/atomic/AtomicBoolean;Lkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function2;)Ljava/lang/Object;
 	public static final fun updateAndGet (Larrow/atomic/AtomicBoolean;Lkotlin/jvm/functions/Function1;)Z
 }
 
@@ -22,7 +24,9 @@ public final class arrow/atomic/AtomicIntKt {
 	public static final fun loop (Ljava/util/concurrent/atomic/AtomicInteger;Lkotlin/jvm/functions/Function1;)Ljava/lang/Void;
 	public static final fun setValue (Ljava/util/concurrent/atomic/AtomicInteger;I)V
 	public static final fun tryUpdate (Ljava/util/concurrent/atomic/AtomicInteger;Lkotlin/jvm/functions/Function1;)Z
+	public static final fun tryUpdate (Ljava/util/concurrent/atomic/AtomicInteger;Lkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function2;)Z
 	public static final fun update (Ljava/util/concurrent/atomic/AtomicInteger;Lkotlin/jvm/functions/Function1;)V
+	public static final fun update (Ljava/util/concurrent/atomic/AtomicInteger;Lkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function2;)Ljava/lang/Object;
 	public static final fun updateAndGet (Ljava/util/concurrent/atomic/AtomicInteger;Lkotlin/jvm/functions/Function1;)I
 }
 
@@ -32,7 +36,9 @@ public final class arrow/atomic/AtomicKt {
 	public static final fun loop (Ljava/util/concurrent/atomic/AtomicReference;Lkotlin/jvm/functions/Function1;)Ljava/lang/Void;
 	public static final fun setValue (Ljava/util/concurrent/atomic/AtomicReference;Ljava/lang/Object;)V
 	public static final fun tryUpdate (Ljava/util/concurrent/atomic/AtomicReference;Lkotlin/jvm/functions/Function1;)Z
+	public static final fun tryUpdate (Ljava/util/concurrent/atomic/AtomicReference;Lkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function2;)Z
 	public static final fun update (Ljava/util/concurrent/atomic/AtomicReference;Lkotlin/jvm/functions/Function1;)V
+	public static final fun update (Ljava/util/concurrent/atomic/AtomicReference;Lkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function2;)Ljava/lang/Object;
 	public static final fun updateAndGet (Ljava/util/concurrent/atomic/AtomicReference;Lkotlin/jvm/functions/Function1;)Ljava/lang/Object;
 }
 
@@ -42,7 +48,9 @@ public final class arrow/atomic/AtomicLongKt {
 	public static final fun loop (Ljava/util/concurrent/atomic/AtomicLong;Lkotlin/jvm/functions/Function1;)Ljava/lang/Void;
 	public static final fun setValue (Ljava/util/concurrent/atomic/AtomicLong;J)V
 	public static final fun tryUpdate (Ljava/util/concurrent/atomic/AtomicLong;Lkotlin/jvm/functions/Function1;)Z
+	public static final fun tryUpdate (Ljava/util/concurrent/atomic/AtomicLong;Lkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function2;)Z
 	public static final fun update (Ljava/util/concurrent/atomic/AtomicLong;Lkotlin/jvm/functions/Function1;)V
+	public static final fun update (Ljava/util/concurrent/atomic/AtomicLong;Lkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function2;)Ljava/lang/Object;
 	public static final fun updateAndGet (Ljava/util/concurrent/atomic/AtomicLong;Lkotlin/jvm/functions/Function1;)J
 }
 
diff --git a/arrow-libs/core/arrow-atomic/build.gradle.kts b/arrow-libs/core/arrow-atomic/build.gradle.kts
@@ -1,7 +1,5 @@
 @file:Suppress("DSL_SCOPE_VIOLATION")
 
-import org.jetbrains.kotlin.gradle.tasks.KotlinCompile
-
 plugins {
   id(libs.plugins.kotlin.multiplatform.get().pluginId)
   alias(libs.plugins.arrowGradleConfig.kotlin)
@@ -19,6 +17,12 @@ spotless {
 apply(from = property("ANIMALSNIFFER_MPP"))
 
 kotlin {
+  targets.all {
+    compilations.all {
+      kotlinOptions.freeCompilerArgs += "-Xexpect-actual-classes"
+    }
+  }
+
   sourceSets {
     commonMain {
       dependencies {
@@ -57,9 +61,3 @@ kotlin {
     }
   }
 }
-
-tasks.withType<KotlinCompile>().configureEach {
-  kotlinOptions {
-    freeCompilerArgs = freeCompilerArgs + "-Xexpect-actual-classes"
-  }
-}
diff --git a/arrow-libs/core/arrow-atomic/src/commonMain/kotlin/arrow/atomic/Atomic.kt b/arrow-libs/core/arrow-atomic/src/commonMain/kotlin/arrow/atomic/Atomic.kt
@@ -1,7 +1,7 @@
 package arrow.atomic
 
 /**
- * [Atomic] value of [A].
+ * [Atomic] value of [V].
  *
  * ```kotlin
  * import arrow.atomic.AtomicInt
@@ -44,44 +44,32 @@ public var <T> Atomic<T>.value: T
 /**
  * Infinite loop that reads this atomic variable and performs the specified [action] on its value.
  */
-public inline fun <V> Atomic<V>.loop(action: (V) -> Unit): Nothing {
-  while (true) {
-    action(value)
-  }
-}
+public inline fun <V> Atomic<V>.loop(action: (V) -> Unit): Nothing { while(true) { action(value) } }
 
-public fun <V> Atomic<V>.tryUpdate(function: (V) -> V): Boolean {
-  val cur = value
-  val upd = function(cur)
-  return compareAndSet(cur, upd)
-}
+public inline fun <V> Atomic<V>.tryUpdate(function: (V) -> V): Boolean = tryUpdate(function) { _, _ -> }
 
-public inline fun <V> Atomic<V>.update(function: (V) -> V) {
-  while (true) {
-    val cur = value
-    val upd = function(cur)
-    if (compareAndSet(cur, upd)) return
-  }
-}
+public inline fun <V> Atomic<V>.update(function: (V) -> V): Unit = update(function) { _, _ -> }
 
 /**
  * Updates variable atomically using the specified [function] of its value and returns its old value.
  */
-public inline fun <V> Atomic<V>.getAndUpdate(function: (V) -> V): V {
-  while (true) {
-    val cur = value
-    val upd = function(cur)
-    if (compareAndSet(cur, upd)) return cur
-  }
-}
+public inline fun <V> Atomic<V>.getAndUpdate(function: (V) -> V): V = update(function) { old, _ -> old }
 
 /**
  * Updates variable atomically using the specified [function] of its value and returns its new value.
  */
-public inline fun <V> Atomic<V>.updateAndGet(function: (V) -> V): V {
+public inline fun <V> Atomic<V>.updateAndGet(function: (V) -> V): V = update(function) { _, new -> new }
+
+@PublishedApi
+internal inline fun <V, U: V, R> Atomic<V>.update(function: (V) -> U, transform: (old: V, new: U) -> R): R {
   while (true) {
-    val cur = value
-    val upd = function(cur)
-    if (compareAndSet(cur, upd)) return upd
+    tryUpdate(function) { old, new -> return transform(old, new) }
   }
 }
+
+@PublishedApi
+internal inline fun <V, U: V> Atomic<V>.tryUpdate(function: (V) -> U, onUpdated: (old: V, new: U) -> Unit): Boolean {
+  val cur = value
+  val upd = function(cur)
+  return compareAndSet(cur, upd).also { if (it) onUpdated(cur, upd) }
+}
diff --git a/arrow-libs/core/arrow-atomic/src/commonMain/kotlin/arrow/atomic/AtomicBoolean.kt b/arrow-libs/core/arrow-atomic/src/commonMain/kotlin/arrow/atomic/AtomicBoolean.kt
@@ -14,7 +14,7 @@ public class AtomicBoolean(value: Boolean) {
 
   public fun get(): Boolean = value
   public fun set(value: Boolean) {
-    inner.value = value.toInt()
+     this.value = value
   }
 
   public fun getAndSet(value: Boolean): Boolean =
@@ -28,45 +28,32 @@ public class AtomicBoolean(value: Boolean) {
 /**
  * Infinite loop that reads this atomic variable and performs the specified [action] on its value.
  */
-public inline fun AtomicBoolean.loop(action: (Boolean) -> Unit): Nothing {
-  while (true) {
-    action(value)
-  }
-}
+public inline fun AtomicBoolean.loop(action: (Boolean) -> Unit): Nothing { while(true) { action(value) } }
 
-public fun AtomicBoolean.tryUpdate(function: (Boolean) -> Boolean): Boolean {
-  val cur = value
-  val upd = function(cur)
-  return compareAndSet(cur, upd)
-}
+public inline fun AtomicBoolean.tryUpdate(function: (Boolean) -> Boolean): Boolean = tryUpdate(function) { _, _ -> }
 
-public inline fun AtomicBoolean.update(function: (Boolean) -> Boolean) {
-  while (true) {
-    val cur = value
-    val upd = function(cur)
-    if (compareAndSet(cur, upd)) return
-  }
-}
+public inline fun AtomicBoolean.update(function: (Boolean) -> Boolean): Unit = update(function) { _, _ -> }
 
 /**
  * Updates variable atomically using the specified [function] of its value and returns its old value.
  */
-public inline fun AtomicBoolean.getAndUpdate(function: (Boolean) -> Boolean): Boolean {
-  while (true) {
-    val cur = value
-    val upd = function(cur)
-    if (compareAndSet(cur, upd)) return cur
-  }
-}
+public inline fun AtomicBoolean.getAndUpdate(function: (Boolean) -> Boolean): Boolean = update(function) { old, _ -> old }
 
 /**
  * Updates variable atomically using the specified [function] of its value and returns its new value.
  */
-public inline fun AtomicBoolean.updateAndGet(function: (Boolean) -> Boolean): Boolean {
+public inline fun AtomicBoolean.updateAndGet(function: (Boolean) -> Boolean): Boolean = update(function) { _, new -> new }
+
+@PublishedApi
+internal inline fun <R> AtomicBoolean.update(function: (Boolean) -> Boolean, transform: (old: Boolean, new: Boolean) -> R): R {
   while (true) {
-    val cur = value
-    val upd = function(cur)
-    if (compareAndSet(cur, upd)) return upd
+    tryUpdate(function) { old, new -> return transform(old, new) }
   }
 }
 
+@PublishedApi
+internal inline fun AtomicBoolean.tryUpdate(function: (Boolean) -> Boolean, onUpdated: (old: Boolean, new: Boolean) -> Unit): Boolean {
+  val cur = value
+  val upd = function(cur)
+  return compareAndSet(cur, upd).also { if (it) onUpdated(cur, upd) }
+}
diff --git a/arrow-libs/core/arrow-atomic/src/commonMain/kotlin/arrow/atomic/AtomicInt.kt b/arrow-libs/core/arrow-atomic/src/commonMain/kotlin/arrow/atomic/AtomicInt.kt
@@ -23,45 +23,32 @@ public var AtomicInt.value: Int
 /**
  * Infinite loop that reads this atomic variable and performs the specified [action] on its value.
  */
-public inline fun AtomicInt.loop(action: (Int) -> Unit): Nothing {
-  while (true) {
-    action(value)
-  }
-}
+public inline fun AtomicInt.loop(action: (Int) -> Unit): Nothing { while(true) { action(value) } }
 
-public fun AtomicInt.tryUpdate(function: (Int) -> Int): Boolean {
-  val cur = value
-  val upd = function(cur)
-  return compareAndSet(cur, upd)
-}
+public inline fun AtomicInt.tryUpdate(function: (Int) -> Int): Boolean = tryUpdate(function) { _, _ -> }
 
-public inline fun AtomicInt.update(function: (Int) -> Int) {
-  while (true) {
-    val cur = value
-    val upd = function(cur)
-    if (compareAndSet(cur, upd)) return
-  }
-}
+public inline fun AtomicInt.update(function: (Int) -> Int) : Unit = update(function) { _, _ -> }
 
 /**
  * Updates variable atomically using the specified [function] of its value and returns its old value.
  */
-public inline fun AtomicInt.getAndUpdate(function: (Int) -> Int): Int {
-  while (true) {
-    val cur = value
-    val upd = function(cur)
-    if (compareAndSet(cur, upd)) return cur
-  }
-}
+public inline fun AtomicInt.getAndUpdate(function: (Int) -> Int): Int = update(function) { old, _ -> old }
 
 /**
  * Updates variable atomically using the specified [function] of its value and returns its new value.
  */
-public inline fun AtomicInt.updateAndGet(function: (Int) -> Int): Int {
+public inline fun AtomicInt.updateAndGet(function: (Int) -> Int): Int = update(function) { _, new -> new }
+
+@PublishedApi
+internal inline fun <R> AtomicInt.update(function: (Int) -> Int, transform: (old: Int, new: Int) -> R): R {
   while (true) {
-    val cur = value
-    val upd = function(cur)
-    if (compareAndSet(cur, upd)) return upd
+    tryUpdate(function) { old, new -> return transform(old, new) }
   }
 }
 
+@PublishedApi
+internal inline fun AtomicInt.tryUpdate(function: (Int) -> Int, onUpdated: (old: Int, new: Int) -> Unit): Boolean {
+  val cur = value
+  val upd = function(cur)
+  return compareAndSet(cur, upd).also { if (it) onUpdated(cur, upd) }
+}
diff --git a/arrow-libs/core/arrow-atomic/src/commonMain/kotlin/arrow/atomic/AtomicLong.kt b/arrow-libs/core/arrow-atomic/src/commonMain/kotlin/arrow/atomic/AtomicLong.kt
@@ -23,44 +23,32 @@ public var AtomicLong.value: Long
 /**
  * Infinite loop that reads this atomic variable and performs the specified [action] on its value.
  */
-public inline fun AtomicLong.loop(action: (Long) -> Unit): Nothing {
-  while (true) {
-    action(value)
-  }
-}
+public inline fun AtomicLong.loop(action: (Long) -> Unit): Nothing { while(true) { action(value) } }
 
-public fun AtomicLong.tryUpdate(function: (Long) -> Long): Boolean {
-  val cur = value
-  val upd = function(cur)
-  return compareAndSet(cur, upd)
-}
+public inline fun AtomicLong.tryUpdate(function: (Long) -> Long): Boolean = tryUpdate(function) { _, _ -> }
 
-public inline fun AtomicLong.update(function: (Long) -> Long) {
-  while (true) {
-    val cur = value
-    val upd = function(cur)
-    if (compareAndSet(cur, upd)) return
-  }
-}
+public inline fun AtomicLong.update(function: (Long) -> Long): Unit = update(function) { _, _ -> }
 
 /**
  * Updates variable atomically using the specified [function] of its value and returns its old value.
  */
-public inline fun AtomicLong.getAndUpdate(function: (Long) -> Long): Long {
-  while (true) {
-    val cur = value
-    val upd = function(cur)
-    if (compareAndSet(cur, upd)) return cur
-  }
-}
+public inline fun AtomicLong.getAndUpdate(function: (Long) -> Long): Long = update(function) { old, _ -> old }
 
 /**
  * Updates variable atomically using the specified [function] of its value and returns its new value.
  */
-public inline fun AtomicLong.updateAndGet(function: (Long) -> Long): Long {
+public inline fun AtomicLong.updateAndGet(function: (Long) -> Long): Long = update(function) { _, new -> new }
+
+@PublishedApi
+internal inline fun <R> AtomicLong.update(function: (Long) -> Long, transform: (old: Long, new: Long) -> R): R {
   while (true) {
-    val cur = value
-    val upd = function(cur)
-    if (compareAndSet(cur, upd)) return upd
+    tryUpdate(function) { old, new -> return transform(old, new) }
   }
 }
+
+@PublishedApi
+internal inline fun AtomicLong.tryUpdate(function: (Long) -> Long, onUpdated: (old: Long, new: Long) -> Unit): Boolean {
+  val cur = value
+  val upd = function(cur)
+  return compareAndSet(cur, upd).also { if (it) onUpdated(cur, upd) }
+}
diff --git a/arrow-libs/core/arrow-atomic/src/jvmMain/kotlin/arrow/atomic/AtomicLong.kt b/arrow-libs/core/arrow-atomic/src/jvmMain/kotlin/arrow/atomic/AtomicLong.kt
@@ -4,4 +4,4 @@ package arrow.atomic
 
 import java.util.concurrent.atomic.AtomicLong
 
-public actual typealias AtomicLong= AtomicLong
+public actual typealias AtomicLong = AtomicLong
diff --git a/arrow-libs/core/arrow-core/src/commonMain/kotlin/arrow/core/MemoizedDeepRecursiveFunction.kt b/arrow-libs/core/arrow-core/src/commonMain/kotlin/arrow/core/MemoizedDeepRecursiveFunction.kt
@@ -1,7 +1,7 @@
 package arrow.core
 
 import arrow.atomic.Atomic
-import arrow.atomic.loop
+import arrow.atomic.updateAndGet
 import kotlin.jvm.JvmInline
 import kotlin.jvm.JvmOverloads
 
@@ -15,16 +15,9 @@ public value class AtomicMemoizationCache<K, V>(
   private val cache: Atomic<Map<K, V>> = Atomic(emptyMap())
 ): MemoizationCache<K, V> {
   override fun get(key: K): V? = cache.get()[key]
-  override fun set(key: K, value: V): V = cache.loop { old ->
-    when (key) {
-      in old ->
-        return@set old.getValue(key)
-      else -> {
-        if (cache.compareAndSet(old, old + Pair(key, value)))
-          return@set value
-      }
-    }
-  }
+  override fun set(key: K, value: V): V = cache.updateAndGet { old ->
+    if (key in old) old else old + (key to value)
+  }.getValue(key)
 }
 
 /**
@@ -41,11 +34,9 @@ public value class AtomicMemoizationCache<K, V>(
 public fun <T, R> MemoizedDeepRecursiveFunction(
   cache: MemoizationCache<T, R> = AtomicMemoizationCache(),
   block: suspend DeepRecursiveScope<T, R>.(T) -> R
-): DeepRecursiveFunction<T, R> {
-  return DeepRecursiveFunction { x ->
-    when (val v = cache.get(x)) {
-      null -> cache.set(x, block(x))
-      else -> v
-    }
+): DeepRecursiveFunction<T, R> = DeepRecursiveFunction { x ->
+  when (val v = cache.get(x)) {
+    null -> cache.set(x, block(x))
+    else -> v
   }
 }