From 04b8bc95392033a08bfdbb998ea67b5f2b43060b Mon Sep 17 00:00:00 2001
From: Anastasiya Shadrina <a.shadrina5@mail.ru>
Date: Mon, 20 Jan 2020 01:36:15 +0700
Subject: [PATCH] [Coverage] Handle return inside nested loop

---
 .../llvm/coverage/CoverageInformation.kt      | 19 ++---------------
 .../llvm/coverage/CoverageRegionCollector.kt  | 21 +++++++++++--------
 backend.native/tests/build.gradle             |  4 ++--
 .../tests/coverage/basic/jumps/main.kt        | 16 ++++++++++++++
 4 files changed, 32 insertions(+), 28 deletions(-)

diff --git a/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/llvm/coverage/CoverageInformation.kt b/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/llvm/coverage/CoverageInformation.kt
index 338bac123fb..be9d4b98868 100644
--- a/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/llvm/coverage/CoverageInformation.kt
+++ b/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/llvm/coverage/CoverageInformation.kt
@@ -20,26 +20,11 @@ import org.jetbrains.kotlin.ir.declarations.name
  * Besides the obvious [file] and line/column borders, it has [RegionKind] which is described later.
  */
 class Region(
-        startOffset: Int,
-        endOffset: Int,
+        val startOffset: Int,
+        val endOffset: Int,
         val file: IrFile,
         val kind: RegionKind
 ) {
-
-    var startOffset: Int = startOffset
-        set(value) {
-            if (value != UNDEFINED_OFFSET) {
-                field = value
-            }
-        }
-
-    var endOffset: Int = endOffset
-        set(value) {
-            if (value != UNDEFINED_OFFSET) {
-                field = value
-            }
-        }
-
     val startLine: Int
         get() = file.fileEntry.line(startOffset)
 
diff --git a/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/llvm/coverage/CoverageRegionCollector.kt b/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/llvm/coverage/CoverageRegionCollector.kt
index f4fbc6c89ca..9b2fd13d317 100644
--- a/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/llvm/coverage/CoverageRegionCollector.kt
+++ b/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/llvm/coverage/CoverageRegionCollector.kt
@@ -165,19 +165,21 @@ private class IrFunctionRegionsCollector(
         body.acceptVoid(this)
     }
 
-    override fun visitReturn(expression: IrReturn) {
-        val next = irStatementsStack.lastOrNull()?.next ?: return
-        val nextRegion = recordRegion(next, expression.endOffset, currentRegion.endOffset) ?: return
-        currentRegion.endOffset = expression.endOffset
-        regionStack.pop()
-        regionStack.push(nextRegion)
-    }
-
     override fun visitBreakContinue(jump: IrBreakContinue) {
         val (current, next) = irStatementsStack.lastOrNull() ?: return
         recordRegion(next ?: return, current.endOffset, jump.loop.endOffset)
     }
 
+    override fun visitReturn(expression: IrReturn) {
+        irStatementsStack.subList(0, irStatementsStack.lastIndex)
+                .filter { (current, next) -> next != null && current.endOffset > expression.endOffset }
+                .forEach { (current, next) -> recordRegion(next!!, expression.endOffset, current.endOffset) }
+        val next = irStatementsStack.lastOrNull()?.next ?: return
+        val nextRegion = recordRegion(next, expression.endOffset, currentRegion.endOffset) ?: return
+        regionStack.pop()
+        regionStack.push(nextRegion)
+    }
+
     private fun visitInFileContext(file: IrFile, visit: () -> Unit) {
         irFileStack.push(file)
         visit()
@@ -199,7 +201,8 @@ private class IrFunctionRegionsCollector(
             if (!current.hasValidOffsets()) {
                 continue
             }
-            val next = if (i < statements.lastIndex && statements[i + 1].hasValidOffsets()) statements[i + 1] else null
+            val nextInContext = irStatementsStack.lastOrNull()?.next
+            val next = if (i < statements.lastIndex && statements[i + 1].hasValidOffsets()) statements[i + 1] else nextInContext
             irStatementsStack.push(StatementContext(current, next))
             visit(current)
             irStatementsStack.pop()
diff --git a/backend.native/tests/build.gradle b/backend.native/tests/build.gradle
index 349494f36c6..f700e9f685f 100644
--- a/backend.native/tests/build.gradle
+++ b/backend.native/tests/build.gradle
@@ -4021,8 +4021,8 @@ if (UtilsKt.getTestTargetSupportsCodeCoverage(project)) {
 
     task coverage_jumps(type: CoverageTest) {
         binaryName = "CoverageJumps"
-        numberOfCoveredFunctions = 7
-        numberOfCoveredLines = 66
+        numberOfCoveredFunctions = 8
+        numberOfCoveredLines = 74
 
         konanArtifacts {
             program(binaryName, targets: [ target ]) {
diff --git a/backend.native/tests/coverage/basic/jumps/main.kt b/backend.native/tests/coverage/basic/jumps/main.kt
index 249f3c75090..a0306316b00 100644
--- a/backend.native/tests/coverage/basic/jumps/main.kt
+++ b/backend.native/tests/coverage/basic/jumps/main.kt
@@ -55,6 +55,21 @@ fun singleReturn() {
     return
 }
 
+fun nestedReturn() {
+    while (true) {
+        while (true) {
+            while (true) {
+                if (1 < 2) {
+                    return
+                }
+                println()
+            }
+            println()
+        }
+    }
+    println()
+}
+
 fun main() {
     simpleReturn(0)
     simpleReturn(1)
@@ -71,4 +86,5 @@ fun main() {
     breakFromWhile()
     continueFromDoWhile()
     singleReturn()
+    nestedReturn()
 }
\ No newline at end of file