Special handling of inline function to track values from individual c…

…alls

idea/src/org/jetbrains/kotlin/idea/slicer/ArgumentSliceProducer.kt

@@ -32,10 +32,10 @@ data class ArgumentSliceProducer private constructor(
         parameterDescriptor.containingDeclaration.isExtension
     )
 
-    override fun produce(usage: UsageInfo, behaviour: KotlinSliceUsage.SpecialBehaviour?, parent: SliceUsage): Collection<SliceUsage>? {
+    override fun produce(usage: UsageInfo, mode: KotlinSliceAnalysisMode, parent: SliceUsage): Collection<SliceUsage>? {
         val element = usage.element ?: return emptyList()
         val argumentExpression = extractArgumentExpression(element) ?: return emptyList()
-        return listOf(KotlinSliceUsage(argumentExpression, parent, behaviour, forcedExpressionMode = true))
+        return listOf(KotlinSliceUsage(argumentExpression, parent, mode, forcedExpressionMode = true))
     }
 
     private fun extractArgumentExpression(refElement: PsiElement): PsiElement? {

idea/src/org/jetbrains/kotlin/idea/slicer/CallSliceProducer.kt

@@ -15,24 +15,31 @@ import org.jetbrains.kotlin.psi.psiUtil.getParentOfTypeAndBranch
 import org.jetbrains.kotlin.psi.psiUtil.isAncestor
 
 object CallSliceProducer : SliceProducer {
-    override fun produce(usage: UsageInfo, behaviour: KotlinSliceUsage.SpecialBehaviour?, parent: SliceUsage): Collection<SliceUsage>? {
-        if ((parent as? KotlinSliceUsage)?.behaviour is LambdaCallsBehaviour) {
+    override fun produce(usage: UsageInfo, mode: KotlinSliceAnalysisMode, parent: SliceUsage): Collection<SliceUsage>? {
+        if ((parent as? KotlinSliceUsage)?.mode?.currentBehaviour is LambdaCallsBehaviour) {
             // UsageInfo produced by LambdaCallsBehaviour has full call-element and does not require any processing
             return null
         }
 
         when (val refElement = usage.element) {
             null -> {
                 val element = (usage.reference as? LightMemberReference)?.element ?: return emptyList()
-                return listOf(KotlinSliceUsage(element, parent, behaviour, false))
+                return listOf(KotlinSliceUsage(element, parent, mode, false))
             }
 
             is KtExpression -> {
                 return mutableListOf<SliceUsage>().apply {
                     refElement.getCallElementForExactCallee()
-                        ?.let { this += KotlinSliceUsage(it, parent, behaviour, false) }
+                        ?.let { this += KotlinSliceUsage(it, parent, mode, false) }
                     refElement.getCallableReferenceForExactCallee()
-                        ?.let { this += KotlinSliceUsage(it, parent, LambdaCallsBehaviour(SliceProducer.Trivial, behaviour), false) }
+                        ?.let {
+                            this += KotlinSliceUsage(
+                                it,
+                                parent,
+                                mode.withBehaviour(LambdaCallsBehaviour(SliceProducer.Trivial)),
+                                forcedExpressionMode = true
+                            )
+                        }
                 }
             }
 

idea/src/org/jetbrains/kotlin/idea/slicer/InflowSlicer.kt

@@ -31,6 +31,7 @@ import org.jetbrains.kotlin.idea.references.readWriteAccessWithFullExpression
 import org.jetbrains.kotlin.idea.search.declarationsSearch.HierarchySearchRequest
 import org.jetbrains.kotlin.idea.search.declarationsSearch.searchOverriders
 import org.jetbrains.kotlin.idea.util.actualsForExpected
+import org.jetbrains.kotlin.idea.util.hasInlineModifier
 import org.jetbrains.kotlin.idea.util.isExpectDeclaration
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi.*
@@ -141,8 +142,8 @@ class InflowSlicer(
             else -> null
         }
         if (lambda != null) {
-            if (behaviour is LambdaResultInflowBehaviour) {
-                lambda.passToProcessor(behaviour.originalBehaviour)
+            if (mode.currentBehaviour is LambdaResultInflowBehaviour) {
+                lambda.passToProcessor(mode.dropBehaviour())
             }
             return
         }
@@ -176,7 +177,7 @@ class InflowSlicer(
                         val callable = accessedDescriptor.containingDeclaration as? CallableDescriptor ?: return
                         when (val declaration = callable.originalSource.getPsi()) {
                             is KtFunctionLiteral -> {
-                                declaration.passToProcessorAsValue(LambdaCallsBehaviour(ReceiverSliceProducer, behaviour))
+                                declaration.passToProcessorAsValue(mode.withBehaviour(LambdaCallsBehaviour(ReceiverSliceProducer)))
                             }
 
                             is KtCallableDeclaration -> {
@@ -194,12 +195,12 @@ class InflowSlicer(
                                 processCalls(functionLiteral, false, ArgumentSliceProducer(parameterDescriptor))
                             }
                         } else {
-                            accessedDeclaration.passDeclarationToProcessorWithOverriders()
+                            accessedDeclaration.passDeclarationToProcessorWithOverriders(createdAt.element)
                         }
                     }
 
                     else -> {
-                        accessedDeclaration?.passDeclarationToProcessorWithOverriders()
+                        accessedDeclaration?.passDeclarationToProcessorWithOverriders(createdAt.element)
                     }
                 }
             }
@@ -215,8 +216,8 @@ class InflowSlicer(
                     val referencedDescriptor = bindingContext[BindingContext.REFERENCE_TARGET, callableRefExpr.callableReference] ?: return
                     val referencedDeclaration = (referencedDescriptor as? DeclarationDescriptorWithSource)?.originalSource?.getPsi()
                         ?: return
-                    if (behaviour is LambdaResultInflowBehaviour) {
-                        referencedDeclaration.passToProcessor(behaviour.originalBehaviour)
+                    if (mode.currentBehaviour is LambdaResultInflowBehaviour) {
+                        referencedDeclaration.passToProcessor(mode.dropBehaviour())
                     }
                 }
 
@@ -228,9 +229,9 @@ class InflowSlicer(
                 val resultingDescriptor = resolvedCall.resultingDescriptor
                 if (resultingDescriptor is FunctionInvokeDescriptor) {
                     (resolvedCall.dispatchReceiver as? ExpressionReceiver)?.expression
-                        ?.passToProcessorAsValue(LambdaResultInflowBehaviour(behaviour))
+                        ?.passToProcessorAsValue(mode.withBehaviour(LambdaResultInflowBehaviour))
                 } else {
-                    resultingDescriptor.originalSource.getPsi()?.passDeclarationToProcessorWithOverriders()
+                    resultingDescriptor.originalSource.getPsi()?.passDeclarationToProcessorWithOverriders(createdAt.element)
                 }
             }
         }
@@ -305,19 +306,37 @@ class InflowSlicer(
         }
     }
 
-    private fun PsiElement.passDeclarationToProcessorWithOverriders() {
-        passToProcessor()
+    private fun PsiElement.passDeclarationToProcessorWithOverriders(callElement: KtElement) {
+        val newMode = if (this is KtNamedFunction && hasInlineModifier())
+            mode.withInlineFunctionCall(callElement, this)
+        else
+            mode
+
+        passToProcessor(newMode)
 
         HierarchySearchRequest(this, analysisScope)
             .searchOverriders()
-            .forEach { it.namedUnwrappedElement?.passToProcessor() }
+            .forEach { it.namedUnwrappedElement?.passToProcessor(newMode) }
 
         if (this is KtCallableDeclaration && isExpectDeclaration()) {
             resolveToDescriptorIfAny(BodyResolveMode.FULL)
                 ?.actualsForExpected()
                 ?.forEach {
-                    (it as? DeclarationDescriptorWithSource)?.originalSource?.getPsi()?.passToProcessor()
+                    (it as? DeclarationDescriptorWithSource)?.originalSource?.getPsi()?.passToProcessor(newMode)
                 }
         }
     }
+
+    override fun processCalls(callable: KtCallableDeclaration, includeOverriders: Boolean, sliceProducer: SliceProducer) {
+        if (callable is KtNamedFunction) {
+            val (newMode, callElement) = mode.popInlineFunctionCall(callable)
+            if (newMode != null && callElement != null) {
+                val sliceUsage = KotlinSliceUsage(callElement, parentUsage, newMode, false)
+                sliceProducer.produceAndProcess(sliceUsage, newMode, parentUsage, processor)
+                return
+            }
+        }
+
+        super.processCalls(callable, includeOverriders, sliceProducer)
+    }
 }

idea/src/org/jetbrains/kotlin/idea/slicer/KotlinSliceAnalysisMode.kt

@@ -0,0 +1,67 @@
+/*
+ * Copyright 2010-2020 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.idea.slicer
+
+import org.jetbrains.kotlin.idea.findUsages.handlers.SliceUsageProcessor
+import org.jetbrains.kotlin.psi.KtElement
+import org.jetbrains.kotlin.psi.KtNamedFunction
+import org.jetbrains.kotlin.psi.psiUtil.createSmartPointer
+
+data class KotlinSliceAnalysisMode(val behaviourStack: List<Behaviour>, val inlineCallStack: List<InlineFunctionCall>) {
+    fun withBehaviour(behaviour: Behaviour) = copy(behaviourStack = behaviourStack + behaviour)
+
+    fun withInlineFunctionCall(
+        callElement: KtElement,
+        function: KtNamedFunction
+    ) = copy(inlineCallStack = inlineCallStack + InlineFunctionCall(callElement, function))
+
+    fun dropBehaviour(): KotlinSliceAnalysisMode {
+        check(behaviourStack.isNotEmpty())
+        return copy(behaviourStack = behaviourStack.dropLast(1))
+    }
+
+    fun popInlineFunctionCall(function: KtNamedFunction): Pair<KotlinSliceAnalysisMode?, KtElement?> {
+        val last = inlineCallStack.lastOrNull()
+        if (last?.function != function) return null to null
+        val newMode = copy(inlineCallStack = inlineCallStack.dropLast(1))
+        return newMode to last.callElement
+    }
+
+    val currentBehaviour: Behaviour?
+        get() = behaviourStack.lastOrNull()
+
+    interface Behaviour {
+        fun processUsages(element: KtElement, parent: KotlinSliceUsage, uniqueProcessor: SliceUsageProcessor)
+
+        val slicePresentationPrefix: String
+        val testPresentationPrefix: String
+
+        override fun equals(other: Any?): Boolean
+        override fun hashCode(): Int
+    }
+
+    class InlineFunctionCall(callElement: KtElement, function: KtNamedFunction) {
+        private val callElementPointer = callElement.createSmartPointer()
+        private val functionPointer = function.createSmartPointer()
+
+        val callElement: KtElement?
+            get() = callElementPointer.element
+
+        val function: KtNamedFunction?
+            get() = functionPointer.element
+
+        override fun equals(other: Any?): Boolean {
+            if (this === other) return true
+            return other is InlineFunctionCall && other.callElement == callElement && other.function == function
+        }
+
+        override fun hashCode() = 0
+    }
+
+    companion object {
+        val Default = KotlinSliceAnalysisMode(emptyList(), emptyList())
+    }
+}

idea/src/org/jetbrains/kotlin/idea/slicer/KotlinSliceDereferenceUsage.kt

@@ -14,8 +14,8 @@ import org.jetbrains.kotlin.idea.KotlinBundle
 class KotlinSliceDereferenceUsage(
     element: PsiElement,
     parent: KotlinSliceUsage,
-    behaviour: SpecialBehaviour?
-) : KotlinSliceUsage(element, parent, behaviour, false) {
+    mode: KotlinSliceAnalysisMode
+) : KotlinSliceUsage(element, parent, mode, false) {
     override fun processChildren(processor: Processor<in SliceUsage>) {
         // no children
     }

idea/src/org/jetbrains/kotlin/idea/slicer/KotlinSliceDereferenceUsage.kt.191

@@ -25,8 +25,8 @@ import org.jetbrains.kotlin.idea.KotlinBundle
 class KotlinSliceDereferenceUsage(
     element: PsiElement,
     parent: KotlinSliceUsage,
-    behaviour: SpecialBehaviour?
-) : KotlinSliceUsage(element, parent, behaviour, false) {
+    mode: KotlinSliceAnalysisMode
+) : KotlinSliceUsage(element, parent, mode, false) {
     override fun processChildren(processor: Processor<SliceUsage>) {
         // no children
     }

idea/src/org/jetbrains/kotlin/idea/slicer/KotlinSliceProvider.kt

@@ -49,6 +49,7 @@ class KotlinSliceProvider : SliceLanguageSupportProvider, SliceUsageTransformer
     }
 
     val leafAnalyzer by lazy { SliceLeafAnalyzer(LEAF_ELEMENT_EQUALITY, this) }
+
     val nullnessAnalyzer: SliceNullnessAnalyzerBase by lazy {
         object : SliceNullnessAnalyzerBase(LEAF_ELEMENT_EQUALITY, this) {
             override fun checkNullability(element: PsiElement?): Nullability {
@@ -72,7 +73,7 @@ class KotlinSliceProvider : SliceLanguageSupportProvider, SliceUsageTransformer
 
     override fun transform(usage: SliceUsage): Collection<SliceUsage>? {
         if (usage is KotlinSliceUsage) return null
-        return listOf(KotlinSliceUsage(usage.element, usage.parent, null, false))
+        return listOf(KotlinSliceUsage(usage.element, usage.parent, KotlinSliceAnalysisMode.Default, false))
     }
 
     override fun getExpressionAtCaret(atCaret: PsiElement, dataFlowToThis: Boolean): KtElement? {

idea/src/org/jetbrains/kotlin/idea/slicer/KotlinSliceUsage.kt

@@ -24,57 +24,61 @@ import org.jetbrains.kotlin.idea.findUsages.handlers.SliceUsageProcessor
 import org.jetbrains.kotlin.psi.KtElement
 
 open class KotlinSliceUsage : SliceUsage {
-    interface SpecialBehaviour {
-        val originalBehaviour: SpecialBehaviour?
-        fun processUsages(element: KtElement, parent: KotlinSliceUsage, uniqueProcessor: SliceUsageProcessor)
 
-        val slicePresentationPrefix: String
-        val testPresentationPrefix: String
-
-        override fun equals(other: Any?): Boolean
-        override fun hashCode(): Int
-    }
-
-    val behaviour: SpecialBehaviour?
+    val mode: KotlinSliceAnalysisMode
     val forcedExpressionMode: Boolean
 
+    private var usageInfo: UsageInfo? = null
+
     constructor(
         element: PsiElement,
         parent: SliceUsage,
-        behaviour: SpecialBehaviour?,
+        mode: KotlinSliceAnalysisMode,
         forcedExpressionMode: Boolean,
     ) : super(element, parent) {
-        this.behaviour = behaviour
+        this.mode = mode
         this.forcedExpressionMode = forcedExpressionMode
+        initializeUsageInfo()
     }
 
     constructor(element: PsiElement, params: SliceAnalysisParams) : super(element, params) {
-        this.behaviour = null
+        this.mode = KotlinSliceAnalysisMode.Default
         this.forcedExpressionMode = false
+        initializeUsageInfo()
+    }
+
+    private fun initializeUsageInfo() {
+        val originalInfo = getUsageInfo()
+        if (mode != KotlinSliceAnalysisMode.Default) {
+            val element = originalInfo.element
+            if (element != null) {
+                usageInfo = UsageInfoWrapper(element, mode)
+            } else {
+                usageInfo = null
+            }
+        } else {
+            usageInfo = originalInfo
+        }
+    }
+
+    // we have to replace UsageInfo with another one whose equality takes into account mode
+    override fun getUsageInfo(): UsageInfo {
+        return usageInfo ?: super.getUsageInfo()
     }
 
     override fun copy(): KotlinSliceUsage {
-        val element = usageInfo.element!!
+        val element = getUsageInfo().element!!
         return if (parent == null)
             KotlinSliceUsage(element, params)
         else
-            KotlinSliceUsage(element, parent, behaviour, forcedExpressionMode)
-    }
-
-    override fun getUsageInfo(): UsageInfo {
-        val originalInfo = super.getUsageInfo()
-        if (behaviour != null) {
-            val element = originalInfo.element ?: return originalInfo
-            // Do not let IDEA consider usages of the same anonymous function as duplicates when their levels differ
-            return UsageInfoWrapper(element, behaviour)
-        }
-        return originalInfo
+            KotlinSliceUsage(element, parent, mode, forcedExpressionMode)
     }
 
-    override fun canBeLeaf() = element != null && behaviour == null
+    override fun canBeLeaf() = element != null && mode == KotlinSliceAnalysisMode.Default
 
     public override fun processUsagesFlownDownTo(element: PsiElement, uniqueProcessor: SliceUsageProcessor) {
         val ktElement = element as? KtElement ?: return
+        val behaviour = mode.currentBehaviour
         if (behaviour != null) {
             behaviour.processUsages(ktElement, this, uniqueProcessor)
         } else {
@@ -84,6 +88,7 @@ open class KotlinSliceUsage : SliceUsage {
 
     public override fun processUsagesFlownFromThe(element: PsiElement, uniqueProcessor: SliceUsageProcessor) {
         val ktElement = element as? KtElement ?: return
+        val behaviour = mode.currentBehaviour
         if (behaviour != null) {
             behaviour.processUsages(ktElement, this, uniqueProcessor)
         } else {
@@ -92,9 +97,9 @@ open class KotlinSliceUsage : SliceUsage {
     }
 
     @Suppress("EqualsOrHashCode")
-    private class UsageInfoWrapper(element: PsiElement, private val behaviour: SpecialBehaviour) : UsageInfo(element) {
+    private class UsageInfoWrapper(element: PsiElement, private val mode: KotlinSliceAnalysisMode) : UsageInfo(element) {
         override fun equals(other: Any?): Boolean {
-            return other is UsageInfoWrapper && super.equals(other) && behaviour == other.behaviour
+            return other is UsageInfoWrapper && super.equals(other) && mode == other.mode
         }
     }
 }

idea/src/org/jetbrains/kotlin/idea/slicer/KotlinSliceUsageCellRenderer.kt

@@ -43,10 +43,8 @@ object KotlinSliceUsageCellRenderer : SliceUsageCellRendererBase() {
             }
         }
 
-        var behaviour = sliceUsage.behaviour
-        while (behaviour != null) {
+        for (behaviour in sliceUsage.mode.behaviourStack.reversed()) {
             append(behaviour.slicePresentationPrefix, SimpleTextAttributes.REGULAR_BOLD_ATTRIBUTES)
-            behaviour = behaviour.originalBehaviour
         }
 
         val declaration = sliceUsage.element?.parents?.firstOrNull {