[K/N] Support escape analysis with new mm

kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/ToplevelPhases.kt

@@ -463,8 +463,6 @@ internal fun PhaseConfig.konanPhasesConfig(config: KonanConfig) {
         disableUnless(buildDFGPhase, getBoolean(KonanConfigKeys.OPTIMIZATION))
         disableUnless(devirtualizationPhase, getBoolean(KonanConfigKeys.OPTIMIZATION))
         disableUnless(escapeAnalysisPhase, getBoolean(KonanConfigKeys.OPTIMIZATION))
-        // TODO: Support escape analysis with experimental MM.
-        disableIf(escapeAnalysisPhase, config.memoryModel == MemoryModel.EXPERIMENTAL)
         // Inline accessors only in optimized builds due to separate compilation and possibility to get broken
         // debug information.
         disableUnless(propertyAccessorInlinePhase, getBoolean(KonanConfigKeys.OPTIMIZATION))

kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/llvm/CodeGenerator.kt

@@ -255,12 +255,16 @@ internal class StackLocalsManagerImpl(
         val bbInitStackLocals: LLVMBasicBlockRef
 ) : StackLocalsManager {
     private class StackLocal(val isArray: Boolean, val irClass: IrClass,
-                             val bodyPtr: LLVMValueRef, val objHeaderPtr: LLVMValueRef)
+                             val bodyPtr: LLVMValueRef, val objHeaderPtr: LLVMValueRef,
+                             val gcRootSetSlot: LLVMValueRef?)
 
     private val stackLocals = mutableListOf<StackLocal>()
 
     fun isEmpty() = stackLocals.isEmpty()
 
+    private fun FunctionGenerationContext.createRootSetSlot() =
+            if (context.memoryModel == MemoryModel.EXPERIMENTAL) alloca(kObjHeaderPtr) else null
+
     override fun alloc(irClass: IrClass, cleanFieldsExplicitly: Boolean): LLVMValueRef = with(functionGenerationContext) {
         val type = context.llvmDeclarations.forClass(irClass).bodyType
         val stackLocal = appendingTo(bbInitStackLocals) {
@@ -271,12 +275,16 @@ internal class StackLocalsManagerImpl(
             val objectHeader = structGep(stackSlot, 0, "objHeader")
             val typeInfo = codegen.typeInfoForAllocation(irClass)
             setTypeInfoForLocalObject(objectHeader, typeInfo)
-            StackLocal(false, irClass, stackSlot, objectHeader)
+            val gcRootSetSlot = createRootSetSlot()
+            StackLocal(false, irClass, stackSlot, objectHeader, gcRootSetSlot)
         }
 
         stackLocals += stackLocal
         if (cleanFieldsExplicitly)
             clean(stackLocal, false)
+        if (stackLocal.gcRootSetSlot != null) {
+            storeStackRef(stackLocal.objHeaderPtr, stackLocal.gcRootSetSlot)
+        }
         stackLocal.objHeaderPtr
     }
 
@@ -323,11 +331,16 @@ internal class StackLocalsManagerImpl(
                     0,
                     constCount * LLVMSizeOfTypeInBits(codegen.llvmTargetData, arrayToElementType[irClass.symbol]).toInt() / 8
             )
-            StackLocal(true, irClass, arraySlot, arrayHeaderSlot)
+            val gcRootSetSlot = createRootSetSlot()
+            StackLocal(true, irClass, arraySlot, arrayHeaderSlot, gcRootSetSlot)
         }
 
         stackLocals += stackLocal
-        bitcast(kObjHeaderPtr, stackLocal.objHeaderPtr)
+        val result = bitcast(kObjHeaderPtr, stackLocal.objHeaderPtr)
+        if (stackLocal.gcRootSetSlot != null) {
+            storeStackRef(result, stackLocal.gcRootSetSlot)
+        }
+        result
     }
 
     override fun clean(refsOnly: Boolean) = stackLocals.forEach { clean(it, refsOnly) }
@@ -360,6 +373,9 @@ internal class StackLocalsManagerImpl(
                 memset(bodyWithSkippedServiceInfoPtr, 0, bodySize - serviceInfoSize)
             }
         }
+        if (stackLocal.gcRootSetSlot != null) {
+            storeStackRef(kNullObjHeaderPtr, stackLocal.gcRootSetSlot)
+        }
     }
 
     private fun setTypeInfoForLocalObject(objectHeader: LLVMValueRef, typeInfoPointer: LLVMValueRef) = with(functionGenerationContext) {

kotlin-native/runtime/src/gc/common/cpp/MarkAndSweepUtils.hpp

@@ -24,7 +24,6 @@ void Mark(KStdVector<ObjHeader*> graySet) noexcept {
         graySet.pop_back();
 
         RuntimeAssert(!isNullOrMarker(top), "Got invalid reference %p in gray set", top);
-        RuntimeAssert(!top->local(), "TODO: Stack objects are not supported yet, top=%p", top);
 
         if (top->heap()) {
             if (!Traits::TryMark(top)) {
@@ -32,9 +31,9 @@ void Mark(KStdVector<ObjHeader*> graySet) noexcept {
             }
         }
 
-        if (!top->permanent()) {
+        if (top->heap() || top->local()) {
             traverseReferredObjects(top, [&graySet](ObjHeader* field) noexcept {
-                if (!isNullOrMarker(field) && !field->permanent() && !Traits::IsMarked(field)) {
+                if (!isNullOrMarker(field) && field->heap() && !Traits::IsMarked(field)) {
                     graySet.push_back(field);
                 }
             });

kotlin-native/runtime/src/gc/common/cpp/MarkAndSweepUtilsMarkTest.cpp

@@ -38,6 +38,7 @@ class BaseObject {
 public:
     enum class Kind {
         kPermanent,
+        kStackLocal,
         kHeapLike // Treated as heap object for the purposes of the test.
     };
 
@@ -60,6 +61,11 @@ protected:
             case Kind::kHeapLike:
                 RuntimeAssert(GetObjHeader()->heap(), "Must be heap");
                 break;
+            case Kind::kStackLocal:
+                GetObjHeader()->typeInfoOrMeta_ = setPointerBits(GetObjHeader()->typeInfoOrMeta_,
+                                                                 OBJECT_TAG_PERMANENT_CONTAINER | OBJECT_TAG_NONTRIVIAL_CONTAINER);
+                RuntimeAssert(GetObjHeader()->local(), "Must be stack local");
+                break;
         }
     }
 
@@ -205,6 +211,31 @@ TEST_F(MarkAndSweepUtilsMarkTest, MarkSinglePermanentCharArray) {
     EXPECT_MARKED();
 }
 
+TEST_F(MarkAndSweepUtilsMarkTest, MarkSingleStackObject) {
+    Object object{BaseObject::Kind::kStackLocal};
+
+    Mark({object});
+
+    EXPECT_MARKED();
+}
+
+TEST_F(MarkAndSweepUtilsMarkTest, MarkSingleStackObjectArray) {
+    ObjectArray array{BaseObject::Kind::kStackLocal};
+
+    Mark({array});
+
+    EXPECT_MARKED();
+}
+
+TEST_F(MarkAndSweepUtilsMarkTest, MarkSingleStackCharArray) {
+    CharArray array{BaseObject::Kind::kStackLocal};
+
+    Mark({array});
+
+    EXPECT_MARKED();
+}
+
+
 TEST_F(MarkAndSweepUtilsMarkTest, MarkSingleObjectWithInvalidFields) {
     Object object;
     object->field1 = kInitializingSingleton;
@@ -416,6 +447,28 @@ TEST_F(MarkAndSweepUtilsMarkTest, MarkTreeWithPermanentLeaf) {
     EXPECT_MARKED(root, root_field1, root_field1_field2, root_field3, root_field3_element1, root_field3_element2, root_field3_element3);
 }
 
+TEST_F(MarkAndSweepUtilsMarkTest, MarkTreeWithStackRoot) {
+    Object root{BaseObject::Kind::kStackLocal};
+    Object root_field1{BaseObject::Kind::kPermanent};
+    Object root_field1_field1{BaseObject::Kind::kPermanent};
+    Object root_field1_field2{BaseObject::Kind::kPermanent};
+    ObjectArray root_field3{BaseObject::Kind::kStackLocal};
+    Object root_field3_element1{BaseObject::Kind::kHeapLike};
+    ObjectArray root_field3_element2{BaseObject::Kind::kHeapLike};
+    CharArray root_field3_element3{BaseObject::Kind::kPermanent};
+    root->field1 = root_field1.header();
+    root_field1->field1 = root_field1_field1.header();
+    root_field1->field2 = root_field1_field2.header();
+    root->field3 = root_field3.header();
+    root_field3.elements()[0] = root_field3_element1.header();
+    root_field3.elements()[1] = root_field3_element2.header();
+    root_field3.elements()[2] = root_field3_element3.header();
+
+    Mark({root, root_field3});
+
+    EXPECT_MARKED(root_field3_element1, root_field3_element2);
+}
+
 TEST_F(MarkAndSweepUtilsMarkTest, MarkRecursiveTree) {
     Object root;
     Object inner1;
@@ -442,6 +495,25 @@ TEST_F(MarkAndSweepUtilsMarkTest, MarkRecursiveTreeWithPermanentRoot) {
     EXPECT_MARKED();
 }
 
+TEST_F(MarkAndSweepUtilsMarkTest, MarkRecursiveTreeWithStackRoot) {
+    Object root{BaseObject::Kind::kStackLocal};
+    Object inner1{BaseObject::Kind::kStackLocal};
+    ObjectArray inner2{BaseObject::Kind::kStackLocal};
+    Object inner3{BaseObject::Kind::kHeapLike};
+    Object inner2_element1{BaseObject::Kind::kHeapLike};
+    root->field1 = inner1.header();
+    inner1->field1 = inner2.header();
+    inner2.elements()[0] = root.header();
+    inner2.elements()[1] = inner2_element1.header();
+    root->field2 = inner3.header();
+    inner3->field1 = inner2.header();
+
+    Mark({root, inner1, inner2});
+
+    EXPECT_MARKED(inner3, inner2_element1);
+}
+
+
 TEST_F(MarkAndSweepUtilsMarkTest, MarkForest) {
     Object root1;
     ObjectArray root2;