kotlinx.atomicfu compiler plugin (JS): tracing supported. (#4691)

* kotlinx.atomicfu compiler plugin (JS): tracing supported.

* Erasing trace invocations from both IrBlock and IrBlockBody.

plugins/atomicfu/atomicfu-compiler/src/org.jetbrains.kotlinx.atomicfu.compiler/extensions/AtomicfuTransformer.kt

@@ -25,10 +25,14 @@ private const val LOCKS = "locks"
 private const val AFU_LOCKS_PKG = "$AFU_PKG.$LOCKS"
 private const val ATOMICFU_RUNTIME_FUNCTION_PREDICATE = "atomicfu_"
 private const val REENTRANT_LOCK_TYPE = "ReentrantLock"
+private const val TRACE_BASE_TYPE = "TraceBase"
 private const val GETTER = "atomicfu\$getter"
 private const val SETTER = "atomicfu\$setter"
 private const val GET = "get"
 private const val ATOMIC_VALUE_FACTORY = "atomic"
+private const val TRACE = "Trace"
+private const val INVOKE = "invoke"
+private const val APPEND = "append"
 private const val ATOMIC_ARRAY_OF_NULLS_FACTORY = "atomicArrayOfNulls"
 private const val REENTRANT_LOCK_FACTORY = "reentrantLock"
 
@@ -101,6 +105,23 @@ class AtomicfuTransformer(private val context: IrPluginContext) {
             return super.visitFunction(declaration, declaration)
         }
 
+        override fun visitBlockBody(body: IrBlockBody, data: IrFunction?): IrBody {
+            // Erase messages added by the Trace object from the function body:
+            // val trace = Trace(size)
+            // Messages may be added via trace invocation:
+            // trace { "Doing something" }
+            // or via multi-append of arguments:
+            // trace.append(index, "CAS", value)
+            body.statements.removeIf { it.isTrace() }
+            return super.visitBlockBody(body, data)
+        }
+
+        override fun visitContainerExpression(expression: IrContainerExpression, data: IrFunction?): IrExpression {
+            // Erase messages added by the Trace object from blocks.
+            expression.statements.removeIf { it.isTrace() }
+            return super.visitContainerExpression(expression, data)
+        }
+
         override fun visitCall(expression: IrCall, data: IrFunction?): IrElement {
             expression.eraseAtomicFactory()?.let { return it.transform(this, data) }
             val isInline = expression.symbol.owner.isInline
@@ -276,6 +297,19 @@ class AtomicfuTransformer(private val context: IrPluginContext) {
                 it.type.isAtomicArrayType() && symbol.owner.name.asString() == GET
             } ?: false
 
+        private fun IrStatement.isTrace() =
+            this is IrCall && (isTraceInvoke() || isTraceAppend())
+
+        private fun IrCall.isTraceInvoke(): Boolean =
+            symbol.isKotlinxAtomicfuPackage() &&
+            symbol.owner.name.asString() == INVOKE &&
+            symbol.owner.dispatchReceiverParameter?.type?.isTraceBaseType() == true
+
+        private fun IrCall.isTraceAppend(): Boolean =
+            symbol.isKotlinxAtomicfuPackage() &&
+            symbol.owner.name.asString() == APPEND &&
+            symbol.owner.dispatchReceiverParameter?.type?.isTraceBaseType() == true
+
         private fun IrCall.getAccessors(): List<IrExpression> {
             val isArrayElement = isArrayElementGetter()
             val valueType = type.atomicToValueType()
@@ -318,6 +352,7 @@ class AtomicfuTransformer(private val context: IrPluginContext) {
                 isAtomicFactory() -> getValueArgument(0) ?: error("Atomic factory should take at least one argument: ${this.render()}")
                 isAtomicArrayFactory() -> buildObjectArray()
                 isReentrantLockFactory() -> context.buildConstNull()
+                isTraceFactory() -> context.buildConstNull()
                 else -> null
             }
 
@@ -344,6 +379,7 @@ class AtomicfuTransformer(private val context: IrPluginContext) {
     private fun IrType.isAtomicValueType() = belongsTo(AFU_PKG, ATOMIC_VALUE_TYPES)
     private fun IrType.isAtomicArrayType() = belongsTo(AFU_PKG, ATOMIC_ARRAY_TYPES)
     private fun IrType.isReentrantLockType() = belongsTo(AFU_LOCKS_PKG, REENTRANT_LOCK_TYPE)
+    private fun IrType.isTraceBaseType() = belongsTo(AFU_PKG, TRACE_BASE_TYPE)
 
     private fun IrType.belongsTo(packageName: String, typeNames: Set<String>) =
         getSignature()?.let { sig ->
@@ -371,6 +407,10 @@ class AtomicfuTransformer(private val context: IrPluginContext) {
         symbol.isKotlinxAtomicfuPackage() && symbol.owner.name.asString() == ATOMIC_VALUE_FACTORY &&
                 type.isAtomicValueType()
 
+    private fun IrCall.isTraceFactory(): Boolean =
+        symbol.isKotlinxAtomicfuPackage() && symbol.owner.name.asString() == TRACE &&
+                type.isTraceBaseType()
+
     private fun IrCall.isAtomicArrayFactory(): Boolean =
         symbol.isKotlinxAtomicfuPackage() && symbol.owner.name.asString() == ATOMIC_ARRAY_OF_NULLS_FACTORY &&
                 type.isAtomicArrayType()

plugins/atomicfu/atomicfu-compiler/test/org/jetbrains/kotlinx/atomicfu/AtomicfuJsIrTestGenerated.java

@@ -145,6 +145,12 @@ public class AtomicfuJsIrTestGenerated extends AbstractAtomicfuJsIrTest {
         runTest("plugins/atomicfu/atomicfu-compiler/testData/box/TopLevelTest.kt");
     }
 
+    @Test
+    @TestMetadata("TraceTest.kt")
+    public void testTraceTest() throws Exception {
+        runTest("plugins/atomicfu/atomicfu-compiler/testData/box/TraceTest.kt");
+    }
+
     @Test
     @TestMetadata("UncheckedCastTest.kt")
     public void testUncheckedCastTest() throws Exception {

plugins/atomicfu/atomicfu-compiler/testData/box/TraceTest.kt

@@ -0,0 +1,106 @@
+import kotlinx.atomicfu.*
+import kotlin.test.*
+
+class TraceTest {
+    private val defaultTrace = Trace()
+    private val a1 = atomic(5, defaultTrace)
+
+    private val traceWithSize = Trace(30)
+    private val a2 = atomic(1, traceWithSize)
+
+    private val traceWithFormat = Trace(format = TraceFormat { i, text -> "[$i: $text]" })
+    private val a = atomic(0, traceWithFormat)
+
+    private val traceWithSizeAndFormat = Trace(30, TraceFormat { id, text -> "$id: $text"})
+    private val a3 = atomic(2)
+
+    private val shortTrace = Trace(4)
+    private val s = atomic(0, shortTrace.named("s"))
+
+    fun testDefaultTrace() {
+        val oldValue = a1.value
+        defaultTrace { "before CAS value = $oldValue" }
+        val res = a1.compareAndSet(oldValue, oldValue * 10)
+        val newValue = a1.value
+        defaultTrace { "after CAS value = $newValue" }
+    }
+
+    fun testTraceWithSize() {
+        val oldValue = a2.value
+        traceWithSize { "before CAS value = $oldValue" }
+        assertTrue(a2.compareAndSet(oldValue, oldValue * 10))
+        traceWithSize { "after CAS value = ${a2.value}" }
+        traceWithSize { "before getAndDecrement value = ${a2.value}" }
+        a2.getAndDecrement()
+        assertEquals(9, a2.value)
+        traceWithSize { "after getAndDecrement value = ${a2.value}" }
+    }
+
+    fun testTraceWithFormat() {
+        val oldValue = a3.value
+        traceWithFormat { "before CAS value = $oldValue" }
+        assertTrue(a3.compareAndSet(oldValue, oldValue * 10))
+        traceWithFormat { "after CAS value = ${a3.value}" }
+        traceWithFormat { "before getAndDecrement value = ${a3.value}" }
+        a3.getAndDecrement()
+        assertEquals(19, a3.value)
+        traceWithFormat { "after getAndDecrement value = ${a3.value}" }
+    }
+
+    fun testNamedTrace() {
+        s.value = 5
+        shortTrace { "before CAS value = ${s.value}" }
+        s.compareAndSet(5, -2)
+        assertEquals(-2, s.value)
+        shortTrace { "after CAS value = ${s.value}" }
+    }
+
+    private enum class Status { START, END }
+
+    fun testMultipleAppend() {
+        val i = 1
+        traceWithFormat.append(i, Status.START)
+        assertEquals(0, a.value)
+        a.incrementAndGet()
+        traceWithFormat.append(i, a.value, "incAndGet")
+        assertEquals(1, a.value)
+        a.lazySet(10)
+        traceWithFormat.append(i, a.value, "lazySet")
+        assertEquals(10, a.value)
+        traceWithFormat.append(i, Status.END)
+    }
+
+    fun testTraceInBlock() {
+        a1.lazySet(5)
+        if (a1.value == 5) {
+            defaultTrace { "Value checked" }
+            if (a1.compareAndSet(5, 10)) {
+                defaultTrace { "CAS succeeded" }
+            }
+        }
+        assertEquals(10, a1.value)
+        while (true) {
+            if (a1.value == 10) {
+                defaultTrace.append("Value checked", a1.value)
+                a1.value = 15
+                break
+            } else {
+                defaultTrace.append("Wrong value", a1.value)
+            }
+        }
+    }
+
+    fun test() {
+        testDefaultTrace()
+        testTraceWithSize()
+        testTraceWithFormat()
+        testNamedTrace()
+        testMultipleAppend()
+        testTraceInBlock()
+    }
+}
+
+fun box(): String {
+    TraceTest().test()
+    return "OK"
+}