Basic support of source mapping in inliner

compiler/backend/src/org/jetbrains/kotlin/codegen/SourceInfo.kt

@@ -16,9 +16,8 @@
 
 package org.jetbrains.kotlin.codegen
 
-import org.jetbrains.kotlin.name.ClassId
-import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.backend.common.CodegenUtil
+import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.psi.KtNamedFunction
 
@@ -26,17 +25,24 @@ data class SourceInfo(val source: String, val pathOrCleanFQN: String, val linesI
 
     companion object {
         fun createInfo(element: KtElement?, internalClassName: String): SourceInfo {
-            assert(element != null) { "Couldn't create source mapper for null element " + internalClassName }
+            assert(element != null) { "Couldn't create source mapper for null element $internalClassName" }
             val lineNumbers = CodegenUtil.getLineNumberForElement(element!!.containingFile, true)
-            assert(lineNumbers != null) { "Couldn't extract line count in " + element.containingFile }
+                    ?: error("Couldn't extract line count in ${element.containingFile}")
 
             //TODO hack condition for package parts cleaning
             val isTopLevel = element is KtFile || (element is KtNamedFunction && element.getParent() is KtFile)
             val cleanedClassFqName = if (!isTopLevel) internalClassName else internalClassName.substringBefore('$')
 
-            return SourceInfo(element.containingKtFile.name, cleanedClassFqName, lineNumbers!!)
+            return SourceInfo(element.containingKtFile.name, cleanedClassFqName, lineNumbers)
+        }
+
+        fun createInfoForIr(lineNumbers: Int, internalClassName: String, containingFileName: String): SourceInfo {
+            //TODO cut topLevel names
+//            val isTopLevel = element is KtFile || (element is KtNamedFunction && element.getParent() is KtFile)
+//            val cleanedClassFqName = if (!isTopLevel) internalClassName else internalClassName.substringBefore('$')
+
+            return SourceInfo(containingFileName, internalClassName, lineNumbers)
         }
     }
-
 }
 

compiler/backend/src/org/jetbrains/kotlin/codegen/inline/SourceCompilerForInline.kt

@@ -274,7 +274,7 @@ class PsiSourceCompilerForInline(private val codegen: ExpressionCodegen, overrid
 
         //for maxLocals calculation
         val maxCalcAdapter = wrapWithMaxLocalCalc(node)
-        val parentContext = context.parentContext ?: error("Context has no parent: $context")
+        val parentContext = context.parentContext ?: error("Context has no parent: " + context)
         val methodContext = parentContext.intoFunction(callableDescriptor)
 
         val smap = if (callDefault) {

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/JvmBackendContext.kt

@@ -40,7 +40,7 @@ import org.jetbrains.kotlin.resolve.scopes.MemberScope
 
 class JvmBackendContext(
     val state: GenerationState,
-    psiSourceManager: PsiSourceManager,
+    val psiSourceManager: PsiSourceManager,
     override val irBuiltIns: IrBuiltIns,
     irModuleFragment: IrModuleFragment, symbolTable: SymbolTable
 ) : CommonBackendContext {

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

@@ -16,12 +16,13 @@
 
 package org.jetbrains.kotlin.backend.jvm.codegen
 
-import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
 import org.jetbrains.kotlin.backend.jvm.JvmLoweredDeclarationOrigin
 import org.jetbrains.kotlin.backend.jvm.descriptors.JvmDescriptorWithExtraFlags
 import org.jetbrains.kotlin.codegen.*
 import org.jetbrains.kotlin.codegen.binding.CodegenBinding
+import org.jetbrains.kotlin.codegen.inline.DefaultSourceMapper
+import org.jetbrains.kotlin.codegen.inline.SourceMapper
 import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.util.dump
@@ -32,7 +33,6 @@ import org.jetbrains.kotlin.resolve.DescriptorUtils
 import org.jetbrains.kotlin.resolve.DescriptorUtils.isTopLevelDeclaration
 import org.jetbrains.kotlin.resolve.jvm.diagnostics.JvmDeclarationOrigin
 import org.jetbrains.kotlin.resolve.jvm.diagnostics.OtherOrigin
-import org.jetbrains.kotlin.resolve.source.PsiSourceElement
 import org.jetbrains.kotlin.resolve.source.getPsi
 import org.jetbrains.kotlin.types.ErrorUtils
 import org.jetbrains.org.objectweb.asm.Opcodes
@@ -60,10 +60,16 @@ class ClassCodegen private constructor(
         descriptor.source.getPsi() as KtElement
     ) else typeMapper.mapType(descriptor)
 
+    private val sourceManager = context.psiSourceManager
+
+    private val fileEntry = sourceManager.getFileEntry(irClass.fileParent)
+
     val psiElement = irClass.descriptor.psiElement!!
 
     val visitor: ClassBuilder = state.factory.newVisitor(OtherOrigin(psiElement, descriptor), type, psiElement.containingFile)
 
+    private var sourceMapper: DefaultSourceMapper? = null
+
     fun generate() {
         val superClassInfo = SuperClassInfo.getSuperClassInfo(descriptor, typeMapper)
         val signature = ImplementationBodyCodegen.signature(descriptor, type, superClassInfo, typeMapper)
@@ -83,8 +89,16 @@ class ClassCodegen private constructor(
             generateDeclaration(it)
         }
 
+        done()
+    }
+
+    private fun done() {
         writeInnerClasses()
 
+        sourceMapper?.let {
+            SourceMapper.flushToClassBuilder(it, visitor)
+        }
+
         visitor.done()
     }
 
@@ -192,6 +206,20 @@ class ClassCodegen private constructor(
         }
     }
 
+
+    fun getOrCreateSourceMapper(): DefaultSourceMapper {
+        if (sourceMapper == null) {
+            sourceMapper = DefaultSourceMapper(
+                SourceInfo.createInfoForIr(
+                    fileEntry.getSourceRangeInfo(irClass.startOffset, irClass.endOffset).endLineNumber + 1,
+                    irClass.name.asString(),
+                    fileEntry.name
+                )
+            )
+        }
+        return sourceMapper!!
+    }
+
 }
 
 fun ClassDescriptor.calculateClassFlags(): Int {
@@ -269,9 +297,6 @@ private val MemberDescriptor.effectiveModality: Modality
         return modality
     }
 
-private val DeclarationDescriptorWithSource.psiElement: PsiElement?
-    get() = (source as? PsiSourceElement)?.psi
-
 private val IrField.OtherOrigin: JvmDeclarationOrigin
     get() = OtherOrigin(descriptor.psiElement, this.descriptor)
 

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

@@ -23,10 +23,7 @@ import org.jetbrains.kotlin.codegen.signature.BothSignatureWriter
 import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.descriptors.impl.TypeAliasConstructorDescriptor
 import org.jetbrains.kotlin.ir.IrElement
-import org.jetbrains.kotlin.ir.declarations.IrClass
+import org.jetbrains.kotlin.ir.declarations.*
-import org.jetbrains.kotlin.ir.declarations.IrFunction
-import org.jetbrains.kotlin.ir.declarations.IrTypeAlias
-import org.jetbrains.kotlin.ir.declarations.IrVariable
 import org.jetbrains.kotlin.ir.expressions.*
 import org.jetbrains.kotlin.ir.symbols.IrSymbol
 import org.jetbrains.kotlin.ir.util.dump
@@ -109,11 +106,18 @@ class ExpressionCodegen(
 
     val state = classCodegen.state
 
+    private val sourceManager = classCodegen.context.psiSourceManager
+
+    private val fileEntry = sourceManager.getFileEntry(irFunction.fileParent)
+
     fun generate() {
         mv.visitCode()
+        irFunction.markLineNumber(true)
         val info = BlockInfo.create()
         val result = irFunction.body!!.accept(this, info)
 
+        irFunction.markLineNumber(false)
+
         val returnType = typeMapper.mapReturnType(irFunction.descriptor)
         if (irFunction.body is IrExpressionBody) {
             mv.areturn(returnType)
@@ -526,11 +530,6 @@ class ExpressionCodegen(
         }
     }
 
-
-    fun markNewLabel(): Label {
-        return Label().apply { mv.visitLabel(this) }
-    }
-
     override fun visitReturn(expression: IrReturn, data: BlockInfo): StackValue {
         val value = expression.value.apply {
             gen(this, returnType, data)
@@ -539,6 +538,7 @@ class ExpressionCodegen(
         val afterReturnLabel = Label()
         generateFinallyBlocksIfNeeded(returnType, afterReturnLabel, data)
 
+        expression.markLineNumber(false)
         mv.areturn(returnType)
         mv.mark(afterReturnLabel)
         mv.nop()/*TODO check RESTORE_STACK_IN_TRY_CATCH processor*/
@@ -1087,6 +1087,12 @@ class ExpressionCodegen(
     override fun markLineNumberAfterInlineIfNeeded() {
         //TODO
     }
+
+    private fun markNewLabel() = Label().apply { mv.visitLabel(this) }
+
+    private fun IrElement.markLineNumber(startOffset: Boolean) {
+        mv.visitLineNumber(fileEntry.getLineNumber(if (startOffset) this.startOffset else endOffset) + 1, markNewLabel())
+    }
 }
 
 private class DefaultArg(val index: Int)

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

@@ -53,8 +53,9 @@ class IrSourceCompilerForInline(
     override val callElementText: String
         get() = callElement.toString()
 
+    //TODO
     override val callsiteFile: PsiFile?
-        get() = TODO("not implemented")
+        get() = callElement.descriptor.psiElement?.containingFile
 
     override val contextKind: OwnerKind
         get() = OwnerKind.getMemberOwnerKind(callElement.descriptor.containingDeclaration)
@@ -63,21 +64,19 @@ class IrSourceCompilerForInline(
         get() = InlineCallSiteInfo("TODO", null, null)
 
     override val lazySourceMapper: DefaultSourceMapper
-        get() = DefaultSourceMapper(SourceInfo("TODO", "TODO", 100))
+        get() = codegen.classCodegen.getOrCreateSourceMapper()
 
     override fun generateLambdaBody(adapter: MethodVisitor, jvmMethodSignature: JvmMethodSignature, lambdaInfo: ExpressionLambda): SMAP {
         lambdaInfo as? IrExpressionLambda ?: error("Expecting ir lambda, but $lambdaInfo")
 
         val functionCodegen = object : FunctionCodegen(lambdaInfo.function, codegen.classCodegen) {
             override fun createMethod(flags: Int, signature: JvmMethodGenericSignature): MethodVisitor {
-                //TODO: to avoid smap assertion
-                adapter.visitLineNumber(1, Label())
                 return adapter
             }
         }
         functionCodegen.generate()
 
-        return SMAP(/*TODO*/listOf(FileMapping("TODO", "TODO").also { it.id = 1; it.addRangeMapping(RangeMapping(1, 1, 1)) }))
+        return SMAP(codegen.classCodegen.getOrCreateSourceMapper().resultMappings)
     }
 
     override fun doCreateMethodNodeFromSource(

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

@@ -5,14 +5,14 @@
 
 package org.jetbrains.kotlin.backend.jvm.codegen
 
+import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.codegen.FrameMapBase
-import org.jetbrains.kotlin.ir.declarations.IrClass
+import org.jetbrains.kotlin.descriptors.DeclarationDescriptorWithSource
-import org.jetbrains.kotlin.ir.declarations.IrConstructor
+import org.jetbrains.kotlin.ir.declarations.*
-import org.jetbrains.kotlin.ir.declarations.IrFunction
-import org.jetbrains.kotlin.ir.declarations.IrSymbolOwner
 import org.jetbrains.kotlin.ir.symbols.IrSymbol
 import org.jetbrains.kotlin.ir.util.*
 import org.jetbrains.kotlin.resolve.annotations.JVM_STATIC_ANNOTATION_FQ_NAME
+import org.jetbrains.kotlin.resolve.source.PsiSourceElement
 import org.jetbrains.org.objectweb.asm.Type
 
 class IrFrameMap : FrameMapBase<IrSymbol>()
@@ -30,3 +30,15 @@ fun IrFrameMap.leave(irDeclaration: IrSymbolOwner): Int {
 }
 
 val IrClass.isJvmInterface get() = isAnnotationClass || isInterface
+
+internal val IrDeclaration.fileParent: IrFile
+    get() {
+        val myParent = parent
+        return when (myParent) {
+            is IrFile -> myParent
+            else -> (myParent as IrDeclaration).fileParent
+        }
+    }
+
+internal val DeclarationDescriptorWithSource.psiElement: PsiElement?
+    get() = (source as? PsiSourceElement)?.psi