[K/N] Fix a race in MemorySharedRefs ^KT-56233

2023-03-31 16:24:56 +02:00
parent 8a8aeed998
commit 4cd1f2ff82
55 changed files with 1693 additions and 464 deletions
@@ -95,7 +95,7 @@ private fun PatchBuilder.addObjCPatches() {
    addProtocolImport("NSCopying")

    addPrivateSelector("toKotlin:")
-    addPrivateSelector("releaseAsAssociatedObject:")
+    addPrivateSelector("releaseAsAssociatedObject")

    addPrivateClass("KIteratorAsNSEnumerator", "iteratorHolder")
    addPrivateClass("KListAsNSArray", "listHolder")
@@ -58,7 +58,6 @@ bool SweepExtraObject(ExtraObjectCell* extraObjectCell, AtomicStack<ExtraObjectC
    }
    extraObject->ClearRegularWeakReferenceImpl();
    if (extraObject->HasAssociatedObject()) {
-        extraObject->DetachAssociatedObject();
        extraObject->setFlag(mm::ExtraObjectData::FLAGS_IN_FINALIZER_QUEUE);
        finalizerQueue.Push(extraObjectCell);
        KeepAlive(baseObject);
@@ -50,6 +50,13 @@ struct SweepTraits {
    }
 };

+struct ProcessWeaksTraits {
+    static bool IsMarked(ObjHeader* obj) noexcept {
+        auto& objectData = mm::ObjectFactory<gc::ConcurrentMarkAndSweep>::NodeRef::From(obj).ObjectData();
+        return objectData.marked();
+    }
+};
+
 } // namespace

 void gc::ConcurrentMarkAndSweep::ThreadData::SafePointAllocation(size_t size) noexcept {
@@ -179,13 +186,17 @@ bool gc::ConcurrentMarkAndSweep::PerformFullGC(int64_t epoch) noexcept {
    auto markStats = gcHandle.getMarked();
    scheduler.gcData().UpdateAliveSetBytes(markStats.totalObjectsSize);

-#ifndef CUSTOM_ALLOCATOR
-    auto extraObjectFactoryIterable = mm::GlobalData::Instance().extraObjectDataFactory().LockForIter();
-    gc::SweepExtraObjects<SweepTraits>(gcHandle, extraObjectFactoryIterable);
+    gc::processWeaks<ProcessWeaksTraits>(gcHandle, mm::SpecialRefRegistry::instance());

+#ifndef CUSTOM_ALLOCATOR
+    // Taking the locks before the pause is completed. So that any destroying thread
+    // would not publish into the global state at an unexpected time.
+    auto extraObjectFactoryIterable = mm::GlobalData::Instance().extraObjectDataFactory().LockForIter();
    auto objectFactoryIterable = objectFactory_.LockForIter();
    mm::ResumeThreads();
    gcHandle.threadsAreResumed();
+
+    gc::SweepExtraObjects<SweepTraits>(gcHandle, extraObjectFactoryIterable);
    auto finalizerQueue = gc::Sweep<SweepTraits>(gcHandle, objectFactoryIterable);
    kotlin::compactObjectPoolInMainThread();
 #else
@@ -22,6 +22,7 @@
 #include "SingleThreadExecutor.hpp"
 #include "TestSupport.hpp"
 #include "ThreadData.hpp"
+#include "WeakRef.hpp"
 #include "std_support/Vector.hpp"

 using namespace kotlin;
@@ -196,6 +197,7 @@ test_support::RegularWeakReferenceImpl& InstallWeakReference(mm::ThreadData& thr
    mm::AllocateObject(&threadData, theRegularWeakReferenceImplTypeInfo, location);
    auto& weakReference = test_support::RegularWeakReferenceImpl::FromObjHeader(*location);
    auto& extraObjectData = mm::ExtraObjectData::GetOrInstall(objHeader);
+    weakReference->weakRef = static_cast<mm::RawSpecialRef*>(mm::WeakRef::create(objHeader));
    weakReference->referred = objHeader;
    auto* setWeakRef = extraObjectData.GetOrSetRegularWeakReferenceImpl(objHeader, weakReference.header());
    EXPECT_EQ(setWeakRef, weakReference.header());
@@ -211,6 +213,7 @@ public:

    ~ConcurrentMarkAndSweepTest() {
        mm::GlobalsRegistry::Instance().ClearForTests();
+        mm::SpecialRefRegistry::instance().clearForTests();
        mm::GlobalData::Instance().extraObjectDataFactory().ClearForTests();
        mm::GlobalData::Instance().gc().ClearForTests();
    }
@@ -337,7 +340,7 @@ TEST_P(ConcurrentMarkAndSweepTest, FreeObjectsWithFinalizers) {
 }

 TEST_P(ConcurrentMarkAndSweepTest, FreeObjectWithFreeWeak) {
-    RunInNewThread([](mm::ThreadData& threadData) {
+    RunInNewThread([this](mm::ThreadData& threadData) {
        auto& object1 = AllocateObject(threadData);
        auto& weak1 = ([&threadData, &object1]() -> test_support::RegularWeakReferenceImpl& {
            ObjHolder holder;
@@ -349,6 +352,7 @@ TEST_P(ConcurrentMarkAndSweepTest, FreeObjectWithFreeWeak) {
        ASSERT_THAT(IsMarked(weak1.header()), false);
        ASSERT_THAT(weak1.get(), object1.header());

+        EXPECT_CALL(finalizerHook(), Call(weak1.header()));
        threadData.gc().ScheduleAndWaitFullGCWithFinalizers();

        EXPECT_THAT(Alive(threadData), testing::UnorderedElementsAre());
@@ -1112,6 +1116,7 @@ TEST_P(ConcurrentMarkAndSweepTest, FreeObjectWithFreeWeakReversedOrder) {

        global1->field1 = nullptr;

+        EXPECT_CALL(finalizerHook(), Call(weak.load()->header()));
        threadData.gc().ScheduleAndWaitFullGC();

        EXPECT_THAT(Alive(threadData), testing::UnorderedElementsAre(global1.header()));
@@ -14,6 +14,7 @@
 #include <cinttypes>
 #include <mutex>

+using namespace kotlin;

 extern "C" {
 void Kotlin_Internal_GC_GCInfoBuilder_setEpoch(KRef thiz, KLong value);
@@ -370,4 +371,14 @@ GCHandle::GCMarkScope::~GCMarkScope() {
               objectsCount, getStageTime(), konan::currentThreadId());
 }

-} // namespace kotlin::gc
+gc::GCHandle::GCProcessWeaksScope::GCProcessWeaksScope(gc::GCHandle& handle) noexcept : handle_(handle) {}
+
+gc::GCHandle::GCProcessWeaksScope::~GCProcessWeaksScope() {
+    GCLogDebug(
+            handle_.getEpoch(),
+            "Processed special refs in %" PRIu64 " microseconds. %" PRIu64 " are undisposed, %" PRIu64 " are alive, %" PRIu64
+            " are nulled out",
+            getStageTime(), undisposedCount_, aliveCount_, nulledCount_);
+}
+
+} // namespace kotlin::gc
@@ -7,12 +7,18 @@

 #include <cstdint>
 #include <pthread.h>
+
 #include "Common.h"
-#include "ThreadData.hpp"
+#include "Porting.h"
+#include "Utils.hpp"

 #define GCLogInfo(epoch, format, ...) RuntimeLogInfo({kTagGC}, "Epoch #%" PRIu64 ": " format, epoch, ##__VA_ARGS__)
 #define GCLogDebug(epoch, format, ...) RuntimeLogDebug({kTagGC}, "Epoch #%" PRIu64 ": " format, epoch, ##__VA_ARGS__)

+namespace kotlin::mm {
+class ThreadData;
+}
+
 namespace kotlin::gc {

 class GCHandle;
@@ -87,6 +93,21 @@ public:
        }
    };

+    class GCProcessWeaksScope : GCStageScopeUsTimer, Pinned {
+        GCHandle& handle_;
+        uint64_t undisposedCount_ = 0;
+        uint64_t aliveCount_ = 0;
+        uint64_t nulledCount_ = 0;
+
+    public:
+        explicit GCProcessWeaksScope(GCHandle& handle) noexcept;
+        ~GCProcessWeaksScope();
+
+        void addUndisposed() noexcept { ++undisposedCount_; }
+        void addAlive() noexcept { ++aliveCount_; }
+        void addNulled() noexcept { ++nulledCount_; }
+    };
+
 private:
    uint64_t epoch_;
    explicit GCHandle(uint64_t epoch) : epoch_(epoch) {}
@@ -117,6 +138,7 @@ public:
    GCGlobalRootSetScope collectGlobalRoots() { return GCGlobalRootSetScope(*this); }
    GCThreadRootSetScope collectThreadRoots(mm::ThreadData& threadData) { return GCThreadRootSetScope(*this, threadData); }
    GCMarkScope mark() { return GCMarkScope(*this); }
+    GCProcessWeaksScope processWeaks() noexcept { return GCProcessWeaksScope(*this); }

    MemoryUsage getMarked();
 };
@@ -16,7 +16,7 @@
 #include "ObjectTraversal.hpp"
 #include "RootSet.hpp"
 #include "Runtime.h"
-#include "StableRefRegistry.hpp"
+#include "SpecialRefRegistry.hpp"
 #include "ThreadData.hpp"
 #include "Types.h"

@@ -116,7 +116,6 @@ void SweepExtraObjects(GCHandle handle, typename Traits::ExtraObjectsFactory::It
        if (!extraObject.getFlag(mm::ExtraObjectData::FLAGS_IN_FINALIZER_QUEUE) && !Traits::IsMarkedByExtraObject(extraObject)) {
            extraObject.ClearRegularWeakReferenceImpl();
            if (extraObject.HasAssociatedObject()) {
-                extraObject.DetachAssociatedObject();
                extraObject.setFlag(mm::ExtraObjectData::FLAGS_IN_FINALIZER_QUEUE);
                ++it;
            } else {
@@ -184,7 +183,6 @@ void collectRootSetForThread(GCHandle gcHandle, typename Traits::MarkQueue& mark
 template <typename Traits>
 void collectRootSetGlobals(GCHandle gcHandle, typename Traits::MarkQueue& markQueue) noexcept {
    auto handle = gcHandle.collectGlobalRoots();
-    mm::StableRefRegistry::Instance().ProcessDeletions();
    // TODO: Remove useless mm::GlobalRootSet abstraction.
    for (auto value : mm::GlobalRootSet()) {
        if (internal::collectRoot<Traits>(markQueue, value.object)) {
@@ -213,6 +211,28 @@ void collectRootSet(GCHandle handle, typename Traits::MarkQueue& markQueue, F&&
    collectRootSetGlobals<Traits>(handle, markQueue);
 }

+template <typename Traits>
+void processWeaks(GCHandle gcHandle, mm::SpecialRefRegistry& registry) noexcept {
+    auto handle = gcHandle.processWeaks();
+    for (auto& object : registry.lockForIter()) {
+        auto* obj = object;
+        if (!obj) {
+            // We already processed it at some point.
+            handle.addUndisposed();
+            continue;
+        }
+        if (obj->permanent() || Traits::IsMarked(obj)) {
+            // TODO: Let's not put permanent objects in here at all?
+            // Object is alive. Nothing to do.
+            handle.addAlive();
+            continue;
+        }
+        // Object is not alive. Clear it out.
+        object = nullptr;
+        handle.addNulled();
+    }
+}
+
 } // namespace gc
 } // namespace kotlin

@@ -14,6 +14,7 @@
 #include "ObjectFactory.hpp"
 #include "ObjectTestSupport.hpp"
 #include "ExtraObjectDataFactory.hpp"
+#include "WeakRef.hpp"

 using namespace kotlin;

@@ -154,6 +155,13 @@ struct SweepTraits {
    }
 };

+struct ProcessWeakTraits {
+    static bool IsMarked(ObjHeader* object) noexcept {
+        auto& objectData = ObjectFactory::NodeRef::From(object).ObjectData();
+        return objectData.state != ObjectFactoryTraits::ObjectData::State::kUnmarked;
+    }
+};
+
 class MarkAndSweepUtilsSweepTest : public ::testing::Test {
 public:
    ~MarkAndSweepUtilsSweepTest() override {
@@ -186,6 +194,7 @@ public:
    }

    std_support::vector<ObjHeader*> Sweep() {
+        gc::processWeaks<ProcessWeakTraits>(gc::GCHandle::getByEpoch(0), specialRefRegistry_);
        gc::SweepExtraObjects<SweepTraits>(gc::GCHandle::getByEpoch(0), extraObjectFactory_);
        auto finalizers = gc::Sweep<SweepTraits>(gc::GCHandle::getByEpoch(0), objectFactory_);
        std_support::vector<ObjHeader*> objects;
@@ -244,7 +253,9 @@ public:
        auto& extraObjectData = InstallExtraData(objHeader);
        auto* setHeader = extraObjectData.GetOrSetRegularWeakReferenceImpl(objHeader, weakReference.header());
        EXPECT_EQ(setHeader, weakReference.header());
+        weakReference->weakRef = static_cast<mm::RawSpecialRef*>(specialRefRegistryThreadQueue_.createWeakRef(objHeader));
        weakReference->referred = objHeader;
+        specialRefRegistryThreadQueue_.publish();
        return weakReference;
    }

@@ -258,6 +269,9 @@ private:
    ObjectFactory::ThreadQueue objectFactoryThreadQueue_{objectFactory_, gc::Allocator()};
    ExtraObjectsDataFactory extraObjectFactory_;
    ExtraObjectsDataFactory::ThreadQueue extraObjectFactoryThreadQueue_{extraObjectFactory_};
+    mm::SpecialRefRegistry specialRefRegistry_;
+    mm::SpecialRefRegistry::ThreadQueue specialRefRegistryThreadQueue_{specialRefRegistry_};
+
    std_support::vector<ObjectFactory::FinalizerQueue> finalizers_;
 };

@@ -445,8 +459,10 @@ TEST_F(MarkAndSweepUtilsSweepTest, SweepSingleObjectWithWeakReference) {

    auto finalizers = Sweep();

-    EXPECT_THAT(finalizers, testing::UnorderedElementsAre());
+    EXPECT_THAT(finalizers, testing::UnorderedElementsAre(weakReference.header()));
    EXPECT_THAT(Alive(), testing::UnorderedElementsAre());
+
+    EXPECT_CALL(finalizerHook(), Call(weakReference.header()));
 }

 TEST_F(MarkAndSweepUtilsSweepTest, SweepSingleObjectArrayWithWeakReference) {
@@ -456,8 +472,10 @@ TEST_F(MarkAndSweepUtilsSweepTest, SweepSingleObjectArrayWithWeakReference) {

    auto finalizers = Sweep();

-    EXPECT_THAT(finalizers, testing::UnorderedElementsAre());
+    EXPECT_THAT(finalizers, testing::UnorderedElementsAre(weakReference.header()));
    EXPECT_THAT(Alive(), testing::UnorderedElementsAre());
+
+    EXPECT_CALL(finalizerHook(), Call(weakReference.header()));
 }

 TEST_F(MarkAndSweepUtilsSweepTest, SweepSingleCharArrayWithWeakReference) {
@@ -467,8 +485,10 @@ TEST_F(MarkAndSweepUtilsSweepTest, SweepSingleCharArrayWithWeakReference) {

    auto finalizers = Sweep();

-    EXPECT_THAT(finalizers, testing::UnorderedElementsAre());
+    EXPECT_THAT(finalizers, testing::UnorderedElementsAre(weakReference.header()));
    EXPECT_THAT(Alive(), testing::UnorderedElementsAre());
+
+    EXPECT_CALL(finalizerHook(), Call(weakReference.header()));
 }

 TEST_F(MarkAndSweepUtilsSweepTest, SweepSingleMarkedObjectWithWeakReference) {
@@ -44,6 +44,13 @@ struct FinalizeTraits {
    using ObjectFactory = mm::ObjectFactory<gc::SameThreadMarkAndSweep>;
 };

+struct ProcessWeaksTraits {
+    static bool IsMarked(ObjHeader* obj) noexcept {
+        auto& objectData = mm::ObjectFactory<gc::SameThreadMarkAndSweep>::NodeRef::From(obj).ObjectData();
+        return objectData.marked();
+    }
+};
+
 // Global, because it's accessed on a hot path: avoid memory load from `this`.
 std::atomic<gc::SameThreadMarkAndSweep::SafepointFlag> gSafepointFlag = gc::SameThreadMarkAndSweep::SafepointFlag::kNone;

@@ -130,6 +137,9 @@ bool gc::SameThreadMarkAndSweep::PerformFullGC() noexcept {
        gc::Mark<internal::MarkTraits>(gcHandle, markQueue_);
        auto markStats = gcHandle.getMarked();
        scheduler.gcData().UpdateAliveSetBytes(markStats.totalObjectsSize);
+
+        gc::processWeaks<ProcessWeaksTraits>(gcHandle, mm::SpecialRefRegistry::instance());
+
        gc::SweepExtraObjects<SweepTraits>(gcHandle, extraObjectsDataFactory);
        finalizerQueue = gc::Sweep<SweepTraits>(gcHandle, objectFactory_);

@@ -22,6 +22,7 @@
 #include "SingleThreadExecutor.hpp"
 #include "TestSupport.hpp"
 #include "ThreadData.hpp"
+#include "WeakRef.hpp"
 #include "std_support/Memory.hpp"
 #include "std_support/Vector.hpp"

@@ -197,6 +198,7 @@ test_support::RegularWeakReferenceImpl& InstallWeakReference(mm::ThreadData& thr
    mm::AllocateObject(&threadData, theRegularWeakReferenceImplTypeInfo, location);
    auto& weakReference = test_support::RegularWeakReferenceImpl::FromObjHeader(*location);
    auto& extraObjectData = mm::ExtraObjectData::GetOrInstall(objHeader);
+    weakReference->weakRef = static_cast<mm::RawSpecialRef*>(mm::WeakRef::create(objHeader));
    weakReference->referred = objHeader;
    auto* setWeakRef = extraObjectData.GetOrSetRegularWeakReferenceImpl(objHeader, weakReference.header());
    EXPECT_EQ(setWeakRef, weakReference.header());
@@ -207,6 +209,7 @@ class SameThreadMarkAndSweepTest : public testing::Test {
 public:
    ~SameThreadMarkAndSweepTest() {
        mm::GlobalsRegistry::Instance().ClearForTests();
+        mm::SpecialRefRegistry::instance().clearForTests();
        mm::GlobalData::Instance().extraObjectDataFactory().ClearForTests();
        mm::GlobalData::Instance().gc().impl().objectFactory().ClearForTests();
    }
@@ -333,7 +336,7 @@ TEST_F(SameThreadMarkAndSweepTest, FreeObjectsWithFinalizers) {
 }

 TEST_F(SameThreadMarkAndSweepTest, FreeObjectWithFreeWeak) {
-    RunInNewThread([](mm::ThreadData& threadData) {
+    RunInNewThread([this](mm::ThreadData& threadData) {
        auto& object1 = AllocateObject(threadData);
        auto& weak1 = ([&threadData, &object1]() -> test_support::RegularWeakReferenceImpl& {
            ObjHolder holder;
@@ -345,6 +348,7 @@ TEST_F(SameThreadMarkAndSweepTest, FreeObjectWithFreeWeak) {
        ASSERT_THAT(IsMarked(weak1.header()), false);
        ASSERT_THAT(weak1.get(), object1.header());

+        EXPECT_CALL(finalizerHook(), Call(weak1.header()));
        threadData.gc().ScheduleAndWaitFullGC();

        EXPECT_THAT(Alive(threadData), testing::UnorderedElementsAre());
@@ -1096,6 +1100,7 @@ TEST_F(SameThreadMarkAndSweepTest, FreeObjectWithFreeWeakReversedOrder) {

        global1->field1 = nullptr;

+        EXPECT_CALL(finalizerHook(), Call(weak.load()->header()));
        threadData.gc().ScheduleAndWaitFullGC();

        EXPECT_THAT(Alive(threadData), testing::UnorderedElementsAre(global1.header()));
@@ -3085,7 +3085,7 @@ void ObjHeader::destroyMetaObject(ObjHeader* object) {
  }

 #ifdef KONAN_OBJC_INTEROP
-  Kotlin_ObjCExport_detachAndReleaseAssociatedObject(meta->associatedObject_);
+  Kotlin_ObjCExport_releaseAssociatedObject(meta->associatedObject_);
 #endif

  std_support::allocator_delete(objectAllocator, meta);
@@ -3730,6 +3730,10 @@ RUNTIME_NOTHROW ALWAYS_INLINE void Kotlin_processEmptyObjectInMark(void* state,
    // no-op, used by the new MM only.
 }

+RUNTIME_NOTHROW void DisposeRegularWeakReferenceImpl(ObjHeader* counter) {
+    RuntimeFail("New MM only");
+}
+
 } // extern "C"

 #if !KONAN_NO_EXCEPTIONS
@@ -183,8 +183,8 @@ void BackRefFromAssociatedObject::detach() {
  obj_ = nullptr; // Handled in addRef/tryAddRef/releaseRef/ref.
 }

-ALWAYS_INLINE void BackRefFromAssociatedObject::assertDetached() {
-    RuntimeAssert(obj_ == nullptr, "Expecting this=%p to be detached, but found obj_=%p", this, obj_);
+void BackRefFromAssociatedObject::dealloc() {
+    RuntimeFail("New MM only");
 }

 template <ErrorPolicy errorPolicy>
@@ -86,6 +86,10 @@ ALWAYS_INLINE inline RuntimeAssertsMode runtimeAssertsMode() noexcept {
    return static_cast<RuntimeAssertsMode>(Kotlin_runtimeAssertsMode);
 }

+ALWAYS_INLINE inline bool runtimeAssertsEnabled() noexcept {
+    return runtimeAssertsMode() != RuntimeAssertsMode::kIgnore;
+}
+
 ALWAYS_INLINE inline std::string_view runtimeLogs() noexcept {
    return Kotlin_runtimeLogs == nullptr ? std::string_view() : std::string_view(Kotlin_runtimeLogs);
 }
@@ -27,6 +27,8 @@ NO_INLINE void RunFinalizerHooksImpl(ObjHeader* object, const TypeInfo* type) no
        DisposeCleaner(object);
    } else if (type == theWorkerBoundReferenceTypeInfo) {
        DisposeWorkerBoundReference(object);
+    } else if (type == theRegularWeakReferenceImplTypeInfo) {
+        DisposeRegularWeakReferenceImpl(object);
    }
 }

@@ -355,6 +355,8 @@ CODEGEN_INLINE_POLICY RUNTIME_NOTHROW void Kotlin_mm_switchThreadStateRunnable()
 CODEGEN_INLINE_POLICY void Kotlin_mm_safePointFunctionPrologue() RUNTIME_NOTHROW;
 CODEGEN_INLINE_POLICY void Kotlin_mm_safePointWhileLoopBody() RUNTIME_NOTHROW;

+RUNTIME_NOTHROW void DisposeRegularWeakReferenceImpl(ObjHeader* counter);
+
 #ifdef __cplusplus
 }
 #endif
@@ -8,7 +8,9 @@

 #include <type_traits>

+#include "ManuallyScoped.hpp"
 #include "Memory.h"
+#include "Mutex.hpp"

 // TODO: Generalize for uses outside this file.
 enum class ErrorPolicy {
@@ -30,16 +32,14 @@ class KRefSharedHolder {

  void dispose();

-  void disposeFromNative() {
-    kotlin::CalledFromNativeGuard guard;
-    dispose();
-  }
-
  OBJ_GETTER0(describe) const;

 private:
  ObjHeader* obj_;
-  ForeignRefContext context_;
+  union {
+    ForeignRefContext context_; // Legacy MM.
+    kotlin::mm::RawSpecialRef* ref_; // New MM.
+  };
 };

 static_assert(std::is_trivially_destructible_v<KRefSharedHolder>, "KRefSharedHolder destructor is not guaranteed to be called.");
@@ -47,6 +47,7 @@ static_assert(std::is_trivially_destructible_v<KRefSharedHolder>, "KRefSharedHol
 class BackRefFromAssociatedObject {
 public:
  void initForPermanentObject(ObjHeader* obj);
+
  void initAndAddRef(ObjHeader* obj);

  // Error if refCount is zero and it's called from the wrong worker with non-frozen obj_.
@@ -59,8 +60,11 @@ class BackRefFromAssociatedObject {

  void releaseRef();

+  // This does nothing with the new MM.
  void detach();
-  void assertDetached();
+
+  // This does nothing with legacy MM.
+  void dealloc();

  // Error if called from the wrong worker with non-frozen obj_.
  template <ErrorPolicy errorPolicy>
@@ -69,9 +73,18 @@ class BackRefFromAssociatedObject {
  ObjHeader* refPermanent() const;

 private:
-  ObjHeader* obj_; // May be null before [initAndAddRef] or after [detach].
-  ForeignRefContext context_;
-  volatile int refCount;
+  union {
+    struct {
+      ObjHeader* obj_; // May be null before [initAndAddRef] or after [detach].
+      ForeignRefContext context_;
+      volatile int refCount;
+    }; // Legacy MM
+    struct {
+      kotlin::mm::RawSpecialRef* ref_;
+      kotlin::ManuallyScoped<kotlin::RWSpinLock<kotlin::MutexThreadStateHandling::kIgnore>> deallocMutex_;
+    }; // New MM. Regular object.
+    ObjHeader* permanentObj_; // New MM. Permanent object.
+  };
 };

 static_assert(
@@ -99,9 +99,9 @@ static Class getOrCreateClass(const TypeInfo* typeInfo);

 namespace {

-ALWAYS_INLINE void send_releaseAsAssociatedObject(void* associatedObject, ReleaseMode mode) {
-  auto msgSend = reinterpret_cast<void (*)(void* self, SEL cmd, ReleaseMode mode)>(&objc_msgSend);
-  msgSend(associatedObject, Kotlin_ObjCExport_releaseAsAssociatedObjectSelector, mode);
+ALWAYS_INLINE void send_releaseAsAssociatedObject(void* associatedObject) {
+  auto msgSend = reinterpret_cast<void (*)(void* self, SEL cmd)>(&objc_msgSend);
+  msgSend(associatedObject, Kotlin_ObjCExport_releaseAsAssociatedObjectSelector);
 }

 } // namespace
@@ -109,22 +109,7 @@ ALWAYS_INLINE void send_releaseAsAssociatedObject(void* associatedObject, Releas
 extern "C" ALWAYS_INLINE void Kotlin_ObjCExport_releaseAssociatedObject(void* associatedObject) {
  if (associatedObject != nullptr) {
    kotlin::ThreadStateGuard guard(kotlin::ThreadState::kNative);
-    send_releaseAsAssociatedObject(associatedObject, ReleaseMode::kRelease);
-  }
-}
-
-extern "C" ALWAYS_INLINE void Kotlin_ObjCExport_detachAndReleaseAssociatedObject(void* associatedObject) {
-  if (associatedObject != nullptr) {
-    kotlin::ThreadStateGuard guard(kotlin::ThreadState::kNative);
-    send_releaseAsAssociatedObject(associatedObject, ReleaseMode::kDetachAndRelease);
-  }
-}
-
-extern "C" ALWAYS_INLINE void Kotlin_ObjCExport_detachAssociatedObject(void* associatedObject) {
-  if (associatedObject != nullptr) {
-    // Switching to Native state is not required, because detach is fast and can't call user code.
-    // Also switching is not possible, because this is called from GC.
-    send_releaseAsAssociatedObject(associatedObject, ReleaseMode::kDetach);
+    send_releaseAsAssociatedObject(associatedObject);
  }
 }

@@ -300,7 +285,7 @@ static OBJ_GETTER(blockToKotlinImp, id self, SEL cmd);
 static OBJ_GETTER(boxedBooleanToKotlinImp, NSNumber* self, SEL cmd);

 static OBJ_GETTER(SwiftObject_toKotlinImp, id self, SEL cmd);
-static void SwiftObject_releaseAsAssociatedObjectImp(id self, SEL cmd, ReleaseMode mode);
+static void SwiftObject_releaseAsAssociatedObjectImp(id self, SEL cmd);

 static void initTypeAdaptersFrom(const ObjCTypeAdapter** adapters, int count) {
  for (int index = 0; index < count; ++index) {
@@ -361,7 +346,7 @@ static void Kotlin_ObjCExport_initializeImpl() {
        swiftRootClass, releaseAsAssociatedObjectSelector,
        (IMP)SwiftObject_releaseAsAssociatedObjectImp, releaseAsAssociatedObjectTypeEncoding
      );
-      RuntimeAssert(added, "Unable to add 'releaseAsAssociatedObject:' method to SwiftObject class");
+      RuntimeAssert(added, "Unable to add 'releaseAsAssociatedObject' method to SwiftObject class");
    }
  }
 }
@@ -381,9 +366,7 @@ static OBJ_GETTER(SwiftObject_toKotlinImp, id self, SEL cmd) {
  RETURN_RESULT_OF(Kotlin_ObjCExport_convertUnmappedObjCObject, self);
 }

-static void SwiftObject_releaseAsAssociatedObjectImp(id self, SEL cmd, ReleaseMode mode) {
-  if (!ReleaseModeHasRelease(mode))
-    return;
+static void SwiftObject_releaseAsAssociatedObjectImp(id self, SEL cmd) {
  objc_release(self);
 }

@@ -18,32 +18,6 @@
 +(instancetype)createRetainedWrapper:(ObjHeader*)obj;
@end

-enum class ReleaseMode {
-    kRelease,
-    kDetachAndRelease,
-    kDetach,
-};
-
-inline bool ReleaseModeHasDetach(ReleaseMode mode) {
-    switch (mode) {
-        case ReleaseMode::kRelease:
-            return false;
-        case ReleaseMode::kDetachAndRelease:
-        case ReleaseMode::kDetach:
-            return true;
-    }
-}
-
-inline bool ReleaseModeHasRelease(ReleaseMode mode) {
-    switch (mode) {
-        case ReleaseMode::kRelease:
-        case ReleaseMode::kDetachAndRelease:
-            return true;
-        case ReleaseMode::kDetach:
-            return false;
-    }
-}
-
 extern "C" void Kotlin_ObjCExport_initializeClass(Class clazz);
 extern "C" const TypeInfo* Kotlin_ObjCExport_getAssociatedTypeInfo(Class clazz);
 extern "C" OBJ_GETTER(Kotlin_ObjCExport_convertUnmappedObjCObject, id obj);
@@ -115,8 +115,18 @@ void releaseImp(id self, SEL _cmd) {
  getBackRef(self)->releaseRef();
 }

-void releaseAsAssociatedObjectImp(id self, SEL _cmd, ReleaseMode mode) {
+void releaseAsAssociatedObjectImp(id self, SEL _cmd) {
  auto* classData = GetKotlinClassData(self);
+  if (CurrentMemoryModel == MemoryModel::kExperimental) {
+    // No need for any special handling. Weak reference handling machinery
+    // has already cleaned up the reference to Kotlin object.
+    // [super release]
+    Class clazz = classData->objcClass;
+    struct objc_super s = {self, clazz};
+    auto messenger = reinterpret_cast<void (*) (struct objc_super*, SEL _cmd)>(objc_msgSendSuper2);
+    messenger(&s, @selector(release));
+    return;
+  }

  // This function is called by the GC. It made a decision to reclaim Kotlin object, and runs
  // deallocation hooks at the moment, including deallocation of the "associated object" ([self])
@@ -132,24 +142,28 @@ void releaseAsAssociatedObjectImp(id self, SEL _cmd, ReleaseMode mode) {
  // Generally retaining and releasing Kotlin object that is being deallocated would lead to
  // use-after-dispose and double-dispose problems (with unpredictable consequences) or to an assertion failure.
  // To workaround this, detach the back ref from the Kotlin object:
-  if (ReleaseModeHasDetach(mode)) {
-    backRef->detach();
-  } else {
-    // With Mark&Sweep this object should already have been detached earlier.
-    backRef->assertDetached();
-  }
+  backRef->detach();

  // So retain/release/etc. on [self] won't affect the Kotlin object, and an attempt to get
  // the reference to it (e.g. when calling Kotlin method on [self]) would crash.
  // The latter is generally ok, because by the time superclass dealloc gets launched, subclass state
  // should already be deinitialized, and Kotlin methods operate on the subclass.
-  if (ReleaseModeHasRelease(mode)) {
-    // [super release]
-    Class clazz = classData->objcClass;
-    struct objc_super s = {self, clazz};
-    auto messenger = reinterpret_cast<void (*) (struct objc_super*, SEL _cmd)>(objc_msgSendSuper2);
-    messenger(&s, @selector(release));
-  }
+  // [super release]
+  Class clazz = classData->objcClass;
+  struct objc_super s = {self, clazz};
+  auto messenger = reinterpret_cast<void (*) (struct objc_super*, SEL _cmd)>(objc_msgSendSuper2);
+  messenger(&s, @selector(release));
+}
+
+void deallocImp(id self, SEL _cmd) {
+  getBackRef(self)->dealloc();
+
+  // [super dealloc]
+  auto* classData = GetKotlinClassData(self);
+  Class clazz = classData->objcClass;
+  struct objc_super s = {self, clazz};
+  auto messenger = reinterpret_cast<void (*) (struct objc_super*, SEL _cmd)>(objc_msgSendSuper2);
+  messenger(&s, @selector(dealloc));
 }

 }
@@ -277,6 +291,9 @@ void* CreateKotlinObjCClass(const KotlinObjCClassInfo* info) {
  AddNSObjectOverride(false, newClass, @selector(release), (void*)&releaseImp);
  AddNSObjectOverride(false, newClass, Kotlin_ObjCExport_releaseAsAssociatedObjectSelector,
      (void*)&releaseAsAssociatedObjectImp);
+  if (CurrentMemoryModel == MemoryModel::kExperimental) {
+    AddNSObjectOverride(false, newClass, @selector(dealloc), (void*)&deallocImp);
+  }

  AddMethods(newClass, info->instanceMethods, info->instanceMethodsNum);
  AddMethods(newMetaclass, info->classMethods, info->classMethodsNum);
@@ -12,8 +12,6 @@
 #if KONAN_OBJC_INTEROP

 extern "C" ALWAYS_INLINE void Kotlin_ObjCExport_releaseAssociatedObject(void* associatedObject);
-extern "C" ALWAYS_INLINE void Kotlin_ObjCExport_detachAndReleaseAssociatedObject(void* associatedObject);
-extern "C" ALWAYS_INLINE void Kotlin_ObjCExport_detachAssociatedObject(void* associatedObject);

 namespace konan {
 class AutoreleasePool : private kotlin::Pinned {
@@ -345,6 +345,7 @@ public:
 };

 struct RegularWeakReferenceImplPayload {
+    void* weakRef;
    void* referred;

    using Field = ObjHeader* RegularWeakReferenceImplPayload::*;
@@ -22,7 +22,7 @@ import kotlin.native.internal.*
 *              [weak1]   [weak2]
 *                 \       /
 *                  V     V
- *     ...... [WeakReferenceImpl] <------
+ *     ... [RegularWeakReferenceImpl] <-
 *     .                                |
 *     .                                |
 *      ->[Object] -> [ExtraObjectData]-
@@ -50,8 +50,10 @@ internal class WeakReferenceCounterLegacyMM(var referred: COpaquePointer?) : Wea

@NoReorderFields
@ExportTypeInfo("theRegularWeakReferenceImplTypeInfo")
+@HasFinalizer // TODO: Consider just using Cleaners.
 internal class RegularWeakReferenceImpl(
-    val referred: COpaquePointer,
+    val weakRef: COpaquePointer,
+    val referred: COpaquePointer, // TODO: This exists only for the ExtraObjectData's sake. Refactor and remove.
 ) : WeakReferenceImpl() {
    @GCUnsafeCall("Konan_RegularWeakReferenceImpl_get")
    external override fun get(): Any?
@@ -73,7 +75,7 @@ internal fun makeWeakReferenceCounterLegacyMM(referred: COpaquePointer) = WeakRe

 // Create a counter object.
@ExportForCppRuntime
-internal fun makeRegularWeakReferenceImpl(referred: COpaquePointer) = RegularWeakReferenceImpl(referred)
+internal fun makeRegularWeakReferenceImpl(weakRef: COpaquePointer, referred: COpaquePointer) = RegularWeakReferenceImpl(weakRef, referred)

 internal class PermanentWeakReferenceImpl(val referred: Any): kotlin.native.ref.WeakReferenceImpl() {
    override fun get(): Any? = referred
@@ -5,7 +5,7 @@

 #include "Memory.h"

-#include "StableRefRegistry.hpp"
+#include "StableRef.hpp"
 #include "ThreadData.hpp"
 #include "ThreadRegistry.hpp"

@@ -14,22 +14,16 @@ using namespace kotlin;
 namespace {

 #if !KONAN_NO_EXCEPTIONS
-class ExceptionObjHolderImpl : public ExceptionObjHolder {
+class ExceptionObjHolderImpl : public ExceptionObjHolder, private Pinned {
 public:
-    explicit ExceptionObjHolderImpl(ObjHeader* obj) noexcept {
-        auto* threadData = mm::ThreadRegistry::Instance().CurrentThreadData();
-        stableRef_ = threadData->stableRefThreadQueue().Insert(obj);
-    }
+    explicit ExceptionObjHolderImpl(ObjHeader* obj) noexcept : stableRef_(mm::StableRef::create(obj)) {}

-    ~ExceptionObjHolderImpl() override {
-        auto* threadData = mm::ThreadRegistry::Instance().CurrentThreadData();
-        threadData->stableRefThreadQueue().Erase(stableRef_);
-    }
+    ~ExceptionObjHolderImpl() override { std::move(stableRef_).dispose(); }

-    ObjHeader* obj() noexcept { return **stableRef_; }
+    ObjHeader* obj() noexcept { return *stableRef_; }

 private:
-    mm::StableRefRegistry::Node* stableRef_;
+    mm::StableRef stableRef_;
 };
 #endif

@@ -21,11 +21,9 @@ namespace {
 class ExceptionObjHolderTest : public ::testing::Test {
 public:
    static std_support::vector<ObjHeader*> Collect(mm::ThreadData& threadData) {
-        auto& stableRefs = mm::StableRefRegistry::Instance();
-        stableRefs.ProcessThread(&threadData);
-        stableRefs.ProcessDeletions();
+        threadData.specialRefRegistry().publish();
        std_support::vector<ObjHeader*> result;
-        for (const auto& obj : stableRefs.LockForIter()) {
+        for (const auto& obj : mm::SpecialRefRegistry::instance().roots()) {
            result.push_back(obj);
        }
        return result;
@@ -66,7 +64,7 @@ TEST_F(ExceptionObjHolderTest, ThrowInsideCatch) {
            try {
                ExceptionObjHolder::Throw(&exception2);
            } catch (...) {
-                EXPECT_THAT(Collect(threadData), testing::ElementsAre(&exception1, &exception2));
+                EXPECT_THAT(Collect(threadData), testing::UnorderedElementsAre(&exception1, &exception2));
            }
            EXPECT_THAT(Collect(threadData), testing::ElementsAre(&exception1));
        }
@@ -94,7 +92,7 @@ TEST_F(ExceptionObjHolderTest, StoreException) {
        } catch (...) {
            storedException2 = std::current_exception();
        }
-        EXPECT_THAT(Collect(threadData), testing::ElementsAre(&exception1, &exception2));
+        EXPECT_THAT(Collect(threadData), testing::UnorderedElementsAre(&exception1, &exception2));

        storedException1 = std::exception_ptr();
        EXPECT_THAT(Collect(threadData), testing::ElementsAre(&exception2));
@@ -54,7 +54,9 @@ mm::ExtraObjectData& mm::ExtraObjectData::Install(ObjHeader* object) noexcept {
 void mm::ExtraObjectData::Uninstall() noexcept {
    auto *object = GetBaseObject();
    atomicSetRelease(const_cast<const TypeInfo**>(&object->typeInfoOrMeta_), typeInfo_);
-    RuntimeAssert(!object->has_meta_object(), "Object has metaobject after removing metaobject");
+    RuntimeAssert(
+            !object->has_meta_object(), "Object %p has metaobject %p after removing metaobject %p", object, object->meta_object_or_null(),
+            this);

 #ifdef KONAN_OBJC_INTEROP
    Kotlin_ObjCExport_releaseAssociatedObject(associatedObject_);
@@ -62,12 +64,6 @@ void mm::ExtraObjectData::Uninstall() noexcept {
 #endif
 }

-void mm::ExtraObjectData::DetachAssociatedObject() noexcept {
-#ifdef KONAN_OBJC_INTEROP
-    Kotlin_ObjCExport_detachAssociatedObject(associatedObject_);
-#endif
-}
-
 bool mm::ExtraObjectData::HasAssociatedObject() noexcept {
 #ifdef KONAN_OBJC_INTEROP
    return associatedObject_ != nullptr;
@@ -77,10 +73,7 @@ bool mm::ExtraObjectData::HasAssociatedObject() noexcept {
 }

 void mm::ExtraObjectData::ClearRegularWeakReferenceImpl() noexcept {
-    if (!HasRegularWeakReferenceImpl()) return;
-
    auto *object = GetBaseObject();
-    disposeRegularWeakReferenceImpl(GetRegularWeakReferenceImpl());
    // 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?
@@ -88,9 +81,16 @@ void mm::ExtraObjectData::ClearRegularWeakReferenceImpl() noexcept {
 }

 mm::ExtraObjectData::~ExtraObjectData() {
-    RuntimeAssert(!HasRegularWeakReferenceImpl(), "Object must have cleared weak references");
+    auto* weakReference = weakReferenceOrBaseObject_.load(std::memory_order_relaxed);
+    if (hasPointerBits(weakReference, WEAK_REF_TAG)) {
+        weakReference = clearPointerBits(weakReference, WEAK_REF_TAG);
+    } else {
+        weakReference = nullptr;
+    }
+    RuntimeAssert(weakReference == nullptr, "ExtraObjectData %p must have cleared weak reference %p", this, weakReference);

 #ifdef KONAN_OBJC_INTEROP
-    RuntimeAssert(associatedObject_ == nullptr, "Object must have cleared associated object");
+    auto* associatedObject = associatedObject_.load(std::memory_order_relaxed);
+    RuntimeAssert(associatedObject == nullptr, "ExtraObjectData %p must have cleared associated object %p", this, associatedObject);
 #endif
 }
@@ -47,13 +47,12 @@ public:
    std::atomic<void*>& AssociatedObject() noexcept { return associatedObject_; }
 #endif
    bool HasAssociatedObject() noexcept;
-    void DetachAssociatedObject() noexcept;

    bool getFlag(Flags value) noexcept { return (flags_.load() & (1u << static_cast<uint32_t>(value))) != 0; }
    void setFlag(Flags value) noexcept { flags_.fetch_or(1u << static_cast<uint32_t>(value)); }

    bool HasRegularWeakReferenceImpl() noexcept { return hasPointerBits(weakReferenceOrBaseObject_.load(), WEAK_REF_TAG); }
-    void ClearRegularWeakReferenceImpl() noexcept;
+    void ClearRegularWeakReferenceImpl() noexcept; // TODO: Only exists for the sake of GetBaseObject. Refactor to remove the need for it.
    ObjHeader* GetRegularWeakReferenceImpl() noexcept {
        auto* pointer = weakReferenceOrBaseObject_.load();
        if (hasPointerBits(pointer, WEAK_REF_TAG)) return clearPointerBits(pointer, WEAK_REF_TAG);
@@ -39,9 +39,5 @@ void mm::ExtraObjectDataFactory::ProcessThread(mm::ThreadData* threadData) noexc
    threadData->extraObjectDataThreadQueue().Publish();
 }

-void mm::ExtraObjectDataFactory::ProcessDeletions() noexcept {
-    extraObjects_.ApplyDeletions();
-}
-
 mm::ExtraObjectDataFactory::ExtraObjectDataFactory() = default;
 mm::ExtraObjectDataFactory::~ExtraObjectDataFactory() = default;
@@ -44,9 +44,6 @@ public:
    // when it's asked by GC to stop.
    void ProcessThread(mm::ThreadData* threadData) noexcept;

-    // Lock registry and apply deletions. Should be called on GC thread after all threads have published, and before `LockForIter`.
-    void ProcessDeletions() noexcept;
-
    // Lock registry for safe iteration.
    Iterable LockForIter() noexcept { return extraObjects_.LockForIter(); }

@@ -10,7 +10,7 @@
 #include "GlobalsRegistry.hpp"
 #include "GC.hpp"
 #include "GCScheduler.hpp"
-#include "StableRefRegistry.hpp"
+#include "SpecialRefRegistry.hpp"
 #include "ThreadRegistry.hpp"
 #include "Utils.hpp"
 #include "ExtraObjectDataFactory.hpp"
@@ -26,7 +26,7 @@ public:

    ThreadRegistry& threadRegistry() noexcept { return threadRegistry_; }
    GlobalsRegistry& globalsRegistry() noexcept { return globalsRegistry_; }
-    StableRefRegistry& stableRefRegistry() noexcept { return stableRefRegistry_; }
+    SpecialRefRegistry& specialRefRegistry() noexcept { return specialRefRegistry_; }
    ExtraObjectDataFactory& extraObjectDataFactory() noexcept { return extraObjectDataFactory_; }
    gc::GC& gc() noexcept { return gc_; }
    AppStateTracking& appStateTracking() noexcept { return appStateTracking_; }
@@ -41,7 +41,7 @@ private:
    ThreadRegistry threadRegistry_;
    AppStateTracking appStateTracking_;
    GlobalsRegistry globalsRegistry_;
-    StableRefRegistry stableRefRegistry_;
+    SpecialRefRegistry specialRefRegistry_;
    ExtraObjectDataFactory extraObjectDataFactory_;
    gc::GC gc_;
 };
@@ -16,7 +16,7 @@
 #include "ObjectOps.hpp"
 #include "Porting.h"
 #include "Runtime.h"
-#include "StableRefRegistry.hpp"
+#include "StableRef.hpp"
 #include "ThreadData.hpp"
 #include "ThreadRegistry.hpp"
 #include "ThreadState.hpp"
@@ -24,29 +24,6 @@

 using namespace kotlin;

-// TODO: This name does not make sense anymore.
-// Delete all means of creating this type directly as it only serves
-// as a typedef for `mm::StableRefRegistry::Node`.
-class ForeignRefManager : Pinned {
-public:
-    ForeignRefManager() = delete;
-    ~ForeignRefManager() = delete;
-};
-
-namespace {
-
-// `reinterpret_cast` to it and back to the same type
-// will yield precisely the same pointer, so it's safe.
-ALWAYS_INLINE ForeignRefManager* ToForeignRefManager(mm::StableRefRegistry::Node* data) {
-    return reinterpret_cast<ForeignRefManager*>(data);
-}
-
-ALWAYS_INLINE mm::StableRefRegistry::Node* FromForeignRefManager(ForeignRefManager* manager) {
-    return reinterpret_cast<mm::StableRefRegistry::Node*>(manager);
-}
-
-} // namespace
-
 ObjHeader* ObjHeader::GetWeakCounter() {
    RuntimeFail("Only for legacy MM");
 }
@@ -464,16 +441,6 @@ extern "C" RUNTIME_NOTHROW void PerformFullGC(MemoryState* memory) {
    threadData->gc().ScheduleAndWaitFullGCWithFinalizers();
 }

-extern "C" RUNTIME_NOTHROW OBJ_GETTER(TryRef, ObjHeader* object) {
-    // TODO: With CMS this needs:
-    //       * during marking phase if `object` is unmarked: barrier (might be automatic because of the stack write)
-    //         and return `object`;
-    //       * during marking phase if `object` is marked: return `object`;
-    //       * during sweeping phase if `object` is unmarked: return nullptr;
-    //       * during sweeping phase if `object` is marked: return `object`;
-    RETURN_OBJ(object);
-}
-
 extern "C" RUNTIME_NOTHROW bool ClearSubgraphReferences(ObjHeader* root, bool checked) {
    // TODO: Remove when legacy MM is gone.
    return true;
@@ -483,24 +450,23 @@ extern "C" RUNTIME_NOTHROW void* CreateStablePointer(ObjHeader* object) {
    if (!object)
        return nullptr;

-    auto* threadData = mm::ThreadRegistry::Instance().CurrentThreadData();
-    AssertThreadState(threadData, ThreadState::kRunnable);
-    return mm::StableRefRegistry::Instance().RegisterStableRef(threadData, object);
+    AssertThreadState(ThreadState::kRunnable);
+    return static_cast<mm::RawSpecialRef*>(mm::StableRef::create(object));
 }

 extern "C" RUNTIME_NOTHROW void DisposeStablePointer(void* pointer) {
-    DisposeStablePointerFor(kotlin::mm::GetMemoryState(), pointer);
+    if (!pointer) return;
+
+    // Can be safely called in any thread state.
+    mm::StableRef(static_cast<mm::RawSpecialRef*>(pointer)).dispose();
 }

 extern "C" RUNTIME_NOTHROW void DisposeStablePointerFor(MemoryState* memoryState, void* pointer) {
    if (!pointer)
        return;

-    auto* threadData = memoryState->GetThreadData();
-    AssertThreadState(threadData, ThreadState::kRunnable);
-
-    auto* node = static_cast<mm::StableRefRegistry::Node*>(pointer);
-    mm::StableRefRegistry::Instance().UnregisterStableRef(threadData, node);
+    // Can be safely called in any thread state.
+    mm::StableRef(static_cast<mm::RawSpecialRef*>(pointer)).disposeOn(*mm::FromMemoryState(memoryState)->Get());
 }

 extern "C" RUNTIME_NOTHROW OBJ_GETTER(DerefStablePointer, void* pointer) {
@@ -508,24 +474,19 @@ extern "C" RUNTIME_NOTHROW OBJ_GETTER(DerefStablePointer, void* pointer) {
        RETURN_OBJ(nullptr);

    AssertThreadState(ThreadState::kRunnable);
-
-    auto* node = static_cast<mm::StableRefRegistry::Node*>(pointer);
-    ObjHeader* object = **node;
-    RETURN_OBJ(object);
+    RETURN_OBJ(*mm::StableRef(static_cast<mm::RawSpecialRef*>(pointer)));
 }

 extern "C" RUNTIME_NOTHROW OBJ_GETTER(AdoptStablePointer, void* pointer) {
    if (!pointer)
        RETURN_OBJ(nullptr);

-    auto* threadData = mm::ThreadRegistry::Instance().CurrentThreadData();
-    AssertThreadState(threadData, ThreadState::kRunnable);
-    auto* node = static_cast<mm::StableRefRegistry::Node*>(pointer);
-    ObjHeader* object = **node;
-    // Make sure `object` stays in the rootset: put it on the stack before removing it from `StableRefRegistry`.
-    mm::SetStackRef(OBJ_RESULT, object);
-    mm::StableRefRegistry::Instance().UnregisterStableRef(threadData, node);
-    return object;
+    AssertThreadState(ThreadState::kRunnable);
+    mm::StableRef stableRef(static_cast<mm::RawSpecialRef*>(pointer));
+    auto* obj = *stableRef;
+    mm::SetStackRef(OBJ_RESULT, obj);
+    std::move(stableRef).dispose();
+    return obj;
 }

 extern "C" void MutationCheck(ObjHeader* obj) {
@@ -554,22 +515,6 @@ extern "C" void EnsureNeverFrozen(ObjHeader* obj) {
    }
 }

-extern "C" ForeignRefContext InitForeignRef(ObjHeader* object) {
-    AssertThreadState(ThreadState::kRunnable);
-    auto* threadData = mm::ThreadRegistry::Instance().CurrentThreadData();
-    auto* node = mm::StableRefRegistry::Instance().RegisterStableRef(threadData, object);
-    return ToForeignRefManager(node);
-}
-
-extern "C" void DeinitForeignRef(ObjHeader* object, ForeignRefContext context) {
-    AssertThreadState(ThreadState::kRunnable);
-    RuntimeAssert(context != nullptr, "DeinitForeignRef must not be called for InitLocalForeignRef");
-    auto* threadData = mm::ThreadRegistry::Instance().CurrentThreadData();
-    auto* node = FromForeignRefManager(context);
-    RuntimeAssert(object == **node, "Must correspond to the same object");
-    mm::StableRefRegistry::Instance().UnregisterStableRef(threadData, node);
-}
-
 extern "C" void CheckGlobalsAccessible() {
    // TODO: Remove when legacy MM is gone.
    // Always accessible
@@ -659,3 +604,7 @@ RUNTIME_NOTHROW extern "C" OBJ_GETTER(Konan_WeakReferenceCounterLegacyMM_get, Ob
 RUNTIME_NOTHROW extern "C" OBJ_GETTER(Konan_RegularWeakReferenceImpl_get, ObjHeader* weakRef) {
    RETURN_RESULT_OF(mm::derefRegularWeakReferenceImpl, weakRef);
 }
+
+RUNTIME_NOTHROW extern "C" void DisposeRegularWeakReferenceImpl(ObjHeader* weakRef) {
+    mm::disposeRegularWeakReferenceImpl(weakRef);
+}
@@ -39,10 +39,4 @@ extern "C" struct MemoryState : kotlin::Pinned {
    }
 };

-extern "C" {
-ForeignRefContext InitForeignRef(ObjHeader* object);
-void DeinitForeignRef(ObjHeader* object, ForeignRefContext context);
-RUNTIME_NOTHROW OBJ_GETTER(TryRef, ObjHeader* object);
-}
-
 #endif //RUNTIME_MEMORYPRIVATE_HPP
@@ -5,19 +5,22 @@

 #include "MemorySharedRefs.hpp"

-#include "MemoryPrivate.hpp"
+#include <shared_mutex>
+
+#include "ObjCBackRef.hpp"
+#include "StableRef.hpp"

 using namespace kotlin;

 void KRefSharedHolder::initLocal(ObjHeader* obj) {
    RuntimeAssert(obj != nullptr, "must not be null");
-    context_ = nullptr;
+    ref_ = nullptr;
    obj_ = obj;
 }

 void KRefSharedHolder::init(ObjHeader* obj) {
    RuntimeAssert(obj != nullptr, "must not be null");
-    context_ = InitForeignRef(obj);
+    ref_ = static_cast<mm::RawSpecialRef*>(mm::StableRef::create(obj));
    obj_ = obj;
 }

@@ -33,46 +36,30 @@ template ObjHeader* KRefSharedHolder::ref<ErrorPolicy::kThrow>() const;
 template ObjHeader* KRefSharedHolder::ref<ErrorPolicy::kTerminate>() const;

 void KRefSharedHolder::dispose() {
-    if (obj_ == nullptr) {
-        // To handle the case when it is not initialized. See [KotlinMutableSet/Dictionary dealloc].
+    // Handles the case when it is not initialized. See [KotlinMutableSet/Dictionary dealloc].
+    if (!ref_) {
        return;
    }
-
-    DeinitForeignRef(obj_, context_);
+    std::move(mm::StableRef::reinterpret(ref_)).dispose();
+    // obj_ is dangling now.
 }

 void BackRefFromAssociatedObject::initForPermanentObject(ObjHeader* obj) {
    RuntimeAssert(obj != nullptr, "must not be null");
    RuntimeAssert(obj->permanent(), "Can only be called with permanent object");
-    obj_ = obj;
+    permanentObj_ = obj;
 }

 void BackRefFromAssociatedObject::initAndAddRef(ObjHeader* obj) {
    RuntimeAssert(obj != nullptr, "must not be null");
    RuntimeAssert(!obj->permanent(), "Can only be called with non-permanent object");
-    obj_ = obj;
-
-    // Generally a specialized addRef below:
-    context_ = InitForeignRef(obj);
-    refCount = 1;
+    ref_ = static_cast<mm::RawSpecialRef*>(mm::ObjCBackRef::create(obj));
+    deallocMutex_.construct();
 }

 template <ErrorPolicy errorPolicy>
 void BackRefFromAssociatedObject::addRef() {
-    static_assert(errorPolicy != ErrorPolicy::kDefaultValue, "Cannot use default return value here");
-
-    // Can be called both from Native state (if ObjC or Swift code adds RC)
-    // and from Runnable state (Kotlin_ObjCExport_refToObjC).
-
-    if (atomicAdd(&refCount, 1) == 1) {
-        if (obj_ == nullptr) return; // E.g. after [detach].
-
-        kotlin::CalledFromNativeGuard guard(/* reentrant */ true);
-
-        // Foreign reference has already been deinitialized (see [releaseRef]).
-        // Create a new one:
-        context_ = InitForeignRef(obj_);
-    }
+    mm::ObjCBackRef::reinterpret(ref_).retain();
 }

 template void BackRefFromAssociatedObject::addRef<ErrorPolicy::kThrow>();
@@ -80,54 +67,36 @@ template void BackRefFromAssociatedObject::addRef<ErrorPolicy::kTerminate>();

 template <ErrorPolicy errorPolicy>
 bool BackRefFromAssociatedObject::tryAddRef() {
-    static_assert(errorPolicy != ErrorPolicy::kDefaultValue, "Cannot use default return value here");
-    kotlin::CalledFromNativeGuard guard;
-
-    if (obj_ == nullptr) return false; // E.g. after [detach].
-
-    ObjHolder holder;
-    ObjHeader* obj = TryRef(obj_, holder.slot());
-    // Failed to lock weak reference.
-    if (obj == nullptr) return false;
-    RuntimeAssert(obj == obj_, "Mismatched locked weak. obj=%p obj_=%p", obj, obj_);
-    // TODO: This is a very weird way to ask for "unsafe" addRef.
-    addRef<ErrorPolicy::kIgnore>();
-    return true;
+    // Only this method can be called in parallel with dealloc.
+    std::shared_lock guard(*deallocMutex_, std::try_to_lock);
+    if (!guard) {
+        // That means `dealloc` is running in parallel, so
+        // cannot possibly retain.
+        return false;
+    }
+    return mm::ObjCBackRef::reinterpret(ref_).tryRetain();
 }

 template bool BackRefFromAssociatedObject::tryAddRef<ErrorPolicy::kThrow>();
 template bool BackRefFromAssociatedObject::tryAddRef<ErrorPolicy::kTerminate>();

 void BackRefFromAssociatedObject::releaseRef() {
-    ForeignRefContext context = context_;
-    if (atomicAdd(&refCount, -1) == 0) {
-        if (obj_ == nullptr) return; // E.g. after [detach].
-
-        kotlin::CalledFromNativeGuard guard;
-
-        // Note: by this moment "subsequent" addRef may have already happened and patched context_.
-        // So use the value loaded before refCount update:
-        DeinitForeignRef(obj_, context);
-        // From this moment [context] is generally a dangling pointer.
-        // This is handled in [IsForeignRefAccessible] and [addRef].
-        // TODO: This probably isn't fine in new MM. Make sure it works.
-    }
+    mm::ObjCBackRef::reinterpret(ref_).release();
 }

 void BackRefFromAssociatedObject::detach() {
-    RuntimeAssert(atomicGet(&refCount) == 0, "unexpected refCount");
-    obj_ = nullptr; // Handled in addRef/tryAddRef/releaseRef/ref.
+    RuntimeFail("Legacy MM only");
 }

-ALWAYS_INLINE void BackRefFromAssociatedObject::assertDetached() {
-    RuntimeAssert(obj_ == nullptr, "Expecting this=%p to be detached, but found obj_=%p", this, obj_);
+void BackRefFromAssociatedObject::dealloc() {
+    // This will wait for all `tryAddRef` to finish.
+    std::unique_lock guard(*deallocMutex_);
+    std::move(mm::ObjCBackRef::reinterpret(ref_)).dispose();
 }

 template <ErrorPolicy errorPolicy>
 ObjHeader* BackRefFromAssociatedObject::ref() const {
-    kotlin::AssertThreadState(kotlin::ThreadState::kRunnable);
-    RuntimeAssert(obj_ != nullptr, "no valid Kotlin object found");
-    return obj_;
+    return *mm::ObjCBackRef::reinterpret(ref_);
 }

 template ObjHeader* BackRefFromAssociatedObject::ref<ErrorPolicy::kDefaultValue>() const;
@@ -135,5 +104,5 @@ template ObjHeader* BackRefFromAssociatedObject::ref<ErrorPolicy::kThrow>() cons
 template ObjHeader* BackRefFromAssociatedObject::ref<ErrorPolicy::kTerminate>() const;

 ObjHeader* BackRefFromAssociatedObject::refPermanent() const {
-  return obj_;
+    return permanentObj_;
 }
@@ -0,0 +1,16 @@
+/*
+ * Copyright 2010-2023 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the LICENSE file.
+ */
+
+#include "ObjCBackRef.hpp"
+
+#include "ThreadData.hpp"
+
+using namespace kotlin;
+
+// static
+mm::ObjCBackRef mm::ObjCBackRef::create(ObjHeader* obj) noexcept {
+    RuntimeAssert(obj != nullptr, "Creating ObjCBackRef for null object");
+    return mm::ThreadRegistry::Instance().CurrentThreadData()->specialRefRegistry().createObjCBackRef(obj);
+}
@@ -0,0 +1,109 @@
+/*
+ * Copyright 2010-2023 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the LICENSE file.
+ */
+
+#pragma once
+
+#include "Memory.h"
+#include "RawPtr.hpp"
+#include "SpecialRefRegistry.hpp"
+#include "Utils.hpp"
+
+namespace kotlin::mm {
+
+// Reference from an ObjC associated object back into a Kotlin object.
+// GC automatically tracks references with refcount > 0 as roots, and invalidates references with refcount = 0 when the Kotlin object is
+// collected. Use `create` and `dispose` to create and destroy the back reference.
+class ObjCBackRef : private MoveOnly {
+public:
+    ObjCBackRef() noexcept = default;
+
+    // Cast raw ref into a back reference.
+    explicit ObjCBackRef(RawSpecialRef* raw) noexcept : node_(SpecialRefRegistry::Node::fromRaw(raw)) {}
+
+    // Cast back reference into a raw ref
+    [[nodiscard("must be manually disposed")]] explicit operator RawSpecialRef*() && noexcept {
+        // Make sure to move out from node_.
+        auto node = std::move(node_);
+        return node->asRaw();
+    }
+
+    // Create new back reference for `obj`.
+    [[nodiscard("must be manually disposed")]] static ObjCBackRef create(ObjHeader* obj) noexcept;
+
+    // Dispose back reference.
+    void dispose() && noexcept {
+        RuntimeAssert(node_, "Disposing null ObjCBackRef");
+        // Make sure to move out from node_.
+        auto node = std::move(node_);
+        // Can be safely called with any thread state.
+        node->dispose();
+    }
+
+    // Increment refcount.
+    void retain() noexcept {
+        // In objc import if KtClass inherits from ObjCClass
+        // calling [self retain] inside [ObjCClass dealloc] will lead to
+        // this->retain() being called after this->dispose()
+        if (!node_) return;
+        // Can be safely called with any thread state.
+        node_->retainRef();
+    }
+
+    // Decrement refcount.
+    void release() noexcept {
+        // In objc import if KtClass inherits from ObjCClass
+        // calling [self release] inside [ObjCClass dealloc] will lead to
+        // this->release() being called after this->dispose()
+        if (!node_) return;
+        // Can be safely called with any thread state.
+        node_->releaseRef();
+    }
+
+    // Try incrementing refcount. Will fail if the underlying object is not alive.
+    [[nodiscard("refcount change must be processed")]] bool tryRetain() noexcept {
+        // In objc export if ObjCClass is objc_setAssociatedObject with KtClass
+        // calling [KtClass _tryRetain] inside [ObjCClass dealloc] will lead to
+        // this->tryRetain() being called after this->dispose()
+        if (!node_) return false;
+        CalledFromNativeGuard guard;
+        return tryRetainIgnoreState();
+    }
+
+    // Get the underlying object.
+    // The result is only safe to use only with refcount > 0.
+    [[nodiscard("expensive pure function")]] ObjHeader* operator*() const noexcept {
+        // In objc import if KtClass inherits from ObjCClass
+        // calling [self retain] inside [ObjCClass dealloc] and then passing the retained
+        // reference back to Kotlin will lead to
+        // this->operator*() being called after this->dispose()
+        if (!node_) return nullptr;
+        return node_->ref();
+    }
+
+    bool tryRetainForTests() noexcept { return tryRetainIgnoreState(); }
+
+    static ObjCBackRef& reinterpret(RawSpecialRef*& raw) noexcept { return reinterpret_cast<ObjCBackRef&>(raw); }
+
+    static const ObjCBackRef& reinterpret(RawSpecialRef* const& raw) noexcept { return reinterpret_cast<const ObjCBackRef&>(raw); }
+
+private:
+    bool tryRetainIgnoreState() noexcept {
+        ObjHolder holder;
+        if (auto* obj = node_->tryRef(holder.slot())) {
+            node_->retainRef();
+            return true;
+        }
+        return false;
+    }
+
+    raw_ptr<SpecialRefRegistry::Node> node_;
+};
+
+static_assert(sizeof(ObjCBackRef) == sizeof(void*), "ObjCBackRef must be a thin wrapper around pointer");
+static_assert(alignof(ObjCBackRef) == alignof(void*), "ObjCBackRef must be a thin wrapper around pointer");
+static_assert(
+        std::is_trivially_destructible_v<ObjCBackRef>, "ObjCBackRef must be trivially destructible. Destruction is manual via dispose()");
+
+} // namespace kotlin::mm
@@ -89,7 +89,7 @@ mm::GlobalRootSet::Value mm::GlobalRootSet::Iterator::operator*() noexcept {
        case Phase::kGlobals:
            return {**globalsIterator_, Source::kGlobal};
        case Phase::kStableRefs:
-            return {*stableRefsIterator_, Source::kStableRef};
+            return {*specialRefsIterator_, Source::kStableRef};
        case Phase::kDone:
            RuntimeFail("Cannot dereference");
    }
@@ -102,7 +102,7 @@ mm::GlobalRootSet::Iterator& mm::GlobalRootSet::Iterator::operator++() noexcept
            Init();
            return *this;
        case Phase::kStableRefs:
-            ++stableRefsIterator_;
+            ++specialRefsIterator_;
            Init();
            return *this;
        case Phase::kDone:
@@ -121,7 +121,7 @@ bool mm::GlobalRootSet::Iterator::operator==(const Iterator& rhs) const noexcept
        case Phase::kGlobals:
            return globalsIterator_ == rhs.globalsIterator_;
        case Phase::kStableRefs:
-            return stableRefsIterator_ == rhs.stableRefsIterator_;
+            return specialRefsIterator_ == rhs.specialRefsIterator_;
    }
 }

@@ -131,10 +131,10 @@ void mm::GlobalRootSet::Iterator::Init() noexcept {
            case Phase::kGlobals:
                if (globalsIterator_ != owner_.globalsIterable_.end()) return;
                phase_ = Phase::kStableRefs;
-                stableRefsIterator_ = owner_.stableRefsIterable_.begin();
+                specialRefsIterator_ = owner_.specialRefsIterable_.begin();
                break;
            case Phase::kStableRefs:
-                if (stableRefsIterator_ != owner_.stableRefsIterable_.end()) return;
+                if (specialRefsIterator_ != owner_.specialRefsIterable_.end()) return;
                phase_ = Phase::kDone;
                break;
            case Phase::kDone:
@@ -146,4 +146,4 @@ void mm::GlobalRootSet::Iterator::Init() noexcept {
 mm::ThreadRootSet::ThreadRootSet(ThreadData& threadData) noexcept : ThreadRootSet(threadData.shadowStack(), threadData.tls()) {}

 mm::GlobalRootSet::GlobalRootSet() noexcept :
-    GlobalRootSet(mm::GlobalData::Instance().globalsRegistry(), mm::GlobalData::Instance().stableRefRegistry()) {}
+    GlobalRootSet(mm::GlobalData::Instance().globalsRegistry(), mm::GlobalData::Instance().specialRefRegistry()) {}
@@ -8,7 +8,7 @@

 #include "GlobalsRegistry.hpp"
 #include "ShadowStack.hpp"
-#include "StableRefRegistry.hpp"
+#include "SpecialRefRegistry.hpp"
 #include "ThreadLocalStorage.hpp"

 struct ObjHeader;
@@ -18,6 +18,7 @@ namespace mm {

 class ThreadData;

+// TODO: Extremely useless class. Remove.
 class ThreadRootSet {
 public:
    enum class Source {
@@ -78,6 +79,7 @@ private:
    ThreadLocalStorage& tls_;
 };

+// TODO: Extremely useless class. Remove.
 class GlobalRootSet {
 public:
    enum class Source {
@@ -86,7 +88,7 @@ public:
    };

    struct Value {
-        ObjHeader*& object;
+        ObjHeader* object;
        Source source;

        bool operator==(const Value& rhs) const noexcept { return object == rhs.object && source == rhs.source; }
@@ -123,12 +125,12 @@ public:
        Phase phase_;
        union {
            GlobalsRegistry::Iterator globalsIterator_;
-            StableRefRegistry::Iterator stableRefsIterator_;
+            SpecialRefRegistry::RootsIterator specialRefsIterator_;
        };
    };

-    GlobalRootSet(GlobalsRegistry& globalsRegistry, StableRefRegistry& stableRefRegistry) noexcept :
-        globalsIterable_(globalsRegistry.LockForIter()), stableRefsIterable_(stableRefRegistry.LockForIter()) {}
+    GlobalRootSet(GlobalsRegistry& globalsRegistry, SpecialRefRegistry& stableRefRegistry) noexcept :
+        globalsIterable_(globalsRegistry.LockForIter()), specialRefsIterable_(stableRefRegistry.roots()) {}
    GlobalRootSet() noexcept;

    Iterator begin() noexcept { return Iterator(Iterator::begin, *this); }
@@ -138,7 +140,7 @@ private:
    // TODO: These use separate locks, which is inefficient, and slightly dangerous. In practice it's
    //       fine, because this is the only place where these two locks are taken simultaneously.
    GlobalsRegistry::Iterable globalsIterable_;
-    StableRefRegistry::Iterable stableRefsIterable_;
+    SpecialRefRegistry::RootsIterable specialRefsIterable_;
 };

 } // namespace mm
@@ -9,6 +9,7 @@
 #include "gtest/gtest.h"

 #include "ShadowStack.hpp"
+#include "StableRef.hpp"
 #include "std_support/Memory.hpp"
 #include "std_support/Vector.hpp"

@@ -97,19 +98,19 @@ TEST(GlobalRootSetTest, Basic) {
    globalsProducer.Insert(&global1);
    globalsProducer.Insert(&global2);

-    mm::StableRefRegistry stableRefs;
-    mm::StableRefRegistry::ThreadQueue stableRefsProducer(stableRefs);
+    mm::SpecialRefRegistry specialRefsRegistry;
+    mm::SpecialRefRegistry::ThreadQueue stableRefsProducer(specialRefsRegistry);
    ObjHeader* stableRef1 = reinterpret_cast<ObjHeader*>(3);
    ObjHeader* stableRef2 = reinterpret_cast<ObjHeader*>(4);
    ObjHeader* stableRef3 = reinterpret_cast<ObjHeader*>(5);
-    stableRefsProducer.Insert(stableRef1);
-    stableRefsProducer.Insert(stableRef2);
-    stableRefsProducer.Insert(stableRef3);
+    auto stableRefHandle1 = stableRefsProducer.createStableRef(stableRef1);
+    auto stableRefHandle2 = stableRefsProducer.createStableRef(stableRef2);
+    auto stableRefHandle3 = stableRefsProducer.createStableRef(stableRef3);

    globalsProducer.Publish();
-    stableRefsProducer.Publish();
+    stableRefsProducer.publish();

-    mm::GlobalRootSet iter(globals, stableRefs);
+    mm::GlobalRootSet iter(globals, specialRefsRegistry);

    std_support::vector<mm::GlobalRootSet::Value> actual;
    for (auto object : iter) {
@@ -120,15 +121,19 @@ TEST(GlobalRootSetTest, Basic) {
    auto asStableRef = [](ObjHeader*& object) -> mm::GlobalRootSet::Value { return {object, mm::GlobalRootSet::Source::kStableRef}; };
    EXPECT_THAT(
            actual,
-            testing::ElementsAre(
+            testing::UnorderedElementsAre(
                    asGlobal(global1), asGlobal(global2), asStableRef(stableRef1), asStableRef(stableRef2), asStableRef(stableRef3)));
+
+    std::move(stableRefHandle1).dispose();
+    std::move(stableRefHandle2).dispose();
+    std::move(stableRefHandle3).dispose();
 }

 TEST(GlobalRootSetTest, Empty) {
    mm::GlobalsRegistry globals;
-    mm::StableRefRegistry stableRefs;
+    mm::SpecialRefRegistry specialRefsRegistry;

-    mm::GlobalRootSet iter(globals, stableRefs);
+    mm::GlobalRootSet iter(globals, specialRefsRegistry);

    std_support::vector<mm::GlobalRootSet::Value> actual;
    for (auto object : iter) {
@@ -28,6 +28,12 @@ class ShadowStack : private Pinned {
 public:
    class Iterator {
    public:
+        using difference_type = ptrdiff_t;
+        using value_type = ObjHeader*;
+        using pointer = value_type*;
+        using reference = value_type&;
+        using iterator_category = std::forward_iterator_tag;
+
        explicit Iterator(FrameOverlay* frame) noexcept : frame_(frame), object_(begin()), end_(end()) { Init(); }

        ObjHeader*& operator*() noexcept { return *object_; }
@@ -0,0 +1,157 @@
+/*
+ * Copyright 2010-2020 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the LICENSE file.
+ */
+
+#include "SpecialRefRegistry.hpp"
+
+#include "GlobalData.hpp"
+#include "MemoryPrivate.hpp"
+#include "ObjCBackRef.hpp"
+#include "StableRef.hpp"
+#include "ThreadData.hpp"
+#include "ThreadState.hpp"
+#include "WeakRef.hpp"
+
+using namespace kotlin;
+
+mm::StableRef mm::SpecialRefRegistry::ThreadQueue::createStableRef(ObjHeader* object) noexcept {
+    return mm::StableRef(registerNode(object, 1, true).asRaw());
+}
+
+mm::WeakRef mm::SpecialRefRegistry::ThreadQueue::createWeakRef(ObjHeader* object) noexcept {
+    return mm::WeakRef(registerNode(object, 0, false).asRaw());
+}
+
+mm::ObjCBackRef mm::SpecialRefRegistry::ThreadQueue::createObjCBackRef(ObjHeader* object) noexcept {
+    return mm::ObjCBackRef(registerNode(object, 1, false).asRaw());
+}
+
+void mm::SpecialRefRegistry::ThreadQueue::deleteNodeIfLocal(Node& node) noexcept {
+    // This is a very weird optimization.
+    // * We're saving some time during root scanning and some memory by
+    //   deleting some short-lived nodes without ever publishing them.
+    // * But in order to do that we have to be in a runnable state, so
+    //   we potentially force a native state thread to go wait for the GC.
+    if (node.owner_ == this) {
+        queue_.erase(node.position_);
+    }
+}
+
+// static
+mm::SpecialRefRegistry& mm::SpecialRefRegistry::instance() noexcept {
+    return GlobalData::Instance().specialRefRegistry();
+}
+
+mm::SpecialRefRegistry::Node* mm::SpecialRefRegistry::nextRoot(Node* current) noexcept {
+    RuntimeAssert(current != nullptr, "current cannot be null");
+    RuntimeAssert(current != rootsTail(), "current cannot be tail");
+    Node* candidate = current->nextRoot_.load(std::memory_order_relaxed);
+    // Not an infinite loop, `candidate` always moves forward and since insertions can only
+    // happen in the head, they will always happen before `candidate`.
+    while (true) {
+        RuntimeAssert(candidate != nullptr, "candidate cannot be null");
+        if (candidate == rootsTail())
+            // Reached tail, nothing to do anymore
+            return candidate;
+        if (candidate->rc_.load(std::memory_order_relaxed) > 0) {
+            // Keeping acquire-release for nextRoot_.
+            std::atomic_thread_fence(std::memory_order_acquire);
+            // Perfectly good node. Stop right there.
+            return candidate;
+        }
+        // Bad node. Let's remove it from the roots.
+        // Racy if someone concurrently inserts in the middle. Or iterates.
+        // But we don't have that here. Inserts are only in the beginning.
+        // Iteration also happens only here.
+        auto [candidatePrev, candidateNext] = eraseFromRoots(current, candidate);
+        // We removed candidate. But should we have?
+        if (candidate->rc_.load(std::memory_order_relaxed) > 0) {
+            RuntimeAssert(candidate->obj_ != nullptr, "candidate cannot have a null obj_");
+            // Ooops. Let's put it back. Okay to put into the head.
+            insertIntoRootsHead(*candidate);
+        }
+        // eraseFromRoots and insertIntoRootsHead are both acquire-release fences.
+        // This means they play nice with each other and we don't need an extra fence
+        // here to ensure synchronization with 0->1 rc_ change:
+        // * We read rc_ after eraseFromRoots.
+        // * retainRef writes rc_ before insertIntoRootsHead.
+        // So the write to rc_ in retainRef happens before the read here.
+        //
+        // Okay, properly deleted. Our new `candidate` is the next of previous candidate,
+        // and our `current` then is our best guess at the previous node of the `candidate`.
+        current = candidatePrev;
+        candidate = candidateNext;
+        // `current` has either moved forward or stayed where it is.
+        // `candidate` has definitely moved forward.
+        // `current` is only used in `eraseFromRoots` which itself ensures that no
+        // infinite loop can happen.
+        // So, this loop is also not infinite.
+    }
+}
+
+std::pair<mm::SpecialRefRegistry::Node*, mm::SpecialRefRegistry::Node*> mm::SpecialRefRegistry::eraseFromRoots(
+        Node* prev, Node* node) noexcept {
+    RuntimeAssert(node != rootsHead(), "node cannot be head");
+    RuntimeAssert(node != rootsTail(), "node cannot be tail");
+    Node* next = node->nextRoot_.load(std::memory_order_acquire);
+    RuntimeAssert(next != nullptr, "node@%p next cannot be null", node);
+    do {
+        Node* prevExpectedNext = node;
+        bool removed =
+                prev->nextRoot_.compare_exchange_strong(prevExpectedNext, next, std::memory_order_release, std::memory_order_acquire);
+        if (removed) {
+            auto* actualNext = node->nextRoot_.exchange(nullptr, std::memory_order_acq_rel);
+            RuntimeAssert(actualNext == next, "node@%p next expected %p actual %p", node, next, actualNext);
+            return {prev, next};
+        }
+        prev = prevExpectedNext;
+        RuntimeAssert(prev != rootsHead(), "prev cannot be head");
+        RuntimeAssert(prev != rootsTail(), "prev cannot be tail");
+        // We moved `prev` forward, nothing can insert after `prev` anymore, this
+        // cannot be an infinite loop, then.
+    } while (true);
+}
+
+void mm::SpecialRefRegistry::insertIntoRootsHead(Node& node) noexcept {
+    Node* next = rootsHead()->nextRoot_.load(std::memory_order_acquire);
+    Node* nodeExpectedNext = nullptr;
+    do {
+        RuntimeAssert(next != nullptr, "head's next cannot be null");
+        if (!node.nextRoot_.compare_exchange_strong(nodeExpectedNext, next, std::memory_order_release, std::memory_order_acquire)) {
+            // So:
+            // * `node` is already in the roots list
+            // * some other thread is inserting it in the roots list
+            // * GC thread may be removing it from the roots list, but
+            //   will recheck rc afterwards and insert it back if needed
+            // In either case, do not touch anything anymore here.
+            return;
+        }
+        // CAS was successfull, so we need to update the expected value of node.nextRoot_
+        nodeExpectedNext = next;
+    } while (!rootsHead()->nextRoot_.compare_exchange_weak(next, &node, std::memory_order_release, std::memory_order_acquire));
+}
+
+std_support::list<mm::SpecialRefRegistry::Node>::iterator mm::SpecialRefRegistry::findAliveNode(
+        std_support::list<Node>::iterator it) noexcept {
+    while (it != all_.end() && it->rc_.load(std::memory_order_relaxed) == Node::disposedMarker) {
+        // Synchronization with `Node::dispose()`
+        std::atomic_thread_fence(std::memory_order_acquire);
+        // Removing disposed nodes.
+        if (it->nextRoot_.load(std::memory_order_relaxed) != nullptr) {
+            // Wait, it's in the roots list. Lets wait until the next GC
+            // for it to get cleaned up from there.
+            ++it;
+            continue;
+        }
+        // If we observe both `nextRoot_ == nullptr` and `rc_ == disposedMarker`, we
+        // can be sure, that no mutator can change `nextRoot_` later.
+        // Proof:
+        // For the mutator to change `nextRoot_`, it would need to perform 0->1 `rc_` transition.
+        // But `rc_` can only be set to `disposedMarker` during `dispose()` call,
+        // which can only happen after any `retainRef()` call. So, it's impossible
+        // to observe any changes in `nextRoot_` after we observe `disposedMarker` in `nextRoot_`.
+        it = all_.erase(it);
+    }
+    return it;
+}
@@ -0,0 +1,367 @@
+/*
+ * Copyright 2010-2023 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the LICENSE file.
+ */
+
+#pragma once
+
+#include <atomic>
+
+#include "Memory.h"
+#include "RawPtr.hpp"
+#include "ThreadRegistry.hpp"
+#include "std_support/List.hpp"
+
+namespace kotlin::mm {
+
+class ObjCBackRef;
+class StableRef;
+class WeakRef;
+
+// Registry for all special references to objects:
+// * stable references (i.e. always part of the root set)
+// * weak references
+// * ObjC back references. A mix between stable and weak references
+//   created for ObjC part of Kotlin objects. Have a count of external
+//   references. When > 0 - stable reference. When = 0 - weak reference.
+//
+// Consists of 2 global lists and 1 thread local list of `Node`s.
+// Each `Node` has a reference `obj_` to an object and an external references counter `rc_`.
+// Invariants:
+// * `rc_ > 0` -> alive externally referenced object, must be in the root set.
+// * `rc_ == 0` -> alive externally unreferenced object, must eventually be out of the root set.
+// * `rc_ == disposedMarker` -> the `Node` itself is no longer externally referenced and can be eventually deleted.
+// * `rc_` can be increased and decreased by any thread in any state.
+//   * In practice 0 -> 1 only happens in mutator threads in runnable state, but the implementation of `SpecialRefRegistry`
+//     does not depend on it.
+// * `Node`s are owned either by a global `std::list` or by thread local `std::list`s. Global list is protected by a mutex.
+// * Insertion into the global list happens by moving from thread local list during STW and when a thread gets destroyed.
+// * Only the GC thread traverses and removes elements from the global list. Removal happens only if `Node` has `rc_ == disposedMarker`
+//   and it's not in the roots list.
+// * During global list traversal `Node`s `obj_` referenced may get nulled out by the GC.
+// * Insertion into thread local lists happens in runnable state.
+// * Removal from thread local list happens during STW, thread destruction, or in the runnable state for `Node`s that can never
+//   go through 0 -> 1 rc transition (created via mm::StableRef).
+// * `Node`s are additionally linked into an intrusive global roots list.
+// * Any thread in any state can insert into the roots list. Insertion only happens into the head.
+// * Only the GC thread can remove from the roots list during root scanning. If after removal
+//   the `rc_` of the `Node` is `> 0`, the GC thread will make sure the node is inserted
+//   into the head
+// * During roots list traversal all nodes to the left are either marked or inserted into the mark queue.
+class SpecialRefRegistry : private Pinned {
+    // TODO: Consider using a real mutex.
+    using Mutex = SpinLock<MutexThreadStateHandling::kIgnore>;
+
+    class Node : private Pinned {
+    public:
+        using Rc = int32_t;
+        inline static constexpr Rc disposedMarker = std::numeric_limits<Rc>::min();
+        static_assert(disposedMarker < 0, "disposedMarker must be an impossible Rc value");
+
+        Node(ObjHeader* obj, Rc rc) noexcept : obj_(obj), rc_(rc) {
+            RuntimeAssert(obj != nullptr, "Creating StableRef for null object");
+            RuntimeAssert(rc >= 0, "Creating StableRef with negative rc %d", rc);
+        }
+
+        ~Node() {
+            if (compiler::runtimeAssertsEnabled()) {
+                auto rc = rc_.load(std::memory_order_relaxed);
+                RuntimeAssert(rc == disposedMarker, "Deleting StableRef@%p with rc %d", this, rc);
+            }
+        }
+
+        void dispose() noexcept {
+            // Synchronization with `findAliveNode()`.
+            // TODO: When assertions are disabled, exchange may pollute the
+            //       generated assembly. Check if this a problem.
+            auto rc = rc_.exchange(disposedMarker, std::memory_order_release);
+            if (compiler::runtimeAssertsEnabled()) {
+                if (rc > 0) {
+                    // In objc export if ObjCClass extends from KtClass
+                    // doing retain+autorelease inside [ObjCClass dealloc] will cause
+                    // this->dispose() be called after this->retain() but before
+                    // subsequent this->release().
+                    // However, since this happens in dealloc, the stored object must
+                    // have been cleared already.
+                    RuntimeAssert(obj_ == nullptr, "Disposing StableRef@%p with rc %d and uncleaned object %p", this, rc, obj_);
+                }
+                RuntimeAssert(rc >= 0, "Disposing StableRef@%p with rc %d", this, rc);
+            }
+        }
+
+        [[nodiscard("expensive pure function")]] ObjHeader* ref() const noexcept {
+            if (compiler::runtimeAssertsEnabled()) {
+                AssertThreadState(ThreadState::kRunnable);
+                auto rc = rc_.load(std::memory_order_relaxed);
+                RuntimeAssert(rc >= 0, "Dereferencing StableRef@%p with rc %d", this, rc);
+            }
+            return obj_;
+        }
+
+        OBJ_GETTER0(tryRef) noexcept {
+            AssertThreadState(ThreadState::kRunnable);
+            // TODO: Weak read barrier with CMS.
+            RETURN_OBJ(obj_);
+        }
+
+        void retainRef() noexcept {
+            auto rc = rc_.fetch_add(1, std::memory_order_relaxed);
+            RuntimeAssert(rc >= 0, "Retaining StableRef@%p with rc %d", this, rc);
+            if (rc == 0) {
+                RuntimeAssert(
+                        position_ == std_support::list<Node>::iterator{},
+                        "Retaining StableRef@%p with fast deletion optimization is disallowed", this);
+
+                if (!obj_) {
+                    // In objc export if ObjCClass extends from KtClass
+                    // calling retain inside [ObjCClass dealloc] will cause
+                    // node.retainRef() be called after node.obj_ was cleared but
+                    // before node.dispose().
+                    // We could place it into the root set, and it'll be removed
+                    // from it at some later point. But let's just skip it.
+                    return;
+                }
+
+                // TODO: With CMS barrier for marking `obj_` should be here.
+                //       Until we have the barrier, the object must already be in the roots.
+                //       If 0->1 happened from `[ObjCClass _tryRetain]`, it would first hold the object
+                //       on the stack via `tryRef`.
+                //       If 0->1 happened during construction:
+                //       * First of all, currently it's impossible because the `Node` is created with rc=1 and not inserted
+                //         into the roots list until publishing.
+                //       * Even if the above changes, for the construction, the object must have been passed in from somewhere,
+                //         so it must be reachable anyway.
+                //       If 0->1 happened because an object is passing through the interop border for the second time (or more)
+                //       (e.g. accessing a non-permanent global a couple of times). Follows the construction case above:
+                //       "the object must have been passed in from somewhere, so it must be reachable anyway".
+
+                // 0->1 changes require putting this node into the root set.
+                SpecialRefRegistry::instance().insertIntoRootsHead(*this);
+            }
+        }
+
+        void releaseRef() noexcept {
+            auto rc = rc_.fetch_sub(1, std::memory_order_relaxed);
+            RuntimeAssert(rc > 0, "Releasing StableRef@%p with rc %d", this, rc);
+        }
+
+        RawSpecialRef* asRaw() noexcept { return reinterpret_cast<RawSpecialRef*>(this); }
+        static Node* fromRaw(RawSpecialRef* ref) noexcept { return reinterpret_cast<Node*>(ref); }
+
+    private:
+        friend class SpecialRefRegistry;
+        friend class SpecialRefRegistryTest;
+
+        // obj_ is set in the constructor and can be nulled out only by the
+        // GC thread when processing weaks. It's the responsibility of the
+        // GC to make sure nulling out obj_ is synchronized with mutators:
+        // * via STW: nulling obj_ only happens when mutators are paused.
+        // * via weak read barriers: when GC enters a weak processing phase,
+        //   it enables weak read barriers which do not read obj_ if obj_ will
+        //   be nulled, and disable the barriers when the phase is completed.
+        //   Synchronization between GC and mutators happens via enabling/disabling
+        //   the barriers.
+        ObjHeader* obj_ = nullptr;
+        // Only ever updated using relaxed memory ordering. Any synchronization
+        // with nextRoot_ is achieved via acquire-release of nextRoot_.
+        std::atomic<Rc> rc_ = 0; // After dispose() will be disposedMarker.
+        // Singly linked lock free list. Using acquire-release throughout.
+        std::atomic<Node*> nextRoot_ = nullptr;
+        // This and the next one only serve fast deletion optimization for shortly lived StableRefs.
+        // TODO: Consider discarding this optimization completely.
+        //       If we were to use custom allocator for these nodes as well they better
+        //       be only deleted in the sweep anyway.
+        //       Alternative: keep stable refs completely separate.
+        void* owner_ = nullptr;
+        std_support::list<Node>::iterator position_{};
+    };
+
+public:
+    class ThreadQueue : private Pinned {
+    public:
+        explicit ThreadQueue(SpecialRefRegistry& registry) : owner_(registry) {}
+
+        ~ThreadQueue() { publish(); }
+
+        void publish() noexcept {
+            for (auto& node : queue_) {
+                // No need to synchronize. These two can only be updated in the runnable state.
+                // TODO: If we were to remove this optimization, we could avoid scanning
+                //       the whole queue here and just have the nodes inserted into the roots
+                //       when they're created.
+                node.owner_ = nullptr;
+                node.position_ = std_support::list<Node>::iterator();
+                RuntimeAssert(node.obj_ != nullptr, "Publishing Node with null obj_");
+                // If the node was created with a positive refcount, we must ensure its put into
+                // the roots.
+                auto rc = node.rc_.load(std::memory_order_relaxed);
+                if (rc > 0) {
+                    // Regular publishing happens before the global root scanning,
+                    // so this insertion will definitely be processed.
+                    // Publishing during the thread destruction is a bit more complicated.
+                    // But the GC makes sure to process all threads before scanning global
+                    // roots. So, it'll either publish the dying thread itself, or
+                    // if the dying thread has already deregistered, it means it published
+                    // itself. In any case, global root scanning happens afterwards.
+                    // TODO: With CMS barrier for marking `node.obj_` should be here.
+                    owner_.insertIntoRootsHead(node);
+                }
+            }
+            std::unique_lock guard(owner_.mutex_);
+            RuntimeAssert(owner_.all_.get_allocator() == queue_.get_allocator(), "allocators must match");
+            owner_.all_.splice(owner_.all_.end(), std::move(queue_));
+        }
+
+        void clearForTests() noexcept { queue_.clear(); }
+
+        [[nodiscard("must be manually disposed")]] StableRef createStableRef(ObjHeader* object) noexcept;
+        [[nodiscard("must be manually disposed")]] WeakRef createWeakRef(ObjHeader* object) noexcept;
+        [[nodiscard("must be manually disposed")]] ObjCBackRef createObjCBackRef(ObjHeader* object) noexcept;
+
+    private:
+        friend class StableRef;
+        friend class SpecialRefRegistryTest;
+
+        [[nodiscard("must be manually disposed")]] Node& registerNode(ObjHeader* obj, Node::Rc rc, bool allowFastDeletion) noexcept {
+            RuntimeAssert(obj != nullptr, "Creating node for null object");
+            queue_.emplace_back(obj, rc);
+            auto& node = queue_.back();
+            if (allowFastDeletion) {
+                node.owner_ = this;
+                node.position_ = std::prev(queue_.end());
+            }
+            return node;
+        }
+
+        void deleteNodeIfLocal(Node& node) noexcept;
+
+        SpecialRefRegistry& owner_;
+        std_support::list<Node> queue_;
+    };
+
+    class RootsIterator {
+    public:
+        ObjHeader* operator*() const noexcept {
+            // Ignoring rc here. If someone nulls out rc during root
+            // scanning, it's okay to be conservative and still make it a root.
+            return node_->obj_;
+        }
+
+        RootsIterator& operator++() noexcept {
+            node_ = owner_->nextRoot(node_);
+            return *this;
+        }
+
+        bool operator==(const RootsIterator& rhs) const noexcept { return node_ == rhs.node_; }
+
+        bool operator!=(const RootsIterator& rhs) const noexcept { return !(*this == rhs); }
+
+    private:
+        friend class SpecialRefRegistry;
+
+        RootsIterator(SpecialRefRegistry& owner, Node* node) noexcept : owner_(&owner), node_(node) {}
+
+        SpecialRefRegistry* owner_;
+        Node* node_;
+    };
+
+    class RootsIterable : private MoveOnly {
+    public:
+        RootsIterator begin() const noexcept { return RootsIterator(*owner_, owner_->nextRoot(owner_->rootsHead())); }
+
+        RootsIterator end() const noexcept { return RootsIterator(*owner_, owner_->rootsTail()); }
+
+    private:
+        friend class SpecialRefRegistry;
+
+        explicit RootsIterable(SpecialRefRegistry& owner) noexcept : owner_(&owner) {}
+
+        raw_ptr<SpecialRefRegistry> owner_;
+    };
+
+    class Iterator {
+    public:
+        ObjHeader*& operator*() noexcept { return iterator_->obj_; }
+
+        Iterator& operator++() noexcept {
+            iterator_ = owner_->findAliveNode(std::next(iterator_));
+            return *this;
+        }
+
+        bool operator==(const Iterator& rhs) const noexcept { return iterator_ == rhs.iterator_; }
+
+        bool operator!=(const Iterator& rhs) const noexcept { return iterator_ != rhs.iterator_; }
+
+    private:
+        friend class SpecialRefRegistry;
+        friend class SpecialRefRegistryTest;
+
+        Iterator(SpecialRefRegistry& owner, std_support::list<Node>::iterator iterator) noexcept : owner_(&owner), iterator_(iterator) {}
+
+        SpecialRefRegistry* owner_;
+        std_support::list<Node>::iterator iterator_;
+    };
+
+    class Iterable : private MoveOnly {
+    public:
+        Iterator begin() noexcept { return Iterator(owner_, owner_.findAliveNode(owner_.all_.begin())); }
+        Iterator end() noexcept { return Iterator(owner_, owner_.all_.end()); }
+
+    private:
+        friend class SpecialRefRegistry;
+
+        Iterable(SpecialRefRegistry& owner) noexcept : owner_(owner), guard_(owner_.mutex_) {}
+
+        SpecialRefRegistry& owner_;
+        std::unique_lock<Mutex> guard_;
+    };
+
+    SpecialRefRegistry() noexcept { rootsHead()->nextRoot_.store(rootsTail(), std::memory_order_relaxed); }
+
+    ~SpecialRefRegistry() = default;
+
+    static SpecialRefRegistry& instance() noexcept;
+
+    void clearForTests() noexcept {
+        rootsHead()->nextRoot_ = rootsTail();
+        for (auto& node : all_) {
+            // Allow the tests not to run the finalizers for weaks.
+            node.rc_ = Node::disposedMarker;
+        }
+        all_.clear();
+    }
+
+    // Should be called on the GC thread after all threads have published.
+    RootsIterable roots() noexcept { return RootsIterable(*this); }
+
+    // Should be called on the GC thread after marking is complete.
+    // Locks the registry and allows safe iteration over it.
+    Iterable lockForIter() noexcept { return Iterable(*this); }
+
+private:
+    friend class ObjCBackRef;
+    friend class StableRef;
+    friend class WeakRef;
+    friend class SpecialRefRegistryTest;
+
+    Node* nextRoot(Node* current) noexcept;
+    // Erase `node` from the roots list. `prev` is the current guess of the node
+    // previous to `node`. Returns two nodes between which `node` was deleted.
+    std::pair<Node*, Node*> eraseFromRoots(Node* prev, Node* node) noexcept;
+    void insertIntoRootsHead(Node& node) noexcept;
+    std_support::list<Node>::iterator findAliveNode(std_support::list<Node>::iterator it) noexcept;
+
+    Node* rootsHead() noexcept { return reinterpret_cast<Node*>(rootsHeadStorage_); }
+    const Node* rootsHead() const noexcept { return reinterpret_cast<const Node*>(rootsHeadStorage_); }
+    static Node* rootsTail() noexcept { return reinterpret_cast<Node*>(rootsTailStorage_); }
+
+    // TODO: Iteration over `all_` will be slow, because it's `std_support::list`
+    //       collected at different times from different threads, and so the nodes
+    //       are all over the memory. Consider using custom allocator for that.
+    std_support::list<Node> all_;
+    Mutex mutex_;
+    alignas(Node) char rootsHeadStorage_[sizeof(Node)] = {0};
+    alignas(Node) static inline char rootsTailStorage_[sizeof(Node)] = {0};
+};
+
+} // namespace kotlin::mm
@@ -0,0 +1,480 @@
+/*
+ * Copyright 2010-2023 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the LICENSE file.
+ */
+
+#include "SpecialRefRegistry.hpp"
+
+#include <condition_variable>
+#include <mutex>
+
+#include "gmock/gmock.h"
+#include "gtest/gtest.h"
+
+#include "ObjCBackRef.hpp"
+#include "StableRef.hpp"
+#include "TestSupport.hpp"
+#include "ThreadRegistry.hpp"
+#include "WeakRef.hpp"
+
+using namespace kotlin;
+
+namespace {
+
+class Waiter : private Pinned {
+public:
+    void allow() noexcept {
+        {
+            std::unique_lock guard(mutex_);
+            allow_ = true;
+        }
+        cv_.notify_all();
+    };
+
+    void wait() noexcept {
+        std::unique_lock guard(mutex_);
+        cv_.wait(guard, [this] { return allow_; });
+    }
+
+private:
+    bool allow_ = false;
+    std::mutex mutex_;
+    std::condition_variable cv_;
+};
+
+} // namespace
+
+class SpecialRefRegistryTest : public testing::Test {
+public:
+    ~SpecialRefRegistryTest() {
+        // Clean up safely.
+        roots();
+        all();
+    }
+
+    void publish() noexcept { mm::ThreadRegistry::Instance().CurrentThreadData()->specialRefRegistry().publish(); }
+
+    template <typename... Invalidated>
+    std::vector<ObjHeader*> all(Invalidated&&... invalidated) noexcept {
+        std::set<ObjHeader*> invalidatedSet({std::forward<Invalidated>(invalidated)...});
+        std::vector<ObjHeader*> result;
+        for (auto& obj : mm::SpecialRefRegistry::instance().lockForIter()) {
+            if (invalidatedSet.find(obj) != invalidatedSet.end()) {
+                obj = nullptr;
+            }
+            result.push_back(obj);
+        }
+        return result;
+    }
+
+    std::vector<ObjHeader*> roots() noexcept {
+        std::vector<ObjHeader*> result;
+        for (auto* obj : mm::SpecialRefRegistry::instance().roots()) {
+            result.push_back(obj);
+        }
+        return result;
+    }
+
+    ObjHeader* tryRef(mm::WeakRef& weakRef) noexcept {
+        ObjHeader* result;
+        return weakRef.tryRef(&result);
+    }
+};
+
+TEST_F(SpecialRefRegistryTest, RegisterStableRefWithoutPublish) {
+    RunInNewThread([this] {
+        ObjHeader* obj = reinterpret_cast<ObjHeader*>(1);
+        ObjHolder holder(obj);
+        auto ref = mm::StableRef::create(obj);
+
+        EXPECT_THAT(roots(), testing::UnorderedElementsAre());
+        EXPECT_THAT(all(), testing::UnorderedElementsAre());
+        EXPECT_THAT(*ref, obj);
+
+        std::move(ref).dispose();
+
+        EXPECT_THAT(roots(), testing::UnorderedElementsAre());
+        EXPECT_THAT(all(), testing::UnorderedElementsAre());
+
+        publish();
+
+        EXPECT_THAT(roots(), testing::UnorderedElementsAre());
+        EXPECT_THAT(all(), testing::UnorderedElementsAre());
+    });
+}
+
+TEST_F(SpecialRefRegistryTest, RegisterStableRef) {
+    RunInNewThread([this] {
+        ObjHeader* obj = reinterpret_cast<ObjHeader*>(1);
+        ObjHolder holder(obj);
+        auto ref = mm::StableRef::create(obj);
+
+        EXPECT_THAT(roots(), testing::UnorderedElementsAre());
+        EXPECT_THAT(all(), testing::UnorderedElementsAre());
+        EXPECT_THAT(*ref, obj);
+
+        publish();
+
+        EXPECT_THAT(roots(), testing::UnorderedElementsAre(obj));
+        EXPECT_THAT(all(), testing::UnorderedElementsAre(obj));
+        EXPECT_THAT(*ref, obj);
+
+        std::move(ref).dispose();
+
+        EXPECT_THAT(roots(), testing::UnorderedElementsAre());
+        EXPECT_THAT(all(), testing::UnorderedElementsAre());
+    });
+}
+
+TEST_F(SpecialRefRegistryTest, RegisterWeakRefWithoutPublish) {
+    RunInNewThread([this] {
+        ObjHeader* obj = reinterpret_cast<ObjHeader*>(1);
+        ObjHolder holder(obj);
+        auto ref = mm::WeakRef::create(obj);
+
+        EXPECT_THAT(roots(), testing::UnorderedElementsAre());
+        EXPECT_THAT(all(), testing::UnorderedElementsAre());
+        EXPECT_THAT(tryRef(ref), obj);
+
+        std::move(ref).dispose();
+
+        EXPECT_THAT(roots(), testing::UnorderedElementsAre());
+        EXPECT_THAT(all(), testing::UnorderedElementsAre());
+
+        publish();
+
+        EXPECT_THAT(roots(), testing::UnorderedElementsAre());
+        EXPECT_THAT(all(), testing::UnorderedElementsAre());
+    });
+}
+
+TEST_F(SpecialRefRegistryTest, RegisterWeakRef) {
+    RunInNewThread([this] {
+        ObjHeader* obj = reinterpret_cast<ObjHeader*>(1);
+        ObjHolder holder(obj);
+        auto ref = mm::WeakRef::create(obj);
+
+        EXPECT_THAT(roots(), testing::UnorderedElementsAre());
+        EXPECT_THAT(all(), testing::UnorderedElementsAre());
+        EXPECT_THAT(tryRef(ref), obj);
+
+        publish();
+
+        EXPECT_THAT(roots(), testing::UnorderedElementsAre());
+        EXPECT_THAT(all(), testing::UnorderedElementsAre(obj));
+        EXPECT_THAT(tryRef(ref), obj);
+
+        std::move(ref).dispose();
+
+        EXPECT_THAT(roots(), testing::UnorderedElementsAre());
+        EXPECT_THAT(all(), testing::UnorderedElementsAre());
+    });
+}
+
+TEST_F(SpecialRefRegistryTest, RegisterObjCRefWithoutPublish) {
+    RunInNewThread([this] {
+        ObjHeader* obj = reinterpret_cast<ObjHeader*>(1);
+        ObjHolder holder(obj);
+        auto ref = mm::ObjCBackRef::create(obj);
+
+        EXPECT_THAT(roots(), testing::UnorderedElementsAre());
+        EXPECT_THAT(all(), testing::UnorderedElementsAre());
+        EXPECT_THAT(*ref, obj);
+
+        ref.release();
+
+        EXPECT_THAT(roots(), testing::UnorderedElementsAre());
+        EXPECT_THAT(all(), testing::UnorderedElementsAre());
+        EXPECT_THAT(*ref, obj);
+
+        std::move(ref).dispose();
+
+        EXPECT_THAT(roots(), testing::UnorderedElementsAre());
+        EXPECT_THAT(all(), testing::UnorderedElementsAre());
+
+        publish();
+
+        EXPECT_THAT(roots(), testing::UnorderedElementsAre());
+        EXPECT_THAT(all(), testing::UnorderedElementsAre());
+    });
+}
+
+TEST_F(SpecialRefRegistryTest, RegisterObjCRef) {
+    RunInNewThread([this] {
+        ObjHeader* obj = reinterpret_cast<ObjHeader*>(1);
+        ObjHolder holder(obj);
+        auto ref = mm::ObjCBackRef::create(obj);
+
+        EXPECT_THAT(roots(), testing::UnorderedElementsAre());
+        EXPECT_THAT(all(), testing::UnorderedElementsAre());
+        EXPECT_THAT(*ref, obj);
+
+        publish();
+
+        EXPECT_THAT(roots(), testing::UnorderedElementsAre(obj));
+        EXPECT_THAT(all(), testing::UnorderedElementsAre(obj));
+        EXPECT_THAT(*ref, obj);
+
+        ref.release();
+
+        EXPECT_THAT(roots(), testing::UnorderedElementsAre());
+        EXPECT_THAT(all(), testing::UnorderedElementsAre(obj));
+        EXPECT_THAT(*ref, obj);
+
+        std::move(ref).dispose();
+
+        EXPECT_THAT(roots(), testing::UnorderedElementsAre());
+        EXPECT_THAT(all(), testing::UnorderedElementsAre());
+    });
+}
+
+TEST_F(SpecialRefRegistryTest, RegisterAllRefsWithoutPublish) {
+    RunInNewThread([this] {
+        ObjHeader* obj = reinterpret_cast<ObjHeader*>(1);
+        ObjHolder holder(obj);
+        auto ref1 = mm::StableRef::create(obj);
+        auto ref2 = mm::WeakRef::create(obj);
+        auto ref3 = mm::ObjCBackRef::create(obj);
+
+        EXPECT_THAT(roots(), testing::UnorderedElementsAre());
+        EXPECT_THAT(all(), testing::UnorderedElementsAre());
+
+        std::move(ref1).dispose();
+        std::move(ref2).dispose();
+        ref3.release();
+        std::move(ref3).dispose();
+
+        EXPECT_THAT(roots(), testing::UnorderedElementsAre());
+        EXPECT_THAT(all(), testing::UnorderedElementsAre());
+
+        publish();
+
+        EXPECT_THAT(roots(), testing::UnorderedElementsAre());
+        EXPECT_THAT(all(), testing::UnorderedElementsAre());
+    });
+}
+
+TEST_F(SpecialRefRegistryTest, RegisterAllRefs) {
+    RunInNewThread([this] {
+        ObjHeader* obj = reinterpret_cast<ObjHeader*>(1);
+        ObjHolder holder(obj);
+        auto ref1 = mm::StableRef::create(obj);
+        auto ref2 = mm::WeakRef::create(obj);
+        auto ref3 = mm::ObjCBackRef::create(obj);
+
+        EXPECT_THAT(roots(), testing::UnorderedElementsAre());
+        EXPECT_THAT(all(), testing::UnorderedElementsAre());
+
+        publish();
+
+        EXPECT_THAT(roots(), testing::UnorderedElementsAre(obj, obj));
+        EXPECT_THAT(all(), testing::UnorderedElementsAre(obj, obj, obj));
+
+        std::move(ref1).dispose();
+        std::move(ref2).dispose();
+        ref3.release();
+        std::move(ref3).dispose();
+
+        EXPECT_THAT(roots(), testing::UnorderedElementsAre());
+        EXPECT_THAT(all(), testing::UnorderedElementsAre());
+    });
+}
+
+TEST_F(SpecialRefRegistryTest, InvalidateWeakRef) {
+    RunInNewThread([this] {
+        ObjHeader* obj = reinterpret_cast<ObjHeader*>(1);
+        ObjHolder holder(obj);
+        auto ref = mm::WeakRef::create(obj);
+        publish();
+
+        EXPECT_THAT(roots(), testing::UnorderedElementsAre());
+        EXPECT_THAT(all(obj), testing::UnorderedElementsAre(nullptr));
+        EXPECT_THAT(tryRef(ref), nullptr);
+
+        EXPECT_THAT(roots(), testing::UnorderedElementsAre());
+        EXPECT_THAT(all(), testing::UnorderedElementsAre(nullptr));
+        EXPECT_THAT(tryRef(ref), nullptr);
+
+        std::move(ref).dispose();
+
+        EXPECT_THAT(roots(), testing::UnorderedElementsAre());
+        EXPECT_THAT(all(), testing::UnorderedElementsAre());
+    });
+}
+
+TEST_F(SpecialRefRegistryTest, InvalidateObjCRef) {
+    RunInNewThread([this] {
+        ObjHeader* obj = reinterpret_cast<ObjHeader*>(1);
+        ObjHolder holder(obj);
+        auto ref = mm::ObjCBackRef::create(obj);
+        ref.release();
+        publish();
+
+        EXPECT_THAT(roots(), testing::UnorderedElementsAre());
+        EXPECT_THAT(all(obj), testing::UnorderedElementsAre(nullptr));
+        EXPECT_FALSE(ref.tryRetainForTests());
+
+        EXPECT_THAT(roots(), testing::UnorderedElementsAre());
+        EXPECT_THAT(all(), testing::UnorderedElementsAre(nullptr));
+
+        std::move(ref).dispose();
+
+        EXPECT_THAT(roots(), testing::UnorderedElementsAre());
+        EXPECT_THAT(all(), testing::UnorderedElementsAre());
+    });
+}
+
+TEST_F(SpecialRefRegistryTest, TryObjCRef) {
+    RunInNewThread([this] {
+        ObjHeader* obj = reinterpret_cast<ObjHeader*>(1);
+        ObjHolder holder(obj);
+        auto ref = mm::ObjCBackRef::create(obj);
+        ref.release();
+        publish();
+
+        EXPECT_TRUE(ref.tryRetainForTests());
+        EXPECT_THAT(*ref, obj);
+
+        EXPECT_THAT(roots(), testing::UnorderedElementsAre(obj));
+        EXPECT_THAT(all(), testing::UnorderedElementsAre(obj));
+        EXPECT_THAT(*ref, obj);
+
+        ref.release();
+        std::move(ref).dispose();
+
+        EXPECT_THAT(roots(), testing::UnorderedElementsAre());
+        EXPECT_THAT(all(), testing::UnorderedElementsAre());
+    });
+}
+
+TEST_F(SpecialRefRegistryTest, ReRetainObjCRefBeforePublish) {
+    RunInNewThread([this] {
+        ObjHeader* obj = reinterpret_cast<ObjHeader*>(1);
+        ObjHolder holder(obj);
+        auto ref = mm::ObjCBackRef::create(obj);
+        ref.release();
+        ref.retain();
+
+        EXPECT_THAT(roots(), testing::UnorderedElementsAre(obj));
+        EXPECT_THAT(all(), testing::UnorderedElementsAre());
+        EXPECT_THAT(*ref, obj);
+
+        publish();
+
+        EXPECT_THAT(roots(), testing::UnorderedElementsAre(obj));
+        EXPECT_THAT(all(), testing::UnorderedElementsAre(obj));
+        EXPECT_THAT(*ref, obj);
+
+        ref.release();
+        std::move(ref).dispose();
+
+        EXPECT_THAT(roots(), testing::UnorderedElementsAre());
+        EXPECT_THAT(all(), testing::UnorderedElementsAre());
+    });
+}
+
+TEST_F(SpecialRefRegistryTest, StressStableRef) {
+    ObjHeader* obj = reinterpret_cast<ObjHeader*>(1);
+    Waiter waiter;
+    std::vector<ScopedThread> mutators;
+    for (int i = 0; i < kDefaultThreadCount; ++i) {
+        mutators.emplace_back([&, this] {
+            ScopedMemoryInit scope;
+            ObjHolder holder(obj);
+            waiter.wait();
+            auto ref = mm::StableRef::create(obj);
+            publish();
+            std::move(ref).dispose();
+        });
+    }
+    waiter.allow();
+    mutators.clear();
+    EXPECT_THAT(roots(), testing::UnorderedElementsAre());
+    EXPECT_THAT(all(), testing::UnorderedElementsAre());
+}
+
+TEST_F(SpecialRefRegistryTest, StressWeakRef) {
+    ObjHeader* obj = reinterpret_cast<ObjHeader*>(1);
+    Waiter waiter;
+    std::vector<ScopedThread> mutators;
+    for (int i = 0; i < kDefaultThreadCount; ++i) {
+        mutators.emplace_back([&, this] {
+            ScopedMemoryInit scope;
+            ObjHolder holder(obj);
+            waiter.wait();
+            auto ref = mm::WeakRef::create(obj);
+            publish();
+            std::move(ref).dispose();
+        });
+    }
+    waiter.allow();
+    mutators.clear();
+    EXPECT_THAT(roots(), testing::UnorderedElementsAre());
+    EXPECT_THAT(all(), testing::UnorderedElementsAre());
+}
+
+TEST_F(SpecialRefRegistryTest, StressObjCRef) {
+    ObjHeader* obj = reinterpret_cast<ObjHeader*>(1);
+    Waiter waiter;
+    std::vector<ScopedThread> mutators;
+    for (int i = 0; i < kDefaultThreadCount; ++i) {
+        mutators.emplace_back([&, this] {
+            ScopedMemoryInit scope;
+            ObjHolder holder(obj);
+            waiter.wait();
+            auto ref = mm::ObjCBackRef::create(obj);
+            publish();
+            ref.release();
+            std::move(ref).dispose();
+        });
+    }
+    waiter.allow();
+    mutators.clear();
+    EXPECT_THAT(roots(), testing::UnorderedElementsAre());
+    EXPECT_THAT(all(), testing::UnorderedElementsAre());
+}
+
+TEST_F(SpecialRefRegistryTest, StressObjCRefRetainRelease) {
+    RunInNewThread([this] {
+        constexpr int kGCCycles = 10000;
+        constexpr int kRefsCount = 3;
+        ObjHeader* obj = reinterpret_cast<ObjHeader*>(1);
+        ObjHolder holder(obj);
+        Waiter waiter;
+        std::atomic<bool> canStop = false;
+        std::vector<mm::ObjCBackRef> refs;
+        for (int i = 0; i < kRefsCount; ++i) {
+            refs.emplace_back(mm::ObjCBackRef::create(obj));
+            refs.back().release();
+        }
+        publish();
+        std::vector<ScopedThread> mutators;
+        mutators.emplace_back([&, this] {
+            waiter.wait();
+            for (int i = 0; i < kGCCycles; ++i) {
+                roots();
+                all();
+            }
+            canStop.store(true, std::memory_order_release);
+        });
+        for (int i = 0; i < kDefaultThreadCount; ++i) {
+            mutators.emplace_back([i, obj, &refs, &waiter, &canStop] {
+                ScopedMemoryInit scope;
+                ObjHolder holder(obj);
+                waiter.wait();
+                auto& ref = refs[i % kRefsCount];
+                while (!canStop.load(std::memory_order_acquire)) {
+                    ref.retain();
+                    ref.release();
+                }
+            });
+        }
+        waiter.allow();
+        mutators.clear();
+        for (auto& ref : refs) {
+            std::move(ref).dispose();
+        }
+        EXPECT_THAT(roots(), testing::UnorderedElementsAre());
+        EXPECT_THAT(all(), testing::UnorderedElementsAre());
+    });
+}
@@ -0,0 +1,34 @@
+/*
+ * Copyright 2010-2023 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the LICENSE file.
+ */
+
+#include "StableRef.hpp"
+
+#include "MemoryPrivate.hpp"
+#include "ThreadData.hpp"
+#include "ThreadRegistry.hpp"
+#include "ThreadState.hpp"
+
+using namespace kotlin;
+
+// static
+mm::StableRef mm::StableRef::create(ObjHeader* obj) noexcept {
+    RuntimeAssert(obj != nullptr, "Creating StableRef for null object");
+    return mm::ThreadRegistry::Instance().CurrentThreadData()->specialRefRegistry().createStableRef(obj);
+}
+
+// static
+void mm::StableRef::tryToDeleteImmediately(raw_ptr<SpecialRefRegistry::Node> node) noexcept {
+    // When we're on the registered thread, perform oportunistic quick deletion.
+    if (auto* threadNode = mm::ThreadRegistry::Instance().CurrentThreadDataNodeOrNull()) {
+        tryToDeleteImmediately(*threadNode->Get(), std::move(node));
+    }
+}
+
+// static
+void mm::StableRef::tryToDeleteImmediately(mm::ThreadData& thread, raw_ptr<SpecialRefRegistry::Node> node) noexcept {
+    auto lastState = SwitchThreadState(&thread, ThreadState::kRunnable, /* reentrant = */ true);
+    thread.specialRefRegistry().deleteNodeIfLocal(*node);
+    SwitchThreadState(&thread, lastState, /* reentrant = */ true);
+}
@@ -0,0 +1,84 @@
+/*
+ * Copyright 2010-2023 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the LICENSE file.
+ */
+
+#pragma once
+
+#include "Memory.h"
+#include "RawPtr.hpp"
+#include "SpecialRefRegistry.hpp"
+#include "Utils.hpp"
+
+namespace kotlin::mm {
+
+class ThreadData;
+
+// Stable reference to a Kotlin object.
+// Every stable reference makes Kotlin object be in the root set.
+// Use `create` and `dispose` to create and destroy the stable reference.
+class StableRef : private MoveOnly {
+public:
+    StableRef() noexcept = default;
+
+    // Cast raw ref into a stable reference.
+    explicit StableRef(RawSpecialRef* raw) noexcept : node_(SpecialRefRegistry::Node::fromRaw(raw)) {}
+
+    // Cast stable reference into raw ref.
+    [[nodiscard("must be manually disposed")]] explicit operator RawSpecialRef*() && noexcept {
+        // Make sure to move out from node_.
+        auto node = std::move(node_);
+        return node->asRaw();
+    }
+
+    // Create new stable reference for `obj`.
+    [[nodiscard("must be manually disposed")]] static StableRef create(ObjHeader* obj) noexcept;
+
+    // Dispose stable reference.
+    void dispose() && noexcept {
+        auto node = std::move(*this).disposeImpl();
+        tryToDeleteImmediately(std::move(node));
+    }
+
+    // Dispose stable reference using `thread` for opportunistic deletion.
+    // Note: `thread` should still be the current thread, and it's used
+    // when the thread is being destroyed and its TLS deallocating.
+    void disposeOn(mm::ThreadData& thread) && noexcept {
+        auto node = std::move(*this).disposeImpl();
+        tryToDeleteImmediately(thread, std::move(node));
+    }
+
+    // Get the underlying object.
+    // Always safe, because the object is guaranteed to be in the root set.
+    [[nodiscard("expensive pure function")]] ObjHeader* operator*() const noexcept {
+        RuntimeAssert(node_, "operator* on null StableRef");
+        return node_->ref();
+    }
+
+    static StableRef& reinterpret(RawSpecialRef*& raw) noexcept { return reinterpret_cast<StableRef&>(raw); }
+
+    static const StableRef& reinterpret(RawSpecialRef* const& raw) noexcept { return reinterpret_cast<const StableRef&>(raw); }
+
+private:
+    raw_ptr<SpecialRefRegistry::Node> disposeImpl() && noexcept {
+        RuntimeAssert(node_, "Disposing null StableRef");
+        // Make sure to move out from node_.
+        auto node = std::move(node_);
+        // Can be safely called with any thread state.
+        node->releaseRef();
+        // Can be safely called with any thread state.
+        node->dispose();
+        return node;
+    }
+
+    static void tryToDeleteImmediately(raw_ptr<SpecialRefRegistry::Node> node) noexcept;
+    static void tryToDeleteImmediately(mm::ThreadData& thread, raw_ptr<SpecialRefRegistry::Node> node) noexcept;
+
+    raw_ptr<SpecialRefRegistry::Node> node_;
+};
+
+static_assert(sizeof(StableRef) == sizeof(void*), "StableRef must be a thin wrapper around pointer");
+static_assert(alignof(StableRef) == alignof(void*), "StableRef must be a thin wrapper around pointer");
+static_assert(std::is_trivially_destructible_v<StableRef>, "StableRef must be trivially destructible. Destruction is manual via dispose()");
+
+} // namespace kotlin::mm
@@ -1,35 +0,0 @@
-/*
- * Copyright 2010-2020 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
- * that can be found in the LICENSE file.
- */
-
-#include "StableRefRegistry.hpp"
-
-#include "GlobalData.hpp"
-#include "ThreadData.hpp"
-
-using namespace kotlin;
-
-// static
-mm::StableRefRegistry& mm::StableRefRegistry::Instance() noexcept {
-    return GlobalData::Instance().stableRefRegistry();
-}
-
-mm::StableRefRegistry::Node* mm::StableRefRegistry::RegisterStableRef(mm::ThreadData* threadData, ObjHeader* object) noexcept {
-    return threadData->stableRefThreadQueue().Insert(object);
-}
-
-void mm::StableRefRegistry::UnregisterStableRef(mm::ThreadData* threadData, Node* node) noexcept {
-    threadData->stableRefThreadQueue().Erase(node);
-}
-
-void mm::StableRefRegistry::ProcessThread(mm::ThreadData* threadData) noexcept {
-    threadData->stableRefThreadQueue().Publish();
-}
-
-void mm::StableRefRegistry::ProcessDeletions() noexcept {
-    stableRefs_.ApplyDeletions();
-}
-
-mm::StableRefRegistry::StableRefRegistry() = default;
-mm::StableRefRegistry::~StableRefRegistry() = default;
@@ -1,75 +0,0 @@
-/*
- * Copyright 2010-2020 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
- * that can be found in the LICENSE file.
- */
-
-#ifndef RUNTIME_MM_STABLE_REF_REGISTRY_H
-#define RUNTIME_MM_STABLE_REF_REGISTRY_H
-
-#include "Memory.h"
-#include "MultiSourceQueue.hpp"
-#include "ThreadRegistry.hpp"
-
-namespace kotlin {
-namespace mm {
-
-// Registry for all objects that have references outside of Kotlin.
-class StableRefRegistry : Pinned {
-    using Mutex = SpinLock<MutexThreadStateHandling::kIgnore>;
-
-public:
-    class ThreadQueue : public MultiSourceQueue<ObjHeader*, Mutex>::Producer {
-    public:
-        explicit ThreadQueue(StableRefRegistry& registry) : Producer(registry.stableRefs_) {}
-        // Do not add fields as this is just a wrapper and Producer does not have virtual destructor.
-    };
-
-    using Iterable = MultiSourceQueue<ObjHeader*, Mutex>::Iterable;
-    using Iterator = MultiSourceQueue<ObjHeader*, Mutex>::Iterator;
-    using Node = MultiSourceQueue<ObjHeader*, Mutex>::Node;
-
-    StableRefRegistry();
-    ~StableRefRegistry();
-
-    static StableRefRegistry& Instance() noexcept;
-
-    Node* RegisterStableRef(mm::ThreadData* threadData, ObjHeader* object) noexcept;
-
-    void UnregisterStableRef(mm::ThreadData* threadData, Node* node) noexcept;
-
-    // Collect stable references from thread corresponding to `threadData`. Must be called by the thread
-    // when it's asked by GC to stop.
-    void ProcessThread(mm::ThreadData* threadData) noexcept;
-
-    // Lock registry and apply deletions. Should be called on GC thread after all threads have published, and before `LockForIter`.
-    void ProcessDeletions() noexcept;
-
-    // Lock registry for safe iteration.
-    // TODO: Iteration over `stableRefs_` will be slow, because it's `std_support::list` collected at different times from
-    // different threads, and so the nodes are all over the memory. Use metrics to understand how
-    // much of a problem is it.
-    Iterable LockForIter() noexcept { return stableRefs_.LockForIter(); }
-
-    void ClearForTests() noexcept { stableRefs_.ClearForTests(); }
-
-private:
-    // Current approach optimizes for creating and disposing of stable refs:
-    // * creation just enqueues ref, disposing either queues or deletes the ref immediately (if it still resides in the current queue).
-    // * when thread is stopped, it'll scan through the local queue (to mark that refs no longer reside in it) and push creation and
-    //   deletion queues to the global registry.
-    // * during marking GC will have to `ProcessDeletions` to actually delete the refs that were enqueued for deletion.
-    // So, we sacrifice memory (to keep deleted queues) and marking time (to process these queues) to improve creation and disposal times.
-    //
-    // Other alternatives:
-    // * Use a single global collection (e.g. lock free doubly linked list).
-    // * Sacrifice disposal time to try to delete as early as possible (e.g. post directly into owning producer, so it processes deletions
-    //   before posting queue to the global registry)
-    //
-    // TODO: Measure to understand, if this approach is problematic.
-    MultiSourceQueue<ObjHeader*, Mutex> stableRefs_;
-};
-
-} // namespace mm
-} // namespace kotlin
-
-#endif // RUNTIME_MM_STABLE_REF_REGISTRY_H
@@ -46,12 +46,12 @@ 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 stableRefs = mm::StableRefRegistry::Instance().LockForIter();
+    auto specialRefs = mm::SpecialRefRegistry::instance().lockForIter();
    auto threads = mm::ThreadRegistry::Instance().LockForIter();

    EXPECT_THAT(collectCopy(globals), testing::UnorderedElementsAre());
    EXPECT_THAT(collectPointers(extraObjects), testing::UnorderedElementsAre());
-    EXPECT_THAT(collectCopy(stableRefs), testing::UnorderedElementsAre());
+    EXPECT_THAT(collectPointers(specialRefs), testing::UnorderedElementsAre());
    EXPECT_THAT(collectPointers(threads), testing::UnorderedElementsAre());
    gc::AssertClear(mm::GlobalData::Instance().gc());
 }
@@ -15,7 +15,7 @@
 #include "ObjectFactory.hpp"
 #include "ExtraObjectDataFactory.hpp"
 #include "ShadowStack.hpp"
-#include "StableRefRegistry.hpp"
+#include "SpecialRefRegistry.hpp"
 #include "ThreadLocalStorage.hpp"
 #include "Utils.hpp"
 #include "ThreadSuspension.hpp"
@@ -33,7 +33,7 @@ public:
    explicit ThreadData(int threadId) noexcept :
        threadId_(threadId),
        globalsThreadQueue_(GlobalsRegistry::Instance()),
-        stableRefThreadQueue_(StableRefRegistry::Instance()),
+        specialRefRegistry_(SpecialRefRegistry::instance()),
        extraObjectDataThreadQueue_(ExtraObjectDataFactory::Instance()),
        gc_(GlobalData::Instance().gc(), *this),
        suspensionData_(ThreadState::kNative, *this) {}
@@ -46,7 +46,7 @@ public:

    ThreadLocalStorage& tls() noexcept { return tls_; }

-    StableRefRegistry::ThreadQueue& stableRefThreadQueue() noexcept { return stableRefThreadQueue_; }
+    SpecialRefRegistry::ThreadQueue& specialRefRegistry() noexcept { return specialRefRegistry_; }

    ExtraObjectDataFactory::ThreadQueue& extraObjectDataThreadQueue() noexcept { return extraObjectDataThreadQueue_; }

@@ -65,14 +65,14 @@ public:
    void Publish() noexcept {
        // TODO: These use separate locks, which is inefficient.
        globalsThreadQueue_.Publish();
-        stableRefThreadQueue_.Publish();
+        specialRefRegistry_.publish();
        extraObjectDataThreadQueue_.Publish();
        gc_.Publish();
    }

    void ClearForTests() noexcept {
        globalsThreadQueue_.ClearForTests();
-        stableRefThreadQueue_.ClearForTests();
+        specialRefRegistry_.clearForTests();
        extraObjectDataThreadQueue_.ClearForTests();
        gc_.ClearForTests();
    }
@@ -81,7 +81,7 @@ private:
    const int threadId_;
    GlobalsRegistry::ThreadQueue globalsThreadQueue_;
    ThreadLocalStorage tls_;
-    StableRefRegistry::ThreadQueue stableRefThreadQueue_;
+    SpecialRefRegistry::ThreadQueue specialRefRegistry_;
    ExtraObjectDataFactory::ThreadQueue extraObjectDataThreadQueue_;
    ShadowStack shadowStack_;
    gc::GC::ThreadData gc_;
@@ -44,6 +44,7 @@ public:
        RuntimeAssert(currentThreadDataNode_ != nullptr, "Thread is not attached to the runtime");
        return currentThreadDataNode_;
    }
+    Node* CurrentThreadDataNodeOrNull() const noexcept { return currentThreadDataNode_; }

    bool IsCurrentThreadRegistered() const noexcept { return currentThreadDataNode_ != nullptr; }

@@ -6,20 +6,21 @@
 #include "Weak.hpp"

 #include "ExtraObjectData.hpp"
-#include "ObjectOps.hpp"
+#include "WeakRef.hpp"
 #include "ThreadState.hpp"
 #include "Types.h"

 using namespace kotlin;

 extern "C" {
-OBJ_GETTER(makeRegularWeakReferenceImpl, void*);
+OBJ_GETTER(makeRegularWeakReferenceImpl, void*, void*);
 }

 namespace {

 struct RegularWeakReferenceImpl {
    ObjHeader header;
+    mm::WeakRef weakRef;
    void* referred;
 };

@@ -38,18 +39,17 @@ OBJ_GETTER(mm::createRegularWeakReferenceImpl, ObjHeader* object) noexcept {
        RETURN_OBJ(weakRef);
    }
    ObjHolder holder;
-    auto* weakRef = makeRegularWeakReferenceImpl(object, holder.slot());
+    auto* weakRef = makeRegularWeakReferenceImpl(static_cast<mm::RawSpecialRef*>(mm::WeakRef::create(object)), object, holder.slot());
    auto* setWeakRef = extraObject.GetOrSetRegularWeakReferenceImpl(object, weakRef);
    RETURN_OBJ(setWeakRef);
 }

 void mm::disposeRegularWeakReferenceImpl(ObjHeader* weakRef) noexcept {
-    asRegularWeakReferenceImpl(weakRef)->referred = nullptr;
+    std::move(asRegularWeakReferenceImpl(weakRef)->weakRef).dispose();
 }

 OBJ_GETTER(mm::derefRegularWeakReferenceImpl, ObjHeader* weakRef) noexcept {
-    ObjHeader** location = reinterpret_cast<ObjHeader**>(&asRegularWeakReferenceImpl(weakRef)->referred);
-    RETURN_RESULT_OF(mm::ReadHeapRefAtomic, location);
+    RETURN_RESULT_OF0(asRegularWeakReferenceImpl(weakRef)->weakRef.tryRef);
 }

 ObjHeader* mm::regularWeakReferenceImplBaseObjectUnsafe(ObjHeader* weakRef) noexcept {
@@ -0,0 +1,16 @@
+/*
+ * Copyright 2010-2023 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the LICENSE file.
+ */
+
+#include "WeakRef.hpp"
+
+#include "ThreadData.hpp"
+
+using namespace kotlin;
+
+// static
+mm::WeakRef mm::WeakRef::create(ObjHeader* obj) noexcept {
+    RuntimeAssert(obj != nullptr, "Creating WeakRef for null object");
+    return mm::ThreadRegistry::Instance().CurrentThreadData()->specialRefRegistry().createWeakRef(obj);
+}
@@ -0,0 +1,64 @@
+/*
+ * Copyright 2010-2023 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the LICENSE file.
+ */
+
+#pragma once
+
+#include "Memory.h"
+#include "RawPtr.hpp"
+#include "SpecialRefRegistry.hpp"
+#include "Utils.hpp"
+
+namespace kotlin::mm {
+
+// Weak reference to a Kotlin object.
+// GC automatically invalidates the reference when the Kotlin object is collected.
+// Use `create` and `dispose` to create and destroy the weak reference.
+class WeakRef : private MoveOnly {
+public:
+    WeakRef() noexcept = default;
+
+    // Cast raw ref into a weak reference.
+    explicit WeakRef(RawSpecialRef* raw) noexcept : node_(SpecialRefRegistry::Node::fromRaw(raw)) {}
+
+    // Cast weak reference into raw ref.
+    [[nodiscard("must be manually disposed")]] explicit operator RawSpecialRef*() && noexcept {
+        // Make sure to move out from node_.
+        auto node = std::move(node_);
+        return node->asRaw();
+    }
+
+    // Create new weak reference for `obj`.
+    [[nodiscard("must be manually disposed")]] static WeakRef create(ObjHeader* obj) noexcept;
+
+    // Dispose weak reference.
+    void dispose() && noexcept {
+        RuntimeAssert(node_, "Disposing null WeakRef");
+        // Make sure to move out from node_.
+        auto node = std::move(node_);
+        // Can be safely called with any thread state.
+        node->dispose();
+    }
+
+    // Safely dereference weak reference. Returns null if the underlying object
+    // is not alive.
+    OBJ_GETTER0(tryRef) const noexcept {
+        RuntimeAssert(node_, "tryRef on null WeakRef");
+        AssertThreadState(ThreadState::kRunnable);
+        RETURN_RESULT_OF0(node_->tryRef);
+    }
+
+    static WeakRef& reinterpret(RawSpecialRef*& raw) noexcept { return reinterpret_cast<WeakRef&>(raw); }
+
+    static const WeakRef& reinterpret(RawSpecialRef* const& raw) noexcept { return reinterpret_cast<const WeakRef&>(raw); }
+
+private:
+    raw_ptr<SpecialRefRegistry::Node> node_;
+};
+
+static_assert(sizeof(WeakRef) == sizeof(void*), "WeakRef must be a thin wrapper around pointer");
+static_assert(alignof(WeakRef) == alignof(void*), "WeakRef must be a thin wrapper around pointer");
+static_assert(std::is_trivially_destructible_v<WeakRef>, "WeakRef must be trivially destructible. Destruction is manual via dispose()");
+
+} // namespace kotlin::mm
@@ -105,7 +105,7 @@ static void injectToRuntime();
        {
          kotlin::ThreadStateGuard guard(kotlin::ThreadState::kNative);
          candidate->refHolder.releaseRef();
-          [candidate releaseAsAssociatedObject:ReleaseMode::kDetachAndRelease];
+          [candidate releaseAsAssociatedObject];
        }
        return objc_retain(old);
      }
@@ -142,9 +142,16 @@ static void injectToRuntime();
  }
 }

-(void)releaseAsAssociatedObject:(ReleaseMode)mode {
+-(void)releaseAsAssociatedObject {
  RuntimeAssert(!permanent, "Cannot be called on permanent objects");

+  if (CurrentMemoryModel == MemoryModel::kExperimental) {
+    // No need for any special handling. Weak reference handling machinery
+    // has already cleaned up the reference to Kotlin object.
+    [super release];
+    return;
+  }
+
  // This function is called by the GC. It made a decision to reclaim Kotlin object, and runs
  // deallocation hooks at the moment, including deallocation of the "associated object" ([self])
  // using the [super release] call below.
@@ -155,20 +162,25 @@ static void injectToRuntime();
  // Generally retaining and releasing Kotlin object that is being deallocated would lead to
  // use-after-dispose and double-dispose problems (with unpredictable consequences) or to an assertion failure.
  // To workaround this, detach the back ref from the Kotlin object:
-  if (ReleaseModeHasDetach(mode)) {
-    refHolder.detach();
-  } else {
-    // With Mark&Sweep this object should already have been detached earlier.
-    refHolder.assertDetached();
-  }
+  refHolder.detach();

  // So retain/release/etc. on [self] won't affect the Kotlin object, and an attempt to get
  // the reference to it (e.g. when calling Kotlin method on [self]) would crash.
  // The latter is generally ok because can be triggered only by user-defined Swift/Obj-C
  // subclasses of Kotlin classes.
-  if (ReleaseModeHasRelease(mode)) {
-    [super release];
+  [super release];
+}
+
+-(void)dealloc {
+  if (CurrentMemoryModel == MemoryModel::kExperimental) {
+    if (!permanent) {
+      refHolder.dealloc();
+    }
+    [super dealloc];
+    return;
  }
+
+  [super dealloc];
 }

 - (instancetype)copyWithZone:(NSZone *)zone {
@@ -186,9 +198,7 @@ static void injectToRuntime();
  RETURN_RESULT_OF(Kotlin_ObjCExport_convertUnmappedObjCObject, self);
 }

-(void)releaseAsAssociatedObject:(ReleaseMode)mode {
-  if (!ReleaseModeHasRelease(mode))
-    return;
+-(void)releaseAsAssociatedObject {
  objc_release(self);
 }
@end
@@ -261,7 +271,7 @@ static void injectToRuntimeImpl() {
  Kotlin_ObjCExport_toKotlinSelector = @selector(toKotlin:);

  RuntimeCheck(Kotlin_ObjCExport_releaseAsAssociatedObjectSelector == nullptr, errorMessage);
-  Kotlin_ObjCExport_releaseAsAssociatedObjectSelector = @selector(releaseAsAssociatedObject:);
+  Kotlin_ObjCExport_releaseAsAssociatedObjectSelector = @selector(releaseAsAssociatedObject);
 }

 static void injectToRuntime() {
@@ -93,9 +93,7 @@ static inline KInt objCIndexToKotlinOrThrow(NSUInteger index) {
  RETURN_RESULT_OF(invokeAndAssociate, Kotlin_NSArrayAsKList_create, objc_retain(self));
 }

-(void)releaseAsAssociatedObject:(ReleaseMode)mode {
-  if (!ReleaseModeHasRelease(mode))
-    return;
+-(void)releaseAsAssociatedObject {
  objc_release(self);
 }
@end
@@ -108,9 +106,7 @@ static inline KInt objCIndexToKotlinOrThrow(NSUInteger index) {
  RETURN_RESULT_OF(invokeAndAssociate, Kotlin_NSMutableArrayAsKMutableList_create, objc_retain(self));
 }

-(void)releaseAsAssociatedObject:(ReleaseMode)mode {
-  if (!ReleaseModeHasRelease(mode))
-    return;
+-(void)releaseAsAssociatedObject {
  objc_release(self);
 }
@end
@@ -124,9 +120,7 @@ static inline KInt objCIndexToKotlinOrThrow(NSUInteger index) {
  RETURN_RESULT_OF(invokeAndAssociate, Kotlin_NSSetAsKSet_create, objc_retain(self));
 }

-(void)releaseAsAssociatedObject:(ReleaseMode)mode {
-  if (!ReleaseModeHasRelease(mode))
-    return;
+-(void)releaseAsAssociatedObject {
  objc_release(self);
 }

@@ -140,9 +134,7 @@ static inline KInt objCIndexToKotlinOrThrow(NSUInteger index) {
  RETURN_RESULT_OF(invokeAndAssociate, Kotlin_NSDictionaryAsKMap_create, objc_retain(self));
 }

-(void)releaseAsAssociatedObject:(ReleaseMode)mode {
-  if (!ReleaseModeHasRelease(mode))
-    return;
+-(void)releaseAsAssociatedObject {
  objc_release(self);
 }

@@ -156,7 +148,7 @@ static inline KInt objCIndexToKotlinOrThrow(NSUInteger index) {
 }

 -(void)dealloc {
-  iteratorHolder.disposeFromNative();
+  iteratorHolder.dispose();
  [super dealloc];
 }

@@ -186,7 +178,7 @@ static inline KInt objCIndexToKotlinOrThrow(NSUInteger index) {
 }

 -(void)dealloc {
-  listHolder.disposeFromNative();
+  listHolder.dispose();
  [super dealloc];
 }

@@ -222,7 +214,7 @@ static inline KInt objCIndexToKotlinOrThrow(NSUInteger index) {
 }

 -(void)dealloc {
-  listHolder.disposeFromNative();
+  listHolder.dispose();
  [super dealloc];
 }

@@ -302,7 +294,7 @@ static inline id KSet_getElement(KRef set, id object) {
 }

 -(void)dealloc {
-  setHolder.disposeFromNative();
+  setHolder.dispose();
  [super dealloc];
 }

@@ -390,7 +382,7 @@ static inline id KSet_getElement(KRef set, id object) {
  // Note: since setHolder initialization is not performed directly with alloc,
  // it is possible that it wasn't initialized properly.
  // Fortunately setHolder.dispose() handles the zero-initialized case too.
-  setHolder.disposeFromNative();
+  setHolder.dispose();
  [super dealloc];
 }

@@ -464,7 +456,7 @@ static inline id KMap_get(KRef map, id aKey) {
 }

 -(void)dealloc {
-  mapHolder.disposeFromNative();
+  mapHolder.dispose();
  [super dealloc];
 }

@@ -510,7 +502,7 @@ static inline id KMap_get(KRef map, id aKey) {
  // Note: since mapHolder initialization is not performed directly with alloc,
  // it is possible that it wasn't initialized properly.
  // Fortunately mapHolder.dispose() handles the zero-initialized case too.
-  mapHolder.disposeFromNative();
+  mapHolder.dispose();
  [super dealloc];
 }

@@ -599,9 +591,7 @@ static inline id KMap_get(KRef map, id aKey) {
@end

@implementation NSEnumerator (NSEnumeratorAsAssociatedObject)
-(void)releaseAsAssociatedObject:(ReleaseMode)mode {
-  if (!ReleaseModeHasRelease(mode))
-    return;
+-(void)releaseAsAssociatedObject {
  objc_release(self);
 }
@end
@@ -35,7 +35,7 @@
 }

 -(void)dealloc {
-  refHolder.disposeFromNative();
+  refHolder.dispose();
  [super dealloc];
 }

@@ -74,9 +74,7 @@ void objc_release(id obj);
 }

 // Called when removing Kotlin object.
-(void)releaseAsAssociatedObject:(ReleaseMode)mode {
-  if (!ReleaseModeHasRelease(mode))
-    return;
+-(void)releaseAsAssociatedObject {
  objc_destroyWeak(&referred);
  objc_release(self);
 }
@@ -51,7 +51,8 @@ kotlin::test_support::TypeInfoHolder theWorkerBoundReferenceTypeInfoHolder{
        kotlin::test_support::TypeInfoHolder::ObjectBuilder<EmptyPayload>()};
 kotlin::test_support::TypeInfoHolder theCleanerImplTypeInfoHolder{kotlin::test_support::TypeInfoHolder::ObjectBuilder<EmptyPayload>()};
 kotlin::test_support::TypeInfoHolder theRegularWeakReferenceImplTypeInfoHolder{
-        kotlin::test_support::TypeInfoHolder::ObjectBuilder<kotlin::test_support::RegularWeakReferenceImplPayload>()};
+        kotlin::test_support::TypeInfoHolder::ObjectBuilder<kotlin::test_support::RegularWeakReferenceImplPayload>().addFlag(
+                TF_HAS_FINALIZER)};

 ArrayHeader theEmptyStringImpl = {theStringTypeInfoHolder.typeInfo(), /* element count */ 0};