Refactoring: abstraction of idea with "lambdaLevel" in slices

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

@@ -55,8 +55,8 @@ class InflowSlicer(
     parentUsage: KotlinSliceUsage
 ) : Slicer(element, processor, parentUsage) {
 
-    override fun processChildren() {
+    override fun processChildren(forcedExpressionMode: Boolean) {
-        if (parentUsage.forcedExpressionMode) {
+        if (forcedExpressionMode) {
             (element as? KtExpression)?.let { processExpression(it) }
             return
         }
@@ -199,8 +199,8 @@ class InflowSlicer(
             else -> null
         }
         if (lambda != null) {
-            if (parentUsage.lambdaLevel > 0) {
+            if (behaviour is LambdaResultInflowBehaviour) {
-                lambda.passToProcessor(parentUsage.lambdaLevel - 1)
+                lambda.passToProcessor(behaviour.originalBehaviour)
             }
             return
         }
@@ -253,7 +253,9 @@ class InflowSlicer(
                     val referencedDescriptor = bindingContext[BindingContext.REFERENCE_TARGET, callableRefExpr.callableReference] ?: return
                     val referencedDeclaration = (referencedDescriptor as? DeclarationDescriptorWithSource)?.originalSource?.getPsi()
                         ?: return
-                    referencedDeclaration.passToProcessor(parentUsage.lambdaLevel - 1)
+                    if (behaviour is LambdaResultInflowBehaviour) {
+                        referencedDeclaration.passToProcessor(behaviour.originalBehaviour)
+                    }
                 }
 
                 else -> return
@@ -263,7 +265,8 @@ class InflowSlicer(
                 val resolvedCall = createdAt.resolvedCall
                 val resultingDescriptor = resolvedCall.resultingDescriptor
                 if (resultingDescriptor is FunctionInvokeDescriptor) {
-                    (resolvedCall.dispatchReceiver as? ExpressionReceiver)?.expression?.passToProcessorAsValue(parentUsage.lambdaLevel + 1)
+                    (resolvedCall.dispatchReceiver as? ExpressionReceiver)?.expression
+                        ?.passToProcessorAsValue(LambdaResultInflowBehaviour(behaviour))
                 } else {
                     resultingDescriptor.originalSource.getPsi()?.passDeclarationToProcessorWithOverriders()
                 }
@@ -340,8 +343,6 @@ class InflowSlicer(
         }
     }
 
-    private fun PsiElement.passToProcessorAsValue(lambdaLevel: Int = parentUsage.lambdaLevel) = passToProcessor(lambdaLevel, true)
-
     private fun PsiElement.passDeclarationToProcessorWithOverriders() {
         passToProcessor()
 

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,
-    lambdaLevel: Int
+    behaviour: SpecialBehaviour?
-) : KotlinSliceUsage(element, parent, lambdaLevel, false) {
+) : KotlinSliceUsage(element, parent, behaviour, false) {
     override fun processChildren(processor: Processor<in SliceUsage>) {
         // no children
     }

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

@@ -23,10 +23,10 @@ import com.intellij.util.Processor
 import org.jetbrains.kotlin.idea.KotlinBundle
 
 class KotlinSliceDereferenceUsage(
-        element: PsiElement,
+    element: PsiElement,
-        parent: KotlinSliceUsage,
+    parent: KotlinSliceUsage,
-        lambdaLevel: Int
+    behaviour: SpecialBehaviour?
-) : KotlinSliceUsage(element, parent, lambdaLevel, false) {
+) : KotlinSliceUsage(element, parent, behaviour, false) {
     override fun processChildren(processor: Processor<SliceUsage>) {
         // no children
     }

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

@@ -72,7 +72,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, 0, false))
+        return listOf(KotlinSliceUsage(usage.element, usage.parent, null, false))
     }
 
     override fun getExpressionAtCaret(atCaret: PsiElement, dataFlowToThis: Boolean): KtElement? {

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

@@ -20,49 +20,112 @@ import com.intellij.psi.PsiElement
 import com.intellij.slicer.SliceAnalysisParams
 import com.intellij.slicer.SliceUsage
 import com.intellij.usageView.UsageInfo
+import org.jetbrains.kotlin.idea.KotlinBundle
 import org.jetbrains.kotlin.idea.findUsages.handlers.SliceUsageProcessor
 import org.jetbrains.kotlin.psi.KtElement
-import org.jetbrains.kotlin.psi.KtExpression
 
 open class KotlinSliceUsage : SliceUsage {
-    class UsageInfoLambdaWrapper(element: PsiElement) : UsageInfo(element)
+    interface SpecialBehaviour {
+        val originalBehaviour: SpecialBehaviour?
+        fun processUsages(element: KtElement, parent: KotlinSliceUsage, uniqueProcessor: SliceUsageProcessor)
 
-    val lambdaLevel: Int
+        val slicePresentationPrefix: String
+        val testPresentationPrefix: String
+
+        override fun equals(other: Any?): Boolean
+        override fun hashCode(): Int
+    }
+
+    val behaviour: SpecialBehaviour?
     val forcedExpressionMode: Boolean
 
-    constructor(element: PsiElement, parent: SliceUsage, lambdaLevel: Int, forcedExpressionMode: Boolean) : super(element, parent) {
+    constructor(
-        this.lambdaLevel = lambdaLevel
+        element: PsiElement,
+        parent: SliceUsage,
+        behaviour: SpecialBehaviour?,
+        forcedExpressionMode: Boolean,
+    ) : super(element, parent) {
+        this.behaviour = behaviour
         this.forcedExpressionMode = forcedExpressionMode
     }
 
     constructor(element: PsiElement, params: SliceAnalysisParams) : super(element, params) {
-        this.lambdaLevel = 0
+        this.behaviour = null
         this.forcedExpressionMode = false
     }
 
     override fun copy(): KotlinSliceUsage {
         val element = usageInfo.element!!
-        if (parent == null) return KotlinSliceUsage(element, params)
+        return if (parent == null)
-        return KotlinSliceUsage(element, parent, lambdaLevel, forcedExpressionMode)
+            KotlinSliceUsage(element, params)
+        else
+            KotlinSliceUsage(element, parent, behaviour, forcedExpressionMode)
     }
 
     override fun getUsageInfo(): UsageInfo {
         val originalInfo = super.getUsageInfo()
-        if (lambdaLevel > 0 && forcedExpressionMode) {
+        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 UsageInfoLambdaWrapper(element)
+            return UsageInfoWrapper(element, behaviour)
         }
         return originalInfo
     }
 
-    override fun canBeLeaf() = element != null && lambdaLevel == 0
+    override fun canBeLeaf() = element != null && behaviour == null
 
     public override fun processUsagesFlownDownTo(element: PsiElement, uniqueProcessor: SliceUsageProcessor) {
-        InflowSlicer(element as? KtElement ?: return, uniqueProcessor, this).processChildren()
+        val ktElement = element as? KtElement ?: return
+        if (behaviour != null) {
+            behaviour.processUsages(ktElement, this, uniqueProcessor)
+        } else {
+            InflowSlicer(ktElement, uniqueProcessor, this).processChildren(forcedExpressionMode)
+        }
     }
 
     public override fun processUsagesFlownFromThe(element: PsiElement, uniqueProcessor: SliceUsageProcessor) {
-        OutflowSlicer(element as? KtElement ?: return, uniqueProcessor, this).processChildren()
+        val ktElement = element as? KtElement ?: return
+        if (behaviour != null) {
+            behaviour.processUsages(ktElement, this, uniqueProcessor)
+        } else {
+            OutflowSlicer(ktElement, uniqueProcessor, this).processChildren(forcedExpressionMode)
+        }
+    }
+
+    @Suppress("EqualsOrHashCode")
+    private class UsageInfoWrapper(element: PsiElement, private val behaviour: SpecialBehaviour) : UsageInfo(element) {
+        override fun equals(other: Any?): Boolean {
+            return other is UsageInfoWrapper && super.equals(other) && behaviour == other.behaviour
+        }
     }
 }
+
+data class LambdaResultOutflowBehaviour(
+    override val originalBehaviour: KotlinSliceUsage.SpecialBehaviour?
+) : KotlinSliceUsage.SpecialBehaviour {
+
+    override fun processUsages(element: KtElement, parent: KotlinSliceUsage, uniqueProcessor: SliceUsageProcessor) {
+        OutflowSlicer(element, uniqueProcessor, parent).processChildren(parent.forcedExpressionMode)
+    }
+
+    override val slicePresentationPrefix: String
+        get() = KotlinBundle.message("slicer.text.tracking.enclosing.lambda")
+
+    override val testPresentationPrefix: String
+        get() = "[LAMBDA] "
+}
+
+data class LambdaResultInflowBehaviour(
+    override val originalBehaviour: KotlinSliceUsage.SpecialBehaviour?
+) : KotlinSliceUsage.SpecialBehaviour {
+
+    override fun processUsages(element: KtElement, parent: KotlinSliceUsage, uniqueProcessor: SliceUsageProcessor) {
+        InflowSlicer(element, uniqueProcessor, parent).processChildren(parent.forcedExpressionMode)
+    }
+
+    override val slicePresentationPrefix: String
+        get() = KotlinBundle.message("slicer.text.tracking.enclosing.lambda")
+
+    override val testPresentationPrefix: String
+        get() = "[LAMBDA] "
+}

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

@@ -43,10 +43,11 @@ object KotlinSliceUsageCellRenderer : SliceUsageCellRendererBase() {
             }
         }
 
-        append(
+        var behaviour = sliceUsage.behaviour
-            KotlinBundle.message("slicer.text.tracking.enclosing.lambda").repeat(sliceUsage.lambdaLevel),
+        while (behaviour != null) {
-            SimpleTextAttributes.REGULAR_BOLD_ATTRIBUTES
+            append(behaviour.slicePresentationPrefix, SimpleTextAttributes.REGULAR_BOLD_ATTRIBUTES)
-        )
+            behaviour = behaviour.originalBehaviour
+        }
 
         val declaration = sliceUsage.element?.parents?.firstOrNull {
             it is KtClass ||

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

@@ -36,8 +36,8 @@ class OutflowSlicer(
     parentUsage: KotlinSliceUsage
 ) : Slicer(element, processor, parentUsage) {
 
-    override fun processChildren() {
+    override fun processChildren(forcedExpressionMode: Boolean) {
-        if (parentUsage.forcedExpressionMode) {
+        if (forcedExpressionMode) {
             (element as? KtExpression)?.let { processExpression(it) }
             return
         }
@@ -136,7 +136,9 @@ class OutflowSlicer(
                             parameters[parameterIndex + shift].passToProcessor()
                         }
 
-                        else -> {} // not supported
+                        else -> {
+                            // not supported
+                        }
                     }
                     true
                 }
@@ -153,7 +155,8 @@ class OutflowSlicer(
                     null -> (usageInfo.reference as? LightMemberReference)?.element?.passToProcessor()
                     is KtExpression -> {
                         refElement.getCallElementForExactCallee()?.passToProcessor()
-                        refElement.getCallableReferenceForExactCallee()?.passToProcessor(parentUsage.lambdaLevel + 1)
+                        refElement.getCallableReferenceForExactCallee()
+                            ?.passToProcessor(LambdaResultOutflowBehaviour(behaviour))
                     }
                     else -> refElement.passToProcessor()
                 }
@@ -166,7 +169,7 @@ class OutflowSlicer(
             is KtNamedFunction -> function
             else -> null
         } ?: return
-        funExpression.passToProcessor(parentUsage.lambdaLevel + 1, true)
+        funExpression.passToProcessor(LambdaResultOutflowBehaviour(behaviour), forcedExpressionMode = true)
     }
 
     private fun processExtensionReceiver(declaration: KtCallableDeclaration, declarationWithBody: KtDeclarationWithBody) {
@@ -229,8 +232,8 @@ class OutflowSlicer(
             }
 
             Call.CallType.INVOKE -> {
-                if (parentUsage.lambdaLevel > 0 && receiverValue == resolvedCall.dispatchReceiver) {
+                if (receiverValue == resolvedCall.dispatchReceiver && behaviour is LambdaResultOutflowBehaviour) {
-                    instruction.element.passToProcessor(parentUsage.lambdaLevel - 1)
+                    instruction.element.passToProcessor(behaviour.originalBehaviour)
                 }
             }
 
@@ -275,7 +278,7 @@ class OutflowSlicer(
         } ?: return
 
         if (receiver != null && resolvedCall.dispatchReceiver == receiver) {
-            processor.process(KotlinSliceDereferenceUsage(expression, parentUsage, parentUsage.lambdaLevel))
+            processor.process(KotlinSliceDereferenceUsage(expression, parentUsage, behaviour))
         }
     }
 

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

@@ -37,9 +37,10 @@ abstract class Slicer(
     protected val processor: SliceUsageProcessor,
     protected val parentUsage: KotlinSliceUsage
 ) {
-    abstract fun processChildren()
+    abstract fun processChildren(forcedExpressionMode: Boolean)
 
     protected val analysisScope: SearchScope = parentUsage.scope.toSearchScope()
+    protected val behaviour: KotlinSliceUsage.SpecialBehaviour? = parentUsage.behaviour
     protected val project = element.project
 
     protected class PseudocodeCache {
@@ -57,10 +58,14 @@ abstract class Slicer(
     protected val pseudocodeCache = PseudocodeCache()
 
     protected fun PsiElement.passToProcessor(
-        lambdaLevel: Int = parentUsage.lambdaLevel,
+        behaviour: KotlinSliceUsage.SpecialBehaviour? = this@Slicer.behaviour,
-        forcedExpressionMode: Boolean = false
+        forcedExpressionMode: Boolean = false,
     ) {
-        processor.process(KotlinSliceUsage(this, parentUsage, lambdaLevel, forcedExpressionMode))
+        processor.process(KotlinSliceUsage(this, parentUsage, behaviour, forcedExpressionMode))
+    }
+
+    protected fun PsiElement.passToProcessorAsValue(behaviour: KotlinSliceUsage.SpecialBehaviour? = this@Slicer.behaviour) {
+        passToProcessor(behaviour, forcedExpressionMode = true)
     }
 
     protected fun processCalls(

idea/tests/org/jetbrains/kotlin/idea/slicer/SlicerTestUtil.kt

@@ -62,7 +62,11 @@ internal fun buildTreeRepresentation(rootNode: SliceNode): String {
                         append("DEREFERENCE: ")
                     }
                     if (usage is KotlinSliceUsage) {
-                        append("[LAMBDA] ".repeat(usage.lambdaLevel))
+                        var behaviour = usage.behaviour
+                        while (behaviour != null) {
+                            append(behaviour.testPresentationPrefix)
+                            behaviour = behaviour.originalBehaviour
+                        }
                     }
                     chunks.slice(1 until chunks.size).joinTo(
                         this,