Extract new smap creation logic into separate class LineNumberMapper

compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/ProvisionalFunctionExpressionLowering.kt

@@ -68,7 +68,7 @@ class ProvisionalFunctionExpressionLowering :
                     valueArgumentsCount = function.valueParameters.size,
                     reflectionTarget = null,
                     origin = origin
-                ).copyAttributes(expression) // TODO better to copy just name
+                ).copyAttributes(expression)
             )
         )
     }

compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/ReturnableBlockLowering.kt

@@ -129,9 +129,9 @@ class ReturnableBlockTransformer(val context: CommonBackendContext, val containe
 
         val newStatements = expression.statements.mapIndexed { i, s ->
             if (expression.statements.size == 1 && s is IrInlinedFunctionBlock) {
-                val transformed = s.statements.mapIndexed { j, it -> transformSingleStatement(it, j == s.statements.lastIndex) }
-                s.statements.clear()
-                s.statements.addAll(transformed)
+                for ((j, statement) in s.statements.withIndex()) {
+                    s.statements[j] = transformSingleStatement(statement, j == s.statements.lastIndex)
+                }
                 s
             } else {
                 transformSingleStatement(s, i == expression.statements.lastIndex)

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

@@ -8,7 +8,6 @@ package org.jetbrains.kotlin.backend.jvm.codegen
 import org.jetbrains.kotlin.backend.common.ir.*
 import org.jetbrains.kotlin.backend.common.lower.BOUND_RECEIVER_PARAMETER
 import org.jetbrains.kotlin.backend.common.lower.LoweredStatementOrigins
-import org.jetbrains.kotlin.backend.common.lower.inline.isAdaptedFunctionReference
 import org.jetbrains.kotlin.backend.jvm.*
 import org.jetbrains.kotlin.backend.jvm.intrinsics.IntrinsicMethod
 import org.jetbrains.kotlin.backend.jvm.intrinsics.JavaClassProperty
@@ -57,7 +56,6 @@ import org.jetbrains.kotlin.resolve.jvm.jvmSignature.JvmMethodSignature
 import org.jetbrains.kotlin.types.TypeSystemCommonBackendContext
 import org.jetbrains.kotlin.types.computeExpandedTypeForInlineClass
 import org.jetbrains.kotlin.types.model.TypeParameterMarker
-import org.jetbrains.kotlin.util.OperatorNameConventions
 import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstanceOrNull
 import org.jetbrains.org.objectweb.asm.Label
 import org.jetbrains.org.objectweb.asm.Opcodes
@@ -144,21 +142,7 @@ class ExpressionCodegen(
     val smap: SourceMapper,
     val reifiedTypeParametersUsages: ReifiedTypeParametersUsages,
 ) : IrElementVisitor<PromisedValue, BlockInfo>, BaseExpressionCodegen {
-    data class AdditionalIrInlineData(val smap: SourceMapCopier, val inlinedBlock: IrInlinedFunctionBlock, val parentSmap: SourceMapper, val tryInfo: TryWithFinallyInfo?) {
-//        fun isInvokeOnLambda(): Boolean = inlinedBlock.inlineCall.symbol.owner.name == OperatorNameConventions.INVOKE
-    }
-
-    val classToCachedSourceMapper = mutableMapOf<IrDeclaration, SourceMapper>()
-    val localSmapCopiersByClass = mutableListOf<AdditionalIrInlineData>()
-
-    private fun getLocalSmap(): List<AdditionalIrInlineData> = /*context.*/localSmapCopiersByClass
-    private fun dropLastLocalSmap() {
-        /*context.*/localSmapCopiersByClass.removeLast()
-    }
-
-    private fun addToLocalSmap(info: AdditionalIrInlineData) {
-        /*context.*/localSmapCopiersByClass.add(info)
-    }
+    private val lineNumberMapper = LineNumberMapper(this)
 
     override fun toString(): String = signature.toString()
 
@@ -173,8 +157,6 @@ class ExpressionCodegen(
 
     val state = context.state
 
-    private val fileEntry = irFunction.fileParentBeforeInline.fileEntry
-
     override val visitor: InstructionAdapter
         get() = mv
 
@@ -208,42 +190,12 @@ class ExpressionCodegen(
 
     private fun markNewLinkedLabel() = linkedLabel().apply { mv.visitLabel(this) }
 
-    private fun getLineNumberForOffset(offset: Int): Int {
-        if (getLocalSmap().isNotEmpty()) {
-            var previousData: AdditionalIrInlineData? = null
-            var result = -1
-            val iterator = getLocalSmap().reversed().iterator()
-            while (iterator.hasNext()) {
-                if (previousData?.inlinedBlock?.isLambdaInlining() == true) {
-                    while (iterator.hasNext() && iterator.next().inlinedBlock.inlineDeclaration != getInlinedAt(previousData.inlinedBlock.inlinedElement)) {
-                        // after lambda's smap we should skip "frames" that were inlined inside body of inline function that accept given lambda
-                        continue
-                    }
-                    if (!iterator.hasNext()) break
-                }
-                val inlineData = iterator.next()
-
-                previousData = inlineData
-                val localFileEntry = if (inlineData.inlinedBlock.inlinedElement is IrCallableReference<*>)
-                    getLocalSmap().reversed().map { it.inlinedBlock }
-                        .firstOrNull { it.isFunctionInlining() }?.inlineDeclaration?.fileOrNull?.fileEntry ?: fileEntry
-                else
-                    inlineData.inlinedBlock.inlineDeclaration.fileEntry
-                val lineNumber = if (result == -1) localFileEntry.getLineNumber(offset) + 1 else result
-                val mappedLineNumber = inlineData.smap.mapLineNumber(lineNumber)
-                result = mappedLineNumber
-            }
-            return result
-        }
-        return fileEntry.getLineNumber(offset) + 1
-    }
-
     private fun IrElement.markLineNumber(startOffset: Boolean) {
         if (noLineNumberScope) return
         val offset = if (startOffset) this.startOffset else endOffset
         if (offset < 0) return
 
-        val lineNumber = getLineNumberForOffset(offset)
+        val lineNumber = lineNumberMapper.getLineNumberForOffset(offset)
 
         assert(lineNumber > 0)
         if (lastLineNumber != lineNumber) {
@@ -300,12 +252,8 @@ class ExpressionCodegen(
             // If this function has an expression body, return the result of that expression.
             // Otherwise, if it does not end in a return statement, it must be void-returning,
             // and an explicit return instruction at the end is still required to pass validation.
+            setExtraLineNumberForVoidReturningFunction(irFunction)
             if (body !is IrStatementContainer || body.statements.lastOrNull() !is IrReturn) {
-                // Allow setting a breakpoint on the closing brace of a void-returning function
-                // without an explicit return, or the `class Something(` line of a primary constructor.
-                if (irFunction.origin != JvmLoweredDeclarationOrigin.CLASS_STATIC_INITIALIZER) {
-                    irFunction.markLineNumber(startOffset = irFunction is IrConstructor && irFunction.isPrimary)
-                }
                 val (returnType, returnIrType) = irFunction.returnAsmAndIrTypes()
                 result.materializeAt(returnType, returnIrType)
                 mv.areturn(returnType)
@@ -319,6 +267,18 @@ class ExpressionCodegen(
         writeParameterInLocalVariableTable(startLabel, endLabel)
     }
 
+    private fun setExtraLineNumberForVoidReturningFunction(irFunction: IrFunction) {
+        val body = irFunction.body ?: return
+        if (body !is IrStatementContainer || body.statements.lastOrNull() !is IrReturn) {
+            // Allow setting a breakpoint on the closing brace of a void-returning function
+            // without an explicit return, or the `class Something(` line of a primary constructor.
+            if (irFunction.origin != JvmLoweredDeclarationOrigin.CLASS_STATIC_INITIALIZER) {
+                irFunction.markLineNumber(startOffset = irFunction is IrConstructor && irFunction.isPrimary)
+                mv.nop()
+            }
+        }
+    }
+
     private fun generateFakeContinuationConstructorIfNeeded() {
         if (!irFunction.isSuspendCapturingCrossinline()) return
         val continuationClass = irFunction.continuationClass() ?: return
@@ -519,129 +479,112 @@ class ExpressionCodegen(
         }
     }
 
-    private inline fun stashSmapForGivenTry(tryInfo: TryWithFinallyInfo, block: () -> Unit) {
-        val lastLineNumberBeforeFinally = lastLineNumber
-        val localSmap = getLocalSmap()
-        val smapCountToDrop = localSmap.indexOfFirst { it.tryInfo == tryInfo }
-        if (smapCountToDrop == -1) {
-            return block()
-        }
+    private fun visitInlinedFunctionBlock(inlinedBlock: IrInlinedFunctionBlock, data: BlockInfo): PromisedValue {
+        val inlineCall = inlinedBlock.inlineCall
+        val callee = inlinedBlock.inlineDeclaration as? IrFunction
+        val callLineNumber = lineNumberMapper.getLineNumberForOffset(inlineCall.startOffset)
 
-        val smapInTryBlock = localSmap.takeLast(localSmap.size - smapCountToDrop)
-        smapInTryBlock.forEach { _ -> dropLastLocalSmap() }
-        block()
-        smapInTryBlock.forEach { addToLocalSmap(it) }
-        if (smapInTryBlock.isNotEmpty()) {
-            lastLineNumber = lastLineNumberBeforeFinally
-        }
-    }
-
-    private fun visitInlinedFunctionBlock(block: IrInlinedFunctionBlock, data: BlockInfo): PromisedValue {
-        val lineNumberForOffset = getLineNumberForOffset(block.inlineCall.startOffset)
-        val callee = block.inlineDeclaration as? IrFunction
-
-        block.getNonDefaultAdditionalStatementsFromInlinedBlock().forEach { exp ->
+        // 1. Evaluate NON DEFAULT arguments from inline function call
+        inlinedBlock.getNonDefaultAdditionalStatementsFromInlinedBlock().forEach { exp ->
             exp.accept(this, data).discard()
         }
 
-        // TODO start_4: reuse code from org/jetbrains/kotlin/codegen/inline/MethodInliner.kt:267
-        if (block.isLambdaInlining()) {
-            val overrideLineNumber = getLocalSmap().reversed().firstOrNull { !it.inlinedBlock.isLambdaInlining() }
-                ?.inlinedBlock?.inlineDeclaration?.isInlineOnly() == true
-            val currentLineNumber = if (overrideLineNumber) getLocalSmap().last().smap.callSite!!.line else lineNumberForOffset
-
-            val firstLine = callee?.body?.statements?.firstOrNull()?.let {
-                block.inlineDeclaration.fileEntry.getLineNumber(it.startOffset) + 1
-            } ?: -1
-            // TODO DefaultLambda
-            if (/*(info is DefaultLambda != overrideLineNumber) &&*/ currentLineNumber >= 0 && firstLine == currentLineNumber) {
-                val label = Label()
-                val fakeLineNumber = (getLocalSmap().lastOrNull()?.smap?.parent ?: smap)
-                    .mapSyntheticLineNumber(SourceMapper.LOCAL_VARIABLE_INLINE_ARGUMENT_SYNTHETIC_LINE_NUMBER)
-                mv.visitLabel(label)
-                mv.visitLineNumber(fakeLineNumber, label)
-            }
+        if (inlinedBlock.isLambdaInlining()) {
+            lineNumberMapper.setUpAdditionalLineNumbersBeforeLambdaInlining(inlinedBlock)
         }
-        // TODO end_4
 
-        noLineNumberScopeWithCondition(block.inlineDeclaration.isInlineOnly()) {
-            buildSmapFor(block, data)
+        noLineNumberScopeWithCondition(inlinedBlock.inlineDeclaration.isInlineOnly()) {
+            inlineCall.markLineNumber(startOffset = true)
+            mv.nop()
 
-            block.getDefaultAdditionalStatementsFromInlinedBlock().forEach { exp ->
+            lineNumberMapper.buildSmapFor(inlinedBlock, inlinedBlock.buildOrGetClassSMAP(data), data)
+
+            if (inlineCall.usesDefaultArguments()) {
+                // $default function has first LN pointing to original callee
+                callee?.markLineNumber(startOffset = true)
+                mv.nop()
+            }
+
+            // 2. Evaluate DEFAULT arguments from inline function call
+            inlinedBlock.getDefaultAdditionalStatementsFromInlinedBlock().forEach { exp ->
                 exp.accept(this, data).discard()
             }
 
-            if (block.inlineCall.usesDefaultArguments()) {
+            if (inlineCall.usesDefaultArguments()) {
                 // we must reset LN because at this point in original inliner we will inline non default call
                 lastLineNumber = -1
             }
 
-            val result = block.getOriginalStatementsFromInlinedBlock().fold(unitValue) { prev, exp ->
+            // 3. Evaluate statements from inline function body
+            val result = inlinedBlock.getOriginalStatementsFromInlinedBlock().fold(unitValue) { prev, exp ->
                 prev.discard()
                 exp.accept(this, data)
             }
 
-            val calleeBody = callee?.body
-            if (block.inlinedElement !is IrCallableReference<*> || callee?.isInline == true) {
-                if (calleeBody !is IrStatementContainer || calleeBody.statements.lastOrNull() !is IrReturn) {
-                    // TODO start_1: reuse code
-                    // Allow setting a breakpoint on the closing brace of a void-returning function
-                    // without an explicit return, or the `class Something(` line of a primary constructor.
-                    if (callee?.origin != JvmLoweredDeclarationOrigin.CLASS_STATIC_INITIALIZER) {
-                        callee?.let {
-                            it.markLineNumber(startOffset = callee is IrConstructor && callee.isPrimary)
-                            mv.nop()
-                        }
-                    }
-                }
+            if (callee != null && (inlinedBlock.inlinedElement !is IrCallableReference<*> || callee.isInline)) {
+                setExtraLineNumberForVoidReturningFunction(callee)
             }
-            // TODO end_1
 
-//            if (!(block.origin as InlinedFunction).isLambdaInlining) {
-            // TODO start_3: reuse from visitReturn and see ReturnableBlockLowering
+            // After `ReturnableBlockLowering` last return could transform, but we still need to place new LN there
             val lastStatement = callee?.body?.statements?.lastOrNull()
             if (lastStatement is IrReturn) {
                 val returnTarget = lastStatement.returnTargetSymbol.owner
-                val originalTarget = (returnTarget as? IrAttributeContainer)?.attributeOwnerId ?: returnTarget
-                if (originalTarget == block.inlineDeclaration) {
+                val originalReturnTarget = (returnTarget as? IrAttributeContainer)?.attributeOwnerId ?: returnTarget
+                if (originalReturnTarget == inlinedBlock.inlineDeclaration) {
                     // if return is implicit we must put new LN at the end of expression
-                    block.statements.last().markLineNumber(startOffset = lastStatement.startOffset != lastStatement.endOffset)
+                    inlinedBlock.statements.last().markLineNumber(startOffset = lastStatement.startOffset != lastStatement.endOffset)
                     mv.nop()
                 }
             }
-            // TODO end_3
-//            }
 
-            dropLastLocalSmap()
+            lineNumberMapper.dropCurrentSmap()
 
-
-            if (block.isLambdaInlining()) {
-                val overrideLineNumber = getLocalSmap().reversed().firstOrNull { !it.inlinedBlock.isLambdaInlining() }
-                    ?.inlinedBlock?.inlineDeclaration?.isInlineOnly() == true
-                val currentLineNumber = if (overrideLineNumber) getLocalSmap().last().smap.callSite!!.line else lineNumberForOffset
-//        currentLineNumber = getLineNumberForOffset(marker.inlineCall.startOffset)
-                // TODO start_2: reuse code
-                if (currentLineNumber != -1) {
-                    if (overrideLineNumber) {
-                        // This is from the function we're inlining into, so no need to remap.
-                        mv.visitLineNumber(currentLineNumber, markNewLabel())
-                    } else {
-                        // Need to go through the superclass here to properly remap the line number via `sourceMapper`.
-                        block.inlineCall.markLineNumber(true)
-                    }
-                    mv.nop()
-                }
-                // TODO end_2
+            if (inlinedBlock.isLambdaInlining()) {
+                lineNumberMapper.setUpAdditionalLineNumbersAfterLambdaInlining(inlinedBlock)
+            } else {
+                // takeUnless is required to avoid markLineNumberAfterInlineIfNeeded for inline only
+                lastLineNumber = callLineNumber.takeUnless { noLineNumberScope } ?: -1
             }
 
-            // takeUnless is required to avoid markLineNumberAfterInlineIfNeeded for inline only
-            if (block.isFunctionInlining()) {
-                lastLineNumber = lineNumberForOffset.takeUnless { noLineNumberScope } ?: -1
-            }
             return result
         }
     }
 
+    private fun IrDeclaration.getClassWithDeclaredFunction(): IrClass? {
+        val parent = this.parentClassOrNull ?: return null
+        if (!parent.isInterface || (this is IrFunction && this.hasJvmDefault())) return parent
+        return parent.declarations.singleOrNull { it.origin == JvmLoweredDeclarationOrigin.DEFAULT_IMPLS } as IrClass
+    }
+
+    private fun IrInlinedFunctionBlock.buildOrGetClassSMAP(data: BlockInfo): SMAP {
+        if (this.isLambdaInlining()) {
+            return context.typeToCachedSMAP[context.getLocalClassType(this.inlinedElement as IrAttributeContainer)]!!
+        }
+
+        val callee = this.inlineDeclaration
+        val calleeFromActualClass = callee.getClassWithDeclaredFunction()!!.declarations
+            .asSequence()
+            .filterIsInstance<IrSimpleFunction>()
+            .filter { it.attributeOwnerId == callee } // original callee could be transformed after lowerings, so we must get correct one
+            .filter {
+                if (inlineCall.usesDefaultArguments()) it.origin == IrDeclarationOrigin.FUNCTION_FOR_DEFAULT_PARAMETER
+                else it.origin != IrDeclarationOrigin.FUNCTION_FOR_DEFAULT_PARAMETER
+            }
+            .filter { it.origin != JvmLoweredDeclarationOrigin.FOR_INLINE_STATE_MACHINE_TEMPLATE } // filter functions with $$forInline postfix
+            .single()
+
+        val nodeAndSmap = calleeFromActualClass.let { actualCallee ->
+            val callToActualCallee = IrCallImpl.fromSymbolOwner(
+                inlineCall.startOffset, inlineCall.endOffset, inlineCall.type, actualCallee.symbol
+            )
+            val callable = methodSignatureMapper.mapToCallableMethod(callToActualCallee, null)
+            val callGenerator = getOrCreateCallGenerator(callToActualCallee, data, callable.signature)
+            (callGenerator as IrInlineCodegen).compileInline()
+        }
+
+        return nodeAndSmap.classSMAP
+    }
+
     private fun visitStatementContainer(container: IrStatementContainer, data: BlockInfo): PromisedValue {
         if (container is IrInlinedFunctionBlock) {
             return visitInlinedFunctionBlock(container, data)
@@ -1049,108 +992,6 @@ class ExpressionCodegen(
                     element.render()
         )
 
-    private fun IrDeclaration.getClassWithDeclaredFunction(): IrClass? {
-        val parent = this.parentClassOrNull ?: return null
-        if (!parent.isInterface || (this is IrFunction && this.hasJvmDefault())) return parent
-        return parent.declarations.singleOrNull { it.origin == JvmLoweredDeclarationOrigin.DEFAULT_IMPLS } as IrClass
-    }
-
-    private fun IrFunctionAccessExpression.isInvokeOnDefaultArg(expected: IrDeclaration): Boolean {
-        if (this.symbol.owner.name != OperatorNameConventions.INVOKE) return false
-
-        val dispatch = this.dispatchReceiver as? IrGetValue
-        val parameter = dispatch?.symbol?.owner as? IrValueParameter
-        val default = parameter?.defaultValue?.expression as? IrFunctionExpression
-
-        return default?.function == expected
-    }
-
-    private fun getInlinedAt(originalExpression: IrElement): IrDeclaration? {
-        for (block in getLocalSmap().map { it.inlinedBlock }.filter { it.isFunctionInlining() }) {
-            block.inlineCall.getAllArgumentsWithIr().forEach {
-                val actualArg = (it.second?.attributeOwnerId ?: ((it.first.defaultValue?.expression?.attributeOwnerId as? IrBlock)?.statements?.firstOrNull() as? IrClass)?.attributeOwnerId) as? IrExpression
-                val extractedAnonymousFunction = if (actualArg?.isAdaptedFunctionReference() == true) ((actualArg as IrBlock).statements.last() as IrFunctionReference).attributeOwnerId else actualArg
-                if (extractedAnonymousFunction == ((originalExpression as? IrAttributeContainer)?.attributeOwnerId ?: originalExpression)) {
-                    return block.inlineDeclaration
-                }
-            }
-        }
-
-        return null
-    }
-
-    private fun buildSmapFor(declaration: IrInlinedFunctionBlock, data: BlockInfo): PromisedValue {
-        val inlineCall = declaration.inlineCall
-        val callee = declaration.inlineDeclaration
-
-        declaration.inlineCall.markLineNumber(startOffset = true)
-        mv.nop()
-
-        val localSmaps = getLocalSmap()
-        if (declaration.isLambdaInlining()) {
-            val callSite = getLocalSmap().lastOrNull()?.smap?.callSite?.takeIf { inlineCall.isInvokeOnDefaultArg(callee) }
-            val classSMAP = context.typeToCachedSMAP[context.getLocalClassType(declaration.inlinedElement as IrAttributeContainer)]!!
-
-            val sourceMapper = if (localSmaps.isEmpty()) {
-                smap
-            } else {
-                localSmaps.reversed().firstOrNull { it.inlinedBlock.inlineDeclaration == getInlinedAt(declaration.inlinedElement) }?.smap?.parent
-                    ?: localSmaps.last().smap.parent // if we are in anonymous inlined class and lambda was declared outside
-            }
-            addToLocalSmap(
-                AdditionalIrInlineData(
-                    SourceMapCopier(sourceMapper, classSMAP, callSite),
-                    declaration,
-                    context.getSourceMapper(if (declaration.inlinedElement is IrCallableReference<*>) irFunction.parentAsClass else callee.parentClassOrNull!!),
-                    data.infos.filterIsInstance<TryWithFinallyInfo>().lastOrNull()
-                )
-            )
-        } else {
-            val calleeFromActualClass = callee.getClassWithDeclaredFunction()!!.declarations
-                .asSequence()
-                .filterIsInstance<IrSimpleFunction>()
-                .filter { it.attributeOwnerId == callee } // original callee could be transformed after lowerings, so we must get correct one
-                .filter { if (inlineCall.usesDefaultArguments()) it.origin == IrDeclarationOrigin.FUNCTION_FOR_DEFAULT_PARAMETER else it.origin != IrDeclarationOrigin.FUNCTION_FOR_DEFAULT_PARAMETER }
-                .filter { it.origin != JvmLoweredDeclarationOrigin.FOR_INLINE_STATE_MACHINE_TEMPLATE } // filter functions with $$forInline postfix
-                .single()
-            val nodeAndSmap = calleeFromActualClass.let { actualCallee ->
-                val callToActualCallee = IrCallImpl.fromSymbolOwner(inlineCall.startOffset, inlineCall.endOffset, inlineCall.type, actualCallee.symbol)
-                val callable = methodSignatureMapper.mapToCallableMethod(callToActualCallee, null)
-                val callGenerator = getOrCreateCallGenerator(callToActualCallee, data, callable.signature)
-                (callGenerator as InlineCodegen<*>).compileInline()
-            }
-
-            val key = callee.parentClassOrNull!!
-            val newSmap = if (localSmaps.isEmpty()) {
-                smap
-            } else {
-                localSmaps.last().parentSmap
-            }
-            val sourcePosition = let {
-                val sourceInfo = newSmap.sourceInfo!!
-                val localFileEntry = getLocalSmap().lastOrNull()?.inlinedBlock?.inlineDeclaration?.fileEntry ?: fileEntry
-                val line = if (inlineCall.startOffset < 0) lastLineNumber else localFileEntry.getLineNumber(inlineCall.startOffset) + 1
-//                val file = fileEntry.name.drop(1)
-                SourcePosition(line, sourceInfo.sourceFileName!!, sourceInfo.pathOrCleanFQN)
-            }
-            val sourceMapCopier = SourceMapCopier(newSmap, nodeAndSmap.classSMAP, sourcePosition)
-            addToLocalSmap(
-                AdditionalIrInlineData(
-                    sourceMapCopier, declaration, context.getSourceMapper(key),
-                    data.infos.filterIsInstance<TryWithFinallyInfo>().lastOrNull()
-                )
-            )
-
-        }
-        if (inlineCall.usesDefaultArguments()) {
-            // $default function has first LN pointing to original callee
-            callee.markLineNumber(startOffset = true)
-            mv.nop()
-        }
-
-        return unitValue
-    }
-
     override fun visitClass(declaration: IrClass, data: BlockInfo): PromisedValue {
         if (declaration.origin != JvmLoweredDeclarationOrigin.CONTINUATION_CLASS) {
             val childCodegen = ClassCodegen.getOrCreate(declaration, context, enclosingFunctionForLocalObjects)
@@ -1609,7 +1450,7 @@ class ExpressionCodegen(
     ) {
         val gapStart = markNewLinkedLabel()
         data.localGapScope(tryWithFinallyInfo) {
-            stashSmapForGivenTry(tryWithFinallyInfo) {
+            lineNumberMapper.stashSmapForGivenTry(tryWithFinallyInfo) {
                 finallyDepth++
                 if (isFinallyMarkerRequired) {
                     generateFinallyMarker(mv, finallyDepth, true)

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

@@ -0,0 +1,232 @@
+/*
+ * Copyright 2010-2022 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.backend.jvm.codegen
+
+import org.jetbrains.kotlin.backend.common.ir.inlineDeclaration
+import org.jetbrains.kotlin.backend.common.ir.isFunctionInlining
+import org.jetbrains.kotlin.backend.common.ir.isLambdaInlining
+import org.jetbrains.kotlin.backend.common.lower.inline.isAdaptedFunctionReference
+import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
+import org.jetbrains.kotlin.backend.jvm.ir.fileParentBeforeInline
+import org.jetbrains.kotlin.backend.jvm.ir.isInlineOnly
+import org.jetbrains.kotlin.codegen.inline.SMAP
+import org.jetbrains.kotlin.codegen.inline.SourceMapCopier
+import org.jetbrains.kotlin.codegen.inline.SourceMapper
+import org.jetbrains.kotlin.codegen.inline.SourcePosition
+import org.jetbrains.kotlin.ir.IrElement
+import org.jetbrains.kotlin.ir.declarations.*
+import org.jetbrains.kotlin.ir.expressions.*
+import org.jetbrains.kotlin.ir.util.*
+import org.jetbrains.kotlin.util.OperatorNameConventions
+import org.jetbrains.org.objectweb.asm.Label
+
+/**
+ * This class basically is 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.
+//  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 val smapStack = mutableListOf<DataForIrInlinedFunction>()
+
+    private data class DataForIrInlinedFunction(
+        val smap: SourceMapCopier,
+        val inlinedBlock: IrInlinedFunctionBlock,
+        val parentSmap: SourceMapper, // this property is used to simulate change in smap but without ruining previous one
+        val tryInfo: TryWithFinallyInfo?
+    )
+
+    fun dropCurrentSmap() {
+        smapStack.removeFirst()
+    }
+
+    fun getLineNumberForOffset(offset: Int): Int {
+        if (smapStack.isEmpty()) {
+            return fileEntry.getLineNumber(offset) + 1
+        }
+
+        var previousData: DataForIrInlinedFunction? = null
+        var result = -1
+        val iterator = smapStack.iterator()
+        while (iterator.hasNext()) {
+            if (previousData != null && previousData.inlinedBlock.isLambdaInlining()) {
+                val inlinedAt = getDeclarationWhereGivenElementWasInlined(previousData.inlinedBlock.inlinedElement)
+                while (iterator.hasNext() && iterator.next().inlinedBlock.inlineDeclaration != inlinedAt) {
+                    // after lambda's smap we should skip "frames" that were inlined inside body of inline function that accept given lambda
+                    continue
+                }
+                if (!iterator.hasNext()) break
+            }
+            val inlineData = iterator.next()
+
+            previousData = inlineData
+            val localFileEntry = inlineData.inlinedBlock.getClassThatContainsDeclaration().fileParentBeforeInline.fileEntry
+            val lineNumber = if (result == -1) localFileEntry.getLineNumber(offset) + 1 else result
+            val mappedLineNumber = inlineData.smap.mapLineNumber(lineNumber)
+            result = mappedLineNumber
+        }
+
+        return result
+    }
+
+    fun buildSmapFor(inlinedBlock: IrInlinedFunctionBlock, classSMAP: SMAP, data: BlockInfo) {
+        val inlineCall = inlinedBlock.inlineCall
+
+        val newData = if (inlinedBlock.isLambdaInlining()) {
+            val callSite = smapStack.firstOrNull()?.smap?.callSite?.takeIf { inlinedBlock.isInvokeOnDefaultArg() }
+
+            val sourceMapper = if (smapStack.isEmpty()) {
+                smap
+            } else {
+                val inlinedAt = getDeclarationWhereGivenElementWasInlined(inlinedBlock.inlinedElement)
+                val inlineData = smapStack.firstOrNull { it.inlinedBlock.inlineDeclaration == inlinedAt }
+                    ?: smapStack.first() // if we are in anonymous inlined class and lambda was declared outside
+                inlineData.smap.parent
+            }
+            DataForIrInlinedFunction(
+                SourceMapCopier(sourceMapper, classSMAP, callSite),
+                inlinedBlock,
+                context.getSourceMapper(inlinedBlock.getClassThatContainsDeclaration()),
+                data.infos.filterIsInstance<TryWithFinallyInfo>().lastOrNull(),
+            )
+        } else {
+            val sourceMapper = if (smapStack.isEmpty()) smap else smapStack.first().parentSmap
+            val sourcePosition = let {
+                val sourceInfo = sourceMapper.sourceInfo!!
+                val localFileEntry = smapStack.firstOrNull()?.inlinedBlock?.inlineDeclaration?.fileParentBeforeInline?.fileEntry ?: fileEntry
+                val line = if (inlineCall.startOffset < 0) lastLineNumber else localFileEntry.getLineNumber(inlineCall.startOffset) + 1
+                SourcePosition(line, sourceInfo.sourceFileName!!, sourceInfo.pathOrCleanFQN)
+            }
+
+            DataForIrInlinedFunction(
+                SourceMapCopier(sourceMapper, classSMAP, sourcePosition),
+                inlinedBlock,
+                context.getSourceMapper(inlinedBlock.getClassThatContainsDeclaration()),
+                data.infos.filterIsInstance<TryWithFinallyInfo>().lastOrNull(),
+            )
+        }
+
+        smapStack.add(0, newData)
+    }
+
+    fun stashSmapForGivenTry(tryInfo: TryWithFinallyInfo, block: () -> Unit) {
+        val lastLineNumberBeforeFinally = lastLineNumber
+        val smapCountToDrop = smapStack.indexOfLast { it.tryInfo == tryInfo } + 1
+        if (smapCountToDrop == 0) {
+            return block()
+        }
+
+        val smapInTryBlock = smapStack.take(smapCountToDrop)
+        smapInTryBlock.forEach { _ -> dropCurrentSmap() }
+        block()
+        smapInTryBlock.reversed().forEach { smapStack.add(0, it) }
+        if (smapInTryBlock.isNotEmpty()) {
+            expressionCodegen.lastLineNumber = lastLineNumberBeforeFinally
+        }
+    }
+
+    fun setUpAdditionalLineNumbersBeforeLambdaInlining(inlinedBlock: IrInlinedFunctionBlock) {
+        val lineNumberForOffset = getLineNumberForOffset(inlinedBlock.inlineCall.startOffset)
+        val callee = inlinedBlock.inlineDeclaration as? IrFunction
+
+        // TODO: reuse code from org/jetbrains/kotlin/codegen/inline/MethodInliner.kt:267
+        val overrideLineNumber = smapStack
+            .map { it.inlinedBlock }
+            .firstOrNull { !it.isLambdaInlining() }?.inlineDeclaration?.isInlineOnly() == true
+        val currentLineNumber = if (overrideLineNumber) smapStack.first().smap.callSite!!.line else lineNumberForOffset
+
+        val firstLine = callee?.body?.statements?.firstOrNull()?.let {
+            inlinedBlock.inlineDeclaration.fileEntry.getLineNumber(it.startOffset) + 1
+        } ?: -1
+        if ((inlinedBlock.isInvokeOnDefaultArg() != overrideLineNumber) && currentLineNumber >= 0 && firstLine == currentLineNumber) {
+            val label = Label()
+            val fakeLineNumber = (smapStack.firstOrNull()?.smap?.parent ?: smap)
+                .mapSyntheticLineNumber(SourceMapper.LOCAL_VARIABLE_INLINE_ARGUMENT_SYNTHETIC_LINE_NUMBER)
+            expressionCodegen.mv.visitLabel(label)
+            expressionCodegen.mv.visitLineNumber(fakeLineNumber, label)
+        }
+    }
+
+    fun setUpAdditionalLineNumbersAfterLambdaInlining(inlinedBlock: IrInlinedFunctionBlock) {
+        val lineNumberForOffset = getLineNumberForOffset(inlinedBlock.inlineCall.startOffset)
+
+        // TODO: reuse code from org/jetbrains/kotlin/codegen/inline/MethodInliner.kt:316
+        val overrideLineNumber = smapStack
+            .map { it.inlinedBlock }
+            .firstOrNull { !it.isLambdaInlining() }?.inlineDeclaration?.isInlineOnly() == true
+        val currentLineNumber = if (overrideLineNumber) smapStack.first().smap.callSite!!.line else lineNumberForOffset
+        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) })
+            } else {
+                // Need to go through the superclass here to properly remap the line number via `sourceMapper`.
+                expressionCodegen.markLineNumber(inlinedBlock.inlineCall)
+            }
+            expressionCodegen.mv.nop()
+        }
+    }
+
+    private fun IrInlinedFunctionBlock.getClassThatContainsDeclaration(): IrClass {
+        val firstFunctionInlineBlock = if (this.inlinedElement is IrCallableReference<*>)
+            smapStack.map { it.inlinedBlock }.firstOrNull { it.isFunctionInlining() }
+        else
+            this
+
+        return firstFunctionInlineBlock?.inlineDeclaration?.parentClassOrNull ?: irFunction.parentAsClass
+    }
+
+    private fun getDeclarationWhereGivenElementWasInlined(inlinedElement: IrElement): IrDeclaration? {
+        val originalInlinedElement = ((inlinedElement as? IrAttributeContainer)?.attributeOwnerId ?: inlinedElement)
+        for (block in smapStack.map { it.inlinedBlock }.filter { it.isFunctionInlining() }) {
+            block.inlineCall.getAllArgumentsWithIr().forEach {
+                // pretty messed up thing, this is needed to get the original expression that was inlined
+                // it was changed a couple of times after all lowerings, so we must get `attributeOwnerId` to ensure that this is original
+                val actualArg = if (it.second == null) {
+                    val blockWithClass = it.first.defaultValue?.expression?.attributeOwnerId as? IrBlock
+                    blockWithClass?.statements?.firstOrNull() as? IrClass
+                } else {
+                    it.second
+                }
+
+                val originalActualArg = actualArg?.attributeOwnerId as? IrExpression
+                val extractedAnonymousFunction = if (originalActualArg?.isAdaptedFunctionReference() == true) {
+                    (originalActualArg as IrBlock).statements.last() as IrFunctionReference
+                } else {
+                    originalActualArg
+                }
+
+                if (extractedAnonymousFunction?.attributeOwnerId == originalInlinedElement) {
+                    return block.inlineDeclaration
+                }
+            }
+        }
+
+        return null
+    }
+
+    private fun IrInlinedFunctionBlock.isInvokeOnDefaultArg(): Boolean {
+        val call = this.inlineCall
+        val expected = this.inlineDeclaration
+        if (call.symbol.owner.name != OperatorNameConventions.INVOKE) return false
+
+        val dispatch = call.dispatchReceiver as? IrGetValue
+        val parameter = dispatch?.symbol?.owner as? IrValueParameter
+        val default = parameter?.defaultValue?.expression as? IrFunctionExpression
+
+        return default?.function == expected
+    }
+}