Debugger: Navigate to the right file in MPP projects (#KT-25152, #KT-25147)

2018-07-08 17:01:02 +03:00
parent 02ee3aef3f
commit a4b3653e1c
27 changed files with 1013 additions and 53 deletions
@@ -1,6 +1,7 @@
 <component name="ProjectDictionaryState">
  <dictionary name="yan">
    <words>
+      <w>debuggee</w>
      <w>deserializes</w>
      <w>impls</w>
      <w>kapt</w>
@@ -1,43 +1,28 @@
 /*
- * Copyright 2010-2016 JetBrains s.r.o.
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
+ * Copyright 2010-2018 JetBrains s.r.o. 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.kapt3
+package org.jetbrains.kotlin.codegen

-import org.jetbrains.kotlin.codegen.AbstractClassBuilder
-import org.jetbrains.kotlin.codegen.ClassBuilder
-import org.jetbrains.kotlin.codegen.ClassBuilderFactory
-import org.jetbrains.kotlin.codegen.ClassBuilderMode
 import org.jetbrains.kotlin.resolve.jvm.diagnostics.JvmDeclarationOrigin
 import org.jetbrains.org.objectweb.asm.*
 import org.jetbrains.org.objectweb.asm.tree.*

-internal class Kapt3BuilderFactory : ClassBuilderFactory {
-    internal val compiledClasses = mutableListOf<ClassNode>()
-    internal val origins = mutableMapOf<Any, JvmDeclarationOrigin>()
+class OriginCollectingClassBuilderFactory(private val builderMode: ClassBuilderMode) : ClassBuilderFactory {
+    val compiledClasses = mutableListOf<ClassNode>()
+    val origins = mutableMapOf<Any, JvmDeclarationOrigin>()

-    override fun getClassBuilderMode(): ClassBuilderMode = ClassBuilderMode.KAPT3
+    override fun getClassBuilderMode(): ClassBuilderMode = builderMode

    override fun newClassBuilder(origin: JvmDeclarationOrigin): AbstractClassBuilder.Concrete {
        val classNode = ClassNode()
        compiledClasses += classNode
-        origins.put(classNode, origin)
-        return Kapt3ClassBuilder(classNode)
+        origins[classNode] = origin
+        return OriginCollectingClassBuilder(classNode)
    }

-    private inner class Kapt3ClassBuilder(val classNode: ClassNode) : AbstractClassBuilder.Concrete(classNode) {
+    private inner class OriginCollectingClassBuilder(val classNode: ClassNode) : AbstractClassBuilder.Concrete(classNode) {
        override fun newField(
                origin: JvmDeclarationOrigin,
                access: Int,
@@ -47,7 +32,7 @@ internal class Kapt3BuilderFactory : ClassBuilderFactory {
                value: Any?
        ): FieldVisitor {
            val fieldNode = super.newField(origin, access, name, desc, signature, value) as FieldNode
-            origins.put(fieldNode, origin)
+            origins[fieldNode] = origin
            return fieldNode
        }

@@ -60,7 +45,7 @@ internal class Kapt3BuilderFactory : ClassBuilderFactory {
                exceptions: Array<out String>?
        ): MethodVisitor {
            val methodNode = super.newMethod(origin, access, name, desc, signature, exceptions) as MethodNode
-            origins.put(methodNode, origin)
+            origins[methodNode] = origin

            // ASM doesn't read information about local variables for the `abstract` methods so we need to get it manually
            if ((access and Opcodes.ACC_ABSTRACT) != 0 && methodNode.localVariables == null) {
@@ -73,7 +58,7 @@ internal class Kapt3BuilderFactory : ClassBuilderFactory {

    override fun asBytes(builder: ClassBuilder): ByteArray {
        val classWriter = ClassWriter(0)
-        (builder as Kapt3ClassBuilder).classNode.accept(classWriter)
+        (builder as OriginCollectingClassBuilder).classNode.accept(classWriter)
        return classWriter.toByteArray()
    }

@@ -94,6 +94,10 @@ fun PsiElement.nextLeaf(filter: (PsiElement) -> Boolean): PsiElement? {
    return leaf
 }

+fun <T : PsiElement> PsiElement.getParentOfTypes(strict: Boolean = false, vararg parentClasses: Class<out T>): T? {
+    return getParentOfTypesAndPredicate(strict, *parentClasses) { true }
+}
+
 fun <T : PsiElement> PsiElement.getParentOfTypesAndPredicate(
    strict: Boolean = false, vararg parentClasses: Class<out T>, predicate: (T) -> Boolean
 ): T? {
@@ -121,6 +125,14 @@ inline fun <reified T : PsiElement> PsiElement.getParentOfType(strict: Boolean):
    return PsiTreeUtil.getParentOfType(this, T::class.java, strict)
 }

+inline fun <reified T : PsiElement, reified V : PsiElement> PsiElement.getParentOfTypes2(): PsiElement? {
+    return PsiTreeUtil.getParentOfType(this, T::class.java, V::class.java)
+}
+
+inline fun <reified T : PsiElement, reified V : PsiElement, reified U : PsiElement> PsiElement.getParentOfTypes3(): PsiElement? {
+    return PsiTreeUtil.getParentOfType(this, T::class.java, V::class.java, U::class.java)
+}
+
 inline fun <reified T : PsiElement> PsiElement.getStrictParentOfType(): T? {
    return PsiTreeUtil.getParentOfType(this, T::class.java, true)
 }
@@ -76,10 +76,7 @@ import org.jetbrains.kotlin.idea.conversion.copy.AbstractLiteralKotlinToKotlinCo
 import org.jetbrains.kotlin.idea.conversion.copy.AbstractLiteralTextToKotlinCopyPasteTest
 import org.jetbrains.kotlin.idea.conversion.copy.AbstractTextJavaToKotlinCopyPasteConversionTest
 import org.jetbrains.kotlin.idea.coverage.AbstractKotlinCoverageOutputFilesTest
-import org.jetbrains.kotlin.idea.debugger.AbstractBeforeExtractFunctionInsertionTest
-import org.jetbrains.kotlin.idea.debugger.AbstractKotlinSteppingTest
-import org.jetbrains.kotlin.idea.debugger.AbstractPositionManagerTest
-import org.jetbrains.kotlin.idea.debugger.AbstractSmartStepIntoTest
+import org.jetbrains.kotlin.idea.debugger.*
 import org.jetbrains.kotlin.idea.debugger.evaluate.*
 import org.jetbrains.kotlin.idea.decompiler.navigation.AbstractNavigateToDecompiledLibraryTest
 import org.jetbrains.kotlin.idea.decompiler.navigation.AbstractNavigateToLibrarySourceTest
@@ -640,6 +637,10 @@ fun main(args: Array<String>) {
            model("debugger/tinyApp/src/evaluate/multipleBreakpoints", testMethod = "doMultipleBreakpointsTest")
        }

+        testClass<AbstractFileRankingTest> {
+            model("debugger/fileRanking")
+        }
+
        testClass<AbstractStubBuilderTest> {
            model("stubs", extension = "kt")
        }
@@ -20,6 +20,7 @@ import com.intellij.openapi.project.DumbService
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.util.io.FileUtilRt
 import com.intellij.psi.search.GlobalSearchScope
+import com.sun.jdi.Location
 import org.jetbrains.kotlin.asJava.finder.JavaElementFinder
 import org.jetbrains.kotlin.descriptors.CallableDescriptor
 import org.jetbrains.kotlin.idea.KotlinFileTypeFactory
@@ -45,13 +46,16 @@ object DebuggerUtils {
            project: Project,
            scope: GlobalSearchScope,
            className: JvmClassName,
-            fileName: String): KtFile? {
+            fileName: String,
+            location: Location? = null
+    ): KtFile? {
        return runReadAction {
            findSourceFileForClass(
                    project,
                    listOf(scope, KotlinSourceFilterScope.librarySources(GlobalSearchScope.allScope(project), project)),
                    className,
-                    fileName)
+                    fileName,
+                    location)
        }
    }

@@ -59,7 +63,9 @@ object DebuggerUtils {
            project: Project,
            scopes: List<GlobalSearchScope>,
            className: JvmClassName,
-            fileName: String): KtFile? {
+            fileName: String,
+            location: Location?
+    ): KtFile? {
        if (!isKotlinSourceFile(fileName)) return null
        if (DumbService.getInstance(project).isDumb) return null

@@ -84,6 +90,8 @@ object DebuggerUtils {
            return null
        }

+        location?.let { return FileRankingCalculatorForIde.findMostAppropriateSource(filesWithExactName, it) }
+
        return filesWithExactName.first()
    }

@@ -0,0 +1,421 @@
+/*
+ * Copyright 2010-2018 JetBrains s.r.o. 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.idea.debugger
+
+import com.intellij.psi.PsiElement
+import com.sun.jdi.*
+import org.jetbrains.kotlin.builtins.KotlinBuiltIns
+import org.jetbrains.kotlin.codegen.ClassBuilderMode
+import org.jetbrains.kotlin.codegen.state.IncompatibleClassTracker
+import org.jetbrains.kotlin.codegen.state.KotlinTypeMapper
+import org.jetbrains.kotlin.descriptors.*
+import org.jetbrains.kotlin.diagnostics.DiagnosticUtils
+import org.jetbrains.kotlin.idea.caches.resolve.analyze
+import org.jetbrains.kotlin.idea.debugger.FileRankingCalculator.Ranking.Companion.LOW
+import org.jetbrains.kotlin.idea.debugger.FileRankingCalculator.Ranking.Companion.ZERO
+import org.jetbrains.kotlin.idea.debugger.FileRankingCalculator.Ranking.Companion.MINOR
+import org.jetbrains.kotlin.idea.debugger.FileRankingCalculator.Ranking.Companion.NORMAL
+import org.jetbrains.kotlin.idea.debugger.FileRankingCalculator.Ranking.Companion.MAJOR
+import org.jetbrains.kotlin.idea.refactoring.getLineStartOffset
+import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.*
+import org.jetbrains.kotlin.resolve.BindingContext
+import org.jetbrains.kotlin.resolve.descriptorUtil.varargParameterPosition
+import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
+import org.jetbrains.kotlin.types.KotlinType
+import org.jetbrains.kotlin.types.typeUtil.builtIns
+import org.jetbrains.kotlin.utils.keysToMap
+import kotlin.jvm.internal.FunctionBase
+import org.jetbrains.org.objectweb.asm.Type as AsmType
+
+object FileRankingCalculatorForIde : FileRankingCalculator() {
+    override fun analyze(element: KtElement) = element.analyze(BodyResolveMode.PARTIAL)
+}
+
+abstract class FileRankingCalculator(
+    private val checkClassFqName: Boolean = true,
+    private val strictMode: Boolean = false
+) {
+    abstract fun analyze(element: KtElement): BindingContext
+
+    fun findMostAppropriateSource(files: Collection<KtFile>, location: Location): KtFile {
+        assert(files.isNotEmpty())
+
+        val fileWithRankings = files.keysToMap { fileRankingSafe(it, location) }
+        val fileWithMaxScore = fileWithRankings.maxBy { it.value }!!
+
+        if (strictMode) {
+            require(fileWithMaxScore.value.value >= 0) { "Max score is negative" }
+
+            // Allow only one element with max ranking
+            require(fileWithRankings.count { it.value == fileWithMaxScore.value } == 1) {
+                "Score is the same for several files"
+            }
+        }
+
+        return fileWithMaxScore.key
+    }
+
+    private class Ranking(val value: Int) : Comparable<Ranking> {
+        companion object {
+            val LOW = Ranking(-1000)
+            val ZERO = Ranking(0)
+            val MINOR = Ranking(1)
+            val NORMAL = Ranking(5)
+            val MAJOR = Ranking(10)
+
+            fun minor(condition: Boolean) = if (condition) MINOR else ZERO
+        }
+
+        operator fun unaryMinus() = Ranking(-value)
+        operator fun plus(other: Ranking) = Ranking(value + other.value)
+        override fun compareTo(other: Ranking) = this.value - other.value
+        override fun toString() = value.toString()
+    }
+
+    private fun collect(vararg conditions: Any): Ranking {
+        return conditions
+            .map { condition ->
+                when (condition) {
+                    is Boolean -> Ranking.minor(condition)
+                    is Int -> Ranking(condition)
+                    is Ranking -> condition
+                    else -> error("Invalid condition type ${condition.javaClass.name}")
+                }
+            }.fold(ZERO) { sum, r -> sum + r }
+    }
+
+    private fun rankingForClass(clazz: KtClassOrObject, fqName: String, virtualMachine: VirtualMachine): Ranking {
+        val bindingContext = analyze(clazz)
+        val descriptor = bindingContext[BindingContext.CLASS, clazz] ?: return ZERO
+
+        val jdiType = virtualMachine.classesByName(fqName).firstOrNull() ?: run {
+            // Check at least the class name if not found
+            return rankingForClassName(fqName, descriptor, bindingContext)
+        }
+
+        return rankingForClass(clazz, jdiType)
+    }
+
+    private fun rankingForClass(clazz: KtClassOrObject, type: ReferenceType): Ranking {
+        val bindingContext = analyze(clazz)
+        val descriptor = bindingContext[BindingContext.CLASS, clazz] ?: return ZERO
+
+        return collect(
+            rankingForClassName(type.name(), descriptor, bindingContext),
+            Ranking.minor(type.isAbstract && descriptor.modality == Modality.ABSTRACT),
+            Ranking.minor(type.isFinal && descriptor.modality == Modality.FINAL),
+            Ranking.minor(type.isStatic && !descriptor.isInner),
+            rankingForVisibility(descriptor, type)
+        )
+    }
+
+    private fun rankingForClassName(fqName: String, descriptor: ClassDescriptor, bindingContext: BindingContext): Ranking {
+        val expectedFqName = makeTypeMapper(bindingContext).mapType(descriptor).className
+        return when {
+            checkClassFqName -> if (expectedFqName == fqName) MAJOR else LOW
+            else -> if (expectedFqName.simpleName() == fqName.simpleName()) MAJOR else LOW
+        }
+    }
+
+    private fun rankingForMethod(function: KtFunction, method: Method): Ranking {
+        val bindingContext = analyze(function)
+        val descriptor = bindingContext[BindingContext.DECLARATION_TO_DESCRIPTOR, function] as? CallableMemberDescriptor ?: return ZERO
+
+        if (function !is KtConstructor<*> && method.name() != descriptor.name.asString())
+            return LOW
+
+        val typeMapper = makeTypeMapper(bindingContext)
+
+        return collect(
+            method.isConstructor && function is KtConstructor<*>,
+            method.isAbstract && descriptor.modality == Modality.ABSTRACT,
+            method.isFinal && descriptor.modality == Modality.FINAL,
+            method.isVarArgs && descriptor.varargParameterPosition() >= 0,
+            rankingForVisibility(descriptor, method),
+            rankingForType(descriptor.returnType, method.safeReturnType(), typeMapper),
+            rankingForValueParameters(descriptor.valueParameters, method.safeArguments(), typeMapper)
+        )
+    }
+
+    private fun rankingForAccessor(accessor: KtPropertyAccessor, method: Method): Ranking {
+        val methodName = method.name()
+        val expectedPropertyName = accessor.property.name ?: return ZERO
+
+        if (accessor.isSetter) {
+            if (!methodName.startsWith("set") || method.returnType() !is VoidType || method.argumentTypes().size != 1)
+                return -MAJOR
+        }
+
+        if (accessor.isGetter) {
+            if (!methodName.startsWith("get") && !methodName.startsWith("is"))
+                return -MAJOR
+            else if (method.returnType() is VoidType || method.argumentTypes().isNotEmpty())
+                return -NORMAL
+        }
+
+        val actualPropertyName = getPropertyName(methodName, accessor.isSetter)
+        if (expectedPropertyName != actualPropertyName)
+            return -NORMAL
+
+        val bindingContext = analyze(accessor.property)
+        val descriptor = bindingContext[BindingContext.PROPERTY_ACCESSOR, accessor] ?: return ZERO
+
+        val jdiPropertyType = when {
+            accessor.isGetter -> method.safeReturnType()
+            else -> method.safeArguments()?.map { it.type() }?.singleOrNull()
+        }
+
+        return rankingForType(descriptor.returnType, jdiPropertyType, makeTypeMapper(bindingContext))
+    }
+
+    private fun getPropertyName(accessorMethodName: String, isSetter: Boolean): String {
+        if (isSetter) {
+            return accessorMethodName.drop(3)
+        }
+
+        return accessorMethodName.drop(if (accessorMethodName.startsWith("is")) 2 else 3)
+    }
+
+    private fun rankingForProperty(property: KtProperty, method: Method): Ranking {
+        val methodName = method.name()
+        val propertyName = property.name ?: return ZERO
+
+        if (property.isTopLevel && method.name() == "<clinit>") {
+            // For top-level property initializers
+            return MINOR
+        }
+
+        if (!methodName.startsWith("get") && !methodName.startsWith("set"))
+            return -MAJOR
+
+        // boolean is
+        return if (methodName.drop(3) == propertyName.capitalize()) MAJOR else -NORMAL
+    }
+
+    private fun rankingForValueParameters(
+        ktParameters: List<ValueParameterDescriptor>,
+        jdiParameters: List<LocalVariable>?,
+        typeMapper: KotlinTypeMapper
+    ): Ranking {
+        if (jdiParameters == null)
+            return ZERO
+
+        if (ktParameters.size != jdiParameters.size)
+            return -NORMAL
+
+        return ktParameters.zip(jdiParameters).fold(ZERO) { sum, (ktParameter, jdiParameter) ->
+            sum + rankingForValueParameter(ktParameter, jdiParameter, typeMapper)
+        }
+    }
+
+    private fun rankingForValueParameter(
+        ktParameter: ValueParameterDescriptor,
+        jdiParameter: LocalVariable,
+        typeMapper: KotlinTypeMapper
+    ): Ranking {
+        return collect(
+            ktParameter.name.asString() == jdiParameter.name(),
+            rankingForType(ktParameter.type, jdiParameter.type(), typeMapper)
+        )
+    }
+
+    private fun rankingForVisibility(descriptor: DeclarationDescriptorWithVisibility, accessible: Accessible): Ranking {
+        return collect(
+            accessible.isPublic && descriptor.visibility == Visibilities.PUBLIC,
+            accessible.isProtected && descriptor.visibility == Visibilities.PROTECTED,
+            accessible.isPrivate && descriptor.visibility == Visibilities.PRIVATE
+        )
+    }
+
+    private fun rankingForType(ktType: KotlinType?, jdiType: Type?, typeMapper: KotlinTypeMapper): Ranking {
+        if (jdiType == null)
+            return ZERO
+
+        if (ktType == null)
+            return Ranking.minor(jdiType is VoidType)
+
+        val asmType = typeMapper.mapType(ktType)
+        return Ranking.minor(
+            when (jdiType) {
+                is VoidType -> asmType == AsmType.VOID_TYPE
+                is LongType -> asmType == AsmType.LONG_TYPE
+                is DoubleType -> asmType == AsmType.DOUBLE_TYPE
+                is CharType -> asmType == AsmType.CHAR_TYPE
+                is FloatType -> asmType == AsmType.FLOAT_TYPE
+                is ByteType -> asmType == AsmType.BYTE_TYPE
+                is IntegerType -> asmType == AsmType.INT_TYPE
+                is BooleanType -> asmType == AsmType.BOOLEAN_TYPE
+                is ShortType -> asmType == AsmType.SHORT_TYPE
+                is ArrayType -> {
+                    asmType.sort == AsmType.ARRAY && KotlinBuiltIns.isArrayOrPrimitiveArray(ktType) && run {
+                        val ktElementType = ktType.builtIns.getArrayElementType(ktType)
+                        val jdiElementType = jdiType.componentType()
+                        rankingForType(ktElementType, jdiElementType, typeMapper) > ZERO
+                    }
+                }
+                is ReferenceType -> asmType.className == jdiType.name()
+                else -> false
+            }
+        )
+    }
+
+    private fun fileRankingSafe(file: KtFile, location: Location): Ranking {
+        return try {
+            fileRanking(file, location)
+        } catch (e: AbsentInformationException) {
+            ZERO
+        } catch (e: InternalException) {
+            ZERO
+        }
+    }
+
+    private fun fileRanking(file: KtFile, location: Location): Ranking {
+        val locationLineNumber = location.lineNumber() - 1
+        val lineStartOffset = file.getLineStartOffset(locationLineNumber) ?: return LOW
+        val elementAt = file.findElementAt(lineStartOffset) ?: return ZERO
+
+        var overallRanking = ZERO
+        val method = location.method()
+
+        if (method.isLambda()) {
+            val (className, methodName) = method.getContainingClassAndMethodNameForLambda() ?: return ZERO
+            if (method.isBridge && method.isSynthetic) {
+                // It might be a static lambda field accessor
+                val containingClass = elementAt.getParentOfType<KtClassOrObject>(false) ?: return LOW
+                return rankingForClass(containingClass, className, location.virtualMachine())
+            } else {
+                val containingFunctionLiteral = findFunctionLiteralOnLine(elementAt) ?: return LOW
+
+                val containingCallable = findNonLocalCallableParent(containingFunctionLiteral) ?: return LOW
+                when (containingCallable) {
+                    is KtFunction -> if (containingCallable.name == methodName) overallRanking += MAJOR
+                    is KtProperty -> if (containingCallable.name == methodName) overallRanking += MAJOR
+                    is KtPropertyAccessor -> if (containingCallable.property.name == methodName) overallRanking += MAJOR
+                }
+
+                val containingClass = containingCallable.getParentOfType<KtClassOrObject>(false)
+                if (containingClass != null) {
+                    overallRanking += rankingForClass(containingClass, className, location.virtualMachine())
+                }
+
+                return overallRanking
+            }
+        }
+
+        // TODO support <clinit>
+        if (method.name() == "<init>") {
+            val containingClass = elementAt.getParentOfType<KtClassOrObject>(false) ?: return LOW
+            val constructorOrInitializer =
+                elementAt.getParentOfTypes2<KtConstructor<*>, KtClassInitializer>()?.takeIf { containingClass.isAncestor(it) }
+                ?: containingClass.primaryConstructor?.takeIf { it.getLine() == containingClass.getLine() }
+
+            if (constructorOrInitializer == null
+                && locationLineNumber < containingClass.getLine()
+                && locationLineNumber > containingClass.lastChild.getLine()
+            ) {
+                return LOW
+            }
+
+            overallRanking += rankingForClass(containingClass, location.declaringType())
+
+            if (constructorOrInitializer is KtConstructor<*>)
+                overallRanking += rankingForMethod(constructorOrInitializer, method)
+        } else {
+            val callable = findNonLocalCallableParent(elementAt) ?: return LOW
+            overallRanking += when (callable) {
+                is KtFunction -> rankingForMethod(callable, method)
+                is KtPropertyAccessor -> rankingForAccessor(callable, method)
+                is KtProperty -> rankingForProperty(callable, method)
+                else -> return LOW
+            }
+
+            val containingClass = elementAt.getParentOfType<KtClassOrObject>(false)
+            if (containingClass != null)
+                overallRanking += rankingForClass(containingClass, location.declaringType())
+        }
+
+        return overallRanking
+    }
+
+    private fun findFunctionLiteralOnLine(element: PsiElement): KtFunctionLiteral? {
+        val literal = element.getParentOfType<KtFunctionLiteral>(false)
+        if (literal != null) {
+            return literal
+        }
+
+        val callExpression = element.getParentOfType<KtCallExpression>(false) ?: return null
+
+        for (lambdaArgument in callExpression.lambdaArguments) {
+            if (element.getLine() == lambdaArgument.getLine()) {
+                val functionLiteral = lambdaArgument.getLambdaExpression()?.functionLiteral
+                if (functionLiteral != null) {
+                    return functionLiteral
+                }
+            }
+        }
+
+        return null
+    }
+
+    private tailrec fun findNonLocalCallableParent(element: PsiElement): PsiElement? {
+        fun PsiElement.isCallableDeclaration() = this is KtProperty || this is KtFunction || this is KtAnonymousInitializer
+
+        // org.jetbrains.kotlin.psi.KtPsiUtil.isLocal
+        fun PsiElement.isLocalDeclaration(): Boolean {
+            val containingDeclaration = getParentOfType<KtDeclaration>(true)
+            return containingDeclaration is KtCallableDeclaration || containingDeclaration is KtPropertyAccessor
+        }
+
+        if (element.isCallableDeclaration() && !element.isLocalDeclaration()) {
+            return element
+        }
+
+        val containingCallable = element.getParentOfTypes3<KtProperty, KtFunction, KtAnonymousInitializer>()
+            ?: return null
+
+        if (containingCallable.isLocalDeclaration()) {
+            return findNonLocalCallableParent(containingCallable)
+        }
+
+        return containingCallable
+    }
+
+    private fun Method.getContainingClassAndMethodNameForLambda(): Pair<String, String>? {
+        // TODO this breaks nested classes
+        val declaringClass = declaringType() as ClassType
+        val (className, methodName) = declaringClass.name().split('$', limit = 3)
+            .takeIf { it.size == 3 }
+            ?: return null
+
+        return Pair(className, methodName)
+    }
+
+    private fun Method.isLambda(): Boolean {
+        val declaringClass = declaringType() as? ClassType ?: return false
+
+        tailrec fun ClassType.isLambdaClass(): Boolean {
+            if (interfaces().any { it.name() == FunctionBase::class.java.name }) {
+                return true
+            }
+
+            val superClass = superclass() ?: return false
+            return superClass.isLambdaClass()
+        }
+
+        return declaringClass.superclass().isLambdaClass()
+    }
+
+    private fun makeTypeMapper(bindingContext: BindingContext): KotlinTypeMapper {
+        return KotlinTypeMapper(bindingContext, ClassBuilderMode.LIGHT_CLASSES, IncompatibleClassTracker.DoNothing, "debugger", false)
+    }
+}
+
+private fun String.simpleName() = substringAfterLast('.').substringAfterLast('$')
+
+private fun PsiElement.getLine(): Int {
+    return DiagnosticUtils.getLineAndColumnInPsiFile(containingFile, textRange).line
+}
@@ -121,7 +121,9 @@ class KotlinPositionManager(private val myDebugProcess: DebugProcess) : MultiReq
                    val javaClassName = JvmClassName.byInternalName(defaultInternalName(location))
                    val project = myDebugProcess.project

-                    val defaultPsiFile = DebuggerUtils.findSourceFileForClass(project, sourceSearchScopes, javaClassName, javaSourceFileName)
+                    val defaultPsiFile = DebuggerUtils.findSourceFileForClass(
+                        project, sourceSearchScopes, javaClassName, javaSourceFileName, location)
+
                    if (defaultPsiFile != null) {
                        return SourcePosition.createFromLine(defaultPsiFile, 0)
                    }
@@ -273,7 +275,7 @@ class KotlinPositionManager(private val myDebugProcess: DebugProcess) : MultiReq

        val project = myDebugProcess.project

-        return DebuggerUtils.findSourceFileForClass(project, sourceSearchScopes, className, sourceName)
+        return DebuggerUtils.findSourceFileForClass(project, sourceSearchScopes, className, sourceName, location)
    }

    private fun defaultInternalName(location: Location): String {
@@ -46,11 +46,27 @@ fun Method.safeLocationsOfLine(line: Int): List<Location> {
    }
 }

-fun Method.safeVariables(): List<LocalVariable> {
+fun Method.safeVariables(): List<LocalVariable>? {
    return try {
        variables()
    } catch (e: AbsentInformationException) {
-        emptyList()
+        null
+    }
+}
+
+fun Method.safeArguments(): List<LocalVariable>? {
+    return try {
+        arguments()
+    } catch (e: AbsentInformationException) {
+        null
+    }
+}
+
+fun Method.safeReturnType(): Type? {
+    return try {
+        returnType()
+    } catch (e: ClassNotLoadedException) {
+        null
    }
 }

@@ -156,7 +172,7 @@ private class MockStackFrame(private val location: Location, private val vm: Vir

    private fun createVisibleVariables() {
        if (visibleVariables == null) {
-            val allVariables = location.method().safeVariables()
+            val allVariables = location.method().safeVariables() ?: emptyList()
            val map = HashMap<String, LocalVariable>(allVariables.size)

            for (allVariable in allVariables) {
@@ -110,7 +110,7 @@ class FrameVisitor(val context: EvaluationContextImpl) {
        val declaringTypeName = location?.declaringType()?.name()?.replace('.', '/')?.let { JvmClassName.byInternalName(it) }

        val sourceFile = if (sourceName != null && declaringTypeName != null) {
-            DebuggerUtils.findSourceFileForClassIncludeLibrarySources(context.project, scope, declaringTypeName, sourceName)
+            DebuggerUtils.findSourceFileForClassIncludeLibrarySources(context.project, scope, declaringTypeName, sourceName, location)
        } else {
            null
        }
@@ -316,7 +316,7 @@ class SelectionAwareScopeHighlighter(val editor: Editor) {
 }

 fun PsiFile.getLineStartOffset(line: Int): Int? {
-    val doc = PsiDocumentManager.getInstance(project).getDocument(this)
+    val doc = viewProvider.document ?: PsiDocumentManager.getInstance(project).getDocument(this)
    if (doc != null && line >= 0 && line < doc.lineCount) {
        val startOffset = doc.getLineStartOffset(line)
        val element = findElementAt(startOffset) ?: return startOffset
@@ -331,11 +331,13 @@ fun PsiFile.getLineStartOffset(line: Int): Int? {
 }

 fun PsiFile.getLineEndOffset(line: Int): Int? {
-    return PsiDocumentManager.getInstance(project).getDocument(this)?.getLineEndOffset(line)
+    val document = viewProvider.document ?: PsiDocumentManager.getInstance(project).getDocument(this)
+    return document?.getLineEndOffset(line)
 }

 fun PsiElement.getLineNumber(start: Boolean = true): Int {
-    return PsiDocumentManager.getInstance(project).getDocument(this.containingFile)?.getLineNumber(if (start) this.startOffset else this.endOffset) ?: 0
+    val document = containingFile.viewProvider.document ?: PsiDocumentManager.getInstance(project).getDocument(containingFile)
+    return document?.getLineNumber(if (start) this.startOffset else this.endOffset) ?: 0
 }

 fun PsiElement.getLineCount(): Int {
@@ -316,7 +316,7 @@ class SelectionAwareScopeHighlighter(val editor: Editor) {
 }

 fun PsiFile.getLineStartOffset(line: Int): Int? {
-    val doc = PsiDocumentManager.getInstance(project).getDocument(this)
+    val doc = viewProvider.document ?: PsiDocumentManager.getInstance(project).getDocument(this)
    if (doc != null && line >= 0 && line < doc.lineCount) {
        val startOffset = doc.getLineStartOffset(line)
        val element = findElementAt(startOffset) ?: return startOffset
@@ -331,11 +331,13 @@ fun PsiFile.getLineStartOffset(line: Int): Int? {
 }

 fun PsiFile.getLineEndOffset(line: Int): Int? {
-    return PsiDocumentManager.getInstance(project).getDocument(this)?.getLineEndOffset(line)
+    val document = viewProvider.document ?: PsiDocumentManager.getInstance(project).getDocument(this)
+    return document?.getLineEndOffset(line)
 }

 fun PsiElement.getLineNumber(start: Boolean = true): Int {
-    return PsiDocumentManager.getInstance(project).getDocument(this.containingFile)?.getLineNumber(if (start) this.startOffset else this.endOffset) ?: 0
+    val document = containingFile.viewProvider.document ?: PsiDocumentManager.getInstance(project).getDocument(containingFile)
+    return document?.getLineNumber(if (start) this.startOffset else this.endOffset) ?: 0
 }

 fun PsiElement.getLineCount(): Int {
@@ -0,0 +1,21 @@
+// DO_NOT_CHECK_CLASS_FQNAME
+
+//FILE: a/a.kt
+package a
+
+class A private constructor() {
+    protected fun a() {
+        val a = 5
+    }
+}
+
+//FILE: b/a.kt
+
+//significant whitespace
+package b
+
+class A public constructor() {
+    private fun a() {
+        val a = 5
+    }
+}
@@ -0,0 +1,14 @@
+//FILE: a/a.kt
+class A {
+    init {
+        val a = 5
+        val b = 3
+    }
+}
+
+//FILE: b/a.kt
+class B {
+    init {
+        val x = 1
+    }
+}
@@ -0,0 +1,34 @@
+// DO_NOT_CHECK_CLASS_FQNAME
+
+//FILE: a/a.kt
+package a
+
+fun block(l: () -> Unit) {}
+
+class A {
+    fun a() {
+        block {
+            val a = 5
+            block {
+                val b = 4
+            }
+        }
+    }
+}
+
+//FILE: b/a.kt
+package b
+
+import a.block
+
+class A {
+    fun b() {
+        val g = 5
+        val x = 1
+        block { val y = 2 }
+        block {
+            val a = 5
+            block { block { val x = 4 }}
+        }
+    }
+}
@@ -0,0 +1,13 @@
+//FILE: a/a.kt
+class A(
+    val firstName: String,
+    val lastName: String,
+    val age: Int
+)
+
+//FILE: b/a.kt
+class B(
+    val firstName: String,
+    val lastName: String,
+    val age: Int
+)
@@ -0,0 +1,21 @@
+//FILE: a/a.kt
+class A(
+    val firstName: String,
+    val lastName: String,
+    val age: Int
+) {
+    val c = 1
+    val d = "A"
+}
+
+//FILE: b/a.kt
+class B(
+    val firstName: String,
+    val lastName: String,
+    val age: Int
+) {
+    init {
+        val a = 5
+        val b = 6
+    }
+}
@@ -0,0 +1,21 @@
+// WITH_RUNTIME
+
+//FILE: a/a.kt
+package foo
+
+class A {
+    val a by lazy {
+        val a = 5
+        val b = 2
+        ""
+    }
+}
+
+//FILE: b/a.kt
+package bar
+
+class B {
+    val a by lazy {
+        val b = 0
+    }
+}
@@ -0,0 +1,19 @@
+//FILE: a/a.kt
+package foo
+
+class A {
+    fun a() {
+        val a = 5
+    }
+}
+
+//FILE: b/a.kt
+package bar
+
+class A {
+    fun a() {
+        val c = 1
+        val d = 3
+        val g = ""
+    }
+}
@@ -0,0 +1,19 @@
+//FILE: a/a.kt
+package bar
+
+class A {
+    fun a() {
+        val a = 5
+    }
+}
+
+//FILE: b/a.kt
+package foo
+
+class A {
+    fun b() {
+        val c = 1
+        val d = 3
+        val g = ""
+    }
+}
@@ -0,0 +1,13 @@
+//FILE: a/a.kt
+class A {
+    fun a() {
+        System.out.println("a")
+    }
+}
+
+//FILE: b/a.kt
+class B {
+    fun b() {
+        System.out.println("b")
+    }
+}
@@ -0,0 +1,25 @@
+// WITH_RUNTIME
+
+//FILE: a/a.kt
+package a
+
+val a = run {
+    val a = 5
+    val b = run {
+        val c = 2
+    }
+    5
+}
+
+fun x() {
+    println("")
+}
+
+//FILE: b/a.kt
+package b
+
+val b = 5
+
+fun y() {
+    println("")
+}
@@ -0,0 +1,57 @@
+package org.jetbrains.kotlin.idea.debugger
+
+import com.sun.jdi.ThreadReference
+import org.jetbrains.kotlin.codegen.OriginCollectingClassBuilderFactory
+import org.jetbrains.kotlin.codegen.state.GenerationState
+import org.jetbrains.kotlin.psi.KtElement
+import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.org.objectweb.asm.tree.ClassNode
+
+abstract class AbstractFileRankingTest : LowLevelDebuggerTestBase() {
+    override fun doTest(
+        options: Set<String>,
+        mainThread: ThreadReference,
+        factory: OriginCollectingClassBuilderFactory,
+        state: GenerationState
+    ) {
+        val allKtFiles = factory.origins.mapNotNull { it.value.element?.containingFile as? KtFile }.distinct()
+        fun getKtFiles(name: String) = allKtFiles.filter { it.name == name }
+
+        val doNotCheckClassFqName = "DO_NOT_CHECK_CLASS_FQNAME" in options
+
+        val calculator = object : FileRankingCalculator(checkClassFqName = !doNotCheckClassFqName, strictMode = true) {
+            override fun analyze(element: KtElement) = state.bindingContext
+        }
+
+        val problems = mutableListOf<String>()
+
+        for ((node, origin) in factory.origins) {
+            val classNode = node as? ClassNode ?: continue
+            val expectedFile = origin.element?.containingFile as? KtFile ?: continue
+            val className = classNode.name.replace('/', '.')
+            val jdiClass = mainThread.virtualMachine().classesByName(className).singleOrNull()
+                ?: error("Class '$className' was not found in the debuggee process class loader")
+
+            val locations = jdiClass.allLineLocations()
+            assert(locations.isNotEmpty()) { "There are no locations for class $className" }
+
+            val allFilesWithSameName = getKtFiles(expectedFile.name)
+
+            for (location in locations) {
+                val actualFile = calculator.findMostAppropriateSource(allFilesWithSameName, location)
+                if (actualFile != expectedFile) {
+                    problems += "Location ${location.sourceName()}:${location.lineNumber() - 1} is associated with a wrong KtFile:\n" +
+                            "    - expected: ${expectedFile.virtualFilePath}\n" +
+                            "    - actual: ${actualFile.virtualFilePath}"
+                }
+            }
+        }
+
+        if (problems.isNotEmpty()) {
+            throw AssertionError(buildString {
+                appendln("There were association errors:").appendln()
+                problems.joinTo(this, "\n\n")
+            })
+        }
+    }
+}
@@ -0,0 +1,81 @@
+/*
+ * Copyright 2010-2018 JetBrains s.r.o. 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.idea.debugger;
+
+import com.intellij.testFramework.TestDataPath;
+import org.jetbrains.kotlin.test.JUnit3RunnerWithInners;
+import org.jetbrains.kotlin.test.KotlinTestUtils;
+import org.jetbrains.kotlin.test.TargetBackend;
+import org.jetbrains.kotlin.test.TestMetadata;
+import org.junit.runner.RunWith;
+
+import java.io.File;
+import java.util.regex.Pattern;
+
+/** This class is generated by {@link org.jetbrains.kotlin.generators.tests.TestsPackage}. DO NOT MODIFY MANUALLY */
+@SuppressWarnings("all")
+@TestMetadata("idea/testData/debugger/fileRanking")
+@TestDataPath("$PROJECT_ROOT")
+@RunWith(JUnit3RunnerWithInners.class)
+public class FileRankingTestGenerated extends AbstractFileRankingTest {
+    private void runTest(String testDataFilePath) throws Exception {
+        KotlinTestUtils.runTest(this::doTest, TargetBackend.ANY, testDataFilePath);
+    }
+
+    public void testAllFilesPresentInFileRanking() throws Exception {
+        KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("idea/testData/debugger/fileRanking"), Pattern.compile("^(.+)\\.kt$"), TargetBackend.ANY, true);
+    }
+
+    @TestMetadata("differentFlags.kt")
+    public void testDifferentFlags() throws Exception {
+        runTest("idea/testData/debugger/fileRanking/differentFlags.kt");
+    }
+
+    @TestMetadata("init.kt")
+    public void testInit() throws Exception {
+        runTest("idea/testData/debugger/fileRanking/init.kt");
+    }
+
+    @TestMetadata("lambdas.kt")
+    public void testLambdas() throws Exception {
+        runTest("idea/testData/debugger/fileRanking/lambdas.kt");
+    }
+
+    @TestMetadata("multilinePrimaryConstructor.kt")
+    public void testMultilinePrimaryConstructor() throws Exception {
+        runTest("idea/testData/debugger/fileRanking/multilinePrimaryConstructor.kt");
+    }
+
+    @TestMetadata("multilinePrimaryConstructorWithBody.kt")
+    public void testMultilinePrimaryConstructorWithBody() throws Exception {
+        runTest("idea/testData/debugger/fileRanking/multilinePrimaryConstructorWithBody.kt");
+    }
+
+    @TestMetadata("propertyDelegates.kt")
+    public void testPropertyDelegates() throws Exception {
+        runTest("idea/testData/debugger/fileRanking/propertyDelegates.kt");
+    }
+
+    @TestMetadata("sameClassName.kt")
+    public void testSameClassName() throws Exception {
+        runTest("idea/testData/debugger/fileRanking/sameClassName.kt");
+    }
+
+    @TestMetadata("sameClassNameDifferentMethodNames.kt")
+    public void testSameClassNameDifferentMethodNames() throws Exception {
+        runTest("idea/testData/debugger/fileRanking/sameClassNameDifferentMethodNames.kt");
+    }
+
+    @TestMetadata("simple.kt")
+    public void testSimple() throws Exception {
+        runTest("idea/testData/debugger/fileRanking/simple.kt");
+    }
+
+    @TestMetadata("topLevel.kt")
+    public void testTopLevel() throws Exception {
+        runTest("idea/testData/debugger/fileRanking/topLevel.kt");
+    }
+}
@@ -0,0 +1,169 @@
+/*
+ * Copyright 2010-2018 JetBrains s.r.o. 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.idea.debugger
+
+import com.intellij.util.PathUtil
+import com.intellij.util.SystemProperties
+import com.sun.jdi.ThreadReference
+import com.sun.jdi.VirtualMachine
+import com.sun.tools.jdi.SocketAttachingConnector
+import org.jetbrains.kotlin.backend.common.output.OutputFile
+import org.jetbrains.kotlin.codegen.*
+import org.jetbrains.kotlin.codegen.state.GenerationState
+import org.jetbrains.kotlin.test.ConfigurationKind
+import org.jetbrains.org.objectweb.asm.tree.ClassNode
+import org.jetbrains.org.objectweb.asm.Type as AsmType
+import java.io.File
+import java.io.IOException
+import java.net.Socket
+import java.nio.file.Files
+import kotlin.properties.Delegates
+
+abstract class LowLevelDebuggerTestBase : CodegenTestCase() {
+    private companion object {
+        private const val DEBUG_ADDRESS = "127.0.0.1"
+        private const val DEBUG_PORT = 5115
+    }
+
+    override fun doMultiFileTest(wholeFile: File, files: List<TestFile>, javaFilesDir: File?) {
+        val javaSources = javaFilesDir?.let { arrayOf(it) } ?: emptyArray()
+        createEnvironmentWithMockJdkAndIdeaAnnotations(ConfigurationKind.ALL, *javaSources)
+
+        val options = wholeFile.readLines()
+            .filter { it.matches("^// ?[\\w_]+$".toRegex()) }
+            .mapTo(mutableSetOf()) { it.drop(2).trim() }
+
+        loadMultiFiles(files)
+        val classBuilderFactory = OriginCollectingClassBuilderFactory(ClassBuilderMode.FULL)
+        val generationState = GenerationUtils.compileFiles(myFiles.psiFiles, myEnvironment, classBuilderFactory)
+        classFileFactory = generationState.factory
+
+        val tempDirForTest = Files.createTempDirectory("debuggerTest").toFile()
+
+        try {
+            val classesDir = File(tempDirForTest, "classes").apply {
+                writeMainClass(this)
+                for (classFile in classFileFactory.getClassFiles()) {
+                    File(this, classFile.relativePath).mkdirAndWriteBytes(classFile.asByteArray())
+                }
+            }
+
+            val process = startDebuggeeProcess(classesDir)
+            waitUntil { isPortOpen() }
+
+            val virtualMachine = attachDebugger()
+
+            try {
+                val mainThread = virtualMachine.allThreads().single { it.name() == "main" }
+                waitUntil { areCompiledClassesLoaded(mainThread, classBuilderFactory) }
+                doTest(options, mainThread, classBuilderFactory, generationState)
+            } finally {
+                virtualMachine.exit(0)
+                process.destroy()
+            }
+        } finally {
+            tempDirForTest.deleteRecursively()
+        }
+    }
+
+    protected abstract fun doTest(
+        options: Set<String>,
+        mainThread: ThreadReference,
+        factory: OriginCollectingClassBuilderFactory,
+        state: GenerationState
+    )
+
+    private fun isPortOpen(): Boolean {
+        return try {
+            Socket(DEBUG_ADDRESS, DEBUG_PORT).close()
+            true
+        } catch (e: IOException) {
+            false
+        }
+    }
+
+    private fun areCompiledClassesLoaded(mainThread: ThreadReference, factory: OriginCollectingClassBuilderFactory): Boolean {
+        for ((node, _) in factory.origins) {
+            val classNode = node as? ClassNode ?: continue
+            mainThread.virtualMachine().classesByName(classNode.name.replace('/', '.')).firstOrNull() ?: return false
+        }
+        return true
+    }
+
+    private fun startDebuggeeProcess(classesDir: File): Process {
+        val classesToLoad = this.classFileFactory.getClassFiles().joinToString(",") { it.qualifiedName }
+
+        val classpath = listOf(
+            classesDir.absolutePath,
+            PathUtil.getJarPathForClass(Delegates::class.java) // Add Kotlin runtime JAR
+        )
+
+        val command = arrayOf(
+            findJavaExecutable().absolutePath,
+            "-agentlib:jdwp=transport=dt_socket,server=y,suspend=n,address=$DEBUG_PORT",
+            "-ea",
+            "-classpath", classpath.joinToString(File.pathSeparator),
+            "-D${DebuggerMain.CLASSES_TO_LOAD}=$classesToLoad",
+            DebuggerMain::class.java.name
+        )
+
+        return ProcessBuilder(*command).inheritIO().start()
+    }
+
+    private fun attachDebugger(): VirtualMachine {
+        val connector = SocketAttachingConnector()
+        return connector.attach(connector.defaultArguments().toMutableMap().apply {
+            getValue("port").setValue("$DEBUG_PORT")
+            getValue("hostname").setValue(DEBUG_ADDRESS)
+        })
+    }
+
+    private fun findJavaExecutable(): File {
+        val javaBin = File(SystemProperties.getJavaHome(), "bin")
+        return File(javaBin, "java.exe").takeIf { it.exists() }
+            ?: File(javaBin, "java").also { assert(it.exists()) }
+    }
+
+    private fun writeMainClass(classesDir: File) {
+        val mainClassResourceName = DebuggerMain::class.java.name.replace('.', '/') + ".class"
+        val mainClassBytes = javaClass.classLoader.getResource(mainClassResourceName).readBytes()
+        File(classesDir, mainClassResourceName).mkdirAndWriteBytes(mainClassBytes)
+    }
+
+    private val OutputFile.internalName
+        get() = relativePath.substringBeforeLast(".class")
+
+    private val OutputFile.qualifiedName
+        get() = internalName.replace('/', '.')
+}
+
+private fun File.mkdirAndWriteBytes(array: ByteArray) {
+    parentFile.mkdirs()
+    writeBytes(array)
+}
+
+private fun waitUntil(condition: () -> Boolean) {
+    while (!condition()) {
+        Thread.sleep(30)
+    }
+}
+
+private object DebuggerMain {
+    const val CLASSES_TO_LOAD = "classes.to.load"
+
+    @JvmField
+    val lock = Any()
+
+    @JvmStatic
+    fun main(args: Array<String>) {
+        System.getProperty(CLASSES_TO_LOAD).split(',').forEach { Class.forName(it) }
+        synchronized(lock) {
+            // Wait until debugger is attached
+            @Suppress("PLATFORM_CLASS_MAPPED_TO_KOTLIN")
+            (lock as java.lang.Object).wait()
+        }
+    }
+}
@@ -28,8 +28,10 @@ import org.jetbrains.kotlin.cli.common.messages.CompilerMessageSeverity.OUTPUT
 import org.jetbrains.kotlin.cli.common.messages.MessageCollector
 import org.jetbrains.kotlin.cli.common.messages.OutputMessageUtil
 import org.jetbrains.kotlin.cli.common.output.outputUtils.writeAll
+import org.jetbrains.kotlin.codegen.ClassBuilderMode
 import org.jetbrains.kotlin.codegen.CompilationErrorHandler
 import org.jetbrains.kotlin.codegen.KotlinCodegenFacade
+import org.jetbrains.kotlin.codegen.OriginCollectingClassBuilderFactory
 import org.jetbrains.kotlin.codegen.state.GenerationState
 import org.jetbrains.kotlin.config.CommonConfigurationKeys
 import org.jetbrains.kotlin.config.CompilerConfiguration
@@ -219,7 +221,7 @@ abstract class AbstractKapt3Extension(
            bindingContext: BindingContext,
            files: List<KtFile>
    ): KaptContextForStubGeneration {
-        val builderFactory = Kapt3BuilderFactory()
+        val builderFactory = OriginCollectingClassBuilderFactory(ClassBuilderMode.KAPT3)

        val targetId = TargetId(
                name = compilerConfiguration[CommonConfigurationKeys.MODULE_NAME] ?: module.name.asString(),
@@ -19,8 +19,10 @@ package org.jetbrains.kotlin.kapt3.test
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.util.text.StringUtil
 import org.jetbrains.kotlin.cli.common.messages.MessageCollector
+import org.jetbrains.kotlin.codegen.ClassBuilderMode
 import org.jetbrains.kotlin.codegen.CodegenTestCase
 import org.jetbrains.kotlin.codegen.GenerationUtils
+import org.jetbrains.kotlin.codegen.OriginCollectingClassBuilderFactory
 import org.jetbrains.kotlin.config.CompilerConfiguration
 import org.jetbrains.kotlin.kapt3.*
 import org.jetbrains.kotlin.kapt3.AptMode.STUBS_AND_APT
@@ -140,7 +142,8 @@ abstract class AbstractKotlinKapt3IntegrationTest : CodegenTestCase() {
        try {
            loadMultiFiles(files)

-            GenerationUtils.compileFiles(myFiles.psiFiles, myEnvironment, Kapt3BuilderFactory()).factory
+            val classBuilderFactory = OriginCollectingClassBuilderFactory(ClassBuilderMode.KAPT3)
+            GenerationUtils.compileFiles(myFiles.psiFiles, myEnvironment, classBuilderFactory).factory

            val actualRaw = kapt3Extension.savedStubs ?: error("Stubs were not saved")
            val actual = StringUtil.convertLineSeparators(actualRaw.trim({ it <= ' ' }))
@@ -29,9 +29,7 @@ import org.jetbrains.kotlin.cli.common.messages.MessageRenderer
 import org.jetbrains.kotlin.cli.common.messages.PrintingMessageCollector
 import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreEnvironment
 import org.jetbrains.kotlin.cli.jvm.config.JvmClasspathRoot
-import org.jetbrains.kotlin.codegen.CodegenTestCase
-import org.jetbrains.kotlin.codegen.CodegenTestFiles
-import org.jetbrains.kotlin.codegen.GenerationUtils
+import org.jetbrains.kotlin.codegen.*
 import org.jetbrains.kotlin.kapt3.*
 import org.jetbrains.kotlin.kapt3.base.KaptContext
 import org.jetbrains.kotlin.kapt3.base.KaptPaths
@@ -112,7 +110,7 @@ abstract class AbstractKotlinKapt3Test : CodegenTestCase() {
        loadMultiFiles(files)

        val txtFile = File(wholeFile.parentFile, wholeFile.nameWithoutExtension + ".txt")
-        val classBuilderFactory = Kapt3BuilderFactory()
+        val classBuilderFactory = OriginCollectingClassBuilderFactory(ClassBuilderMode.KAPT3)
        val generationState = GenerationUtils.compileFiles(myFiles.psiFiles, myEnvironment, classBuilderFactory)

        val logger = MessageCollectorBackedKaptLogger(isVerbose = true, messageCollector = messageCollector)