Analyze Data Flow: Support lambdas/anonymous functions

#KT-11994 In Progress

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

@@ -23,8 +23,9 @@ import com.intellij.util.Processor
 
 class KotlinSliceDereferenceUsage(
         element: PsiElement,
-        parent: KotlinSliceUsage
+        parent: KotlinSliceUsage,
-) : KotlinSliceUsage(element, parent) {
+        lambdaLevel: Int
+) : KotlinSliceUsage(element, parent, lambdaLevel, false) {
     override fun processChildren(processor: Processor<SliceUsage>) {
         // no children
     }

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

@@ -23,13 +23,23 @@ import com.intellij.util.Processor
 import org.jetbrains.kotlin.psi.KtExpression
 
 open class KotlinSliceUsage : SliceUsage {
-    constructor(element: PsiElement, parent: SliceUsage) : super(element, parent)
+    val lambdaLevel: Int
-    constructor(element: PsiElement, params: SliceAnalysisParams) : super(element, params)
+    val forcedExpressionMode: Boolean
+
+    constructor(element: PsiElement, parent: SliceUsage, lambdaLevel: Int, forcedExpressionMode: Boolean) : super(element, parent) {
+        this.lambdaLevel = lambdaLevel
+        this.forcedExpressionMode = forcedExpressionMode
+    }
+
+    constructor(element: PsiElement, params: SliceAnalysisParams) : super(element, params) {
+        this.lambdaLevel = 0
+        this.forcedExpressionMode = false
+    }
 
     override fun copy(): KotlinSliceUsage {
         val element = usageInfo.element!!
         if (parent == null) return KotlinSliceUsage(element, params)
-        return KotlinSliceUsage(element, parent)
+        return KotlinSliceUsage(element, parent, lambdaLevel, forcedExpressionMode)
     }
 
     public override fun processUsagesFlownDownTo(element: PsiElement, uniqueProcessor: Processor<SliceUsage>) {

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

@@ -53,6 +53,8 @@ object KotlinSliceUsageCellRenderer : SliceUsageCellRendererBase() {
             }
         }
 
+        append(" (Tracking enclosing lambda)".repeat(sliceUsage.lambdaLevel), SimpleTextAttributes.REGULAR_BOLD_ATTRIBUTES)
+
         val declaration = sliceUsage.element?.parents?.firstOrNull {
             it is KtClass ||
             it is KtObjectDeclaration && !it.isObjectLiteral() ||

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

@@ -16,13 +16,13 @@
 
 package org.jetbrains.kotlin.idea.slicer
 
-import com.intellij.codeInsight.highlighting.ReadWriteAccessDetector
 import com.intellij.codeInsight.highlighting.ReadWriteAccessDetector.Access
 import com.intellij.psi.PsiElement
 import com.intellij.psi.search.SearchScope
 import com.intellij.slicer.SliceUsage
 import com.intellij.usageView.UsageInfo
 import com.intellij.util.Processor
+import org.jetbrains.kotlin.builtins.functions.FunctionInvokeDescriptor
 import org.jetbrains.kotlin.cfg.pseudocode.PseudoValue
 import org.jetbrains.kotlin.cfg.pseudocode.Pseudocode
 import org.jetbrains.kotlin.cfg.pseudocode.containingDeclarationForPseudocode
@@ -46,6 +46,7 @@ import org.jetbrains.kotlin.psi.psiUtil.*
 import org.jetbrains.kotlin.resolve.calls.callUtil.getResolvedCall
 import org.jetbrains.kotlin.resolve.calls.model.DefaultValueArgument
 import org.jetbrains.kotlin.resolve.calls.model.ExpressionValueArgument
+import org.jetbrains.kotlin.resolve.scopes.receivers.ExpressionReceiver
 import org.jetbrains.kotlin.resolve.source.getPsi
 import java.util.*
 
@@ -111,8 +112,11 @@ abstract class Slicer(
 
     protected val pseudocodeCache = PseudocodeCache()
 
-    protected fun PsiElement.passToProcessor() {
+    protected fun PsiElement.passToProcessor(
-        processor.process(KotlinSliceUsage(this, parentUsage))
+            lambdaLevel: Int = parentUsage.lambdaLevel,
+            forcedExpressionMode: Boolean = false
+    ) {
+        processor.process(KotlinSliceUsage(this, parentUsage, lambdaLevel, forcedExpressionMode))
     }
 
     abstract fun processChildren()
@@ -123,14 +127,16 @@ class InflowSlicer(
         processor: Processor<SliceUsage>,
         parentUsage: KotlinSliceUsage
 ) : Slicer(element, processor, parentUsage) {
+    private fun PsiElement.passToProcessorAsValue(lambdaLevel: Int = parentUsage.lambdaLevel) = passToProcessor(lambdaLevel, true)
+
     private fun KtProperty.processProperty() {
         if (!canProcess()) return
 
-        initializer?.passToProcessor()
+        initializer?.passToProcessorAsValue()
         processVariableAccesses(parentUsage.scope.toSearchScope(), Access.Write) body@ {
             val refExpression = it.element as? KtExpression ?: return@body
             val rhs = KtPsiUtil.safeDeparenthesize(refExpression).getAssignmentByLHS()?.right ?: return@body
-            rhs.passToProcessor()
+            rhs.passToProcessorAsValue()
         }
     }
 
@@ -152,7 +158,7 @@ class InflowSlicer(
                                       is ExpressionValueArgument -> resolvedArgument.valueArgument?.getArgumentExpression()
                                       else -> null
                                   } ?: return@body
-            flownExpression.passToProcessor()
+            flownExpression.passToProcessorAsValue()
         }
     }
 
@@ -160,21 +166,46 @@ class InflowSlicer(
         val pseudocode = pseudocodeCache[bodyExpression ?: return] ?: return
         pseudocode.traverse(TraversalOrder.FORWARD) { instr ->
             if (instr is ReturnValueInstruction && instr.subroutine == this) {
-                (instr.returnExpressionIfAny?.returnedExpression ?: instr.element as? KtExpression)?.passToProcessor()
+                (instr.returnExpressionIfAny?.returnedExpression ?: instr.element as? KtExpression)?.passToProcessorAsValue()
             }
         }
     }
 
     private fun Instruction.passInputsToProcessor() {
-        inputValues.forEach { it.element?.passToProcessor() }
+        inputValues.forEach {
+            if (it.createdAt != null) {
+                it.element?.passToProcessorAsValue()
+            }
+        }
     }
 
     private fun KtExpression.processExpression() {
+        val lambda = when (this) {
+                         is KtLambdaExpression -> functionLiteral
+                         is KtNamedFunction -> if (name == null) this else null
+                         else -> null
+                     }
+        if (lambda != null) {
+            if (parentUsage.lambdaLevel > 0) {
+                lambda.passToProcessor(parentUsage.lambdaLevel - 1)
+            }
+            return
+        }
+
         val pseudocode = pseudocodeCache[this] ?: return
         val expressionValue = pseudocode.getElementValue(this) ?: return
         val createdAt = expressionValue.createdAt
         when (createdAt) {
-            is ReadValueInstruction -> (createdAt.target.accessedDescriptor?.source?.getPsi() as? KtDeclaration)?.passToProcessor()
+            is ReadValueInstruction -> {
+                if (createdAt.target == AccessTarget.BlackBox) {
+                    val originalElement = expressionValue.element as? KtExpression ?: return
+                    if (originalElement != this) {
+                        originalElement.processExpression()
+                    }
+                    return
+                }
+                (createdAt.target.accessedDescriptor?.source?.getPsi() as? KtDeclaration)?.passToProcessor()
+            }
 
             is MergeInstruction -> createdAt.passInputsToProcessor()
 
@@ -183,11 +214,22 @@ class InflowSlicer(
                 else -> return
             }
 
-            is CallInstruction -> (createdAt.resolvedCall.resultingDescriptor.source.getPsi() as? KtDeclarationWithBody)?.passToProcessor()
+            is CallInstruction -> {
+                val resolvedCall = createdAt.resolvedCall
+                val resultingDescriptor = resolvedCall.resultingDescriptor
+                if (resultingDescriptor is FunctionInvokeDescriptor) {
+                    (resolvedCall.dispatchReceiver as? ExpressionReceiver)?.expression?.passToProcessorAsValue(parentUsage.lambdaLevel + 1)
+                }
+                else {
+                    resultingDescriptor.source.getPsi()?.passToProcessor()
+                }
+            }
         }
     }
 
     override fun processChildren() {
+        if (parentUsage.forcedExpressionMode) return element.processExpression()
+
         when (element) {
             is KtProperty -> element.processProperty()
             is KtParameter -> element.processParameter()
@@ -235,10 +277,19 @@ class OutflowSlicer(
         return null
     }
 
-    private fun KtFunction.processFunctionCalls() {
+    private fun KtFunction.processFunction() {
-        processCalls(parentUsage.scope.toSearchScope()) {
+        if (this is KtConstructor<*> || this is KtNamedFunction && name != null) {
-            it.element?.getCallElementForExactCallee()?.passToProcessor()
+            processCalls(parentUsage.scope.toSearchScope()) {
+                it.element?.getCallElementForExactCallee()?.passToProcessor()
+            }
+            return
         }
+        val funExpression = when (this) {
+            is KtFunctionLiteral -> parent as? KtLambdaExpression
+            is KtNamedFunction -> this
+            else -> null
+        } ?: return
+        (funExpression as PsiElement).passToProcessor(parentUsage.lambdaLevel + 1, true)
     }
 
     private fun processDereferenceIsNeeded(
@@ -256,7 +307,7 @@ class OutflowSlicer(
         } ?: return
 
         if (receiver != null && resolvedCall.dispatchReceiver == receiver) {
-            processor.process(KotlinSliceDereferenceUsage(expression, parentUsage))
+            processor.process(KotlinSliceDereferenceUsage(expression, parentUsage, parentUsage.lambdaLevel))
         }
     }
 
@@ -279,8 +330,15 @@ class OutflowSlicer(
         processPseudocodeUsages { pseudoValue, instr ->
             when (instr) {
                 is WriteValueInstruction -> (instr.target.accessedDescriptor?.source?.getPsi() as? KtDeclaration)?.passToProcessor()
-                is CallInstruction -> instr.arguments[pseudoValue]?.source?.getPsi()?.passToProcessor()
+                is CallInstruction -> {
-                is ReturnValueInstruction -> (instr.owner.correspondingElement as? KtFunction)?.passToProcessor()
+                    if (parentUsage.lambdaLevel > 0 && instr.receiverValues[pseudoValue] != null) {
+                        instr.element.passToProcessor(parentUsage.lambdaLevel - 1)
+                    }
+                    else {
+                        instr.arguments[pseudoValue]?.source?.getPsi()?.passToProcessor()
+                    }
+                }
+                is ReturnValueInstruction -> instr.subroutine.passToProcessor()
                 is MagicInstruction -> when (instr.kind) {
                     MagicKind.NOT_NULL_ASSERTION, MagicKind.CAST -> instr.outputValue.element?.passToProcessor()
                     else -> {}
@@ -290,9 +348,11 @@ class OutflowSlicer(
     }
 
     override fun processChildren() {
+        if (parentUsage.forcedExpressionMode) return element.processExpression()
+
         when (element) {
             is KtProperty, is KtParameter -> (element as KtDeclaration).processVariable()
-            is KtFunction -> element.processFunctionCalls()
+            is KtFunction -> element.processFunction()
             else -> element.processExpression()
         }
     }

idea/testData/slicer/inflow/anonymousFunBodyExpression.kt

@@ -0,0 +1,9 @@
+// FLOW: IN
+
+fun foo(f: (Int) -> Int): Int {
+    return f(1)
+}
+
+fun test() {
+    val <caret>x = foo(fun(n: Int) = n)
+}

idea/testData/slicer/inflow/anonymousFunBodyExpression.results.txt

@@ -0,0 +1,10 @@
+8 val <bold>x = foo(fun(n: Int) = n)</bold>
+8 	val x = <bold>foo(fun(n: Int) = n)</bold>
+3 		fun <bold>foo(f: (Int) -> Int): Int {</bold>
+4 			return <bold>f(1)</bold>
+4 				[LAMBDA] return <bold>f</bold>(1)
+3 					[LAMBDA] fun foo(<bold>f: (Int) -> Int</bold>): Int {
+8 						[LAMBDA] val x = foo(<bold>fun(n: Int) = n</bold>)
+8 							val x = foo(<bold>fun(n: Int) = n</bold>)
+8 								val x = foo(fun(n: Int) = <bold>n</bold>)
+8 									val x = foo(fun(<bold>n: Int</bold>) = n)

idea/testData/slicer/inflow/anonymousFunReturnExpression.kt

@@ -0,0 +1,9 @@
+// FLOW: IN
+
+fun foo(f: (Int) -> Int): Int {
+    return f(1)
+}
+
+fun test() {
+    val <caret>x = foo(fun(n: Int): Int { return n })
+}

idea/testData/slicer/inflow/anonymousFunReturnExpression.results.txt

@@ -0,0 +1,10 @@
+8 val <bold>x = foo(fun(n: Int): Int { return n })</bold>
+8 	val x = <bold>foo(fun(n: Int): Int { return n })</bold>
+3 		fun <bold>foo(f: (Int) -> Int): Int {</bold>
+4 			return <bold>f(1)</bold>
+4 				[LAMBDA] return <bold>f</bold>(1)
+3 					[LAMBDA] fun foo(<bold>f: (Int) -> Int</bold>): Int {
+8 						[LAMBDA] val x = foo(<bold>fun(n: Int): Int { return n }</bold>)
+8 							val x = foo(<bold>fun(n: Int): Int { return n }</bold>)
+8 								val x = foo(fun(n: Int): Int { return <bold>n</bold> })
+8 									val x = foo(fun(<bold>n: Int</bold>): Int { return n })

idea/testData/slicer/inflow/doubleLambdaResult.kt

@@ -0,0 +1,9 @@
+// FLOW: IN
+
+fun foo(a: Int, b: Int, f: (Int) -> (Int) -> Int): Int {
+    return f(a)(b)
+}
+
+fun test() {
+    val <caret>x = foo(1, 2) { { it } }
+}

idea/testData/slicer/inflow/doubleLambdaResult.results.txt

@@ -0,0 +1,12 @@
+8 val <bold>x = foo(1, 2) { { it } }</bold>
+8 	val x = <bold>foo(1, 2) { { it } }</bold>
+3 		fun <bold>foo(a: Int, b: Int, f: (Int) -> (Int) -> Int): Int {</bold>
+4 			return <bold>f(a)(b)</bold>
+4 				[LAMBDA] return <bold>f(a)</bold>(b)
+4 					[LAMBDA] [LAMBDA] return <bold>f</bold>(a)(b)
+3 						[LAMBDA] [LAMBDA] fun foo(a: Int, b: Int, <bold>f: (Int) -> (Int) -> Int</bold>): Int {
+8 							[LAMBDA] [LAMBDA] val x = foo(1, 2) <bold>{ { it } }</bold>
+8 								[LAMBDA] val x = foo(1, 2) <bold>{ { it } }</bold>
+8 									[LAMBDA] val x = foo(1, 2) { <bold>{ it }</bold> }
+8 										val x = foo(1, 2) { <bold>{ it }</bold> }
+8 											val x = foo(1, 2) { { <bold>it</bold> } }

idea/testData/slicer/inflow/lambdaResult.kt

@@ -0,0 +1,9 @@
+// FLOW: IN
+
+fun foo(f: (Int) -> Int): Int {
+    return f(1)
+}
+
+fun test() {
+    val <caret>x = foo { it }
+}

idea/testData/slicer/inflow/lambdaResult.results.txt

@@ -0,0 +1,9 @@
+8 val <bold>x = foo { it }</bold>
+8 	val x = <bold>foo { it }</bold>
+3 		fun <bold>foo(f: (Int) -> Int): Int {</bold>
+4 			return <bold>f(1)</bold>
+4 				[LAMBDA] return <bold>f</bold>(1)
+3 					[LAMBDA] fun foo(<bold>f: (Int) -> Int</bold>): Int {
+8 						[LAMBDA] val x = foo <bold>{ it }</bold>
+8 							val x = foo <bold>{ it }</bold>
+8 								val x = foo { <bold>it</bold> }

idea/testData/slicer/inflow/lambdaResultWithAssignments.kt

@@ -0,0 +1,10 @@
+// FLOW: IN
+
+fun foo(f: (Int) -> Int): Int {
+    val x = f
+    return x(1)
+}
+
+fun test() {
+    val <caret>y = foo { it }
+}

idea/testData/slicer/inflow/lambdaResultWithAssignments.results.txt

@@ -0,0 +1,11 @@
+9 val <bold>y = foo { it }</bold>
+9 	val y = <bold>foo { it }</bold>
+3 		fun <bold>foo(f: (Int) -> Int): Int {</bold>
+5 			return <bold>x(1)</bold>
+5 				[LAMBDA] return <bold>x</bold>(1)
+4 					[LAMBDA] val <bold>x = f</bold>
+4 						[LAMBDA] val x = <bold>f</bold>
+3 							[LAMBDA] fun foo(<bold>f: (Int) -> Int</bold>): Int {
+9 								[LAMBDA] val y = foo <bold>{ it }</bold>
+9 									val y = foo <bold>{ it }</bold>
+9 										val y = foo { <bold>it</bold> }

idea/testData/slicer/inflow/lambdaResultWithDirectCall.kt

@@ -0,0 +1,5 @@
+// FLOW: IN
+
+fun test() {
+    val <caret>x = { 1 }()
+}

idea/testData/slicer/inflow/lambdaResultWithDirectCall.results.txt

@@ -0,0 +1,5 @@
+4 val <bold>x = { 1 }()</bold>
+4 	val x = <bold>{ 1 }()</bold>
+4 		[LAMBDA] val x = <bold>{ 1 }</bold>()
+4 			val x = <bold>{ 1 }</bold>()
+4 				val x = { <bold>1</bold> }()

idea/testData/slicer/inflow/lambdaResultWithDirectCallViaAssignment.kt

@@ -0,0 +1,6 @@
+// FLOW: IN
+
+fun test() {
+    val f = { 1 }
+    val <caret>x = f()
+}

idea/testData/slicer/inflow/lambdaResultWithDirectCallViaAssignment.results.txt

@@ -0,0 +1,7 @@
+5 val <bold>x = f()</bold>
+5 	val x = <bold>f()</bold>
+5 		[LAMBDA] val x = <bold>f</bold>()
+4 			[LAMBDA] val <bold>f = { 1 }</bold>
+4 				[LAMBDA] val f = <bold>{ 1 }</bold>
+4 					val f = <bold>{ 1 }</bold>
+4 						val f = { <bold>1</bold> }

idea/testData/slicer/inflow/nonLocalReturn.kt

@@ -0,0 +1,5 @@
+// FLOW: IN
+
+inline fun foo(a: Int, f: (Int) -> Int) = f(a)
+
+fun <caret>bar(a: Int): Int = foo(a) { if (it > 0) it else return 0 }

idea/testData/slicer/inflow/nonLocalReturn.results.txt

@@ -0,0 +1,11 @@
+5 fun <bold>bar(a: Int): Int = foo(a) { if (it > 0) it else return 0 }</bold>
+5 	fun bar(a: Int): Int = foo(a) { if (it > 0) it else return <bold>0</bold> }
+5 	fun bar(a: Int): Int = <bold>foo(a) { if (it > 0) it else return 0 }</bold>
+3 		inline fun <bold>foo(a: Int, f: (Int) -> Int) = f(a)</bold>
+3 			inline fun foo(a: Int, f: (Int) -> Int) = <bold>f(a)</bold>
+3 				[LAMBDA] inline fun foo(a: Int, f: (Int) -> Int) = <bold>f</bold>(a)
+3 					[LAMBDA] inline fun foo(a: Int, <bold>f: (Int) -> Int</bold>) = f(a)
+5 						[LAMBDA] fun bar(a: Int): Int = foo(a) <bold>{ if (it > 0) it else return 0 }</bold>
+5 							fun bar(a: Int): Int = foo(a) <bold>{ if (it > 0) it else return 0 }</bold>
+5 								fun bar(a: Int): Int = foo(a) { <bold>if (it > 0) it else return 0</bold> }
+5 									fun bar(a: Int): Int = foo(a) { if (it > 0) <bold>it</bold> else return 0 }

idea/testData/slicer/outflow/anonymousFunBodyExpression.kt

@@ -0,0 +1,9 @@
+// FLOW: OUT
+
+fun foo(f: (Int) -> Int): Int {
+    return f(1)
+}
+
+fun test() {
+    val x = foo(fun(n: Int) = <caret>n)
+}

idea/testData/slicer/outflow/anonymousFunBodyExpression.results.txt

@@ -0,0 +1,8 @@
+8 val x = foo(fun(n: Int) = <bold>n</bold>)
+8 	val x = foo(<bold>fun(n: Int) = n</bold>)
+8 		[LAMBDA] val x = foo(<bold>fun(n: Int) = n</bold>)
+3 			[LAMBDA] fun foo(<bold>f: (Int) -> Int</bold>): Int {
+4 				return <bold>f(1)</bold>
+3 					fun <bold>foo(f: (Int) -> Int): Int {</bold>
+8 						val x = <bold>foo(fun(n: Int) = n)</bold>
+8 							val <bold>x = foo(fun(n: Int) = n)</bold>

idea/testData/slicer/outflow/anonymousFunReturnExpression.kt

@@ -0,0 +1,9 @@
+// FLOW: OUT
+
+fun foo(f: (Int) -> Int): Int {
+    return f(1)
+}
+
+fun test() {
+    val x = foo(fun(n: Int): Int { return <caret>n })
+}

idea/testData/slicer/outflow/anonymousFunReturnExpression.results.txt

@@ -0,0 +1,8 @@
+8 val x = foo(fun(n: Int): Int { return <bold>n</bold> })
+8 	val x = foo(<bold>fun(n: Int): Int { return n }</bold>)
+8 		[LAMBDA] val x = foo(<bold>fun(n: Int): Int { return n }</bold>)
+3 			[LAMBDA] fun foo(<bold>f: (Int) -> Int</bold>): Int {
+4 				return <bold>f(1)</bold>
+3 					fun <bold>foo(f: (Int) -> Int): Int {</bold>
+8 						val x = <bold>foo(fun(n: Int): Int { return n })</bold>
+8 							val <bold>x = foo(fun(n: Int): Int { return n })</bold>

idea/testData/slicer/outflow/doubleLambdaResult.kt

@@ -0,0 +1,9 @@
+// FLOW: OUT
+
+fun foo(a: Int, b: Int, f: (Int) -> (Int) -> Int): Int {
+    return f(a)(b)
+}
+
+fun test() {
+    val x = foo(1, 2) { { <caret>it } }
+}

idea/testData/slicer/outflow/doubleLambdaResult.results.txt

@@ -0,0 +1,11 @@
+8 val x = foo(1, 2) { { <bold>it</bold> } }
+8 	val x = foo(1, 2) { <bold>{ it }</bold> }
+8 		[LAMBDA] val x = foo(1, 2) { <bold>{ it }</bold> }
+8 			[LAMBDA] val x = foo(1, 2) <bold>{ { it } }</bold>
+8 				[LAMBDA] [LAMBDA] val x = foo(1, 2) <bold>{ { it } }</bold>
+3 					[LAMBDA] [LAMBDA] fun foo(a: Int, b: Int, <bold>f: (Int) -> (Int) -> Int</bold>): Int {
+4 						[LAMBDA] return <bold>f(a)</bold>(b)
+4 							return <bold>f(a)(b)</bold>
+3 								fun <bold>foo(a: Int, b: Int, f: (Int) -> (Int) -> Int): Int {</bold>
+8 									val x = <bold>foo(1, 2) { { it } }</bold>
+8 										val <bold>x = foo(1, 2) { { it } }</bold>

idea/testData/slicer/outflow/explicitLambdaReturnExpression.kt

@@ -0,0 +1,12 @@
+// FLOW: OUT
+
+fun foo(f: (Int) -> Int): Int {
+    return f(1)
+}
+
+fun test() {
+    val x = foo {
+        if (it < 0) return@foo <caret>0
+        it
+    }
+}

idea/testData/slicer/outflow/explicitLambdaReturnExpression.results.txt

@@ -0,0 +1,8 @@
+9 if (it < 0) return@foo <bold>0</bold>
+8 	val x = foo <bold>{</bold>
+8 		[LAMBDA] val x = foo <bold>{</bold>
+3 			[LAMBDA] fun foo(<bold>f: (Int) -> Int</bold>): Int {
+4 				return <bold>f(1)</bold>
+3 					fun <bold>foo(f: (Int) -> Int): Int {</bold>
+8 						val x = <bold>foo {</bold>
+8 							val <bold>x = foo {</bold>

idea/testData/slicer/outflow/lambdaResult.kt

@@ -0,0 +1,9 @@
+// FLOW: OUT
+
+fun foo(f: (Int) -> Int): Int {
+    return f(1)
+}
+
+fun test() {
+    val x = foo { <caret>it }
+}

idea/testData/slicer/outflow/lambdaResult.results.txt

@@ -0,0 +1,8 @@
+8 val x = foo { <bold>it</bold> }
+8 	val x = foo <bold>{ it }</bold>
+8 		[LAMBDA] val x = foo <bold>{ it }</bold>
+3 			[LAMBDA] fun foo(<bold>f: (Int) -> Int</bold>): Int {
+4 				return <bold>f(1)</bold>
+3 					fun <bold>foo(f: (Int) -> Int): Int {</bold>
+8 						val x = <bold>foo { it }</bold>
+8 							val <bold>x = foo { it }</bold>

idea/testData/slicer/outflow/lambdaResultWithAssignments.kt

@@ -0,0 +1,10 @@
+// FLOW: OUT
+
+fun foo(f: (Int) -> Int): Int {
+    val x = f
+    return x(1)
+}
+
+fun test() {
+    val y = foo { <caret>it }
+}

idea/testData/slicer/outflow/lambdaResultWithAssignments.results.txt

@@ -0,0 +1,9 @@
+9 val y = foo { <bold>it</bold> }
+9 	val y = foo <bold>{ it }</bold>
+9 		[LAMBDA] val y = foo <bold>{ it }</bold>
+3 			[LAMBDA] fun foo(<bold>f: (Int) -> Int</bold>): Int {
+4 				[LAMBDA] val <bold>x = f</bold>
+5 					return <bold>x(1)</bold>
+3 						fun <bold>foo(f: (Int) -> Int): Int {</bold>
+9 							val y = <bold>foo { it }</bold>
+9 								val <bold>y = foo { it }</bold>

idea/testData/slicer/outflow/lambdaResultWithDirectCall.kt

@@ -0,0 +1,5 @@
+// FLOW: OUT
+
+fun test() {
+    val x = { <caret>1 }()
+}

idea/testData/slicer/outflow/lambdaResultWithDirectCall.results.txt

@@ -0,0 +1,5 @@
+4 val x = { <bold>1</bold> }()
+4 	val x = <bold>{ 1 }</bold>()
+4 		[LAMBDA] val x = <bold>{ 1 }</bold>()
+4 			val x = <bold>{ 1 }()</bold>
+4 				val <bold>x = { 1 }()</bold>

idea/testData/slicer/outflow/lambdaResultWithDirectCallViaAssignment.kt

@@ -0,0 +1,6 @@
+// FLOW: OUT
+
+fun test() {
+    val f = { <caret>1 }
+    val x = f()
+}

idea/testData/slicer/outflow/lambdaResultWithDirectCallViaAssignment.results.txt

@@ -0,0 +1,6 @@
+4 val f = { <bold>1</bold> }
+4 	val f = <bold>{ 1 }</bold>
+4 		[LAMBDA] val f = <bold>{ 1 }</bold>
+4 			[LAMBDA] val <bold>f = { 1 }</bold>
+5 				val x = <bold>f()</bold>
+5 					val <bold>x = f()</bold>

idea/testData/slicer/outflow/nonLocalReturn.kt

@@ -0,0 +1,9 @@
+// FLOW: OUT
+
+inline fun foo(a: Int, f: (Int) -> Int) = f(a)
+
+fun bar(a: Int): Int = foo(a) { if (it > 0) it else return <caret>0 }
+
+fun test() {
+    val x = bar(1)
+}

idea/testData/slicer/outflow/nonLocalReturn.results.txt

@@ -0,0 +1,4 @@
+5 fun bar(a: Int): Int = foo(a) { if (it > 0) it else return <bold>0</bold> }
+5 	fun <bold>bar(a: Int): Int = foo(a) { if (it > 0) it else return 0 }</bold>
+8 		val x = <bold>bar(1)</bold>
+8 			val <bold>x = bar(1)</bold>

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

@@ -18,11 +18,12 @@ package org.jetbrains.kotlin.idea.slicer
 
 import com.intellij.analysis.AnalysisScope
 import com.intellij.openapi.util.io.FileUtil
-import com.intellij.psi.PsiElement
 import com.intellij.slicer.SliceAnalysisParams
 import com.intellij.slicer.SliceUsage
 import com.intellij.usages.TextChunk
 import com.intellij.util.CommonProcessors
+import com.intellij.util.Processor
+import gnu.trove.THashSet
 import gnu.trove.TObjectHashingStrategy
 import org.jetbrains.kotlin.idea.test.KotlinLightCodeInsightFixtureTestCase
 import org.jetbrains.kotlin.idea.util.application.runReadAction
@@ -33,9 +34,14 @@ import java.awt.Font
 import java.io.File
 
 abstract class AbstractSlicerTest : KotlinLightCodeInsightFixtureTestCase() {
-    object SliceUsageHashingStrategy : TObjectHashingStrategy<SliceUsage> {
+    object SliceUsageHashingStrategy : TObjectHashingStrategy<KotlinSliceUsage> {
-        override fun computeHashCode(`object`: SliceUsage) = `object`.usageInfo.hashCode()
+        override fun computeHashCode(`object`: KotlinSliceUsage): Int {
-        override fun equals(o1: SliceUsage, o2: SliceUsage) = o1.usageInfo == o2.usageInfo
+            return `object`.usageInfo.hashCode() * 31 + `object`.forcedExpressionMode.hashCode()
+        }
+
+        override fun equals(o1: KotlinSliceUsage, o2: KotlinSliceUsage): Boolean {
+            return o1.usageInfo == o2.usageInfo && o1.forcedExpressionMode == o2.forcedExpressionMode
+        }
     }
 
     // Based on SliceUsage.processChildren
@@ -44,13 +50,14 @@ abstract class AbstractSlicerTest : KotlinLightCodeInsightFixtureTestCase() {
 
         val element = runReadAction { element }
 
+        @Suppress("UNCHECKED_CAST")
         val uniqueProcessor = CommonProcessors.UniqueProcessor(
                 {
                     processor(it as KotlinSliceUsage)
                     true
                 },
                 SliceUsageHashingStrategy
-        )
+        ) as Processor<SliceUsage>
 
         runReadAction {
             if (params.dataFlowToThis) {
@@ -63,7 +70,7 @@ abstract class AbstractSlicerTest : KotlinLightCodeInsightFixtureTestCase() {
     }
 
     private fun buildTreeRepresentation(rootUsage: KotlinSliceUsage): String {
-        val visitedElements = HashSet<PsiElement>()
+        val visitedUsages = THashSet<KotlinSliceUsage>(SliceUsageHashingStrategy)
 
         fun TextChunk.render(): String {
             var text = text
@@ -74,7 +81,7 @@ abstract class AbstractSlicerTest : KotlinLightCodeInsightFixtureTestCase() {
         }
 
         fun process(usage: KotlinSliceUsage, indent: Int): String {
-            val isDuplicated = !visitedElements.add(usage.element)
+            val isDuplicated = !visitedUsages.add(usage)
 
             return buildString {
                 val chunks = usage.text
@@ -83,6 +90,7 @@ abstract class AbstractSlicerTest : KotlinLightCodeInsightFixtureTestCase() {
                 if (usage is KotlinSliceDereferenceUsage) {
                     append("DEREFERENCE: ")
                 }
+                append("[LAMBDA] ".repeat(usage.lambdaLevel))
                 chunks.slice(1..chunks.size - 1).joinTo(
                         this,
                         separator = "",

idea/tests/org/jetbrains/kotlin/idea/slicer/SlicerTestGenerated.java

@@ -36,12 +36,30 @@ public class SlicerTestGenerated extends AbstractSlicerTest {
         KotlinTestUtils.assertAllTestsPresentInSingleGeneratedClass(this.getClass(), new File("idea/testData/slicer"), Pattern.compile("^(.+)\\.kt$"), TargetBackend.ANY);
     }
 
+    @TestMetadata("inflow/anonymousFunBodyExpression.kt")
+    public void testInflow_AnonymousFunBodyExpression() throws Exception {
+        String fileName = KotlinTestUtils.navigationMetadata("idea/testData/slicer/inflow/anonymousFunBodyExpression.kt");
+        doTest(fileName);
+    }
+
+    @TestMetadata("inflow/anonymousFunReturnExpression.kt")
+    public void testInflow_AnonymousFunReturnExpression() throws Exception {
+        String fileName = KotlinTestUtils.navigationMetadata("idea/testData/slicer/inflow/anonymousFunReturnExpression.kt");
+        doTest(fileName);
+    }
+
     @TestMetadata("inflow/cast.kt")
     public void testInflow_Cast() throws Exception {
         String fileName = KotlinTestUtils.navigationMetadata("idea/testData/slicer/inflow/cast.kt");
         doTest(fileName);
     }
 
+    @TestMetadata("inflow/doubleLambdaResult.kt")
+    public void testInflow_DoubleLambdaResult() throws Exception {
+        String fileName = KotlinTestUtils.navigationMetadata("idea/testData/slicer/inflow/doubleLambdaResult.kt");
+        doTest(fileName);
+    }
+
     @TestMetadata("inflow/funParamerer.kt")
     public void testInflow_FunParamerer() throws Exception {
         String fileName = KotlinTestUtils.navigationMetadata("idea/testData/slicer/inflow/funParamerer.kt");
@@ -72,6 +90,30 @@ public class SlicerTestGenerated extends AbstractSlicerTest {
         doTest(fileName);
     }
 
+    @TestMetadata("inflow/lambdaResult.kt")
+    public void testInflow_LambdaResult() throws Exception {
+        String fileName = KotlinTestUtils.navigationMetadata("idea/testData/slicer/inflow/lambdaResult.kt");
+        doTest(fileName);
+    }
+
+    @TestMetadata("inflow/lambdaResultWithAssignments.kt")
+    public void testInflow_LambdaResultWithAssignments() throws Exception {
+        String fileName = KotlinTestUtils.navigationMetadata("idea/testData/slicer/inflow/lambdaResultWithAssignments.kt");
+        doTest(fileName);
+    }
+
+    @TestMetadata("inflow/lambdaResultWithDirectCall.kt")
+    public void testInflow_LambdaResultWithDirectCall() throws Exception {
+        String fileName = KotlinTestUtils.navigationMetadata("idea/testData/slicer/inflow/lambdaResultWithDirectCall.kt");
+        doTest(fileName);
+    }
+
+    @TestMetadata("inflow/lambdaResultWithDirectCallViaAssignment.kt")
+    public void testInflow_LambdaResultWithDirectCallViaAssignment() throws Exception {
+        String fileName = KotlinTestUtils.navigationMetadata("idea/testData/slicer/inflow/lambdaResultWithDirectCallViaAssignment.kt");
+        doTest(fileName);
+    }
+
     @TestMetadata("inflow/localVal.kt")
     public void testInflow_LocalVal() throws Exception {
         String fileName = KotlinTestUtils.navigationMetadata("idea/testData/slicer/inflow/localVal.kt");
@@ -108,6 +150,12 @@ public class SlicerTestGenerated extends AbstractSlicerTest {
         doTest(fileName);
     }
 
+    @TestMetadata("inflow/nonLocalReturn.kt")
+    public void testInflow_NonLocalReturn() throws Exception {
+        String fileName = KotlinTestUtils.navigationMetadata("idea/testData/slicer/inflow/nonLocalReturn.kt");
+        doTest(fileName);
+    }
+
     @TestMetadata("inflow/notNullAssertion.kt")
     public void testInflow_NotNullAssertion() throws Exception {
         String fileName = KotlinTestUtils.navigationMetadata("idea/testData/slicer/inflow/notNullAssertion.kt");
@@ -162,6 +210,18 @@ public class SlicerTestGenerated extends AbstractSlicerTest {
         doTest(fileName);
     }
 
+    @TestMetadata("outflow/anonymousFunBodyExpression.kt")
+    public void testOutflow_AnonymousFunBodyExpression() throws Exception {
+        String fileName = KotlinTestUtils.navigationMetadata("idea/testData/slicer/outflow/anonymousFunBodyExpression.kt");
+        doTest(fileName);
+    }
+
+    @TestMetadata("outflow/anonymousFunReturnExpression.kt")
+    public void testOutflow_AnonymousFunReturnExpression() throws Exception {
+        String fileName = KotlinTestUtils.navigationMetadata("idea/testData/slicer/outflow/anonymousFunReturnExpression.kt");
+        doTest(fileName);
+    }
+
     @TestMetadata("outflow/callArgument.kt")
     public void testOutflow_CallArgument() throws Exception {
         String fileName = KotlinTestUtils.navigationMetadata("idea/testData/slicer/outflow/callArgument.kt");
@@ -174,6 +234,18 @@ public class SlicerTestGenerated extends AbstractSlicerTest {
         doTest(fileName);
     }
 
+    @TestMetadata("outflow/doubleLambdaResult.kt")
+    public void testOutflow_DoubleLambdaResult() throws Exception {
+        String fileName = KotlinTestUtils.navigationMetadata("idea/testData/slicer/outflow/doubleLambdaResult.kt");
+        doTest(fileName);
+    }
+
+    @TestMetadata("outflow/explicitLambdaReturnExpression.kt")
+    public void testOutflow_ExplicitLambdaReturnExpression() throws Exception {
+        String fileName = KotlinTestUtils.navigationMetadata("idea/testData/slicer/outflow/explicitLambdaReturnExpression.kt");
+        doTest(fileName);
+    }
+
     @TestMetadata("outflow/extensionIndexingDereferences.kt")
     public void testOutflow_ExtensionIndexingDereferences() throws Exception {
         String fileName = KotlinTestUtils.navigationMetadata("idea/testData/slicer/outflow/extensionIndexingDereferences.kt");
@@ -222,6 +294,30 @@ public class SlicerTestGenerated extends AbstractSlicerTest {
         doTest(fileName);
     }
 
+    @TestMetadata("outflow/lambdaResult.kt")
+    public void testOutflow_LambdaResult() throws Exception {
+        String fileName = KotlinTestUtils.navigationMetadata("idea/testData/slicer/outflow/lambdaResult.kt");
+        doTest(fileName);
+    }
+
+    @TestMetadata("outflow/lambdaResultWithAssignments.kt")
+    public void testOutflow_LambdaResultWithAssignments() throws Exception {
+        String fileName = KotlinTestUtils.navigationMetadata("idea/testData/slicer/outflow/lambdaResultWithAssignments.kt");
+        doTest(fileName);
+    }
+
+    @TestMetadata("outflow/lambdaResultWithDirectCall.kt")
+    public void testOutflow_LambdaResultWithDirectCall() throws Exception {
+        String fileName = KotlinTestUtils.navigationMetadata("idea/testData/slicer/outflow/lambdaResultWithDirectCall.kt");
+        doTest(fileName);
+    }
+
+    @TestMetadata("outflow/lambdaResultWithDirectCallViaAssignment.kt")
+    public void testOutflow_LambdaResultWithDirectCallViaAssignment() throws Exception {
+        String fileName = KotlinTestUtils.navigationMetadata("idea/testData/slicer/outflow/lambdaResultWithDirectCallViaAssignment.kt");
+        doTest(fileName);
+    }
+
     @TestMetadata("outflow/localVariableUsages.kt")
     public void testOutflow_LocalVariableUsages() throws Exception {
         String fileName = KotlinTestUtils.navigationMetadata("idea/testData/slicer/outflow/localVariableUsages.kt");
@@ -234,6 +330,12 @@ public class SlicerTestGenerated extends AbstractSlicerTest {
         doTest(fileName);
     }
 
+    @TestMetadata("outflow/nonLocalReturn.kt")
+    public void testOutflow_NonLocalReturn() throws Exception {
+        String fileName = KotlinTestUtils.navigationMetadata("idea/testData/slicer/outflow/nonLocalReturn.kt");
+        doTest(fileName);
+    }
+
     @TestMetadata("outflow/notNullAssertion.kt")
     public void testOutflow_NotNullAssertion() throws Exception {
         String fileName = KotlinTestUtils.navigationMetadata("idea/testData/slicer/outflow/notNullAssertion.kt");