[Coverage] Handle return inside nested loop

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,
+        val startOffset: Int,
-        endOffset: 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)
 

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()

backend.native/tests/build.gradle

@@ -4021,8 +4021,8 @@ if (UtilsKt.getTestTargetSupportsCodeCoverage(project)) {
 
     task coverage_jumps(type: CoverageTest) {
         binaryName = "CoverageJumps"
-        numberOfCoveredFunctions = 7
+        numberOfCoveredFunctions = 8
-        numberOfCoveredLines = 66
+        numberOfCoveredLines = 74
 
         konanArtifacts {
             program(binaryName, targets: [ target ]) {

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()
 }