atomicfu-compiler-plugin: generate atomic extension functions on call site
diff --git a/plugins/atomicfu/atomicfu-compiler/src/org/jetbrains/kotlinx/atomicfu/compiler/backend/jvm/AtomicfuJvmIrTransformer.kt b/plugins/atomicfu/atomicfu-compiler/src/org/jetbrains/kotlinx/atomicfu/compiler/backend/jvm/AtomicfuJvmIrTransformer.kt
index bfff872..e8c6772 100644
--- a/plugins/atomicfu/atomicfu-compiler/src/org/jetbrains/kotlinx/atomicfu/compiler/backend/jvm/AtomicfuJvmIrTransformer.kt
+++ b/plugins/atomicfu/atomicfu-compiler/src/org/jetbrains/kotlinx/atomicfu/compiler/backend/jvm/AtomicfuJvmIrTransformer.kt
@@ -59,8 +59,8 @@
 
     fun transform(moduleFragment: IrModuleFragment) {
         transformAtomicFields(moduleFragment)
-        transformAtomicExtensions(moduleFragment)
         transformAtomicfuDeclarations(moduleFragment)
+        transformAtomicExtensions(moduleFragment)
         for (irFile in moduleFragment.files) {
             irFile.patchDeclarationParents()
         }
@@ -398,72 +398,19 @@
 
     private inner class AtomicExtensionTransformer : IrElementTransformerVoid() {
         override fun visitFile(declaration: IrFile): IrFile {
-            declaration.transformAllAtomicExtensions()
+            declaration.removeAllAtomicExtensions()
             return super.visitFile(declaration)
         }
 
         override fun visitClass(declaration: IrClass): IrStatement {
-            declaration.transformAllAtomicExtensions()
+            declaration.removeAllAtomicExtensions()
             return super.visitClass(declaration)
         }
 
-        private fun IrDeclarationContainer.transformAllAtomicExtensions() {
-            // Transform the signature of kotlinx.atomicfu.Atomic* class extension functions:
+        private fun IrDeclarationContainer.removeAllAtomicExtensions() {
+            // Remove the original signature of kotlinx.atomicfu.Atomic* class extension functions:
             // inline fun AtomicInt.foo(arg: T)
-            // For every signature there are 2 new declarations generated (because of different types of atomic handlers):
-            // 1. for the case of atomic value receiver at the invocation:
-            // inline fun foo$atomicfu(dispatchReceiver: Any?, handler: j.u.c.a.AtomicIntegerFieldUpdater, arg': T)
-            // 2. for the case of atomic array element receiver at the invocation:
-            // inline fun foo$atomicfu$array(dispatchReceiver: Any?, handler: j.u.c.a.AtomicIntegerArray, index: Int, arg': T)
-            declarations.filter { it is IrFunction && it.isAtomicExtension() }.forEach { atomicExtension ->
-                atomicExtension as IrFunction
-                declarations.add(generateAtomicExtension(atomicExtension, this, false))
-                declarations.add(generateAtomicExtension(atomicExtension, this, true))
-                declarations.remove(atomicExtension)
-            }
-        }
-
-        private fun generateAtomicExtension(
-            atomicExtension: IrFunction,
-            parent: IrDeclarationParent,
-            isArrayReceiver: Boolean
-        ): IrFunction {
-            val mangledName = mangleFunctionName(atomicExtension.name.asString(), isArrayReceiver)
-            val valueType = atomicExtension.extensionReceiverParameter!!.type.atomicToValueType()
-            return context.irFactory.buildFun {
-                name = Name.identifier(mangledName)
-                isInline = true
-            }.apply {
-                val newDeclaration = this
-                extensionReceiverParameter = null
-                dispatchReceiverParameter = atomicExtension.dispatchReceiverParameter?.deepCopyWithSymbols(this)
-                if (isArrayReceiver) {
-                    addValueParameter(DISPATCH_RECEIVER, irBuiltIns.anyNType)
-                    addValueParameter(ATOMIC_HANDLER, atomicSymbols.getAtomicArrayClassByValueType(valueType).defaultType)
-                    addValueParameter(INDEX, irBuiltIns.intType)
-                } else {
-                    addValueParameter(DISPATCH_RECEIVER, irBuiltIns.anyNType)
-                    addValueParameter(ATOMIC_HANDLER, atomicSymbols.getFieldUpdaterType(valueType))
-                }
-                atomicExtension.valueParameters.forEach { addValueParameter(it.name, it.type) }
-                // the body will be transformed later by `AtomicFUTransformer`
-                body = atomicExtension.body?.deepCopyWithSymbols(this)
-                body?.transform(
-                    object : IrElementTransformerVoid() {
-                        override fun visitReturn(expression: IrReturn): IrExpression = super.visitReturn(
-                            if (expression.returnTargetSymbol == atomicExtension.symbol) {
-                                with(atomicSymbols.createBuilder(newDeclaration.symbol)) {
-                                    irReturn(expression.value)
-                                }
-                            } else {
-                                expression
-                            }
-                        )
-                    }, null
-                )
-                returnType = atomicExtension.returnType
-                this.parent = parent
-            }
+            declarations.removeIf { it is IrFunction && it.isAtomicExtension() }
         }
     }
 
@@ -555,7 +502,7 @@
                                 }
                                 return super.visitExpression(irCall, data)
                             }
-                            if (expression.symbol.owner.isInline && expression.extensionReceiver != null) {
+                            if (expression.extensionReceiver != null) {
                                 // Transform invocation of the kotlinx.atomicfu.Atomic* class extension functions,
                                 // delegating them to the corresponding transformed atomic extensions:
                                 // for atomic property recevers:
@@ -568,12 +515,11 @@
                                 // The invocation on the atomic array element will be transformed:
                                 // a.foo(arg) -> a.foo$atomicfu$array(dispatchReceiver, atomicHandler, index, arg)
                                 val declaration = expression.symbol.owner
-                                val parent = declaration.parent as IrDeclarationContainer
-                                val transformedAtomicExtension = parent.getTransformedAtomicExtension(declaration, isArrayReceiver)
                                 require(data != null) { "Function containing invocation of the extension function ${expression.render()} is null" }
+                                val transformedAtomicExtension = data.parentDeclarationContainer.getOrBuildTransformedExtension(declaration, isArrayReceiver)
                                 val irCall = callAtomicExtension(
                                     symbol = transformedAtomicExtension.symbol,
-                                    dispatchReceiver = expression.dispatchReceiver,
+                                    dispatchReceiver = data.containingFunction.dispatchReceiverParameter?.capture(),
                                     syntheticValueArguments = if (isArrayReceiver) {
                                         listOf(dispatchReceiver, atomicHandler, receiver.getArrayElementIndex(data))
                                     } else {
@@ -761,13 +707,59 @@
             }
         }
 
-        private fun IrDeclarationContainer.getTransformedAtomicExtension(
+        private fun IrDeclarationContainer.getOrBuildTransformedExtension(
             declaration: IrSimpleFunction,
             isArrayReceiver: Boolean
-        ): IrSimpleFunction = findDeclaration {
-                it.name.asString() == mangleFunctionName(declaration.name.asString(), isArrayReceiver) &&
-                        it.isTransformedAtomicExtension()
-            } ?: error("Could not find corresponding transformed declaration for the atomic extension ${declaration.render()}")
+        ): IrSimpleFunction {
+            // Transform the signature of kotlinx.atomicfu.Atomic* class extension functions:
+            // inline fun AtomicInt.foo(arg: T)
+            // 1. for the case of atomic value receiver at the invocation:
+            // inline fun foo$atomicfu(dispatchReceiver: Any?, handler: j.u.c.a.AtomicIntegerFieldUpdater, arg': T)
+            // 2. for the case of atomic array element receiver at the invocation:
+            // inline fun foo$atomicfu$array(dispatchReceiver: Any?, handler: j.u.c.a.AtomicIntegerArray, index: Int, arg': T)
+            val parent = this
+            val mangledName = mangleFunctionName(declaration.name.asString(), isArrayReceiver)
+            findDeclaration<IrSimpleFunction> {
+                it.name.asString() == mangledName && it.isTransformedAtomicExtension()
+            }?.let { return it }
+            val valueType = declaration.extensionReceiverParameter!!.type.atomicToValueType()
+            return context.irFactory.buildFun {
+                name = Name.identifier(mangledName)
+                isInline = declaration.isInline
+            }.apply {
+                val newDeclaration = this
+                extensionReceiverParameter = null
+                dispatchReceiverParameter = (parent as? IrClass)?.thisReceiver?.deepCopyWithSymbols(this)
+                if (isArrayReceiver) {
+                    addValueParameter(DISPATCH_RECEIVER, irBuiltIns.anyNType)
+                    addValueParameter(ATOMIC_HANDLER, atomicSymbols.getAtomicArrayClassByValueType(valueType).defaultType)
+                    addValueParameter(INDEX, irBuiltIns.intType)
+                } else {
+                    addValueParameter(DISPATCH_RECEIVER, irBuiltIns.anyNType)
+                    addValueParameter(ATOMIC_HANDLER, atomicSymbols.getFieldUpdaterType(valueType))
+                }
+                declaration.valueParameters.forEach { addValueParameter(it.name, it.type) }
+                returnType = declaration.returnType
+
+                this.parent = parent
+                parent.declarations.add(this)
+
+                body = declaration.body?.deepCopyWithSymbols(this)?.transform(this@AtomicfuTransformer, this)
+                body?.transform(
+                    object : IrElementTransformerVoid() {
+                        override fun visitReturn(expression: IrReturn): IrExpression = super.visitReturn(
+                            if (expression.returnTargetSymbol == declaration.symbol) {
+                                with(atomicSymbols.createBuilder(newDeclaration.symbol)) {
+                                    irReturn(expression.value)
+                                }
+                            } else {
+                                expression
+                            }
+                        )
+                    }, null
+                )
+            }
+        }
 
         private fun IrSimpleFunction.generateAtomicfuLoop(valueType: IrType) {
             addValueParameter(ATOMIC_HANDLER, atomicSymbols.getFieldUpdaterType(valueType))
@@ -858,7 +850,7 @@
                 type.isAtomicValueType()
 
     private fun IrFunction.isAtomicExtension(): Boolean =
-        extensionReceiverParameter?.let { it.type.isAtomicValueType() && this.isInline } ?: false
+        extensionReceiverParameter?.type?.isAtomicValueType() ?: false
 
     private fun IrStatement.isTraceCall() = this is IrCall && (isTraceInvoke() || isTraceAppend())
 
diff --git a/plugins/atomicfu/atomicfu-compiler/testData/box/ArrayInlineExtensionTest.txt b/plugins/atomicfu/atomicfu-compiler/testData/box/ArrayInlineExtensionTest.txt
index 8537aff..7b24a51 100644
--- a/plugins/atomicfu/atomicfu-compiler/testData/box/ArrayInlineExtensionTest.txt
+++ b/plugins/atomicfu/atomicfu-compiler/testData/box/ArrayInlineExtensionTest.txt
@@ -18,19 +18,14 @@
     static method <clinit>(): void
     public method <init>(): void
     public final method bar$atomicfu$array(@org.jetbrains.annotations.Nullable p0: java.lang.Object, @org.jetbrains.annotations.NotNull p1: java.util.concurrent.atomic.AtomicIntegerArray, p2: int, p3: int): int
-    public final method bar$atomicfu(@org.jetbrains.annotations.Nullable p0: java.lang.Object, @org.jetbrains.annotations.NotNull p1: java.util.concurrent.atomic.AtomicIntegerFieldUpdater, p2: int): int
     private final method casLoop(p0: int): int
     private final method casLoopExpression(p0: long): long
     public final method extensionLoop$atomicfu$array(@org.jetbrains.annotations.Nullable p0: java.lang.Object, @org.jetbrains.annotations.NotNull p1: java.util.concurrent.atomic.AtomicIntegerArray, p2: int, p3: int): int
     public final method extensionLoop$atomicfu(@org.jetbrains.annotations.Nullable p0: java.lang.Object, @org.jetbrains.annotations.NotNull p1: java.util.concurrent.atomic.AtomicIntegerFieldUpdater, p2: int): int
     public final method extensionLoopExpression$atomicfu$array(@org.jetbrains.annotations.Nullable p0: java.lang.Object, @org.jetbrains.annotations.NotNull p1: java.util.concurrent.atomic.AtomicIntegerArray, p2: int, p3: int): int
-    public final method extensionLoopExpression$atomicfu(@org.jetbrains.annotations.Nullable p0: java.lang.Object, @org.jetbrains.annotations.NotNull p1: java.util.concurrent.atomic.AtomicIntegerFieldUpdater, p2: int): int
     public final method extensionLoopMixedReceivers$atomicfu$array(@org.jetbrains.annotations.Nullable p0: java.lang.Object, @org.jetbrains.annotations.NotNull p1: java.util.concurrent.atomic.AtomicIntegerArray, p2: int, p3: int, p4: int, p5: int): int
-    public final method extensionLoopMixedReceivers$atomicfu(@org.jetbrains.annotations.Nullable p0: java.lang.Object, @org.jetbrains.annotations.NotNull p1: java.util.concurrent.atomic.AtomicIntegerFieldUpdater, p2: int, p3: int, p4: int): int
     public final method extensionLoopRecursive$atomicfu$array(@org.jetbrains.annotations.Nullable p0: java.lang.Object, @org.jetbrains.annotations.NotNull p1: java.util.concurrent.atomic.AtomicIntegerArray, p2: int, p3: int): int
-    public final method extensionLoopRecursive$atomicfu(@org.jetbrains.annotations.Nullable p0: java.lang.Object, @org.jetbrains.annotations.NotNull p1: java.util.concurrent.atomic.AtomicIntegerFieldUpdater, p2: int): int
     public final method foo$atomicfu$array(@org.jetbrains.annotations.Nullable p0: java.lang.Object, @org.jetbrains.annotations.NotNull p1: java.util.concurrent.atomic.AtomicIntegerArray, p2: int, p3: int): int
-    public final method foo$atomicfu(@org.jetbrains.annotations.Nullable p0: java.lang.Object, @org.jetbrains.annotations.NotNull p1: java.util.concurrent.atomic.AtomicIntegerFieldUpdater, p2: int): int
     public final static @org.jetbrains.annotations.NotNull method getA$FU(): java.util.concurrent.atomic.AtomicIntegerFieldUpdater
     public final method getA(): int
     public final @org.jetbrains.annotations.NotNull method getIntArr(): java.util.concurrent.atomic.AtomicIntegerArray
diff --git a/plugins/atomicfu/atomicfu-compiler/testData/box/ComplexLoopTest.txt b/plugins/atomicfu/atomicfu-compiler/testData/box/ComplexLoopTest.txt
index 05c2ab1..23c2349 100644
--- a/plugins/atomicfu/atomicfu-compiler/testData/box/ComplexLoopTest.txt
+++ b/plugins/atomicfu/atomicfu-compiler/testData/box/ComplexLoopTest.txt
@@ -22,7 +22,6 @@
     private final method embeddedUpdate(p0: int): int
     public final method extensionEmbeddedLoops$atomicfu$array(@org.jetbrains.annotations.Nullable p0: java.lang.Object, @org.jetbrains.annotations.NotNull p1: java.util.concurrent.atomic.AtomicIntegerArray, p2: int, p3: int): int
     public final method extensionEmbeddedLoops$atomicfu(@org.jetbrains.annotations.Nullable p0: java.lang.Object, @org.jetbrains.annotations.NotNull p1: java.util.concurrent.atomic.AtomicIntegerFieldUpdater, p2: int): int
-    public final @org.jetbrains.annotations.NotNull method extesntionEmbeddedRefUpdate$atomicfu$array(@org.jetbrains.annotations.Nullable p0: java.lang.Object, @org.jetbrains.annotations.NotNull p1: java.util.concurrent.atomic.AtomicReferenceArray, p2: int, @org.jetbrains.annotations.NotNull p3: java.lang.String): java.lang.String
     public final @org.jetbrains.annotations.NotNull method extesntionEmbeddedRefUpdate$atomicfu(@org.jetbrains.annotations.Nullable p0: java.lang.Object, @org.jetbrains.annotations.NotNull p1: java.util.concurrent.atomic.AtomicReferenceFieldUpdater, @org.jetbrains.annotations.NotNull p2: java.lang.String): java.lang.String
     public final static @org.jetbrains.annotations.NotNull method getA$FU(): java.util.concurrent.atomic.AtomicIntegerFieldUpdater
     public final method getA(): int
@@ -34,11 +33,9 @@
     public final static @org.jetbrains.annotations.NotNull method getR$FU(): java.util.concurrent.atomic.AtomicReferenceFieldUpdater
     public final @org.jetbrains.annotations.Nullable method getR(): java.lang.Object
     public final method loop$atomicfu$array(@org.jetbrains.annotations.NotNull p0: java.util.concurrent.atomic.AtomicIntegerArray, p1: int, @org.jetbrains.annotations.NotNull p2: kotlin.jvm.functions.Function1): void
-    public final method loop$atomicfu$array(@org.jetbrains.annotations.NotNull p0: java.util.concurrent.atomic.AtomicReferenceArray, p1: int, @org.jetbrains.annotations.NotNull p2: kotlin.jvm.functions.Function1): void
     public final method loop$atomicfu(@org.jetbrains.annotations.NotNull p0: java.util.concurrent.atomic.AtomicIntegerFieldUpdater, @org.jetbrains.annotations.NotNull p1: kotlin.jvm.functions.Function1, @org.jetbrains.annotations.Nullable p2: java.lang.Object): void
     public final method loop$atomicfu(@org.jetbrains.annotations.NotNull p0: java.util.concurrent.atomic.AtomicReferenceFieldUpdater, @org.jetbrains.annotations.NotNull p1: kotlin.jvm.functions.Function1, @org.jetbrains.annotations.Nullable p2: java.lang.Object): void
     public final method test(): void
-    public final @org.jetbrains.annotations.Nullable method updateAndGet$atomicfu$array(@org.jetbrains.annotations.NotNull p0: java.util.concurrent.atomic.AtomicReferenceArray, p1: int, @org.jetbrains.annotations.NotNull p2: kotlin.jvm.functions.Function1): java.lang.Object
     public final method updateAndGet$atomicfu(@org.jetbrains.annotations.NotNull p0: java.util.concurrent.atomic.AtomicIntegerFieldUpdater, @org.jetbrains.annotations.NotNull p1: kotlin.jvm.functions.Function1, @org.jetbrains.annotations.Nullable p2: java.lang.Object): int
     public final @org.jetbrains.annotations.Nullable method updateAndGet$atomicfu(@org.jetbrains.annotations.NotNull p0: java.util.concurrent.atomic.AtomicReferenceFieldUpdater, @org.jetbrains.annotations.NotNull p1: kotlin.jvm.functions.Function1, @org.jetbrains.annotations.Nullable p2: java.lang.Object): java.lang.Object
 }
diff --git a/plugins/atomicfu/atomicfu-compiler/testData/box/ExtensionLoopTest.txt b/plugins/atomicfu/atomicfu-compiler/testData/box/ExtensionLoopTest.txt
index 9d12d51..fa6e411 100644
--- a/plugins/atomicfu/atomicfu-compiler/testData/box/ExtensionLoopTest.txt
+++ b/plugins/atomicfu/atomicfu-compiler/testData/box/ExtensionLoopTest.txt
@@ -4,10 +4,8 @@
     public final static @org.jetbrains.annotations.NotNull field ref$ExtensionLoopTest$VolatileWrapper: Ref$ExtensionLoopTest$VolatileWrapper
     static method <clinit>(): void
     public final static @org.jetbrains.annotations.NotNull method box(): java.lang.String
-    public final static method loop$atomicfu$array(@org.jetbrains.annotations.NotNull p0: java.util.concurrent.atomic.AtomicReferenceArray, p1: int, @org.jetbrains.annotations.NotNull p2: kotlin.jvm.functions.Function1): void
     public final static method loop$atomicfu(@org.jetbrains.annotations.NotNull p0: java.util.concurrent.atomic.AtomicReferenceFieldUpdater, @org.jetbrains.annotations.NotNull p1: kotlin.jvm.functions.Function1, @org.jetbrains.annotations.Nullable p2: java.lang.Object): void
     public final static method testTopLevelExtensionLoop(): void
-    public final static @org.jetbrains.annotations.NotNull method topLevelExtensionLoop$atomicfu$array(@org.jetbrains.annotations.Nullable p0: java.lang.Object, @org.jetbrains.annotations.NotNull p1: java.util.concurrent.atomic.AtomicReferenceArray, p2: int, @org.jetbrains.annotations.NotNull p3: java.lang.String): java.lang.String
     public final static @org.jetbrains.annotations.NotNull method topLevelExtensionLoop$atomicfu(@org.jetbrains.annotations.Nullable p0: java.lang.Object, @org.jetbrains.annotations.NotNull p1: java.util.concurrent.atomic.AtomicReferenceFieldUpdater, @org.jetbrains.annotations.NotNull p2: java.lang.String): java.lang.String
 }
 
@@ -37,19 +35,13 @@
     private volatile @kotlin.jvm.Volatile @org.jetbrains.annotations.Nullable field rs: java.lang.Object
     static method <clinit>(): void
     public method <init>(): void
-    public final method bar$atomicfu$array(@org.jetbrains.annotations.Nullable p0: java.lang.Object, @org.jetbrains.annotations.NotNull p1: java.util.concurrent.atomic.AtomicIntegerArray, p2: int, p3: int): int
     public final method bar$atomicfu(@org.jetbrains.annotations.Nullable p0: java.lang.Object, @org.jetbrains.annotations.NotNull p1: java.util.concurrent.atomic.AtomicIntegerFieldUpdater, p2: int): int
     private final method casLoop(p0: int): int
     private final method casLoopExpression(p0: int): int
-    public final method extensionLoop$atomicfu$array(@org.jetbrains.annotations.Nullable p0: java.lang.Object, @org.jetbrains.annotations.NotNull p1: java.util.concurrent.atomic.AtomicIntegerArray, p2: int, p3: int): int
     public final method extensionLoop$atomicfu(@org.jetbrains.annotations.Nullable p0: java.lang.Object, @org.jetbrains.annotations.NotNull p1: java.util.concurrent.atomic.AtomicIntegerFieldUpdater, p2: int): int
-    public final method extensionLoopExpression$atomicfu$array(@org.jetbrains.annotations.Nullable p0: java.lang.Object, @org.jetbrains.annotations.NotNull p1: java.util.concurrent.atomic.AtomicIntegerArray, p2: int, p3: int): int
     public final method extensionLoopExpression$atomicfu(@org.jetbrains.annotations.Nullable p0: java.lang.Object, @org.jetbrains.annotations.NotNull p1: java.util.concurrent.atomic.AtomicIntegerFieldUpdater, p2: int): int
-    public final method extensionLoopMixedReceivers$atomicfu$array(@org.jetbrains.annotations.Nullable p0: java.lang.Object, @org.jetbrains.annotations.NotNull p1: java.util.concurrent.atomic.AtomicIntegerArray, p2: int, p3: int, p4: int): int
     public final method extensionLoopMixedReceivers$atomicfu(@org.jetbrains.annotations.Nullable p0: java.lang.Object, @org.jetbrains.annotations.NotNull p1: java.util.concurrent.atomic.AtomicIntegerFieldUpdater, p2: int, p3: int): int
-    public final method extensionLoopRecursive$atomicfu$array(@org.jetbrains.annotations.Nullable p0: java.lang.Object, @org.jetbrains.annotations.NotNull p1: java.util.concurrent.atomic.AtomicIntegerArray, p2: int, p3: int): int
     public final method extensionLoopRecursive$atomicfu(@org.jetbrains.annotations.Nullable p0: java.lang.Object, @org.jetbrains.annotations.NotNull p1: java.util.concurrent.atomic.AtomicIntegerFieldUpdater, p2: int): int
-    public final method foo$atomicfu$array(@org.jetbrains.annotations.Nullable p0: java.lang.Object, @org.jetbrains.annotations.NotNull p1: java.util.concurrent.atomic.AtomicIntegerArray, p2: int, p3: int): int
     public final method foo$atomicfu(@org.jetbrains.annotations.Nullable p0: java.lang.Object, @org.jetbrains.annotations.NotNull p1: java.util.concurrent.atomic.AtomicIntegerFieldUpdater, p2: int): int
     public final static @org.jetbrains.annotations.NotNull method getA$FU(): java.util.concurrent.atomic.AtomicIntegerFieldUpdater
     public final method getA(): int
@@ -63,7 +55,6 @@
     public final @org.jetbrains.annotations.Nullable method getR(): java.lang.Object
     public final static @org.jetbrains.annotations.NotNull method getRs$FU(): java.util.concurrent.atomic.AtomicReferenceFieldUpdater
     public final @org.jetbrains.annotations.Nullable method getRs(): java.lang.Object
-    public final method loop$atomicfu$array(@org.jetbrains.annotations.NotNull p0: java.util.concurrent.atomic.AtomicIntegerArray, p1: int, @org.jetbrains.annotations.NotNull p2: kotlin.jvm.functions.Function1): void
     public final method loop$atomicfu(@org.jetbrains.annotations.NotNull p0: java.util.concurrent.atomic.AtomicIntegerFieldUpdater, @org.jetbrains.annotations.NotNull p1: kotlin.jvm.functions.Function1, @org.jetbrains.annotations.Nullable p2: java.lang.Object): void
     public final method testIntExtensionLoops(): void
     public final inner class LoopTest$A
diff --git a/plugins/atomicfu/atomicfu-compiler/testData/box/ExtensionsTest.kt b/plugins/atomicfu/atomicfu-compiler/testData/box/ExtensionsTest.kt
index 510f90d..20ea1e8 100644
--- a/plugins/atomicfu/atomicfu-compiler/testData/box/ExtensionsTest.kt
+++ b/plugins/atomicfu/atomicfu-compiler/testData/box/ExtensionsTest.kt
@@ -112,4 +112,4 @@
     testClass.testScopedFieldGetters()
     testClass.testExtension()
     return "OK"
-}
\ No newline at end of file
+}
diff --git a/plugins/atomicfu/atomicfu-compiler/testData/box/ExtensionsTest.txt b/plugins/atomicfu/atomicfu-compiler/testData/box/ExtensionsTest.txt
index b866de4..9d66682 100644
--- a/plugins/atomicfu/atomicfu-compiler/testData/box/ExtensionsTest.txt
+++ b/plugins/atomicfu/atomicfu-compiler/testData/box/ExtensionsTest.txt
@@ -11,7 +11,6 @@
     private volatile @kotlin.jvm.Volatile @org.jetbrains.annotations.Nullable field s: java.lang.Object
     static method <clinit>(): void
     public method <init>(): void
-    public final method booleanExtensionArithmetic$atomicfu$array(@org.jetbrains.annotations.Nullable p0: java.lang.Object, @org.jetbrains.annotations.NotNull p1: java.util.concurrent.atomic.AtomicIntegerArray, p2: int): void
     public final method booleanExtensionArithmetic$atomicfu(@org.jetbrains.annotations.Nullable p0: java.lang.Object, @org.jetbrains.annotations.NotNull p1: java.util.concurrent.atomic.AtomicIntegerFieldUpdater): void
     public final static @org.jetbrains.annotations.NotNull method getA$FU(): java.util.concurrent.atomic.AtomicIntegerFieldUpdater
     public final method getA(): int
@@ -21,11 +20,8 @@
     public final method getL(): long
     public final static @org.jetbrains.annotations.NotNull method getS$FU(): java.util.concurrent.atomic.AtomicReferenceFieldUpdater
     public final @org.jetbrains.annotations.Nullable method getS(): java.lang.Object
-    public final method intExtensionArithmetic$atomicfu$array(@org.jetbrains.annotations.Nullable p0: java.lang.Object, @org.jetbrains.annotations.NotNull p1: java.util.concurrent.atomic.AtomicIntegerArray, p2: int): void
     public final method intExtensionArithmetic$atomicfu(@org.jetbrains.annotations.Nullable p0: java.lang.Object, @org.jetbrains.annotations.NotNull p1: java.util.concurrent.atomic.AtomicIntegerFieldUpdater): void
-    public final method longExtensionArithmetic$atomicfu$array(@org.jetbrains.annotations.Nullable p0: java.lang.Object, @org.jetbrains.annotations.NotNull p1: java.util.concurrent.atomic.AtomicLongArray, p2: int): void
     public final method longExtensionArithmetic$atomicfu(@org.jetbrains.annotations.Nullable p0: java.lang.Object, @org.jetbrains.annotations.NotNull p1: java.util.concurrent.atomic.AtomicLongFieldUpdater): void
-    public final method refExtension$atomicfu$array(@org.jetbrains.annotations.Nullable p0: java.lang.Object, @org.jetbrains.annotations.NotNull p1: java.util.concurrent.atomic.AtomicReferenceArray, p2: int): void
     public final method refExtension$atomicfu(@org.jetbrains.annotations.Nullable p0: java.lang.Object, @org.jetbrains.annotations.NotNull p1: java.util.concurrent.atomic.AtomicReferenceFieldUpdater): void
     public final method testExtension(): void
     public final method testScopedFieldGetters(): void
diff --git a/plugins/atomicfu/atomicfu-compiler/testData/box/FieldInObjectTest.txt b/plugins/atomicfu/atomicfu-compiler/testData/box/FieldInObjectTest.txt
new file mode 100644
index 0000000..9c5e576
--- /dev/null
+++ b/plugins/atomicfu/atomicfu-compiler/testData/box/FieldInObjectTest.txt
@@ -0,0 +1,111 @@
+@kotlin.Metadata
+public final class DelegatedProvider$_a$DelegatedProvider$VolatileWrapper {
+    // source: 'FieldInObjectTest.kt'
+    private final static @org.jetbrains.annotations.NotNull field _a$FU: java.util.concurrent.atomic.AtomicIntegerFieldUpdater
+    private volatile @kotlin.jvm.Volatile field _a: int
+    static method <clinit>(): void
+    public method <init>(): void
+    public synthetic final static method access$get_a$p(p0: DelegatedProvider$_a$DelegatedProvider$VolatileWrapper): int
+    public synthetic final static method access$set_a$p(p0: DelegatedProvider$_a$DelegatedProvider$VolatileWrapper, p1: int): void
+    public final static @org.jetbrains.annotations.NotNull method get_a$FU(): java.util.concurrent.atomic.AtomicIntegerFieldUpdater
+    public final method get_a(): int
+    public final inner class DelegatedProvider$_a$DelegatedProvider$VolatileWrapper
+}
+
+@kotlin.Metadata
+public final class DelegatedProvider {
+    // source: 'FieldInObjectTest.kt'
+    public final static @org.jetbrains.annotations.NotNull field INSTANCE: DelegatedProvider
+    public final static @org.jetbrains.annotations.NotNull field _a$DelegatedProvider$VolatileWrapper: DelegatedProvider$_a$DelegatedProvider$VolatileWrapper
+    private volatile @kotlin.jvm.Volatile field vInt: int
+    static method <clinit>(): void
+    private method <init>(): void
+    public final method getA(): int
+    public final method getVInt(): int
+    public final static @org.jetbrains.annotations.NotNull method get_a$DelegatedProvider$VolatileWrapper(): DelegatedProvider$_a$DelegatedProvider$VolatileWrapper
+    public final method setA(p0: int): void
+    public final method setVInt(p0: int): void
+    public final inner class DelegatedProvider$_a$DelegatedProvider$VolatileWrapper
+}
+
+@kotlin.Metadata
+public final class FieldInObjectTestKt {
+    // source: 'FieldInObjectTest.kt'
+    public final static @org.jetbrains.annotations.NotNull method box(): java.lang.String
+    private final static method testDelegatedPropertiesInObject(): void
+    private final static method testFieldInObject(): void
+}
+
+@kotlin.Metadata
+public final class Provider$Port$Provider$VolatileWrapper {
+    // source: 'FieldInObjectTest.kt'
+    private final static @org.jetbrains.annotations.NotNull field port$FU: java.util.concurrent.atomic.AtomicIntegerFieldUpdater
+    private volatile @kotlin.jvm.Volatile field port: int
+    static method <clinit>(): void
+    public method <init>(): void
+    public final static @org.jetbrains.annotations.NotNull method getPort$FU(): java.util.concurrent.atomic.AtomicIntegerFieldUpdater
+    public final method getPort(): int
+    public final inner class Provider$Port$Provider$VolatileWrapper
+}
+
+@kotlin.Metadata
+public final class Provider$_l$Provider$VolatileWrapper {
+    // source: 'FieldInObjectTest.kt'
+    private final static @org.jetbrains.annotations.NotNull field _l$FU: java.util.concurrent.atomic.AtomicLongFieldUpdater
+    private volatile @kotlin.jvm.Volatile field _l: long
+    static method <clinit>(): void
+    public method <init>(): void
+    public synthetic final static method access$get_l$FU$p(): java.util.concurrent.atomic.AtomicLongFieldUpdater
+    public final inner class Provider$_l$Provider$VolatileWrapper
+}
+
+@kotlin.Metadata
+public final class Provider$_ref$Provider$VolatileWrapper {
+    // source: 'FieldInObjectTest.kt'
+    private final static @org.jetbrains.annotations.NotNull field _ref$FU: java.util.concurrent.atomic.AtomicReferenceFieldUpdater
+    private volatile @kotlin.jvm.Volatile @org.jetbrains.annotations.Nullable field _ref: java.lang.Object
+    static method <clinit>(): void
+    public method <init>(): void
+    public final static @org.jetbrains.annotations.NotNull method get_ref$FU(): java.util.concurrent.atomic.AtomicReferenceFieldUpdater
+    public final @org.jetbrains.annotations.Nullable method get_ref(): java.lang.Object
+    public final inner class Provider$_ref$Provider$VolatileWrapper
+}
+
+@kotlin.Metadata
+public final class Provider$_x$Provider$VolatileWrapper {
+    // source: 'FieldInObjectTest.kt'
+    private final static @org.jetbrains.annotations.NotNull field _x$FU: java.util.concurrent.atomic.AtomicIntegerFieldUpdater
+    private volatile @kotlin.jvm.Volatile field _x: int
+    static method <clinit>(): void
+    public method <init>(): void
+    public final static @org.jetbrains.annotations.NotNull method get_x$FU(): java.util.concurrent.atomic.AtomicIntegerFieldUpdater
+    public final method get_x(): int
+    public final inner class Provider$_x$Provider$VolatileWrapper
+}
+
+@kotlin.Metadata
+public final class Provider {
+    // source: 'FieldInObjectTest.kt'
+    public final static @org.jetbrains.annotations.NotNull field INSTANCE: Provider
+    public final static @org.jetbrains.annotations.NotNull field _l$Provider$VolatileWrapper: Provider$_l$Provider$VolatileWrapper
+    public final static @org.jetbrains.annotations.NotNull field _ref$Provider$VolatileWrapper: Provider$_ref$Provider$VolatileWrapper
+    public final static @org.jetbrains.annotations.NotNull field _x$Provider$VolatileWrapper: Provider$_x$Provider$VolatileWrapper
+    private final static @org.jetbrains.annotations.NotNull field intArr: java.util.concurrent.atomic.AtomicIntegerArray
+    private final static @org.jetbrains.annotations.NotNull field longArr: java.util.concurrent.atomic.AtomicLongArray
+    public final static @org.jetbrains.annotations.NotNull field port$Provider$VolatileWrapper: Provider$Port$Provider$VolatileWrapper
+    private final static @org.jetbrains.annotations.NotNull field refArr: java.util.concurrent.atomic.AtomicReferenceArray
+    static method <clinit>(): void
+    private method <init>(): void
+    public final static @org.jetbrains.annotations.NotNull method getIntArr(): java.util.concurrent.atomic.AtomicIntegerArray
+    public final method getL(): long
+    public final static @org.jetbrains.annotations.NotNull method getLongArr(): java.util.concurrent.atomic.AtomicLongArray
+    public final static @org.jetbrains.annotations.NotNull method getPort$Provider$VolatileWrapper(): Provider$Port$Provider$VolatileWrapper
+    public final static @org.jetbrains.annotations.NotNull method getRefArr(): java.util.concurrent.atomic.AtomicReferenceArray
+    public final static @org.jetbrains.annotations.NotNull method get_ref$Provider$VolatileWrapper(): Provider$_ref$Provider$VolatileWrapper
+    public final static @org.jetbrains.annotations.NotNull method get_x$Provider$VolatileWrapper(): Provider$_x$Provider$VolatileWrapper
+    public final method next(): int
+    public final inner class Provider$Port$Provider$VolatileWrapper
+    public final inner class Provider$_l$Provider$VolatileWrapper
+    public final inner class Provider$_ref$Provider$VolatileWrapper
+    public final inner class Provider$_x$Provider$VolatileWrapper
+}
diff --git a/plugins/atomicfu/atomicfu-compiler/testData/box/InlineExtensionWithTypeParameterTest.txt b/plugins/atomicfu/atomicfu-compiler/testData/box/InlineExtensionWithTypeParameterTest.txt
index 7e8096b..1cecdb2 100644
--- a/plugins/atomicfu/atomicfu-compiler/testData/box/InlineExtensionWithTypeParameterTest.txt
+++ b/plugins/atomicfu/atomicfu-compiler/testData/box/InlineExtensionWithTypeParameterTest.txt
@@ -21,7 +21,6 @@
     private volatile @kotlin.jvm.Volatile @org.jetbrains.annotations.Nullable field sref: java.lang.Object
     static method <clinit>(): void
     public method <init>(): void
-    public final method foo$atomicfu$array(@org.jetbrains.annotations.Nullable p0: java.lang.Object, @org.jetbrains.annotations.NotNull p1: java.util.concurrent.atomic.AtomicReferenceArray, p2: int, p3: int, @org.jetbrains.annotations.NotNull p4: InlineExtensionWithTypeParameterTest$Segment): int
     public final method foo$atomicfu(@org.jetbrains.annotations.Nullable p0: java.lang.Object, @org.jetbrains.annotations.NotNull p1: java.util.concurrent.atomic.AtomicReferenceFieldUpdater, p2: int, @org.jetbrains.annotations.NotNull p3: InlineExtensionWithTypeParameterTest$Segment): int
     private final method getSegmentId(p0: InlineExtensionWithTypeParameterTest$Segment): int
     public final static @org.jetbrains.annotations.NotNull method getSref$FU(): java.util.concurrent.atomic.AtomicReferenceFieldUpdater
diff --git a/plugins/atomicfu/atomicfu-compiler/testData/box/LockTest.txt b/plugins/atomicfu/atomicfu-compiler/testData/box/LockTest.txt
index 0684146..90ce95b 100644
--- a/plugins/atomicfu/atomicfu-compiler/testData/box/LockTest.txt
+++ b/plugins/atomicfu/atomicfu-compiler/testData/box/LockTest.txt
@@ -6,6 +6,7 @@
     static method <clinit>(): void
     public method <init>(): void
     public final method testLock(): void
+    public final method tryAcquire$atomicfu(@org.jetbrains.annotations.Nullable p0: java.lang.Object, @org.jetbrains.annotations.NotNull p1: java.util.concurrent.atomic.AtomicIntegerFieldUpdater): boolean
 }
 
 @kotlin.Metadata
@@ -13,6 +14,4 @@
     // source: 'LockTest.kt'
     public final static @org.jetbrains.annotations.NotNull method box(): java.lang.String
     public final static @org.jetbrains.annotations.NotNull method reflectionTest(@org.jetbrains.annotations.NotNull p0: java.lang.String, @org.jetbrains.annotations.NotNull p1: java.util.Map): java.util.List
-    public final static method tryAcquire$atomicfu$array(@org.jetbrains.annotations.Nullable p0: java.lang.Object, @org.jetbrains.annotations.NotNull p1: java.util.concurrent.atomic.AtomicIntegerArray, p2: int): boolean
-    public final static method tryAcquire$atomicfu(@org.jetbrains.annotations.Nullable p0: java.lang.Object, @org.jetbrains.annotations.NotNull p1: java.util.concurrent.atomic.AtomicIntegerFieldUpdater): boolean
 }
diff --git a/plugins/atomicfu/atomicfu-compiler/testData/box/ParameterizedInlineFunExtensionTest.txt b/plugins/atomicfu/atomicfu-compiler/testData/box/ParameterizedInlineFunExtensionTest.txt
index a16a127..1b4de7d 100644
--- a/plugins/atomicfu/atomicfu-compiler/testData/box/ParameterizedInlineFunExtensionTest.txt
+++ b/plugins/atomicfu/atomicfu-compiler/testData/box/ParameterizedInlineFunExtensionTest.txt
@@ -5,9 +5,7 @@
     private volatile @kotlin.jvm.Volatile @org.jetbrains.annotations.Nullable field tail: java.lang.Object
     static method <clinit>(): void
     public method <init>(): void
-    public final method bar$atomicfu$array(@org.jetbrains.annotations.Nullable p0: java.lang.Object, @org.jetbrains.annotations.NotNull p1: java.util.concurrent.atomic.AtomicReferenceArray, p2: int, p3: java.lang.Object, p4: java.lang.Object): java.lang.Object
     public final method bar$atomicfu(@org.jetbrains.annotations.Nullable p0: java.lang.Object, @org.jetbrains.annotations.NotNull p1: java.util.concurrent.atomic.AtomicReferenceFieldUpdater, p2: java.lang.Object, p3: java.lang.Object): java.lang.Object
-    public final method foo$atomicfu$array(@org.jetbrains.annotations.Nullable p0: java.lang.Object, @org.jetbrains.annotations.NotNull p1: java.util.concurrent.atomic.AtomicReferenceArray, p2: int, p3: java.lang.Object, p4: java.lang.Object, @org.jetbrains.annotations.NotNull p5: kotlin.jvm.functions.Function1): java.lang.Object
     public final method foo$atomicfu(@org.jetbrains.annotations.Nullable p0: java.lang.Object, @org.jetbrains.annotations.NotNull p1: java.util.concurrent.atomic.AtomicReferenceFieldUpdater, p2: java.lang.Object, p3: java.lang.Object, @org.jetbrains.annotations.NotNull p4: kotlin.jvm.functions.Function1): java.lang.Object
     public final method testClose(): void
 }