Supported unchecked cast erasure (including array elements)

Author: SokolovaMaria

atomicfu-transformer/src/main/kotlin/kotlinx/atomicfu/transformer/AtomicFUTransformer.kt

@@ -28,6 +28,8 @@ private val INT_ARRAY_TYPE = getType("[I")
 private val LONG_ARRAY_TYPE = getType("[J")
 private val BOOLEAN_ARRAY_TYPE = getType("[Z")
 private val REF_ARRAY_TYPE = getType("[Ljava/lang/Object;")
+private val REF_TYPE = getType("L$AFU_PKG/AtomicRef;")
+private val ATOMIC_ARRAY_TYPE = getType("L$AFU_PKG/AtomicArray;")
 
 private val AFU_CLASSES: Map<String, TypeInfo> = mapOf(
     "$AFU_PKG/AtomicInt" to TypeInfo(getObjectType("$JUCA_PKG/AtomicIntegerFieldUpdater"), INT_TYPE, INT_TYPE),
@@ -798,12 +800,29 @@ class AtomicFUTransformer(
             iv.desc = getMethodDescriptor(ret, OBJECT_TYPE, *args)
         }
 
+        private fun tryEraseUncheckedCast(getter: AbstractInsnNode) {
+            if (getter.next.opcode == DUP && getter.next.next.opcode == IFNONNULL) {
+                // unchecked cast upon AtomicRef var is performed
+                // erase compiler check for this var being not null:
+                // (remove all insns from ld till the non null branch label)
+                val ifnonnull = (getter.next.next as JumpInsnNode)
+                var i: AbstractInsnNode = getter.next
+                while (!(i is LabelNode && i.label == ifnonnull.label.label)) {
+                    val next = i.next
+                    instructions.remove(i)
+                    i = next
+                }
+            }
+        }
+
         private fun fixupLoadedAtomicVar(f: FieldInfo, ld: FieldInsnNode): AbstractInsnNode? {
+            if (f.fieldType == REF_TYPE) tryEraseUncheckedCast(ld)
             val j = FlowAnalyzer(ld.next).execute()
             return fixupOperationOnAtomicVar(j, f, ld, null)
         }
 
         private fun fixupLoadedArrayElement(f: FieldInfo, ld: FieldInsnNode, getter: MethodInsnNode): AbstractInsnNode? {
+            if (f.fieldType == ATOMIC_ARRAY_TYPE) tryEraseUncheckedCast(getter)
             // contains array field load (in vh case: + swap and pure type array load) and array element index
             // this array element information is only used in case the reference to this element is stored (copied and inserted at the point of loading)
             val arrayElementInfo = mutableListOf<AbstractInsnNode>()
@@ -880,6 +899,7 @@ class AtomicFUTransformer(
         }
 
         private fun fixupArrayElementLoad(f: FieldInfo, ld: FieldInsnNode, otherLd: VarInsnNode, arrayElementInfo: List<AbstractInsnNode>): AbstractInsnNode? {
+            if (f.fieldType == ATOMIC_ARRAY_TYPE) tryEraseUncheckedCast(otherLd)
             // index instructions from array element info: drop owner class load instruction (in vh case together with preceding getting VH + swap)
             val index = arrayElementInfo.drop(if (vh) 3 else 1)
             // previously stored array element reference is loaded -> arrayElementInfo should be cloned and inserted at the point of this load

atomicfu-transformer/src/main/kotlin/kotlinx/atomicfu/transformer/AtomicFUTransformerJS.kt

@@ -20,6 +20,8 @@ private const val SCOPE = "scope"
 private const val FACTORY = "factory"
 private const val REQUIRE = "require"
 private const val KOTLINX_ATOMICFU = "'kotlinx-atomicfu'"
+private const val KOTLIN_TYPE_CHECK = "Kotlin.isType"
+private const val ATOMIC_REF = "AtomicRef"
 private const val MODULE_KOTLINX_ATOMICFU = "\$module\$kotlinx_atomicfu"
 private const val ARRAY = "Array"
 private const val FILL = "fill"
@@ -266,8 +268,10 @@ class AtomicFUTransformerJS(
                 }
             }
             // remove value property call
-            if (node.type == Token.GETPROP) {
-                if ((node as PropertyGet).property.toSource() == MANGLED_VALUE_PROP) {
+            if (node is PropertyGet) {
+                if (node.property.toSource() == MANGLED_VALUE_PROP) {
+                    // check whether atomic operation is performed on the type casted atomic field
+                    node.target.eraseAtomicFieldFromUncheckedCast()?.let { node.target = it }
                     // A.a.value
                     if (node.target.type == Token.GETPROP) {
                         val clearField = node.target as PropertyGet
@@ -335,6 +339,7 @@ class AtomicFUTransformerJS(
                             field = rr.receiver
                         }
                     }
+                    field.eraseAtomicFieldFromUncheckedCast()?.let { field = it }
                     val args = node.arguments
                     val inlined = node.inlineAtomicOperation(funcName.toSource(), field, args)
                     return !inlined
@@ -344,6 +349,21 @@ class AtomicFUTransformerJS(
         }
     }
 
+    private fun AstNode.eraseAtomicFieldFromUncheckedCast(): AstNode? {
+        if (this is ParenthesizedExpression && expression is ConditionalExpression) {
+            val testExpression = (expression as ConditionalExpression).testExpression
+            if (testExpression is FunctionCall && testExpression.target.toSource() == KOTLIN_TYPE_CHECK) {
+                // type check
+                val typeToCast = testExpression.arguments[1]
+                if ((typeToCast as Name).identifier == ATOMIC_REF) {
+                    // unchecked type cast -> erase atomic field itself
+                    return (testExpression.arguments[0] as Assignment).right
+                }
+            }
+        }
+        return null
+    }
+
     private fun AstNode.isThisNode(): Boolean {
         return when(this) {
             is PropertyGet -> {

atomicfu/src/commonTest/kotlin/kotlinx/atomicfu/test/UncheckedCastTest.kt

@@ -0,0 +1,51 @@
+package kotlinx.atomicfu.test
+
+import kotlinx.atomicfu.AtomicRef
+import kotlinx.atomicfu.atomic
+import kotlinx.atomicfu.atomicArrayOfNulls
+import kotlin.test.Test
+import kotlin.test.assertEquals
+import kotlin.test.assertTrue
+
+private val topLevelS = atomic<Any>(listOf("A", "B"))
+
+class UncheckedCastTest {
+    private val s = atomic<Any>("AAA")
+    private val bs = atomic<Any?>(null)
+
+    @Test
+    fun testAtomicValUncheckedCast() {
+        assertEquals((s as AtomicRef<String>).value, "AAA")
+        bs.lazySet(mapOf(1 to listOf(Box(1), Box(2))))
+        assertEquals((bs as AtomicRef<Map<Int, List<Box>>>).value[1]!![0].b * 10, 10)
+    }
+
+    @Test
+    fun testTopLevelValUnchekedCast() {
+        assertEquals((topLevelS as AtomicRef<List<String>>).value[1], "B")
+    }
+
+    private data class Box(val b: Int)
+
+    private inline fun <T> AtomicRef<T>.getString(): String =
+        (this as AtomicRef<String>).value
+
+    @Test
+    fun testInlineFunc() {
+        assertEquals("AAA", s.getString())
+    }
+
+    private val a = atomicArrayOfNulls<Any?>(10)
+
+    @Test
+    fun testArrayValueUncheckedCast() {
+        a[0].value = "OK"
+        assertEquals("OK", (a[0] as AtomicRef<String>).value)
+    }
+
+    @Test
+    fun testArrayValueUncheckedCastInlineFunc() {
+        a[0].value = "OK"
+        assertEquals("OK", a[0].getString())
+    }
+}