From c778e0454d0767a4364438a249871370fcca56df Mon Sep 17 00:00:00 2001
From: Dmitry Jemerov <yole@jetbrains.com>
Date: Tue, 9 May 2017 10:26:42 +0200
Subject: [PATCH] Show line markers for suspending iteration

 #KT-16803 Fixed
---
 .../KotlinSuspendCallLineMarkerProvider.kt    | 26 ++++++++++++-------
 .../suspendCall/suspendIteration.kt           | 17 ++++++++++++
 .../codeInsight/LineMarkersTestGenerated.java |  6 +++++
 3 files changed, 40 insertions(+), 9 deletions(-)
 create mode 100644 idea/testData/codeInsight/lineMarker/suspendCall/suspendIteration.kt

diff --git a/idea/src/org/jetbrains/kotlin/idea/highlighter/KotlinSuspendCallLineMarkerProvider.kt b/idea/src/org/jetbrains/kotlin/idea/highlighter/KotlinSuspendCallLineMarkerProvider.kt
index b88014a65ba..9e0f60a61bb 100644
--- a/idea/src/org/jetbrains/kotlin/idea/highlighter/KotlinSuspendCallLineMarkerProvider.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/highlighter/KotlinSuspendCallLineMarkerProvider.kt
@@ -26,20 +26,18 @@ import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.idea.KotlinIcons
 import org.jetbrains.kotlin.idea.caches.resolve.analyze
 import org.jetbrains.kotlin.idea.refactoring.getLineNumber
-import org.jetbrains.kotlin.psi.KtCallExpression
-import org.jetbrains.kotlin.psi.KtElement
-import org.jetbrains.kotlin.psi.KtExpression
-import org.jetbrains.kotlin.psi.KtOperationReferenceExpression
+import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.calls.callUtil.getResolvedCall
 import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
 
 class KotlinSuspendCallLineMarkerProvider : LineMarkerProvider {
-    private class SuspendCallMarkerInfo(callElement: KtElement) : LineMarkerInfo<KtElement>(
+    private class SuspendCallMarkerInfo(callElement: KtElement, message: String) : LineMarkerInfo<KtElement>(
             callElement,
             callElement.textRange,
             KotlinIcons.SUSPEND_CALL,
             Pass.UPDATE_OVERRIDDEN_MARKERS,
-            { "Suspend function call" },
+            { message },
             null,
             GutterIconRenderer.Alignment.RIGHT
     ) {
@@ -53,7 +51,7 @@ class KotlinSuspendCallLineMarkerProvider : LineMarkerProvider {
     override fun getLineMarkerInfo(element: PsiElement) = null
 
     private fun isValidCandidateExpression(expression: KtExpression): Boolean {
-        if (expression is KtOperationReferenceExpression) return true
+        if (expression is KtOperationReferenceExpression || expression is KtForExpression) return true
         val parent = expression.parent
         if (parent is KtCallExpression && parent.calleeExpression == expression) return true
         return false
@@ -74,12 +72,22 @@ class KotlinSuspendCallLineMarkerProvider : LineMarkerProvider {
             val lineNumber = element.getLineNumber()
             if (lineNumber in markedLineNumbers) continue
 
-            val resolvedCall = element.getResolvedCall(element.analyze(BodyResolveMode.PARTIAL)) ?: continue
+            val bindingContext = element.analyze(BodyResolveMode.PARTIAL)
+            val resolvedCall = if (element is KtForExpression) {
+                bindingContext[BindingContext.LOOP_RANGE_NEXT_RESOLVED_CALL, element.loopRange]
+            } else {
+                element.getResolvedCall(bindingContext)
+            } ?: continue
             val calleeDescriptor = resolvedCall.resultingDescriptor as? FunctionDescriptor ?: continue
             if (!calleeDescriptor.isSuspend) continue
 
             markedLineNumbers += lineNumber
-            result += SuspendCallMarkerInfo(element)
+            if (element is KtForExpression) {
+                result += SuspendCallMarkerInfo(element.loopRange!!, "Suspending iteration")
+
+            } else {
+                result += SuspendCallMarkerInfo(element, "Suspend function call")
+            }
         }
     }
 }
diff --git a/idea/testData/codeInsight/lineMarker/suspendCall/suspendIteration.kt b/idea/testData/codeInsight/lineMarker/suspendCall/suspendIteration.kt
new file mode 100644
index 00000000000..157ab3ee099
--- /dev/null
+++ b/idea/testData/codeInsight/lineMarker/suspendCall/suspendIteration.kt
@@ -0,0 +1,17 @@
+fun coroutine(block: suspend () -> Unit) {}
+
+class SIter {
+    operator fun iterator() = this
+    suspend operator fun hasNext(): Boolean = false
+    suspend operator fun next(): Int = 0
+    suspend fun test() {}
+}
+
+fun foo() {
+    val iter = SIter()
+    coroutine {
+        iter.<lineMarker descr="Suspend function call">test</lineMarker>() // this line is marked (LINE1)
+        for (x in <lineMarker descr="Suspending iteration">iter</lineMarker>) // this line is not (LINE2)
+            println(x)
+    }
+}
\ No newline at end of file
diff --git a/idea/tests/org/jetbrains/kotlin/idea/codeInsight/LineMarkersTestGenerated.java b/idea/tests/org/jetbrains/kotlin/idea/codeInsight/LineMarkersTestGenerated.java
index e37aa5cdf4b..13a606ddb7f 100644
--- a/idea/tests/org/jetbrains/kotlin/idea/codeInsight/LineMarkersTestGenerated.java
+++ b/idea/tests/org/jetbrains/kotlin/idea/codeInsight/LineMarkersTestGenerated.java
@@ -340,5 +340,11 @@ public class LineMarkersTestGenerated extends AbstractLineMarkersTest {
             String fileName = KotlinTestUtils.navigationMetadata("idea/testData/codeInsight/lineMarker/suspendCall/suspendCall.kt");
             doTest(fileName);
         }
+
+        @TestMetadata("suspendIteration.kt")
+        public void testSuspendIteration() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/codeInsight/lineMarker/suspendCall/suspendIteration.kt");
+            doTest(fileName);
+        }
     }
 }