Support lateinit local vars in redundant null check elimination

Lateinit local vars are guaranteed to be non-null after store.
So we mark such stores as storing non-null value
(could be useful for some other constructs, too),
and optimize null checks accordingly.

compiler/backend/src/org/jetbrains/kotlin/codegen/StackValue.java

@@ -26,6 +26,7 @@ import org.jetbrains.annotations.Nullable;
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns;
 import org.jetbrains.kotlin.builtins.PrimitiveType;
 import org.jetbrains.kotlin.codegen.intrinsics.IntrinsicMethods;
+import org.jetbrains.kotlin.codegen.pseudoInsns.PseudoInsnsKt;
 import org.jetbrains.kotlin.codegen.state.KotlinTypeMapper;
 import org.jetbrains.kotlin.descriptors.*;
 import org.jetbrains.kotlin.descriptors.impl.SyntheticFieldDescriptor;
@@ -672,6 +673,9 @@ public abstract class StackValue {
         public void storeSelector(@NotNull Type topOfStackType, @NotNull InstructionAdapter v) {
             coerceFrom(topOfStackType, v);
             v.store(index, this.type);
+            if (isLateinit) {
+                PseudoInsnsKt.storeNotNull(v);
+            }
         }
     }
 

compiler/backend/src/org/jetbrains/kotlin/codegen/optimization/common/Util.kt

@@ -17,6 +17,8 @@
 package org.jetbrains.kotlin.codegen.optimization.common
 
 import org.jetbrains.kotlin.codegen.inline.insnText
+import org.jetbrains.kotlin.codegen.optimization.removeNodeGetNext
+import org.jetbrains.kotlin.codegen.pseudoInsns.PseudoInsn
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 import org.jetbrains.org.objectweb.asm.Opcodes
 import org.jetbrains.org.objectweb.asm.Opcodes.*
@@ -48,6 +50,8 @@ class InsnSequence(val from: AbstractInsnNode, val to: AbstractInsnNode?) : Sequ
 fun InsnList.asSequence() = InsnSequence(this)
 
 fun MethodNode.prepareForEmitting() {
+    stripOptimizationMarkers()
+
     removeEmptyCatchBlocks()
 
     // local variables with live ranges starting after last meaningful instruction lead to VerifyError
@@ -69,6 +73,21 @@ fun MethodNode.prepareForEmitting() {
     }
 }
 
+fun MethodNode.stripOptimizationMarkers() {
+    var insn = instructions.first
+    while (insn != null) {
+        if (isOptimizationMarker(insn)) {
+            insn = instructions.removeNodeGetNext(insn)
+        }
+        else {
+            insn = insn.next
+        }
+    }
+}
+
+private fun isOptimizationMarker(insn: AbstractInsnNode) =
+        PseudoInsn.STORE_NOT_NULL.isa(insn)
+
 fun MethodNode.removeEmptyCatchBlocks() {
     tryCatchBlocks = tryCatchBlocks.filter { tcb ->
         InsnSequence(tcb.start, tcb.end).any(AbstractInsnNode::isMeaningful)

compiler/backend/src/org/jetbrains/kotlin/codegen/optimization/nullCheck/RedundantNullCheckMethodTransformer.kt

@@ -27,8 +27,11 @@ import org.jetbrains.kotlin.codegen.optimization.common.isInsn
 import org.jetbrains.kotlin.codegen.optimization.fixStack.peek
 import org.jetbrains.kotlin.codegen.optimization.fixStack.top
 import org.jetbrains.kotlin.codegen.optimization.transformer.MethodTransformer
+import org.jetbrains.kotlin.codegen.pseudoInsns.PseudoInsn
+import org.jetbrains.kotlin.codegen.pseudoInsns.isPseudo
 import org.jetbrains.kotlin.resolve.jvm.AsmTypes
 import org.jetbrains.kotlin.utils.SmartList
+import org.jetbrains.kotlin.utils.addToStdlib.cast
 import org.jetbrains.org.objectweb.asm.Label
 import org.jetbrains.org.objectweb.asm.Opcodes
 import org.jetbrains.org.objectweb.asm.Type
@@ -55,18 +58,24 @@ class RedundantNullCheckMethodTransformer : MethodTransformer() {
             val insns = methodNode.instructions.toArray()
             val frames = analyze(internalClassName, methodNode, OptimizationBasicInterpreter())
 
-            val checkedReferenceTypes = HashMap<AbstractInsnNode, Type>()
+            val relevantReferenceTypes = HashMap<AbstractInsnNode, Type>()
-            for (i in insns.indices) {
+            insnLoop@ for (i in insns.indices) {
                 val insn = insns[i]
-                val frame = frames[i]
+                val frame = frames[i] ?: continue
-                if (insn.isInstanceOfOrNullCheck()) {
+                when {
-                    checkedReferenceTypes[insn] = frame?.top()?.type ?: continue
+                    insn.isInstanceOfOrNullCheck() ->
-                }
+                        relevantReferenceTypes[insn] = frame.top()?.type ?: continue@insnLoop
-                else if (insn.isCheckParameterIsNotNull() || insn.isCheckExpressionValueIsNotNull()) {
+                    insn.isCheckParameterIsNotNull() ||
-                    checkedReferenceTypes[insn] = frame?.peek(1)?.type ?: continue
+                    insn.isCheckExpressionValueIsNotNull() ->
-                }
+                        relevantReferenceTypes[insn] = frame.peek(1)?.type ?: continue@insnLoop
-                else if (insn.isThrowNpeIntrinsic()) {
+                    insn.isThrowIntrinsicWithoutArguments() ->
-                    stackOnThrowExceptionsHolder[insn] = frame?.maxStackSize ?: continue
+                        stackOnThrowExceptionsHolder[insn] = frame.maxStackSize
+                    insn.isPseudo(PseudoInsn.STORE_NOT_NULL) -> {
+                        val previous = insn.previous ?: continue@insnLoop
+                        if (previous.opcode != Opcodes.ASTORE) continue@insnLoop
+                        val previousFrame = frames[i - 1] ?: continue@insnLoop
+                        relevantReferenceTypes[insn] = previousFrame.top()?.type ?: continue@insnLoop
+                    }
                 }
             }
 
@@ -75,7 +84,7 @@ class RedundantNullCheckMethodTransformer : MethodTransformer() {
                 changes = true
             }
 
-            return checkedReferenceTypes
+            return relevantReferenceTypes
         }
 
         private fun eliminateRedundantChecks(
@@ -178,7 +187,7 @@ class RedundantNullCheckMethodTransformer : MethodTransformer() {
         }
 
         private inner class NullabilityAssumptionsBuilder(
-                val checkedReferenceTypes: Map<AbstractInsnNode, Type>,
+                val relatedReferenceTypes: Map<AbstractInsnNode, Type>,
                 val stackOnThrowExceptions: MutableMap<AbstractInsnNode, Int>
         ) {
 
@@ -230,7 +239,6 @@ class RedundantNullCheckMethodTransformer : MethodTransformer() {
                             if (aLoadInsn == null) continue@insnLoop
                             addDependentCheck(insn, aLoadInsn)
                         }
-
                     }
                 }
             }
@@ -245,29 +253,34 @@ class RedundantNullCheckMethodTransformer : MethodTransformer() {
                 val nullabilityAssumptions = NullabilityAssumptions()
                 for ((varIndex, dependentChecks) in checksDependingOnVariable) {
                     for (checkInsn in dependentChecks) {
-                        val varType = checkedReferenceTypes[checkInsn]
+                        val varType = relatedReferenceTypes[checkInsn]
                                       ?: AsmTypes.OBJECT_TYPE
-                        nullabilityAssumptions.injectAssumptionsForCheck(varIndex, checkInsn, varType)
+                        nullabilityAssumptions.injectAssumptionsForInsn(varIndex, checkInsn, varType)
                     }
                 }
                 for (insn in methodNode.instructions) {
-                    if (insn.isThrowNpeIntrinsic()) {
+                    if (insn.isThrowIntrinsic()) {
-                        nullabilityAssumptions.injectCodeForThrowNpe(insn)
+                        nullabilityAssumptions.injectCodeForThrowIntrinsic(insn)
                     }
                 }
                 return nullabilityAssumptions
             }
 
-            private fun NullabilityAssumptions.injectAssumptionsForCheck(varIndex: Int, insn: AbstractInsnNode, varType: Type) {
+            private fun NullabilityAssumptions.injectAssumptionsForInsn(varIndex: Int, insn: AbstractInsnNode, varType: Type) {
                 when (insn.opcode) {
                     Opcodes.IFNULL,
                     Opcodes.IFNONNULL ->
                         injectAssumptionsForNullCheck(varIndex, insn as JumpInsnNode, varType)
                     Opcodes.INVOKESTATIC -> {
-                        assert(insn.isCheckParameterIsNotNull() || insn.isCheckExpressionValueIsNotNull()) {
+                        when {
-                            "Expected non-null assertion: ${insn.debugText}"
+                            insn.isCheckParameterIsNotNull() ||
+                            insn.isCheckExpressionValueIsNotNull() ->
+                                injectAssumptionsForNotNullAssertion(varIndex, insn, varType)
+                            insn.isPseudo(PseudoInsn.STORE_NOT_NULL) ->
+                                injectCodeForStoreNotNull(insn, varType)
+                            else ->
+                                throw AssertionError("Expected non-null assertion: ${insn.debugText}")
                         }
-                        injectAssumptionsForNotNullAssertion(varIndex, insn, varType)
                     }
                     Opcodes.INSTANCEOF ->
                         injectAssumptionsForInstanceOfCheck(varIndex, insn, varType)
@@ -367,7 +380,7 @@ class RedundantNullCheckMethodTransformer : MethodTransformer() {
                 }
             }
 
-            private fun NullabilityAssumptions.injectCodeForThrowNpe(insn: AbstractInsnNode) {
+            private fun NullabilityAssumptions.injectCodeForThrowIntrinsic(insn: AbstractInsnNode) {
                 methodNode.instructions.run {
                     insert(insn, listOfSynthetics {
                         aconst(null)
@@ -378,6 +391,21 @@ class RedundantNullCheckMethodTransformer : MethodTransformer() {
                 methodNode.maxStack = Math.max(methodNode.maxStack, (stackOnThrowExceptions[insn] ?: -1) + 1)
             }
 
+            private fun NullabilityAssumptions.injectCodeForStoreNotNull(insn: AbstractInsnNode, varType: Type) {
+                // ASTORE v
+                // [STORE_NOT_NULL]
+                // <...>    -- v is not null here because codegen told us so
+                val previous = insn.previous
+                if (previous.opcode != Opcodes.ASTORE) return
+                val varIndex = previous.cast<VarInsnNode>().`var`
+
+                methodNode.instructions.run {
+                    insert(insn, listOfSynthetics {
+                        anew(varType)
+                        store(varIndex, varType)
+                    })
+                }
+            }
         }
 
         inner class NullabilityAssumptions {
@@ -428,13 +456,31 @@ internal fun AbstractInsnNode.isCheckExpressionValueIsNotNull() =
             desc == "(Ljava/lang/Object;Ljava/lang/String;)V"
         }
 
-internal fun AbstractInsnNode.isThrowNpeIntrinsic() =
+internal fun AbstractInsnNode.isThrowIntrinsic() =
         isInsn<MethodInsnNode>(Opcodes.INVOKESTATIC) {
             owner == IntrinsicMethods.INTRINSICS_CLASS_NAME &&
-            name == "throwNpe" &&
+            name in THROW_INTRINSIC_METHOD_NAMES
+        }
+
+internal fun AbstractInsnNode.isThrowIntrinsicWithoutArguments() =
+        isInsn<MethodInsnNode>(Opcodes.INVOKESTATIC) {
+            owner == IntrinsicMethods.INTRINSICS_CLASS_NAME &&
+            name in THROW_INTRINSIC_METHOD_NAMES &&
             desc == "()V"
         }
 
+internal val THROW_INTRINSIC_METHOD_NAMES =
+        setOf(
+                "throwNpe",
+                "throwUninitializedProperty",
+                "throwUninitializedPropertyAccessException",
+                "throwAssert",
+                "throwIllegalArgument",
+                "throwIllegalState",
+                "throwParameterIsNullException",
+                "throwUndefinedForReified"
+        )
+
 internal fun InsnList.popReferenceValueBefore(insn: AbstractInsnNode) {
     val prev = insn.previous
     when (prev?.opcode) {

compiler/backend/src/org/jetbrains/kotlin/codegen/pseudoInsns/PseudoInsns.kt

@@ -26,11 +26,12 @@ import org.jetbrains.org.objectweb.asm.tree.MethodInsnNode
 val PSEUDO_INSN_CALL_OWNER: String = "kotlin/jvm/internal/\$PseudoInsn"
 
 enum class PseudoInsn(val signature: String = "()V") {
-    FIX_STACK_BEFORE_JUMP(),
+    FIX_STACK_BEFORE_JUMP,
     FAKE_ALWAYS_TRUE_IFEQ("()I"),
     FAKE_ALWAYS_FALSE_IFEQ("()I"),
-    SAVE_STACK_BEFORE_TRY(),
+    SAVE_STACK_BEFORE_TRY,
-    RESTORE_STACK_IN_TRY_CATCH()
+    RESTORE_STACK_IN_TRY_CATCH,
+    STORE_NOT_NULL
     ;
 
     fun emit(iv: InstructionAdapter) {
@@ -66,3 +67,10 @@ fun InstructionAdapter.fakeAlwaysFalseIfeq(label: Label) {
     PseudoInsn.FAKE_ALWAYS_FALSE_IFEQ.emit(this)
     this.ifeq(label)
 }
+
+fun InstructionAdapter.storeNotNull() {
+    PseudoInsn.STORE_NOT_NULL.emit(this)
+}
+
+fun AbstractInsnNode.isPseudo(pseudoInsn: PseudoInsn) =
+        pseudoInsn.isa(this)

compiler/testData/codegen/bytecodeText/nullCheckOptimization/localLateinit/checkedAlways.kt

@@ -0,0 +1,26 @@
+// LANGUAGE_VERSION: 1.2
+
+fun almostAlwaysTrue() = true
+
+fun runNoInline(f: () -> Unit) = f()
+
+fun test() {
+    lateinit var z: String
+
+    runNoInline {
+        // NB this code can be executed in a different thread multiple times, each time with different results.
+        // So, 'z' can be initialized at any moment, and should be checked on every read.
+
+        if (almostAlwaysTrue()) {
+            z = ""
+        }
+    }
+
+    println(z)
+    println(z)
+    println(z)
+}
+
+// 0 IFNULL
+// 3 IFNONNULL
+// 3 throwUninitializedPropertyAccessException

compiler/testData/codegen/bytecodeText/nullCheckOptimization/localLateinit/checkedOnce.kt

@@ -0,0 +1,19 @@
+// LANGUAGE_VERSION: 1.2
+
+fun almostAlwaysTrue() = true
+
+fun test() {
+    lateinit var z: String
+    run {
+        if (almostAlwaysTrue()) {
+            z = ""
+        }
+    }
+    println(z)
+    println(z)
+    println(z)
+}
+
+// 0 IFNULL
+// 1 IFNONNULL
+// 1 throwUninitializedPropertyAccessException

compiler/testData/codegen/bytecodeText/nullCheckOptimization/localLateinit/initialized.kt

@@ -0,0 +1,12 @@
+// LANGUAGE_VERSION: 1.2
+
+fun test() {
+    lateinit var z: String
+    run {
+        z = ""
+    }
+    println(z)
+}
+
+// 0 IFNULL
+// 0 IFNONNULL

compiler/tests/org/jetbrains/kotlin/codegen/BytecodeTextTestGenerated.java

@@ -1891,6 +1891,33 @@ public class BytecodeTextTestGenerated extends AbstractBytecodeTextTest {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/bytecodeText/nullCheckOptimization/trivialInstanceOf.kt");
             doTest(fileName);
         }
+
+        @TestMetadata("compiler/testData/codegen/bytecodeText/nullCheckOptimization/localLateinit")
+        @TestDataPath("$PROJECT_ROOT")
+        @RunWith(JUnit3RunnerWithInners.class)
+        public static class LocalLateinit extends AbstractBytecodeTextTest {
+            public void testAllFilesPresentInLocalLateinit() throws Exception {
+                KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("compiler/testData/codegen/bytecodeText/nullCheckOptimization/localLateinit"), Pattern.compile("^(.+)\\.kt$"), TargetBackend.ANY, true);
+            }
+
+            @TestMetadata("checkedAlways.kt")
+            public void testCheckedAlways() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/bytecodeText/nullCheckOptimization/localLateinit/checkedAlways.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("checkedOnce.kt")
+            public void testCheckedOnce() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/bytecodeText/nullCheckOptimization/localLateinit/checkedOnce.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("initialized.kt")
+            public void testInitialized() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/bytecodeText/nullCheckOptimization/localLateinit/initialized.kt");
+                doTest(fileName);
+            }
+        }
     }
 
     @TestMetadata("compiler/testData/codegen/bytecodeText/ranges")