Inline rematerializable safe call receivers.

compiler/ir/ir.psi2ir/src/org/jetbrains/kotlin/psi2ir/transformations/InlineSafeCallChains.kt

@@ -28,9 +28,11 @@ import org.jetbrains.kotlin.psi2ir.generators.GeneratorContext
 import org.jetbrains.kotlin.psi2ir.generators.constNull
 import org.jetbrains.kotlin.psi2ir.generators.equalsNull
 import org.jetbrains.kotlin.psi2ir.intermediate.OnceExpressionValue
+import org.jetbrains.kotlin.psi2ir.intermediate.createRematerializableValue
 
 fun inlineSafeCallChains(context: GeneratorContext, element: IrElement) {
     element.accept(InlineSafeCallChains(context), null)
+    element.accept(InlineSafeCallStableReceiverValues(context), null)
 }
 
 class InlineSafeCallChains(val context: GeneratorContext) : IrElementVisitor<Unit, Nothing?> {
@@ -55,29 +57,63 @@ class InlineSafeCallChains(val context: GeneratorContext) : IrElementVisitor<Uni
         outer.root.replaceWith {
             val newBlock = IrBlockImpl(it.startOffset, it.endOffset, it.type, it.hasResult, IrOperator.SAFE_CALL)
             newBlock.addStatement(inner.receiverVariable.detach())
-            outer.nestedCall.acceptChildren(
+
-                    ReplaceTemporaryVariable(outer.receiverVariable, OnceExpressionValue(inner.nestedCall.detach())),
+            val replaceWithValue = OnceExpressionValue(inner.nestedCall.detach())
-                    null)
+            outer.nestedCall.acceptChildren(ReplaceTemporaryVariable(outer.receiverVariable, replaceWithValue), null)
             newBlock.addStatement(IrIfThenElseImpl(
                     it.startOffset, it.endOffset, it.type,
                     context.equalsNull(it.startOffset, it.endOffset, inner.receiverVariable.defaultLoad()),
                     context.constNull(it.startOffset, it.endOffset),
                     outer.nestedCall.detach(),
                     IrOperator.SAFE_CALL))
+
             newBlock
         }
     }
 
-    private class SafeCallInfo(val root: IrBlock, val receiverVariable: IrVariable, val nestedCall: IrExpression) {
-        val receiverValue = receiverVariable.initializer
-    }
 
-    private fun getSafeCallInfo(block: IrBlock): SafeCallInfo? {
-        if (block.operator != IrOperator.SAFE_CALL) return null
-        val receiverVariable = block.statements[0] as? IrVariable ?: return null
-        if (receiverVariable.initializer == null) return null
-        val nestedCall = (block.statements[1] as? IrWhen)?.elseBranch ?: return null
-        return SafeCallInfo(block, receiverVariable, nestedCall)
-    }
 }
 
+internal class SafeCallInfo(val root: IrBlock, val receiverVariable: IrVariable, val ifThenElse: IrWhen, val nestedCall: IrExpression) {
+    val receiverValue = receiverVariable.initializer!!
+}
+
+internal fun getSafeCallInfo(block: IrBlock): SafeCallInfo? {
+    if (block.operator != IrOperator.SAFE_CALL) return null
+    val receiverVariable = block.statements[0] as? IrVariable ?: return null
+    if (receiverVariable.initializer == null) return null
+    val irWhen = (block.statements[1] as? IrWhen) ?: return null
+    val nestedCall = irWhen.elseBranch ?: return null
+    return SafeCallInfo(block, receiverVariable, irWhen, nestedCall)
+}
+
+class InlineSafeCallStableReceiverValues(val context: GeneratorContext) : IrElementVisitor<Unit, Nothing?> {
+    override fun visitElement(element: IrElement, data: Nothing?) {
+        element.acceptChildren(this, data)
+    }
+
+    override fun visitBlock(expression: IrBlock, data: Nothing?) {
+        expression.acceptChildren(this, data)
+
+        if (expression.operator == IrOperator.SAFE_CALL) {
+            val safeCall = getSafeCallInfo(expression) ?: return
+            if (isOkToInlineReceiverValue(safeCall.receiverValue)) {
+                expression.replaceWith {
+                    val replaceWithValue = createRematerializableValue(safeCall.receiverValue) ?: return
+                    safeCall.ifThenElse.acceptChildren(ReplaceTemporaryVariable(safeCall.receiverVariable, replaceWithValue), null)
+                    safeCall.ifThenElse.detach()
+                }
+            }
+        }
+    }
+
+    private fun isOkToInlineReceiverValue(receiverValue: IrExpression): Boolean {
+        return if (receiverValue is IrGetVariable) {
+            !receiverValue.descriptor.isVar
+        }
+        else {
+            // For now, IrExpressionWithCopy has stable instances only.
+            receiverValue is IrExpressionWithCopy
+        }
+    }
+}

compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/IrElement.kt

@@ -55,7 +55,7 @@ fun IrElement.replaceWith(otherElement: IrElement) {
     parent.replaceChild(slot, otherElement.detach())
 }
 
-fun <T : IrElement> T.replaceWith(transformation: (T) -> IrElement) {
+inline fun <T : IrElement> T.replaceWith(transformation: (T) -> IrElement) {
     val originalParent = this.parent ?: throw AssertionError("Can't replace a non-root element $this")
     val originalSlot = this.slot
     val transformed = transformation(this)

compiler/testData/ir/irText/chainOfSafeCalls.txt

@@ -6,17 +6,14 @@ IrFile /chainOfSafeCalls.kt
         RETURN type=<no-type>
           BLOCK type=C? hasResult=true operator=SAFE_CALL
             VAR val tmp2_safe_receiver: C?
-              BLOCK type=C? hasResult=true operator=SAFE_CALL
+              WHEN type=C? operator=SAFE_CALL
-                VAR val tmp0_safe_receiver: C?
+                if: CALL .EQEQ type=kotlin.Boolean operator=EQEQ
-                  GET_VAR nc type=C? operator=null
+                  arg0: GET_VAR nc type=C? operator=null
-                WHEN type=C? operator=SAFE_CALL
+                  arg1: CONST Null type=kotlin.Nothing? value='null'
-                  if: CALL .EQEQ type=kotlin.Boolean operator=EQEQ
+                then: CONST Null type=kotlin.Nothing? value='null'
-                    arg0: GET_VAR tmp0_safe_receiver type=C? operator=null
+                else: CALL .bar type=C? operator=null
-                    arg1: CONST Null type=kotlin.Nothing? value='null'
+                  $this: CALL .foo type=C operator=null
-                  then: CONST Null type=kotlin.Nothing? value='null'
+                    $this: GET_VAR nc type=C? operator=null
-                  else: CALL .bar type=C? operator=null
-                    $this: CALL .foo type=C operator=null
-                      $this: GET_VAR tmp0_safe_receiver type=C? operator=null
             WHEN type=C? operator=SAFE_CALL
               if: CALL .EQEQ type=kotlin.Boolean operator=EQEQ
                 arg0: GET_VAR tmp2_safe_receiver type=C? operator=null

compiler/testData/ir/irText/dotQualified.txt

@@ -9,13 +9,10 @@ IrFile /dotQualified.kt
     IrExpressionBody
       BLOCK type=<no-type> hasResult=false operator=null
         RETURN type=<no-type>
-          BLOCK type=kotlin.Int? hasResult=true operator=SAFE_CALL
+          WHEN type=kotlin.Int? operator=SAFE_CALL
-            VAR val tmp0_safe_receiver: kotlin.String?
+            if: CALL .EQEQ type=kotlin.Boolean operator=EQEQ
-              GET_VAR s type=kotlin.String? operator=null
+              arg0: GET_VAR s type=kotlin.String? operator=null
-            WHEN type=kotlin.Int? operator=SAFE_CALL
+              arg1: CONST Null type=kotlin.Nothing? value='null'
-              if: CALL .EQEQ type=kotlin.Boolean operator=EQEQ
+            then: CONST Null type=kotlin.Nothing? value='null'
-                arg0: GET_VAR tmp0_safe_receiver type=kotlin.String? operator=null
+            else: CALL .<get-length> type=kotlin.Int operator=GET_PROPERTY
-                arg1: CONST Null type=kotlin.Nothing? value='null'
+              $this: GET_VAR s type=kotlin.String? operator=null
-              then: CONST Null type=kotlin.Nothing? value='null'
-              else: CALL .<get-length> type=kotlin.Int operator=GET_PROPERTY
-                $this: GET_VAR tmp0_safe_receiver type=kotlin.String? operator=null

compiler/testData/ir/irText/safeCallWithIncrementDecrement.txt

@@ -13,16 +13,13 @@ IrFile /safeCallWithIncrementDecrement.kt
     IrExpressionBody
       BLOCK type=<no-type> hasResult=false operator=null
         RETURN type=<no-type>
-          BLOCK type=kotlin.Int? hasResult=true operator=SAFE_CALL
+          WHEN type=kotlin.Int? operator=SAFE_CALL
-            VAR val tmp0_safe_receiver: kotlin.Int?
+            if: CALL .EQEQ type=kotlin.Boolean operator=EQEQ
-              $RECEIVER of: inc type=kotlin.Int?
+              arg0: $RECEIVER of: inc type=kotlin.Int?
-            WHEN type=kotlin.Int? operator=SAFE_CALL
+              arg1: CONST Null type=kotlin.Nothing? value='null'
-              if: CALL .EQEQ type=kotlin.Boolean operator=EQEQ
+            then: CONST Null type=kotlin.Nothing? value='null'
-                arg0: GET_VAR tmp0_safe_receiver type=kotlin.Int? operator=null
+            else: CALL .inc type=kotlin.Int operator=null
-                arg1: CONST Null type=kotlin.Nothing? value='null'
+              $this: $RECEIVER of: inc type=kotlin.Int?
-              then: CONST Null type=kotlin.Nothing? value='null'
-              else: CALL .inc type=kotlin.Int operator=null
-                $this: GET_VAR tmp0_safe_receiver type=kotlin.Int? operator=null
   IrFunction public operator fun kotlin.Int?.get(/*0*/ index: kotlin.Int): kotlin.Int
     IrExpressionBody
       BLOCK type=<no-type> hasResult=false operator=null
@@ -34,38 +31,32 @@ IrFile /safeCallWithIncrementDecrement.kt
   IrFunction public fun testProperty(/*0*/ nc: test.C?): kotlin.Unit
     IrExpressionBody
       BLOCK type=<no-type> hasResult=false operator=null
-        BLOCK type=kotlin.Int? hasResult=true operator=SAFE_CALL
+        WHEN type=kotlin.Int? operator=SAFE_CALL
-          VAR val tmp0_safe_receiver: test.C?
+          if: CALL .EQEQ type=kotlin.Boolean operator=EQEQ
-            GET_VAR nc type=test.C? operator=null
+            arg0: GET_VAR nc type=test.C? operator=null
-          WHEN type=kotlin.Int? operator=SAFE_CALL
+            arg1: CONST Null type=kotlin.Nothing? value='null'
-            if: CALL .EQEQ type=kotlin.Boolean operator=EQEQ
+          then: CONST Null type=kotlin.Nothing? value='null'
-              arg0: GET_VAR tmp0_safe_receiver type=test.C? operator=null
+          else: BLOCK type=kotlin.Int hasResult=true operator=POSTFIX_INCR
-              arg1: CONST Null type=kotlin.Nothing? value='null'
+            VAR val tmp1: kotlin.Int
-            then: CONST Null type=kotlin.Nothing? value='null'
+              CALL .<get-p> type=kotlin.Int operator=POSTFIX_INCR
-            else: BLOCK type=kotlin.Int hasResult=true operator=POSTFIX_INCR
+                $this: GET_VAR nc type=test.C? operator=null
-              VAR val tmp1: kotlin.Int
+            CALL .<set-p> type=kotlin.Unit operator=POSTFIX_INCR
-                CALL .<get-p> type=kotlin.Int operator=POSTFIX_INCR
+              $this: GET_VAR nc type=test.C? operator=null
-                  $this: GET_VAR tmp0_safe_receiver type=test.C? operator=null
+              value: CALL .inc type=kotlin.Int? operator=POSTFIX_INCR
-              CALL .<set-p> type=kotlin.Unit operator=POSTFIX_INCR
+                $receiver: GET_VAR tmp1 type=kotlin.Int operator=null
-                $this: GET_VAR tmp0_safe_receiver type=test.C? operator=null
+            GET_VAR tmp1 type=kotlin.Int operator=null
-                value: CALL .inc type=kotlin.Int? operator=POSTFIX_INCR
-                  $receiver: GET_VAR tmp1 type=kotlin.Int operator=null
-              GET_VAR tmp1 type=kotlin.Int operator=null
   IrFunction public fun testArrayAccess(/*0*/ nc: test.C?): kotlin.Unit
     IrExpressionBody
       BLOCK type=<no-type> hasResult=false operator=null
         BLOCK type=kotlin.Int hasResult=true operator=POSTFIX_INCR
           VAR val tmp3_array: kotlin.Int?
-            BLOCK type=kotlin.Int? hasResult=true operator=SAFE_CALL
+            WHEN type=kotlin.Int? operator=SAFE_CALL
-              VAR val tmp0_safe_receiver: test.C?
+              if: CALL .EQEQ type=kotlin.Boolean operator=EQEQ
-                GET_VAR nc type=test.C? operator=null
+                arg0: GET_VAR nc type=test.C? operator=null
-              WHEN type=kotlin.Int? operator=SAFE_CALL
+                arg1: CONST Null type=kotlin.Nothing? value='null'
-                if: CALL .EQEQ type=kotlin.Boolean operator=EQEQ
+              then: CONST Null type=kotlin.Nothing? value='null'
-                  arg0: GET_VAR tmp0_safe_receiver type=test.C? operator=null
+              else: CALL .<get-p> type=kotlin.Int operator=GET_PROPERTY
-                  arg1: CONST Null type=kotlin.Nothing? value='null'
+                $this: GET_VAR nc type=test.C? operator=null
-                then: CONST Null type=kotlin.Nothing? value='null'
-                else: CALL .<get-p> type=kotlin.Int operator=GET_PROPERTY
-                  $this: GET_VAR tmp0_safe_receiver type=test.C? operator=null
           VAR val tmp4: kotlin.Int
             CALL .get type=kotlin.Int operator=POSTFIX_INCR
               $receiver: GET_VAR tmp3_array type=kotlin.Int? operator=null

compiler/testData/ir/irText/safeCalls.txt

@@ -5,71 +5,56 @@ IrFile /safeCalls.kt
     IrExpressionBody
       BLOCK type=<no-type> hasResult=false operator=null
         RETURN type=<no-type>
-          BLOCK type=kotlin.Int? hasResult=true operator=SAFE_CALL
+          WHEN type=kotlin.Int? operator=SAFE_CALL
-            VAR val tmp0_safe_receiver: kotlin.String?
+            if: CALL .EQEQ type=kotlin.Boolean operator=EQEQ
-              GET_VAR x type=kotlin.String? operator=null
+              arg0: GET_VAR x type=kotlin.String? operator=null
-            WHEN type=kotlin.Int? operator=SAFE_CALL
+              arg1: CONST Null type=kotlin.Nothing? value='null'
-              if: CALL .EQEQ type=kotlin.Boolean operator=EQEQ
+            then: CONST Null type=kotlin.Nothing? value='null'
-                arg0: GET_VAR tmp0_safe_receiver type=kotlin.String? operator=null
+            else: CALL .<get-length> type=kotlin.Int operator=GET_PROPERTY
-                arg1: CONST Null type=kotlin.Nothing? value='null'
+              $this: GET_VAR x type=kotlin.String? operator=null
-              then: CONST Null type=kotlin.Nothing? value='null'
-              else: CALL .<get-length> type=kotlin.Int operator=GET_PROPERTY
-                $this: GET_VAR tmp0_safe_receiver type=kotlin.String? operator=null
   IrFunction public fun test2(/*0*/ x: kotlin.String?): kotlin.Int?
     IrExpressionBody
       BLOCK type=<no-type> hasResult=false operator=null
         RETURN type=<no-type>
-          BLOCK type=kotlin.Int? hasResult=true operator=SAFE_CALL
+          WHEN type=kotlin.Int? operator=SAFE_CALL
-            VAR val tmp0_safe_receiver: kotlin.String?
+            if: CALL .EQEQ type=kotlin.Boolean operator=EQEQ
-              GET_VAR x type=kotlin.String? operator=null
+              arg0: GET_VAR x type=kotlin.String? operator=null
-            WHEN type=kotlin.Int? operator=SAFE_CALL
+              arg1: CONST Null type=kotlin.Nothing? value='null'
-              if: CALL .EQEQ type=kotlin.Boolean operator=EQEQ
+            then: CONST Null type=kotlin.Nothing? value='null'
-                arg0: GET_VAR tmp0_safe_receiver type=kotlin.String? operator=null
+            else: CALL .hashCode type=kotlin.Int operator=null
-                arg1: CONST Null type=kotlin.Nothing? value='null'
+              $this: GET_VAR x type=kotlin.String? operator=null
-              then: CONST Null type=kotlin.Nothing? value='null'
-              else: CALL .hashCode type=kotlin.Int operator=null
-                $this: GET_VAR tmp0_safe_receiver type=kotlin.String? operator=null
   IrFunction public fun test3(/*0*/ x: kotlin.String?, /*1*/ y: kotlin.Any?): kotlin.Boolean?
     IrExpressionBody
       BLOCK type=<no-type> hasResult=false operator=null
         RETURN type=<no-type>
-          BLOCK type=kotlin.Boolean? hasResult=true operator=SAFE_CALL
+          WHEN type=kotlin.Boolean? operator=SAFE_CALL
-            VAR val tmp0_safe_receiver: kotlin.String?
+            if: CALL .EQEQ type=kotlin.Boolean operator=EQEQ
-              GET_VAR x type=kotlin.String? operator=null
+              arg0: GET_VAR x type=kotlin.String? operator=null
-            WHEN type=kotlin.Boolean? operator=SAFE_CALL
+              arg1: CONST Null type=kotlin.Nothing? value='null'
-              if: CALL .EQEQ type=kotlin.Boolean operator=EQEQ
+            then: CONST Null type=kotlin.Nothing? value='null'
-                arg0: GET_VAR tmp0_safe_receiver type=kotlin.String? operator=null
+            else: CALL .equals type=kotlin.Boolean operator=null
-                arg1: CONST Null type=kotlin.Nothing? value='null'
+              $this: GET_VAR x type=kotlin.String? operator=null
-              then: CONST Null type=kotlin.Nothing? value='null'
+              other: GET_VAR y type=kotlin.Any? operator=null
-              else: CALL .equals type=kotlin.Boolean operator=null
-                $this: GET_VAR tmp0_safe_receiver type=kotlin.String? operator=null
-                other: GET_VAR y type=kotlin.Any? operator=null
   IrFunction public fun test4(/*0*/ x: Ref?): kotlin.Unit
     IrExpressionBody
       BLOCK type=<no-type> hasResult=false operator=null
-        BLOCK type=kotlin.Unit? hasResult=true operator=SAFE_CALL
+        WHEN type=kotlin.Unit? operator=SAFE_CALL
-          VAR val tmp0_safe_receiver: Ref?
+          if: CALL .EQEQ type=kotlin.Boolean operator=EQEQ
-            GET_VAR x type=Ref? operator=null
+            arg0: GET_VAR x type=Ref? operator=null
-          WHEN type=kotlin.Unit? operator=SAFE_CALL
+            arg1: CONST Null type=kotlin.Nothing? value='null'
-            if: CALL .EQEQ type=kotlin.Boolean operator=EQEQ
+          then: CONST Null type=kotlin.Nothing? value='null'
-              arg0: GET_VAR tmp0_safe_receiver type=Ref? operator=null
+          else: CALL .<set-value> type=kotlin.Unit operator=EQ
-              arg1: CONST Null type=kotlin.Nothing? value='null'
+            $this: GET_VAR x type=Ref? operator=null
-            then: CONST Null type=kotlin.Nothing? value='null'
+            <set-?>: CONST Int type=kotlin.Int value='0'
-            else: CALL .<set-value> type=kotlin.Unit operator=EQ
-              $this: GET_VAR tmp0_safe_receiver type=Ref? operator=null
-              <set-?>: CONST Int type=kotlin.Int value='0'
   IrFunction public fun IHost.test5(/*0*/ s: kotlin.String?): kotlin.Int?
     IrExpressionBody
       BLOCK type=<no-type> hasResult=false operator=null
         RETURN type=<no-type>
-          BLOCK type=kotlin.Int? hasResult=true operator=SAFE_CALL
+          WHEN type=kotlin.Int? operator=SAFE_CALL
-            VAR val tmp0_safe_receiver: kotlin.String?
+            if: CALL .EQEQ type=kotlin.Boolean operator=EQEQ
-              GET_VAR s type=kotlin.String? operator=null
+              arg0: GET_VAR s type=kotlin.String? operator=null
-            WHEN type=kotlin.Int? operator=SAFE_CALL
+              arg1: CONST Null type=kotlin.Nothing? value='null'
-              if: CALL .EQEQ type=kotlin.Boolean operator=EQEQ
+            then: CONST Null type=kotlin.Nothing? value='null'
-                arg0: GET_VAR tmp0_safe_receiver type=kotlin.String? operator=null
+            else: CALL .extLength type=kotlin.Int operator=null
-                arg1: CONST Null type=kotlin.Nothing? value='null'
+              $this: $RECEIVER of: test5 type=IHost
-              then: CONST Null type=kotlin.Nothing? value='null'
+              $receiver: GET_VAR s type=kotlin.String? operator=null
-              else: CALL .extLength type=kotlin.Int operator=null
-                $this: $RECEIVER of: test5 type=IHost
-                $receiver: GET_VAR tmp0_safe_receiver type=kotlin.String? operator=null

compiler/testData/ir/irText/variableAsFunctionCall.txt

@@ -28,14 +28,11 @@ IrFile /variableAsFunctionCall.kt
     IrExpressionBody
       BLOCK type=<no-type> hasResult=false operator=null
         RETURN type=<no-type>
-          BLOCK type=kotlin.String? hasResult=true operator=SAFE_CALL
+          WHEN type=kotlin.String? operator=SAFE_CALL
-            VAR val tmp0_safe_receiver: kotlin.String?
+            if: CALL .EQEQ type=kotlin.Boolean operator=EQEQ
-              GET_VAR ns type=kotlin.String? operator=null
+              arg0: GET_VAR ns type=kotlin.String? operator=null
-            WHEN type=kotlin.String? operator=SAFE_CALL
+              arg1: CONST Null type=kotlin.Nothing? value='null'
-              if: CALL .EQEQ type=kotlin.Boolean operator=EQEQ
+            then: CONST Null type=kotlin.Nothing? value='null'
-                arg0: GET_VAR tmp0_safe_receiver type=kotlin.String? operator=null
+            else: CALL .invoke type=kotlin.String operator=null
-                arg1: CONST Null type=kotlin.Nothing? value='null'
+              $this: CALL .k type=() -> kotlin.String operator=null
-              then: CONST Null type=kotlin.Nothing? value='null'
+                $this: GET_VAR ns type=kotlin.String? operator=null
-              else: CALL .invoke type=kotlin.String operator=null
-                $this: CALL .k type=() -> kotlin.String operator=null
-                  $this: GET_VAR tmp0_safe_receiver type=kotlin.String? operator=null