[JVM_IR] Move methods that generate LN into LineNumberMapper class

compiler/ir/backend.jvm/codegen/src/org/jetbrains/kotlin/backend/jvm/codegen/ExpressionCodegen.kt

@@ -62,8 +62,6 @@ import org.jetbrains.org.objectweb.asm.Opcodes
 import org.jetbrains.org.objectweb.asm.Type
 import org.jetbrains.org.objectweb.asm.commons.InstructionAdapter
 import java.util.*
-import kotlin.contracts.ExperimentalContracts
-import kotlin.contracts.contract
 
 sealed class ExpressionInfo {
     var blockInfo: BlockInfo? = null
@@ -142,8 +140,6 @@ class ExpressionCodegen(
     val smap: SourceMapper,
     val reifiedTypeParametersUsages: ReifiedTypeParametersUsages,
 ) : IrElementVisitor<PromisedValue, BlockInfo>, BaseExpressionCodegen {
-    private val lineNumberMapper = LineNumberMapper(this)
-
     override fun toString(): String = signature.toString()
 
     var finallyDepth = 0
@@ -165,10 +161,13 @@ class ExpressionCodegen(
     override val typeSystem: TypeSystemCommonBackendContext
         get() = typeMapper.typeSystem
 
-    override var lastLineNumber: Int = -1
+    private val lineNumberMapper = LineNumberMapper(this)
-    var noLineNumberScope: Boolean = false
 
-    private var isInsideCondition = false
+    override val lastLineNumber: Int
+        get() = lineNumberMapper.getLineNumber()
+
+    var isInsideCondition: Boolean = false
+        private set
 
     private val closureReifiedMarkers = hashMapOf<IrClass, ReifiedTypeParametersUsages>()
 
@@ -190,38 +189,16 @@ class ExpressionCodegen(
 
     private fun markNewLinkedLabel() = linkedLabel().apply { mv.visitLabel(this) }
 
-    private fun IrElement.markLineNumber(startOffset: Boolean) {
+    fun IrElement.markLineNumber(startOffset: Boolean) {
-        if (noLineNumberScope) return
+        lineNumberMapper.markLineNumber(this, startOffset)
-        val offset = if (startOffset) this.startOffset else endOffset
-        if (offset < 0) return
-
-        val lineNumber = lineNumberMapper.getLineNumberForOffset(offset)
-
-        assert(lineNumber > 0)
-        if (lastLineNumber != lineNumber) {
-            lastLineNumber = lineNumber
-            mv.visitLineNumber(lineNumber, markNewLabel())
-        }
-    }
-
-    fun markLineNumber(element: IrElement) = element.markLineNumber(true)
-
-    @OptIn(ExperimentalContracts::class)
-    private inline fun noLineNumberScopeWithCondition(flag: Boolean, block: () -> Unit) {
-        contract {
-            callsInPlace(block, kotlin.contracts.InvocationKind.EXACTLY_ONCE)
-        }
-        val previousState = noLineNumberScope
-        noLineNumberScope = noLineNumberScope || flag
-        block()
-        noLineNumberScope = previousState
     }
 
     fun noLineNumberScope(block: () -> Unit) {
-        val previousState = noLineNumberScope
+        lineNumberMapper.noLineNumberScope(block)
-        noLineNumberScope = true
+    }
-        block()
+
-        noLineNumberScope = previousState
+    override fun markLineNumberAfterInlineIfNeeded(registerLineNumberAfterwards: Boolean) {
+        lineNumberMapper.markLineNumberAfterInlineIfNeeded(registerLineNumberAfterwards)
     }
 
     fun gen(expression: IrExpression, type: Type, irType: IrType, data: BlockInfo) {
@@ -422,16 +399,9 @@ class ExpressionCodegen(
             }
         }
 
+        // This block must be executed after `writeLocalVariablesInTable`
         if (expression is IrInlinedFunctionBlock) {
-            // This block must be executed after `writeLocalVariablesInTable`
+            lineNumberMapper.afterIrInline(expression)
-            if (expression.isFunctionInlining()) {
-                val callLineNumber = lineNumberMapper.getLineNumberForOffset(expression.inlineCall.startOffset)
-                // takeUnless is required to avoid markLineNumberAfterInlineIfNeeded for inline only
-                lastLineNumber = callLineNumber.takeUnless { noLineNumberScope } ?: -1
-                markLineNumberAfterInlineIfNeeded(isInsideCondition)
-            } else {
-                lineNumberMapper.setUpAdditionalLineNumbersAfterLambdaInlining(expression)
-            }
         }
 
         if (isSynthesizedInitBlock) {
@@ -496,11 +466,9 @@ class ExpressionCodegen(
             exp.accept(this, data).discard()
         }
 
-        if (inlinedBlock.isLambdaInlining()) {
+        lineNumberMapper.beforeIrInline(inlinedBlock)
-            lineNumberMapper.setUpAdditionalLineNumbersBeforeLambdaInlining(inlinedBlock)
-        }
 
-        noLineNumberScopeWithCondition(inlinedBlock.inlineDeclaration.isInlineOnly()) {
+        lineNumberMapper.noLineNumberScopeWithCondition(inlinedBlock.inlineDeclaration.isInlineOnly()) {
             inlineCall.markLineNumber(startOffset = true)
             mv.nop()
 
@@ -519,7 +487,7 @@ class ExpressionCodegen(
 
             if (inlineCall.usesDefaultArguments()) {
                 // we must reset LN because at this point in original inliner we will inline non default call
-                lastLineNumber = -1
+                lineNumberMapper.resetLineNumber()
             }
 
             // 3. Evaluate statements from inline function body
@@ -729,7 +697,7 @@ class ExpressionCodegen(
         val owner = typeMapper.mapClass(callee.constructedClass)
         val signature = methodSignatureMapper.mapSignatureSkipGeneric(callee)
 
-        markLineNumber(expression)
+        expression.markLineNumber(startOffset = true)
 
         // In this case the receiver is `this` (not specified in IR) and the return value is discarded anyway.
         mv.load(0, OBJECT_TYPE)
@@ -742,7 +710,7 @@ class ExpressionCodegen(
         }
 
         generateConstructorArguments(expression, signature, data)
-        markLineNumber(expression)
+        expression.markLineNumber(startOffset = true)
 
         mv.visitMethodInsn(Opcodes.INVOKESPECIAL, owner.internalName, signature.asmMethod.name, signature.asmMethod.descriptor, false)
 
@@ -759,13 +727,13 @@ class ExpressionCodegen(
 
         // IR constructors have no receiver and return the new instance, but on JVM they are void-returning
         // instance methods named <init>.
-        markLineNumber(expression)
+        expression.markLineNumber(startOffset = true)
         putNeedClassReificationMarker(callee.constructedClass)
         mv.anew(owner)
         mv.dup()
 
         generateConstructorArguments(expression, signature, data)
-        markLineNumber(expression)
+        expression.markLineNumber(startOffset = true)
 
         mv.visitMethodInsn(Opcodes.INVOKESPECIAL, owner.internalName, signature.asmMethod.name, signature.asmMethod.descriptor, false)
 
@@ -1631,20 +1599,6 @@ class ExpressionCodegen(
         }
     }
 
-    override fun markLineNumberAfterInlineIfNeeded(registerLineNumberAfterwards: Boolean) {
-        if (noLineNumberScope || registerLineNumberAfterwards) {
-            if (lastLineNumber > -1) {
-                val label = Label()
-                mv.visitLabel(label)
-                mv.visitLineNumber(lastLineNumber, label)
-            }
-        } else {
-            // Inline function has its own line number which is in a separate instance of codegen,
-            // therefore we need to reset lastLineNumber to force a line number generation after visiting inline function.
-            lastLineNumber = -1
-        }
-    }
-
     val isFinallyMarkerRequired: Boolean
         get() = irFunction.isInline || irFunction.origin == JvmLoweredDeclarationOrigin.INLINE_LAMBDA
 

compiler/ir/backend.jvm/codegen/src/org/jetbrains/kotlin/backend/jvm/codegen/LineNumberMapper.kt

@@ -31,24 +31,26 @@ import org.jetbrains.kotlin.ir.util.parentClassOrNull
 import org.jetbrains.kotlin.ir.util.statements
 import org.jetbrains.kotlin.util.OperatorNameConventions
 import org.jetbrains.org.objectweb.asm.Label
+import kotlin.contracts.ExperimentalContracts
+import kotlin.contracts.InvocationKind
+import kotlin.contracts.contract
 
 /**
- * This class basically is just another wrapper around SMAP.
+ * This class serves two purposes:
+ * 1. Generate line numbers for given `IrElement`
+ * 2. Keep track of `smap` for functions inlined by IR inliner. We can consider this class as basically just another wrapper around SMAP.
  * It is used to unify smap creation for functions inlined from IR and from bytecode.
  */
-// TODO extract all functions and fields responsible for LN modification from ExpressionCodegen.
+class LineNumberMapper(private val expressionCodegen: ExpressionCodegen) {
-//  This includes `lastLineNumber`, `markLineNumber`, `noLineNumberScope`, `markLineNumberAfterInlineIfNeeded`, ...
-class LineNumberMapper(
-    private val expressionCodegen: ExpressionCodegen,
-) {
     private val context: JvmBackendContext
         get() = expressionCodegen.context
     private val smap = expressionCodegen.smap
     private val irFunction = expressionCodegen.irFunction
-    private val lastLineNumber: Int
-        get() = expressionCodegen.lastLineNumber
     private val fileEntry = irFunction.fileParentBeforeInline.fileEntry
 
+    private var lastLineNumber: Int = -1
+    private var noLineNumberScope: Boolean = false
+
     private val smapStack = mutableListOf<DataForIrInlinedFunction>()
 
     private data class DataForIrInlinedFunction(
@@ -58,11 +60,84 @@ class LineNumberMapper(
         val tryInfo: TryWithFinallyInfo?
     )
 
+    private fun markNewLabel() = Label().apply { expressionCodegen.mv.visitLabel(this) }
+
+    fun markLineNumber(element: IrElement, startOffset: Boolean) {
+        if (noLineNumberScope) return
+        val offset = if (startOffset) element.startOffset else element.endOffset
+        if (offset < 0) return
+
+        val lineNumber = getLineNumberForOffset(offset)
+
+        assert(lineNumber > 0)
+        if (lastLineNumber != lineNumber) {
+            lastLineNumber = lineNumber
+            expressionCodegen.mv.visitLineNumber(lineNumber, markNewLabel())
+        }
+    }
+
+    @OptIn(ExperimentalContracts::class)
+    internal inline fun noLineNumberScopeWithCondition(flag: Boolean, block: () -> Unit) {
+        contract {
+            callsInPlace(block, InvocationKind.EXACTLY_ONCE)
+        }
+        val previousState = noLineNumberScope
+        noLineNumberScope = noLineNumberScope || flag
+        block()
+        noLineNumberScope = previousState
+    }
+
+    fun noLineNumberScope(block: () -> Unit) {
+        val previousState = noLineNumberScope
+        noLineNumberScope = true
+        block()
+        noLineNumberScope = previousState
+    }
+
+    fun markLineNumberAfterInlineIfNeeded(registerLineNumberAfterwards: Boolean) {
+        if (noLineNumberScope || registerLineNumberAfterwards) {
+            if (lastLineNumber > -1) {
+                val label = Label()
+                expressionCodegen.mv.visitLabel(label)
+                expressionCodegen.mv.visitLineNumber(lastLineNumber, label)
+            }
+        } else {
+            // Inline function has its own line number which is in a separate instance of codegen,
+            // therefore we need to reset lastLineNumber to force a line number generation after visiting inline function.
+            lastLineNumber = -1
+        }
+    }
+
+    fun getLineNumber(): Int {
+        return lastLineNumber
+    }
+
+    fun resetLineNumber() {
+        lastLineNumber = -1
+    }
+
+    fun beforeIrInline(inlinedBlock: IrInlinedFunctionBlock) {
+        if (inlinedBlock.isLambdaInlining()) {
+            setUpAdditionalLineNumbersBeforeLambdaInlining(inlinedBlock)
+        }
+    }
+
+    fun afterIrInline(inlinedBlock: IrInlinedFunctionBlock) {
+        if (inlinedBlock.isFunctionInlining()) {
+            val callLineNumber = getLineNumberForOffset(inlinedBlock.inlineCall.startOffset)
+            // `takeUnless` is required to avoid `markLineNumberAfterInlineIfNeeded` for inline only
+            lastLineNumber = callLineNumber.takeUnless { noLineNumberScope } ?: -1
+            markLineNumberAfterInlineIfNeeded(expressionCodegen.isInsideCondition)
+        } else {
+            setUpAdditionalLineNumbersAfterLambdaInlining(inlinedBlock)
+        }
+    }
+
     fun dropCurrentSmap() {
         smapStack.removeFirst()
     }
 
-    fun getLineNumberForOffset(offset: Int): Int {
+    private fun getLineNumberForOffset(offset: Int): Int {
         if (smapStack.isEmpty()) {
             return fileEntry.getLineNumber(offset) + 1
         }
@@ -143,11 +218,11 @@ class LineNumberMapper(
         block()
         smapInTryBlock.reversed().forEach { smapStack.add(0, it) }
         if (smapInTryBlock.isNotEmpty()) {
-            expressionCodegen.lastLineNumber = lastLineNumberBeforeFinally
+            lastLineNumber = lastLineNumberBeforeFinally
         }
     }
 
-    fun setUpAdditionalLineNumbersBeforeLambdaInlining(inlinedBlock: IrInlinedFunctionBlock) {
+    private fun setUpAdditionalLineNumbersBeforeLambdaInlining(inlinedBlock: IrInlinedFunctionBlock) {
         val lineNumberForOffset = getLineNumberForOffset(inlinedBlock.inlineCall.startOffset)
         val callee = inlinedBlock.inlineDeclaration as? IrFunction
 
@@ -169,7 +244,7 @@ class LineNumberMapper(
         }
     }
 
-    fun setUpAdditionalLineNumbersAfterLambdaInlining(inlinedBlock: IrInlinedFunctionBlock) {
+    private fun setUpAdditionalLineNumbersAfterLambdaInlining(inlinedBlock: IrInlinedFunctionBlock) {
         val lineNumberForOffset = getLineNumberForOffset(inlinedBlock.inlineCall.startOffset)
 
         // TODO: reuse code from org/jetbrains/kotlin/codegen/inline/MethodInliner.kt:316
@@ -180,10 +255,10 @@ class LineNumberMapper(
         if (currentLineNumber != -1) {
             if (overrideLineNumber) {
                 // This is from the function we're inlining into, so no need to remap.
-                expressionCodegen.mv.visitLineNumber(currentLineNumber, Label().apply { expressionCodegen.mv.visitLabel(this) })
+                expressionCodegen.mv.visitLineNumber(currentLineNumber, markNewLabel())
             } else {
                 // Need to go through the superclass here to properly remap the line number via `sourceMapper`.
-                expressionCodegen.markLineNumber(inlinedBlock.inlineCall)
+                markLineNumber(inlinedBlock.inlineCall, startOffset = true)
             }
             expressionCodegen.mv.nop()
         }

compiler/ir/backend.jvm/codegen/src/org/jetbrains/kotlin/backend/jvm/intrinsics/IrIllegalArgumentException.kt

@@ -51,7 +51,7 @@ object IrIllegalArgumentException : IntrinsicMethod() {
                 data: BlockInfo,
                 expression: IrFunctionAccessExpression
             ): StackValue {
-                codegen.markLineNumber(expression)
+                with(codegen) { expression.markLineNumber(startOffset = true) }
                 v.anew(exceptionTypeDescriptor)
                 v.dup()
                 return super.invoke(v, codegen, data, expression)

compiler/ir/backend.jvm/codegen/src/org/jetbrains/kotlin/backend/jvm/intrinsics/IrIntrinsicFunction.kt

@@ -71,7 +71,7 @@ open class IrIntrinsicFunction(
         expression: IrFunctionAccessExpression
     ): StackValue {
         loadArguments(codegen, data)
-        codegen.markLineNumber(expression)
+        with(codegen) { expression.markLineNumber(startOffset = true) }
         return StackValue.onStack(genInvokeInstructionWithResult(v))
     }
 

compiler/ir/backend.jvm/codegen/src/org/jetbrains/kotlin/backend/jvm/intrinsics/RangeTo.kt

@@ -35,7 +35,7 @@ object RangeTo : IntrinsicMethod() {
                 data: BlockInfo,
                 expression: IrFunctionAccessExpression
             ): StackValue {
-                codegen.markLineNumber(expression)
+                with(codegen) { expression.markLineNumber(startOffset = true) }
                 v.anew(returnType)
                 v.dup()
                 return super.invoke(v, codegen, data, expression)

plugins/kotlinx-serialization/kotlinx-serialization.backend/src/org/jetbrains/kotlinx/serialization/compiler/backend/ir/SerializationJvmIrIntrinsicSupport.kt

@@ -88,7 +88,7 @@ class SerializationJvmIrIntrinsicSupport(
                     mv.store(storedIndex, materialVal.type)
                     IntrinsicType.WithModule(storedIndex)
                 } else {
-                    codegen.markLineNumber(expression)
+                    expression.markLineNumber(startOffset = true)
                     IntrinsicType.Simple
                 }
                 generateSerializerForType(