Move lateinit assertion for companion property to companion object

Previously, for a property named `x` in the companion object of a class
named `Foo`, we generated:
- `Foo.access$getX$cp`, consisting of `GETFIELD Foo.x` and lateinit
  assertion
- `Foo.Companion.getX`, consisting of `INVOKEVIRTUAL Foo.access$getX$cp`

Now, we generate:
- `Foo.access$getX$cp`, consisting of `GETFIELD Foo.x`
- `Foo.Companion.getX`, consisting of `INVOKEVIRTUAL Foo.access$getX$cp`
  and lateinit assertion

The reason is that this way we can avoid generating another accessor and
reuse `Foo.access$getX$cp` in case `isInitialized` is called on a
lateinit property from companion.

For private properties, getX is not generated, but instead the assertion
is generated on each access to the field (which can be improved, see
KT-28331). The same happens for access to non-private properties from
inside the same context where they're declared.

 #KT-21862 In Progress

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

@@ -1660,8 +1660,14 @@ public abstract class StackValue {
 
                 v.visitFieldInsn(isStaticPut ? GETSTATIC : GETFIELD,
                                  backingFieldOwner.getInternalName(), fieldName, this.type.getDescriptor());
-                if (!skipLateinitAssertion) {
-                    genNotNullAssertionForLateInitIfNeeded(v);
+                if (!skipLateinitAssertion && descriptor.isLateInit()) {
+                    CallableMemberDescriptor contextDescriptor = codegen.context.getContextDescriptor();
+                    boolean isCompanionAccessor =
+                            contextDescriptor instanceof AccessorForPropertyBackingField &&
+                            ((AccessorForPropertyBackingField) contextDescriptor).getAccessorKind() == AccessorKind.IN_CLASS_COMPANION;
+                    if (!isCompanionAccessor) {
+                        genNonNullAssertForLateinit(v, this.descriptor.getName().asString());
+                    }
                 }
                 coerceTo(type, kotlinType, v);
             }
@@ -1694,6 +1700,19 @@ public abstract class StackValue {
 
                 coerce(typeOfValueOnStack, kotlinTypeOfValueOnStack, type, kotlinType, v);
 
+                // For non-private lateinit properties in companion object, the assertion is generated in the public getFoo method
+                // in the companion and _not_ in the synthetic accessor access$getFoo$cp in the outer class. The reason is that this way,
+                // the synthetic accessor can be reused for isInitialized checks, which require there to be no assertion.
+                // For lateinit properties that are accessed via the backing field directly (or via the synthetic accessor, if the access
+                // is from a different context), the assertion will be generated on each access, see KT-28331.
+                if (descriptor instanceof AccessorForPropertyBackingField) {
+                    PropertyDescriptor property = ((AccessorForPropertyBackingField) descriptor).getCalleeDescriptor();
+                    if (!skipLateinitAssertion && property.isLateInit() && JvmAbi.isPropertyWithBackingFieldInOuterClass(property) &&
+                        !JvmCodegenUtil.couldUseDirectAccessToProperty(property, true, false, codegen.context, false)) {
+                        genNonNullAssertForLateinit(v, property.getName().asString());
+                    }
+                }
+
                 KotlinType returnType = descriptor.getReturnType();
                 if (returnType != null && KotlinBuiltIns.isNothing(returnType)) {
                     v.aconst(null);
@@ -1732,12 +1751,6 @@ public abstract class StackValue {
             return true;
         }
 
-        private void genNotNullAssertionForLateInitIfNeeded(@NotNull InstructionAdapter v) {
-            if (!descriptor.isLateInit()) return;
-
-            StackValue.genNonNullAssertForLateinit(v, descriptor.getName().asString());
-        }
-
         @Override
         public void store(@NotNull StackValue rightSide, @NotNull InstructionAdapter v, boolean skipReceiver) {
             PropertySetterDescriptor setterDescriptor = descriptor.getSetter();

compiler/testData/codegen/box/properties/lateinit/accessorException.kt

@@ -1,21 +1,32 @@
 public class A {
-
-    fun getMyStr(): String {
+    fun getFromClass(): Boolean {
         try {
             val a = str
+            return false
         } catch (e: RuntimeException) {
-            return "OK"
+            return true
         }
-
-        return "FAIL"
     }
 
+    fun getFromCompanion() = Companion.getFromCompanion()
+
     private companion object {
         private lateinit var str: String
+
+        fun getFromCompanion(): Boolean {
+            try {
+                val a = str
+                return false
+            } catch (e: RuntimeException) {
+                return true
+            }
+        }
     }
 }
 
 fun box(): String {
-    val a = A()
-    return a.getMyStr()
-}
+    if (!A().getFromClass()) return "Fail getFromClass"
+    if (!A().getFromCompanion()) return "Fail getFromCompanion"
+
+    return "OK"
+}

compiler/testData/codegen/box/properties/lateinit/accessorExceptionPublic.kt

@@ -0,0 +1,40 @@
+public class A {
+    fun getFromClass(): Boolean {
+        try {
+            val a = str
+            return false
+        } catch (e: RuntimeException) {
+            return true
+        }
+    }
+
+    fun getFromLambda(): Boolean {
+        try {
+            val a = { str }()
+            return false
+        } catch (e: RuntimeException) {
+            return true
+        }
+    }
+
+    companion object {
+        lateinit var str: String
+
+        fun getFromCompanion(): Boolean {
+            try {
+                val a = str
+                return false
+            } catch (e: RuntimeException) {
+                return true
+            }
+        }
+    }
+}
+
+fun box(): String {
+    if (!A().getFromClass()) return "Fail getFromClass"
+    if (!A().getFromLambda()) return "Fail getFromLambda"
+    if (!A.getFromCompanion()) return "Fail getFromCompanion"
+
+    return "OK"
+}

compiler/testData/codegen/bytecodeText/properties/lateinit/companionObject.kt

@@ -0,0 +1,24 @@
+class Foo {
+    private companion object {
+        lateinit var x: String
+
+        fun test() {
+            consume(x)
+            consume(x)
+            consume(x)
+            consume(x)
+        }
+    }
+
+    fun test2() {
+        consume(x)
+        consume(x)
+        consume(x)
+        consume(x)
+    }
+}
+
+fun consume(s: String) {}
+
+// There's 1 assertion in Foo.Companion.getX, and 4 in Foo.test2 (see KT-28331)
+// 5 throwUninitializedPropertyAccessException

compiler/testData/codegen/bytecodeText/properties/lateinit/companionObjectFromLambda.kt

@@ -0,0 +1,20 @@
+class Foo {
+    private companion object {
+        lateinit var x: String
+
+        fun test() {
+            consume({ x }());
+            { consume(x) }()
+        }
+    }
+
+    fun test2() {
+        consume({ x }());
+        { consume(x) }()
+    }
+}
+
+fun consume(s: String) {}
+
+// There's only one assertion in Foo.Companion.getX
+// 1 throwUninitializedPropertyAccessException

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

@@ -17794,6 +17794,11 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
                 runTest("compiler/testData/codegen/box/properties/lateinit/accessorException.kt");
             }
 
+            @TestMetadata("accessorExceptionPublic.kt")
+            public void testAccessorExceptionPublic() throws Exception {
+                runTest("compiler/testData/codegen/box/properties/lateinit/accessorExceptionPublic.kt");
+            }
+
             public void testAllFilesPresentInLateinit() throws Exception {
                 KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("compiler/testData/codegen/box/properties/lateinit"), Pattern.compile("^(.+)\\.kt$"), TargetBackend.JVM, true);
             }

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

@@ -2724,6 +2724,42 @@ public class BytecodeTextTestGenerated extends AbstractBytecodeTextTest {
         }
     }
 
+    @TestMetadata("compiler/testData/codegen/bytecodeText/properties")
+    @TestDataPath("$PROJECT_ROOT")
+    @RunWith(JUnit3RunnerWithInners.class)
+    public static class Properties extends AbstractBytecodeTextTest {
+        private void runTest(String testDataFilePath) throws Exception {
+            KotlinTestUtils.runTest(this::doTest, TargetBackend.ANY, testDataFilePath);
+        }
+
+        public void testAllFilesPresentInProperties() throws Exception {
+            KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("compiler/testData/codegen/bytecodeText/properties"), Pattern.compile("^(.+)\\.kt$"), TargetBackend.ANY, true);
+        }
+
+        @TestMetadata("compiler/testData/codegen/bytecodeText/properties/lateinit")
+        @TestDataPath("$PROJECT_ROOT")
+        @RunWith(JUnit3RunnerWithInners.class)
+        public static class Lateinit extends AbstractBytecodeTextTest {
+            private void runTest(String testDataFilePath) throws Exception {
+                KotlinTestUtils.runTest(this::doTest, TargetBackend.ANY, testDataFilePath);
+            }
+
+            public void testAllFilesPresentInLateinit() throws Exception {
+                KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("compiler/testData/codegen/bytecodeText/properties/lateinit"), Pattern.compile("^(.+)\\.kt$"), TargetBackend.ANY, true);
+            }
+
+            @TestMetadata("companionObject.kt")
+            public void testCompanionObject() throws Exception {
+                runTest("compiler/testData/codegen/bytecodeText/properties/lateinit/companionObject.kt");
+            }
+
+            @TestMetadata("companionObjectFromLambda.kt")
+            public void testCompanionObjectFromLambda() throws Exception {
+                runTest("compiler/testData/codegen/bytecodeText/properties/lateinit/companionObjectFromLambda.kt");
+            }
+        }
+    }
+
     @TestMetadata("compiler/testData/codegen/bytecodeText/ranges")
     @TestDataPath("$PROJECT_ROOT")
     @RunWith(JUnit3RunnerWithInners.class)

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

@@ -17794,6 +17794,11 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
                 runTest("compiler/testData/codegen/box/properties/lateinit/accessorException.kt");
             }
 
+            @TestMetadata("accessorExceptionPublic.kt")
+            public void testAccessorExceptionPublic() throws Exception {
+                runTest("compiler/testData/codegen/box/properties/lateinit/accessorExceptionPublic.kt");
+            }
+
             public void testAllFilesPresentInLateinit() throws Exception {
                 KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("compiler/testData/codegen/box/properties/lateinit"), Pattern.compile("^(.+)\\.kt$"), TargetBackend.JVM, true);
             }

compiler/tests/org/jetbrains/kotlin/codegen/ir/IrBlackBoxCodegenTestGenerated.java

@@ -17799,6 +17799,11 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
                 runTest("compiler/testData/codegen/box/properties/lateinit/accessorException.kt");
             }
 
+            @TestMetadata("accessorExceptionPublic.kt")
+            public void testAccessorExceptionPublic() throws Exception {
+                runTest("compiler/testData/codegen/box/properties/lateinit/accessorExceptionPublic.kt");
+            }
+
             public void testAllFilesPresentInLateinit() throws Exception {
                 KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("compiler/testData/codegen/box/properties/lateinit"), Pattern.compile("^(.+)\\.kt$"), TargetBackend.JVM_IR, true);
             }

js/js.tests/test/org/jetbrains/kotlin/js/test/semantics/IrJsCodegenBoxTestGenerated.java

@@ -15289,6 +15289,11 @@ public class IrJsCodegenBoxTestGenerated extends AbstractIrJsCodegenBoxTest {
                 runTest("compiler/testData/codegen/box/properties/lateinit/accessorException.kt");
             }
 
+            @TestMetadata("accessorExceptionPublic.kt")
+            public void testAccessorExceptionPublic() throws Exception {
+                runTest("compiler/testData/codegen/box/properties/lateinit/accessorExceptionPublic.kt");
+            }
+
             public void testAllFilesPresentInLateinit() throws Exception {
                 KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("compiler/testData/codegen/box/properties/lateinit"), Pattern.compile("^(.+)\\.kt$"), TargetBackend.JS_IR, true);
             }

js/js.tests/test/org/jetbrains/kotlin/js/test/semantics/JsCodegenBoxTestGenerated.java

@@ -16349,6 +16349,11 @@ public class JsCodegenBoxTestGenerated extends AbstractJsCodegenBoxTest {
                 runTest("compiler/testData/codegen/box/properties/lateinit/accessorException.kt");
             }
 
+            @TestMetadata("accessorExceptionPublic.kt")
+            public void testAccessorExceptionPublic() throws Exception {
+                runTest("compiler/testData/codegen/box/properties/lateinit/accessorExceptionPublic.kt");
+            }
+
             public void testAllFilesPresentInLateinit() throws Exception {
                 KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("compiler/testData/codegen/box/properties/lateinit"), Pattern.compile("^(.+)\\.kt$"), TargetBackend.JS, true);
             }