From 2c0f93b1335e68976813c72c2a68650d8c6d39cf Mon Sep 17 00:00:00 2001
From: Igor Chevdar <igor.chevdar@jetbrains.com>
Date: Tue, 11 Dec 2018 18:17:33 +0300
Subject: [PATCH] IR: Supported special coroutine intrinsic See
 https://youtrack.jetbrains.com/issue/KT-26317 for details

---
 .../common/lower/SharedVariablesLowering.kt   | 38 ++++++++++++-------
 1 file changed, 25 insertions(+), 13 deletions(-)

diff --git a/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/SharedVariablesLowering.kt b/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/SharedVariablesLowering.kt
index aba5d32561e..9e5883f7f1a 100644
--- a/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/SharedVariablesLowering.kt
+++ b/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/SharedVariablesLowering.kt
@@ -16,21 +16,22 @@
 
 package org.jetbrains.kotlin.backend.common.lower
 
-import org.jetbrains.kotlin.backend.common.BackendContext
-import org.jetbrains.kotlin.backend.common.FunctionLoweringPass
+import org.jetbrains.kotlin.backend.common.*
 import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.IrStatement
-import org.jetbrains.kotlin.ir.declarations.*
-import org.jetbrains.kotlin.ir.expressions.IrExpression
-import org.jetbrains.kotlin.ir.expressions.IrGetValue
-import org.jetbrains.kotlin.ir.expressions.IrSetVariable
-import org.jetbrains.kotlin.ir.expressions.IrValueAccessExpression
+import org.jetbrains.kotlin.ir.declarations.IrDeclaration
+import org.jetbrains.kotlin.ir.declarations.IrDeclarationParent
+import org.jetbrains.kotlin.ir.declarations.IrFunction
+import org.jetbrains.kotlin.ir.declarations.IrVariable
+import org.jetbrains.kotlin.ir.expressions.*
 import org.jetbrains.kotlin.ir.symbols.IrValueSymbol
 import org.jetbrains.kotlin.ir.symbols.IrVariableSymbol
 import org.jetbrains.kotlin.ir.util.dump
 import org.jetbrains.kotlin.ir.visitors.*
 import java.util.*
 
+object CoroutineIntrinsicLambdaOrigin : IrStatementOriginImpl("Coroutine intrinsic lambda")
+
 class SharedVariablesLowering(val context: BackendContext) : FunctionLoweringPass {
     override fun lower(irFunction: IrFunction) {
         SharedVariablesTransformer(irFunction).lowerSharedVariables()
@@ -47,17 +48,28 @@ class SharedVariablesLowering(val context: BackendContext) : FunctionLoweringPas
         }
 
         private fun collectSharedVariables() {
-            irFunction.accept(object : IrElementVisitor<Unit, IrDeclarationParent> {
-                val relevantVars = HashSet<IrVariable>()
+            irFunction.accept(object : IrElementVisitor<Unit, IrDeclarationParent?> {
+                val relevantVars = mutableSetOf<IrVariable>()
 
-                override fun visitElement(element: IrElement, data: IrDeclarationParent) {
+                override fun visitElement(element: IrElement, data: IrDeclarationParent?) {
                     element.acceptChildren(this, data)
                 }
 
-                override fun visitDeclaration(declaration: IrDeclaration, data: IrDeclarationParent) =
+                override fun visitDeclaration(declaration: IrDeclaration, data: IrDeclarationParent?) =
                     super.visitDeclaration(declaration, declaration as? IrDeclarationParent ?: data)
 
-                override fun visitVariable(declaration: IrVariable, data: IrDeclarationParent) {
+                override fun visitContainerExpression(expression: IrContainerExpression, data: IrDeclarationParent?) =
+                    super.visitContainerExpression(
+                        expression,
+                        if (expression is IrReturnableBlock
+                            && expression.origin == CoroutineIntrinsicLambdaOrigin
+                        )
+                            null
+                        else
+                            data
+                    )
+
+                override fun visitVariable(declaration: IrVariable, data: IrDeclarationParent?) {
                     declaration.acceptChildren(this, data)
 
                     if (declaration.isVar) {
@@ -65,7 +77,7 @@ class SharedVariablesLowering(val context: BackendContext) : FunctionLoweringPas
                     }
                 }
 
-                override fun visitValueAccess(expression: IrValueAccessExpression, data: IrDeclarationParent) {
+                override fun visitValueAccess(expression: IrValueAccessExpression, data: IrDeclarationParent?) {
                     expression.acceptChildren(this, data)
 
                     val value = expression.symbol.owner