[K/N] Add ability to get base object from extra data

^KT-49325

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

@@ -40,7 +40,7 @@ void Mark(KStdVector<ObjHeader*> graySet) noexcept {
         }
 
         if (auto* extraObjectData = mm::ExtraObjectData::Get(top)) {
-            auto* weakCounter = *extraObjectData->GetWeakCounterLocation();
+            auto weakCounter = extraObjectData->GetWeakReferenceCounter();
             if (!isNullOrMarker(weakCounter)) {
                 graySet.push_back(weakCounter);
             }

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

@@ -48,7 +48,8 @@ public:
 
     void InstallWeakCounter(BaseObject& counter) {
         auto& extraObjectData = mm::ExtraObjectData::GetOrInstall(GetObjHeader());
-        *extraObjectData.GetWeakCounterLocation() = counter.GetObjHeader();
+        auto *setCounter = extraObjectData.GetOrSetWeakReferenceCounter(GetObjHeader(), counter.GetObjHeader());
+        EXPECT_EQ(setCounter, counter.GetObjHeader());
     }
 
 protected:

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

@@ -225,7 +225,8 @@ public:
         objectFactoryThreadQueue_.Publish();
         auto& weakCounter = WeakCounter::FromObjHeader(weakCounterHeader);
         auto& extraObjectData = mm::ExtraObjectData::GetOrInstall(objHeader);
-        *extraObjectData.GetWeakCounterLocation() = weakCounter.header();
+        auto *setHeader = extraObjectData.GetOrSetWeakReferenceCounter(objHeader, weakCounter.header());
+        EXPECT_EQ(setHeader,  weakCounter.header());
         weakCounter->referred = objHeader;
         return weakCounter;
     }

kotlin-native/runtime/src/gc/stms/cpp/SameThreadMarkAndSweepTest.cpp

@@ -208,7 +208,8 @@ WeakCounter& InstallWeakCounter(mm::ThreadData& threadData, ObjHeader* objHeader
     mm::AllocateObject(&threadData, typeHolderWeakCounter.typeInfo(), location);
     auto& weakCounter = WeakCounter::FromObjHeader(*location);
     auto& extraObjectData = mm::ExtraObjectData::GetOrInstall(objHeader);
-    *extraObjectData.GetWeakCounterLocation() = weakCounter.header();
+    auto *setCounter = extraObjectData.GetOrSetWeakReferenceCounter(objHeader, weakCounter.header());
+    EXPECT_EQ(setCounter, weakCounter.header());
     weakCounter->referred = objHeader;
     return weakCounter;
 }
@@ -222,6 +223,7 @@ public:
 
     ~SameThreadMarkAndSweepTest() {
         mm::GlobalsRegistry::Instance().ClearForTests();
+        mm::GlobalData::Instance().extraObjectDataFactory().ClearForTests();
         mm::GlobalData::Instance().objectFactory().ClearForTests();
         mm::GlobalData::Instance().gcScheduler().ReplaceGCSchedulerDataForTests(
                 [](auto& config, auto scheduleGC) { return gc::internal::MakeGCSchedulerData(config, std::move(scheduleGC)); });

kotlin-native/runtime/src/legacymm/cpp/Memory.cpp

@@ -477,8 +477,13 @@ ALWAYS_INLINE bool isShareable(const ObjHeader* obj) {
     return containerFor(obj)->shareable();
 }
 
-ObjHeader** ObjHeader::GetWeakCounterLocation() {
+ObjHeader* ObjHeader::GetWeakCounter() {
-    return &this->meta_object()->WeakReference.counter_;
+    return this->meta_object()->WeakReference.counter_;
+}
+
+ObjHeader* ObjHeader::GetOrSetWeakCounter(ObjHeader* counter) {
+    UpdateHeapRefIfNull(&meta_object()->WeakReference.counter_, counter);
+    return GetWeakCounter();
 }
 
 #if KONAN_OBJC_INTEROP

kotlin-native/runtime/src/main/cpp/Memory.h

@@ -68,7 +68,9 @@ struct ObjHeader {
 
   MetaObjHeader* meta_object_or_null() const noexcept { return AsMetaObject(typeInfoOrMeta_); }
 
-  ALWAYS_INLINE ObjHeader** GetWeakCounterLocation();
+  ALWAYS_INLINE ObjHeader* GetWeakCounter();
+  ALWAYS_INLINE ObjHeader* GetOrSetWeakCounter(ObjHeader* counter);
+
 
 #ifdef KONAN_OBJC_INTEROP
   ALWAYS_INLINE void* GetAssociatedObject();

kotlin-native/runtime/src/main/cpp/Weak.cpp

@@ -68,14 +68,14 @@ OBJ_GETTER(Konan_getWeakReferenceImpl, ObjHeader* referred) {
   }
 #endif // KONAN_OBJC_INTEROP
 
-  ObjHeader** weakCounterLocation = referred->GetWeakCounterLocation();
+  ObjHeader* weakCounter = referred->GetWeakCounter();
-  if (*weakCounterLocation == nullptr) {
+  if (weakCounter == nullptr) {
       ObjHolder counterHolder;
       // Cast unneeded, just to emphasize we store an object reference as void*.
       ObjHeader* counter = makeWeakReferenceCounter(reinterpret_cast<void*>(referred), counterHolder.slot());
-      UpdateHeapRefIfNull(weakCounterLocation, counter);
+      weakCounter = referred->GetOrSetWeakCounter(counter);
   }
-  RETURN_OBJ(*weakCounterLocation);
+  RETURN_OBJ(weakCounter);
 }
 
 // Materialize a weak reference to either null or the real reference.
@@ -89,6 +89,10 @@ OBJ_GETTER(Konan_WeakReferenceCounter_get, ObjHeader* counter) {
 #endif
 }
 
+ALWAYS_INLINE ObjHeader* UnsafeWeakReferenceCounterGet(ObjHeader* counter) {
+    return asWeakReferenceCounter(counter)->referred;
+}
+
 void WeakReferenceCounterClear(ObjHeader* counter) {
   ObjHeader** referredAddress = &asWeakReferenceCounter(counter)->referred;
   // Note, that we don't do UpdateRef here, as reference is weak.

kotlin-native/runtime/src/main/cpp/Weak.h

@@ -12,6 +12,7 @@ extern "C" {
 
 // Atomically clears counter object reference.
 void WeakReferenceCounterClear(ObjHeader* counter);
+ObjHeader* UnsafeWeakReferenceCounterGet(ObjHeader* counter);
 
 } // extern "C"
 

kotlin-native/runtime/src/mm/cpp/ExtraObjectData.cpp

@@ -44,7 +44,7 @@ mm::ExtraObjectData& mm::ExtraObjectData::Install(ObjHeader* object) noexcept {
     RuntimeCheck(!hasPointerBits(typeInfo, OBJECT_TAG_MASK), "Object must not be tagged");
 
     auto *threadData = mm::ThreadRegistry::Instance().CurrentThreadData();
-    auto& data = mm::ExtraObjectDataFactory::Instance().CreateExtraObjectDataForObject(threadData, typeInfo);
+    auto& data = mm::ExtraObjectDataFactory::Instance().CreateExtraObjectDataForObject(threadData, object, typeInfo);
 
     TypeInfo* old = __sync_val_compare_and_swap(&object->typeInfoOrMeta_, typeInfo, reinterpret_cast<TypeInfo*>(&data));
     if (old != typeInfo) {
@@ -78,18 +78,15 @@ void mm::ExtraObjectData::DetachAssociatedObject() noexcept {
 #endif
 }
 
-bool mm::ExtraObjectData::HasWeakReferenceCounter() noexcept {
-    return weakReferenceCounter_ != nullptr;
-}
-
 void mm::ExtraObjectData::ClearWeakReferenceCounter() noexcept {
     if (!HasWeakReferenceCounter()) return;
 
-    WeakReferenceCounterClear(weakReferenceCounter_);
+    auto *object = GetBaseObject();
+    WeakReferenceCounterClear(GetWeakReferenceCounter());
     // Not using `mm::SetHeapRef here`, because this code is called during sweep phase by the GC thread,
     // and so cannot affect marking.
     // TODO: Asserts on the above?
-    weakReferenceCounter_ = nullptr;
+    weakReferenceCounterOrBaseObject_ = object;
 }
 
 mm::ExtraObjectData::~ExtraObjectData() {

kotlin-native/runtime/src/mm/cpp/ExtraObjectData.hpp

@@ -15,6 +15,7 @@
 #include "TypeInfo.h"
 #include "Utils.hpp"
 #include "MultiSourceQueue.hpp"
+#include "Weak.h"
 
 namespace kotlin {
 namespace mm {
@@ -28,6 +29,8 @@ public:
         FLAGS_NEVER_FROZEN = 1 << 1,
     };
 
+    static constexpr unsigned WEAK_REF_TAG = 1;
+
     MetaObjHeader* AsMetaObjHeader() noexcept { return reinterpret_cast<MetaObjHeader*>(this); }
     static ExtraObjectData& FromMetaObjHeader(MetaObjHeader* header) noexcept { return *reinterpret_cast<ExtraObjectData*>(header); }
 
@@ -45,14 +48,35 @@ public:
 #endif
     void DetachAssociatedObject() noexcept;
 
-    ObjHeader** GetWeakCounterLocation() noexcept { return &weakReferenceCounter_; }
-
     std::atomic<Flags>& flags() noexcept { return flags_; }
 
-    bool HasWeakReferenceCounter() noexcept;
+    bool HasWeakReferenceCounter() noexcept { return hasPointerBits(weakReferenceCounterOrBaseObject_.load(), WEAK_REF_TAG); }
     void ClearWeakReferenceCounter() noexcept;
+    ObjHeader* GetWeakReferenceCounter() noexcept {
+        auto *pointer = weakReferenceCounterOrBaseObject_.load();
+        if (hasPointerBits(pointer, WEAK_REF_TAG)) return clearPointerBits(pointer, WEAK_REF_TAG);
+        return nullptr;
+    }
+    ObjHeader* GetOrSetWeakReferenceCounter(ObjHeader* object, ObjHeader* counter) noexcept {
+        if (weakReferenceCounterOrBaseObject_.compare_exchange_strong(object, setPointerBits(counter, WEAK_REF_TAG))) {
+            return counter;
+        } else {
+            return clearPointerBits(object, WEAK_REF_TAG); // on fail current value of counter is stored to object
+        }
+    }
+    ObjHeader* GetBaseObject() noexcept {
+        auto *header = weakReferenceCounterOrBaseObject_.load();
+        if (hasPointerBits(header, WEAK_REF_TAG)) {
+            return UnsafeWeakReferenceCounterGet(clearPointerBits(header, WEAK_REF_TAG));
+        } else {
+            return header;
+        }
+    }
 
-    explicit ExtraObjectData(const TypeInfo* typeInfo) noexcept : typeInfo_(typeInfo) {}
+    // info must be equal to objHeader->type_info(), but it needs to be loaded in advance to avoid data races
+    explicit ExtraObjectData(ObjHeader* objHeader, const TypeInfo *info) noexcept :
+        typeInfo_(info), weakReferenceCounterOrBaseObject_(objHeader) {
+    }
     ~ExtraObjectData();
 private:
 
@@ -65,7 +89,7 @@ private:
     void* associatedObject_ = nullptr;
 #endif
 
-    ObjHeader* weakReferenceCounter_ = nullptr;
+    std::atomic<ObjHeader*> weakReferenceCounterOrBaseObject_;
 };
 
 } // namespace mm

kotlin-native/runtime/src/mm/cpp/ExtraObjectDataFactory.cpp

@@ -15,12 +15,24 @@ mm::ExtraObjectDataFactory& mm::ExtraObjectDataFactory::Instance() noexcept {
     return GlobalData::Instance().extraObjectDataFactory();
 }
 
-mm::ExtraObjectData& mm::ExtraObjectDataFactory::CreateExtraObjectDataForObject(mm::ThreadData* threadData, TypeInfo* info) noexcept {
+mm::ExtraObjectData& mm::ExtraObjectDataFactory::CreateExtraObjectDataForObject(
-    return **threadData->extraObjectDataThreadQueue().Emplace(info);
+        mm::ThreadData* threadData, ObjHeader* baseObject, const TypeInfo* info
+        ) noexcept {
+    return CreateExtraObjectDataForObject(threadData->extraObjectDataThreadQueue(), baseObject, info);
 }
 
 void mm::ExtraObjectDataFactory::DestroyExtraObjectData(mm::ThreadData* threadData, ExtraObjectData& data) noexcept {
-    threadData->extraObjectDataThreadQueue().Erase(&Queue::Node::fromValue(data));
+        DestroyExtraObjectData(threadData->extraObjectDataThreadQueue(), data);
+}
+
+mm::ExtraObjectData& mm::ExtraObjectDataFactory::CreateExtraObjectDataForObject(
+        ThreadQueue& threadQueue, ObjHeader* baseObject, const TypeInfo* info
+        ) noexcept {
+    return **threadQueue.Emplace(baseObject, info);
+}
+
+void mm::ExtraObjectDataFactory::DestroyExtraObjectData(ThreadQueue& threadQueue, ExtraObjectData& data) noexcept {
+    threadQueue.Erase(&Queue::Node::fromValue(data));
 }
 
 void mm::ExtraObjectDataFactory::ProcessThread(mm::ThreadData* threadData) noexcept {

kotlin-native/runtime/src/mm/cpp/ExtraObjectDataFactory.hpp

@@ -33,9 +33,11 @@ public:
 
     static ExtraObjectDataFactory& Instance() noexcept;
 
-    ExtraObjectData& CreateExtraObjectDataForObject(mm::ThreadData* threadData, TypeInfo* info) noexcept;
+    ExtraObjectData& CreateExtraObjectDataForObject(mm::ThreadData* threadData, ObjHeader* baseObject, const TypeInfo* info) noexcept;
+    ExtraObjectData& CreateExtraObjectDataForObject(ThreadQueue& threadQueue, ObjHeader* baseObject, const TypeInfo* info) noexcept;
 
     void DestroyExtraObjectData(mm::ThreadData* threadData, ExtraObjectData& data) noexcept;
+    void DestroyExtraObjectData(ThreadQueue& threadQueue, ExtraObjectData& data) noexcept;
 
     // Collect extra data objects from thread corresponding to `threadData`. Must be called by the thread
     // when it's asked by GC to stop.

kotlin-native/runtime/src/mm/cpp/ExtraObjectDataTest.cpp

@@ -23,9 +23,20 @@ struct EmptyPayload {
     static constexpr std::array<Field, 0> kFields{};
 };
 
+class ExtraObjectDataTest : public testing::Test {
+public:
+    ExtraObjectDataTest() {}
+
+    ~ExtraObjectDataTest() {
+        mm::GlobalsRegistry::Instance().ClearForTests();
+        mm::GlobalData::Instance().extraObjectDataFactory().ClearForTests();
+        mm::GlobalData::Instance().objectFactory().ClearForTests();
+    }
+};
+
 } // namespace
 
-TEST(ExtraObjectDataTest, Install) {
+TEST_F(ExtraObjectDataTest, Install) {
     ScopedMemoryInit init;
     test_support::TypeInfoHolder type{test_support::TypeInfoHolder::ObjectBuilder<EmptyPayload>()};
     test_support::Object<EmptyPayload> object(type.typeInfo());
@@ -38,6 +49,8 @@ TEST(ExtraObjectDataTest, Install) {
     EXPECT_TRUE(object.header()->has_meta_object());
     EXPECT_THAT(object.header()->meta_object(), extraData.AsMetaObjHeader());
     EXPECT_THAT(object.header()->type_info(), typeInfo);
+    EXPECT_FALSE(extraData.HasWeakReferenceCounter());
+    EXPECT_THAT(extraData.GetBaseObject(), object.header());
 
     mm::ExtraObjectData::Uninstall(object.header());
 
@@ -45,7 +58,7 @@ TEST(ExtraObjectDataTest, Install) {
     EXPECT_THAT(object.header()->type_info(), typeInfo);
 }
 
-TEST(ExtraObjectDataTest, ConcurrentInstall) {
+TEST_F(ExtraObjectDataTest, ConcurrentInstall) {
     ScopedMemoryInit init;
     test_support::TypeInfoHolder type{test_support::TypeInfoHolder::ObjectBuilder<EmptyPayload>()};
     test_support::Object<EmptyPayload> object(type.typeInfo());
@@ -65,13 +78,13 @@ TEST(ExtraObjectDataTest, ConcurrentInstall) {
             }
             auto& extraData = mm::ExtraObjectData::Install(object.header());
             actual[i] = &extraData;
+            mm::GlobalData::Instance().threadRegistry().CurrentThreadData()->Publish();
         });
     }
 
     while (readyCount < kThreadCount) {
     }
     canStart = true;
-
     for (auto& t : threads) {
         t.join();
     }
@@ -79,6 +92,4 @@ TEST(ExtraObjectDataTest, ConcurrentInstall) {
     std::vector<mm::ExtraObjectData*> expected(kThreadCount, actual[0]);
 
     EXPECT_THAT(actual, testing::ElementsAreArray(expected));
-
-    mm::ExtraObjectData::Uninstall(object.header());
 }

kotlin-native/runtime/src/mm/cpp/Memory.cpp

@@ -48,8 +48,12 @@ ALWAYS_INLINE mm::StableRefRegistry::Node* FromForeignRefManager(ForeignRefManag
 
 } // namespace
 
-ObjHeader** ObjHeader::GetWeakCounterLocation() {
+ObjHeader* ObjHeader::GetWeakCounter() {
-    return mm::ExtraObjectData::FromMetaObjHeader(this->meta_object()).GetWeakCounterLocation();
+    return mm::ExtraObjectData::FromMetaObjHeader(this->meta_object()).GetWeakReferenceCounter();
+}
+
+ObjHeader* ObjHeader::GetOrSetWeakCounter(ObjHeader* counter) {
+    return mm::ExtraObjectData::FromMetaObjHeader(this->meta_object()).GetOrSetWeakReferenceCounter(this, counter);
 }
 
 #ifdef KONAN_OBJC_INTEROP

kotlin-native/runtime/src/mm/cpp/TestSupport.cpp

@@ -22,37 +22,38 @@ bool isEmpty(T& iterable) {
     return iterable.begin() == iterable.end();
 }
 
-template <typename E, typename T>
+template <typename T>
-std::vector<E> collect(T& iterable) {
+auto collectCopy(T& iterable) {
-    std::vector<E> result;
+    std::vector<std::remove_reference_t<decltype(*iterable.begin())>> result;
-    for (E element : iterable) {
+    for (const auto &element : iterable) {
         result.push_back(element);
     }
-    return std::move(result);
-}
-
-std::vector<mm::ThreadData*> collect(mm::ThreadRegistry::Iterable& iterable) {
-    std::vector<mm::ThreadData*> result;
-    for (mm::ThreadData& element : iterable) {
-        result.push_back(&element);
-    }
-    // Do not use std::move because clang complains that it prevents copy elision.
     return result;
 }
 
+template <typename T>
+auto collectPointers(T& iterable) {
+    std::vector<const std::remove_reference_t<decltype(*iterable.begin())>*> result;
+    for (const auto &element : iterable) {
+        result.push_back(&element);
+    }
+    return result;
+}
 } // namespace
 
 extern "C" void Kotlin_TestSupport_AssertClearGlobalState() {
     // Validate that global registries are empty.
     auto globals = mm::GlobalsRegistry::Instance().LockForIter();
+    auto extraObjects = mm::GlobalData::Instance().extraObjectDataFactory().LockForIter();
     auto objects = mm::GlobalData::Instance().objectFactory().LockForIter();
     auto stableRefs = mm::StableRefRegistry::Instance().LockForIter();
     auto threads = mm::ThreadRegistry::Instance().LockForIter();
 
-    EXPECT_THAT(collect<ObjHeader**>(globals), testing::UnorderedElementsAre());
+    EXPECT_THAT(collectCopy(globals), testing::UnorderedElementsAre());
-    EXPECT_THAT(collect<mm::ObjectFactory<gc::GC>::NodeRef>(objects), testing::UnorderedElementsAre());
+    EXPECT_THAT(collectPointers(extraObjects), testing::UnorderedElementsAre());
-    EXPECT_THAT(collect<ObjHeader*>(stableRefs), testing::UnorderedElementsAre());
+    EXPECT_THAT(collectCopy(objects), testing::UnorderedElementsAre());
-    EXPECT_THAT(collect(threads), testing::UnorderedElementsAre());
+    EXPECT_THAT(collectCopy(stableRefs), testing::UnorderedElementsAre());
+    EXPECT_THAT(collectPointers(threads), testing::UnorderedElementsAre());
 }
 
 void kotlin::DeinitMemoryForTests(MemoryState* memoryState) {