From a99db601176ed549b3d4a5195f63ee617814293e Mon Sep 17 00:00:00 2001
From: "Aleksei.Glushko" <aleksei.glushko@jetbrains.com>
Date: Mon, 28 Aug 2023 18:33:36 +0200
Subject: [PATCH] [K/N] Reenable mark&sweep statistics match assertion

---
 .../runtime/src/custom_alloc/cpp/GCApi.cpp      |  2 ++
 .../runtime/src/gc/common/cpp/GCStatistics.cpp  | 17 +++++------------
 2 files changed, 7 insertions(+), 12 deletions(-)

diff --git a/kotlin-native/runtime/src/custom_alloc/cpp/GCApi.cpp b/kotlin-native/runtime/src/custom_alloc/cpp/GCApi.cpp
index 43d66b24f9e..47474412b17 100644
--- a/kotlin-native/runtime/src/custom_alloc/cpp/GCApi.cpp
+++ b/kotlin-native/runtime/src/custom_alloc/cpp/GCApi.cpp
@@ -48,11 +48,13 @@ bool SweepObject(uint8_t* object, FinalizerQueue& finalizerQueue, gc::GCHandle::
             CustomAllocDebug("SweepObject: fromExtraObject(%p) = %p", extraObject, ExtraObjectCell::fromExtraObject(extraObject));
             finalizerQueue.Push(ExtraObjectCell::fromExtraObject(extraObject));
             gcHandle.addMarkedObject();
+            gcHandle.addKeptObject();
             return true;
         }
         if (!extraObject->getFlag(mm::ExtraObjectData::FLAGS_FINALIZED)) {
             CustomAllocDebug("SweepObject(%p): already waiting to be finalized", heapObjHeader);
             gcHandle.addMarkedObject();
+            gcHandle.addKeptObject();
             return true;
         }
         extraObject->UnlinkFromBaseObject();
diff --git a/kotlin-native/runtime/src/gc/common/cpp/GCStatistics.cpp b/kotlin-native/runtime/src/gc/common/cpp/GCStatistics.cpp
index 60fbbf0733b..3699fc1511a 100644
--- a/kotlin-native/runtime/src/gc/common/cpp/GCStatistics.cpp
+++ b/kotlin-native/runtime/src/gc/common/cpp/GCStatistics.cpp
@@ -208,6 +208,11 @@ void GCHandle::finished() {
     if (auto* stat = statByEpoch(epoch_)) {
         stat->endTime = static_cast<KLong>(konan::getTimeNanos());
         stat->memoryUsageAfter.heap = currentHeapUsage();
+        if (stat->markStats && stat->sweepStats.heap) {
+            RuntimeAssert(stat->markStats->markedCount == stat->sweepStats.heap->keptCount,
+                          "Mismatch in statistics: marked %" PRId64 " objects, while %" PRId64 " are alive after sweep",
+                          stat->markStats->markedCount, stat->sweepStats.heap->keptCount);
+        }
         if (stat->rootSet) {
             GCLogInfo(
                     epoch_,
@@ -294,16 +299,6 @@ void GCHandle::threadsAreSuspended() {
         if (requestTime) {
             auto time = (startTime - *requestTime) / 1000;
             GCLogDebug(epoch_, "Suspended all threads in %" PRIu64 " microseconds", time);
-            return;
-        }
-    }
-    if (last.epoch) {
-        // Assisted sweeping from the last epoch must be completed before the check can be run.
-        if (last.markStats && last.sweepStats.heap) {
-            RuntimeAssert(
-                    last.markStats->markedCount == last.sweepStats.heap->keptCount,
-                    "Mismatch in statistics: marked %" PRId64 " objects, while %" PRId64 " are alive after sweep",
-                    last.markStats->markedCount, last.sweepStats.heap->keptCount);
         }
     }
 }
@@ -323,7 +318,6 @@ void GCHandle::threadsAreResumed() {
         if (startTime) {
             auto time = (endTime - *startTime) / 1000;
             GCLogDebug(epoch_, "Resume all threads. Total pause time is %" PRId64 " microseconds.", time);
-            return;
         }
     }
 }
@@ -334,7 +328,6 @@ void GCHandle::finalizersDone() {
         if (stat->endTime) {
             auto time = (*stat->finalizersDoneTime - *stat->endTime) / 1000;
             GCLogInfo(epoch_, "Finalization is done in %" PRId64 " microseconds after epoch end.", time);
-            return;
         } else {
             GCLogInfo(epoch_, "Finalization is done.");
         }