From ec891474b0ee14c1ef496977483dca7ca93dc10c Mon Sep 17 00:00:00 2001
From: Troels Bjerre Lund <troels@google.com>
Date: Fri, 6 Jan 2023 12:44:28 +0000
Subject: [PATCH] [K/N] custom-alloc: AtomicStack avoid loop on move ^KT-55364

This change avoids looping through the source stack, looking for the
last element, in the case where the target stack is empty. This could
matter in two places:

* when merging ready_ and used_ into unswept_ in PrepareForGC. The first
  of these will be faster with this change.
* when merging finalizer queues (in CR-592, not merged yet), where we
  expect the target queue to be empty.

Since we can expect used_ to be larger than ready_, since we are about
to do a GC, the order of these two has also been reversed.


Co-authored-by: Troels Lund <troels@google.com>

Merge-request: KOTLIN-MR-599
Merged-by: Alexander Shabalin <alexander.shabalin@jetbrains.com>
---
 .../runtime/src/custom_alloc/cpp/AtomicStack.hpp | 16 ++++++++++------
 .../runtime/src/custom_alloc/cpp/PageStore.hpp   |  2 +-
 2 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/kotlin-native/runtime/src/custom_alloc/cpp/AtomicStack.hpp b/kotlin-native/runtime/src/custom_alloc/cpp/AtomicStack.hpp
index b5ec8972e50..d611b88e832 100644
--- a/kotlin-native/runtime/src/custom_alloc/cpp/AtomicStack.hpp
+++ b/kotlin-native/runtime/src/custom_alloc/cpp/AtomicStack.hpp
@@ -40,17 +40,21 @@ public:
         while (!other.stack_.compare_exchange_weak(otherHead, nullptr, std::memory_order_acq_rel)) {}
         // If the `other` stack was empty, do nothing.
         if (!otherHead) return;
-        // Now find the tail of `other`. If no deletions are performed, this is safe.
+        // If `this` stack is empty, just copy the `other` stack over 
+        T* thisHead = nullptr;
+        if (stack_.compare_exchange_strong(thisHead, otherHead, std::memory_order_acq_rel)) {
+            return;
+        }
+        // `this` stack is not empty. Find the tail of `other`. If no deletions are performed, this is safe.
         T* otherTail = otherHead;
         while (otherTail->next_) otherTail = otherTail->next_;
-        // Now make `otherTail->next_` point to the current head of `this` and
-        // simultaneously make `otherHead` the new current head.
-        T* thisHead = nullptr;
         // can't be because of the loop above
         RuntimeAssert(otherTail->next_ == nullptr, "otherTail->next_ must be a tail");
-        while (!stack_.compare_exchange_weak(thisHead, otherHead, std::memory_order_acq_rel)) {
+        // Now make `otherTail->next_` point to the current head of `this` and
+        // simultaneously make `otherHead` the new current head.
+        do {
             otherTail->next_ = thisHead;
-        }
+        } while (!stack_.compare_exchange_weak(thisHead, otherHead, std::memory_order_acq_rel));
     }
 
     bool isEmpty() noexcept { return stack_.load(std::memory_order_relaxed) == nullptr; }
diff --git a/kotlin-native/runtime/src/custom_alloc/cpp/PageStore.hpp b/kotlin-native/runtime/src/custom_alloc/cpp/PageStore.hpp
index 27808c580e5..14791d38dcb 100644
--- a/kotlin-native/runtime/src/custom_alloc/cpp/PageStore.hpp
+++ b/kotlin-native/runtime/src/custom_alloc/cpp/PageStore.hpp
@@ -16,8 +16,8 @@ template <class T>
 class PageStore {
 public:
     void PrepareForGC() noexcept {
-        unswept_.TransferAllFrom(ready_);
         unswept_.TransferAllFrom(used_);
+        unswept_.TransferAllFrom(ready_);
         T* page;
         while ((page = empty_.Pop())) page->Destroy();
     }