From 53e545aa3f6de25d29ee83a8069f4ed22dd510d6 Mon Sep 17 00:00:00 2001
From: Marco Pennekamp <marco.pennekamp@jetbrains.com>
Date: Wed, 17 Jan 2024 17:14:16 +0100
Subject: [PATCH] [LL] `CleanableSoftValueCache`: Don't perform cleanup when
 putting the same value

- If we have some value `X` in the cache, and we put `X` a second time
  into the same location, `X` shouldn't be cleaned up because it wasn't
  actually removed from the cache.
- With this feature, it's necessary to implement `put` in terms of
  `backingMap.compute`, because we need to compare the old and the new
  value atomically, and based on that, decide whether to create a new
  soft reference.

^KT-61222
---
 .../api/fir/caches/CleanableSoftValueCache.kt | 32 ++++++++++++++++---
 1 file changed, 27 insertions(+), 5 deletions(-)

diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/caches/CleanableSoftValueCache.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/caches/CleanableSoftValueCache.kt
index 22a50f1f337..d86bbe8a138 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/caches/CleanableSoftValueCache.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/caches/CleanableSoftValueCache.kt
@@ -147,15 +147,37 @@ class CleanableSoftValueCache<K : Any, V : Any>(
     /**
      * Adds or replaces [value] to/in the cache at the given [key]. Must be called in a read action.
      *
-     * @return The old value that has been replaced, if any. As replacement constitutes removal, the cleaner associated with the value will
-     * be invoked by [put].
+     * As replacement constitutes removal, cleanup will be performed on the replaced value. When the existing value and the new value are
+     * the same (referentially equal), cleanup will not be performed, because the existing value effectively wasn't removed from the cache.
+     *
+     * @return The old value that has been replaced, if any.
      */
     fun put(key: K, value: V): V? {
-        val oldRef = backingMap.put(key, createSoftReference(key, value))
-        oldRef?.performCleanup()
+        // We implement `put` in terms of `backingMap.compute` to avoid creation of a new soft reference when the old and the new value are
+        // referentially equal. A combination of `backingMap.get` and `backingMap.put` would not be atomic, because the existing value
+        // fetched with `backingMap.get` might be outdated by the time we invoke `backingMap.put` based on the `old === new` comparison.
+        // This function's implementation is different from `CleanableSoftValueCache.compute` because `put` needs to return the old value,
+        // not the new value.
+        var oldValue: V? = null
+        var removedRef: SoftReferenceWithCleanup<K, V>? = null
 
+        // See `compute` for additional comments on the implementation, as it is similar to this implementation.
+        backingMap.compute(key) { _, currentRef ->
+            val currentValue = currentRef?.get()
+            oldValue = currentValue
+
+            if (value === currentValue) {
+                return@compute currentRef
+            }
+
+            removedRef = currentRef
+            createSoftReference(key, value)
+        }
+
+        removedRef?.performCleanup()
         processQueue()
-        return oldRef?.get()
+
+        return oldValue
     }
 
     /**