Add RootSet to unify iterating across the entire rootset (#4685)

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

@@ -80,4 +80,11 @@ extern "C" const int KonanNeedDebugInfo;
         ::internal::TODOImpl(CURRENT_SOURCE_LOCATION, ##__VA_ARGS__); \
     } while (false)
 
+// Use RuntimeFail() to unconditionally fail, signifying compiler/runtime bug.
+// TODO: Consider using `CURRENT_SOURCE_LOCATION` when `KonanNeedDebugInfo` is `true`.
+#define RuntimeFail(format, ...) \
+    do { \
+        RuntimeAssertFailed(nullptr, format, ##__VA_ARGS__); \
+    } while (false)
+
 #endif // RUNTIME_ASSERT_H

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

@@ -26,6 +26,9 @@ public:
 
     using Iterator = MultiSourceQueue<ObjHeader**>::Iterator;
 
+    GlobalsRegistry();
+    ~GlobalsRegistry();
+
     static GlobalsRegistry& Instance() noexcept;
 
     void RegisterStorageForGlobal(mm::ThreadData* threadData, ObjHeader** location) noexcept;
@@ -41,11 +44,6 @@ public:
     Iterable Iter() noexcept { return globals_.Iter(); }
 
 private:
-    friend class GlobalData;
-
-    GlobalsRegistry();
-    ~GlobalsRegistry();
-
     // TODO: Add-only MultiSourceQueue can be made more efficient. Measure, if it's a problem.
     MultiSourceQueue<ObjHeader**> globals_;
 };

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

@@ -0,0 +1,149 @@
+/*
+ * Copyright 2010-2021 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 "RootSet.hpp"
+
+#include "KAssert.h"
+#include "GlobalData.hpp"
+#include "ThreadData.hpp"
+
+using namespace kotlin;
+
+mm::ThreadRootSet::Iterator::Iterator(begin_t, ThreadRootSet& owner) noexcept :
+    owner_(owner), phase_(Phase::kStack), stackIterator_(owner_.stack_.begin()) {
+    Init();
+}
+
+mm::ThreadRootSet::Iterator::Iterator(end_t, ThreadRootSet& owner) noexcept : owner_(owner), phase_(Phase::kDone) {}
+
+ObjHeader*& mm::ThreadRootSet::Iterator::operator*() noexcept {
+    switch (phase_) {
+        case Phase::kStack:
+            return *stackIterator_;
+        case Phase::kTLS:
+            return **tlsIterator_;
+        case Phase::kDone:
+            RuntimeFail("Cannot dereference");
+    }
+}
+
+mm::ThreadRootSet::Iterator& mm::ThreadRootSet::Iterator::operator++() noexcept {
+    switch (phase_) {
+        case Phase::kStack:
+            ++stackIterator_;
+            Init();
+            return *this;
+        case Phase::kTLS:
+            ++tlsIterator_;
+            Init();
+            return *this;
+        case Phase::kDone:
+            return *this;
+    }
+}
+
+bool mm::ThreadRootSet::Iterator::operator==(const Iterator& rhs) const noexcept {
+    if (phase_ != rhs.phase_) {
+        return false;
+    }
+
+    switch (phase_) {
+        case Phase::kDone:
+            return true;
+        case Phase::kStack:
+            return stackIterator_ == rhs.stackIterator_;
+        case Phase::kTLS:
+            return tlsIterator_ == rhs.tlsIterator_;
+    }
+}
+
+void mm::ThreadRootSet::Iterator::Init() noexcept {
+    while (phase_ != Phase::kDone) {
+        switch (phase_) {
+            case Phase::kStack:
+                if (stackIterator_ != owner_.stack_.end()) return;
+                phase_ = Phase::kTLS;
+                tlsIterator_ = owner_.tls_.begin();
+                break;
+            case Phase::kTLS:
+                if (tlsIterator_ != owner_.tls_.end()) return;
+                phase_ = Phase::kDone;
+                break;
+            case Phase::kDone:
+                RuntimeFail("Impossible");
+        }
+    }
+}
+
+mm::GlobalRootSet::Iterator::Iterator(begin_t, GlobalRootSet& owner) noexcept :
+    owner_(owner), phase_(Phase::kGlobals), globalsIterator_(owner_.globalsIterable_.begin()) {
+    Init();
+}
+
+mm::GlobalRootSet::Iterator::Iterator(end_t, GlobalRootSet& owner) noexcept : owner_(owner), phase_(Phase::kDone) {}
+
+ObjHeader*& mm::GlobalRootSet::Iterator::operator*() noexcept {
+    switch (phase_) {
+        case Phase::kGlobals:
+            return **globalsIterator_;
+        case Phase::kStableRefs:
+            return *stableRefsIterator_;
+        case Phase::kDone:
+            RuntimeFail("Cannot dereference");
+    }
+}
+
+mm::GlobalRootSet::Iterator& mm::GlobalRootSet::Iterator::operator++() noexcept {
+    switch (phase_) {
+        case Phase::kGlobals:
+            ++globalsIterator_;
+            Init();
+            return *this;
+        case Phase::kStableRefs:
+            ++stableRefsIterator_;
+            Init();
+            return *this;
+        case Phase::kDone:
+            return *this;
+    }
+}
+
+bool mm::GlobalRootSet::Iterator::operator==(const Iterator& rhs) const noexcept {
+    if (phase_ != rhs.phase_) {
+        return false;
+    }
+
+    switch (phase_) {
+        case Phase::kDone:
+            return true;
+        case Phase::kGlobals:
+            return globalsIterator_ == rhs.globalsIterator_;
+        case Phase::kStableRefs:
+            return stableRefsIterator_ == rhs.stableRefsIterator_;
+    }
+}
+
+void mm::GlobalRootSet::Iterator::Init() noexcept {
+    while (phase_ != Phase::kDone) {
+        switch (phase_) {
+            case Phase::kGlobals:
+                if (globalsIterator_ != owner_.globalsIterable_.end()) return;
+                phase_ = Phase::kStableRefs;
+                stableRefsIterator_ = owner_.stableRefsIterable_.begin();
+                break;
+            case Phase::kStableRefs:
+                if (stableRefsIterator_ != owner_.stableRefsIterable_.end()) return;
+                phase_ = Phase::kDone;
+                break;
+            case Phase::kDone:
+                RuntimeFail("Impossible");
+        }
+    }
+}
+
+mm::ThreadRootSet::ThreadRootSet(ThreadData& threadData) noexcept : ThreadRootSet(threadData.shadowStack(), threadData.tls()) {}
+
+mm::GlobalRootSet::GlobalRootSet() noexcept :
+    GlobalRootSet(mm::GlobalData::Instance().globalsRegistry(), mm::GlobalData::Instance().stableRefRegistry()) {}

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

@@ -0,0 +1,123 @@
+/*
+ * Copyright 2010-2021 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_ROOT_SET_H
+#define RUNTIME_MM_ROOT_SET_H
+
+#include "GlobalsRegistry.hpp"
+#include "ShadowStack.hpp"
+#include "StableRefRegistry.hpp"
+#include "ThreadLocalStorage.hpp"
+
+struct ObjHeader;
+
+namespace kotlin {
+namespace mm {
+
+class ThreadData;
+
+class ThreadRootSet {
+public:
+    class Iterator {
+    public:
+        struct begin_t {};
+        static constexpr inline begin_t begin = begin_t{};
+
+        struct end_t {};
+        static constexpr inline end_t end = end_t{};
+
+        Iterator(begin_t, ThreadRootSet& owner) noexcept;
+        Iterator(end_t, ThreadRootSet& owner) noexcept;
+
+        ObjHeader*& operator*() noexcept;
+
+        Iterator& operator++() noexcept;
+
+        bool operator==(const Iterator& rhs) const noexcept;
+        bool operator!=(const Iterator& rhs) const noexcept { return !(*this == rhs); }
+
+    private:
+        enum class Phase {
+            kStack,
+            kTLS,
+            kDone,
+        };
+
+        void Init() noexcept;
+
+        ThreadRootSet& owner_;
+        Phase phase_;
+        union {
+            ShadowStack::Iterator stackIterator_;
+            ThreadLocalStorage::Iterator tlsIterator_;
+        };
+    };
+
+    ThreadRootSet(ShadowStack& stack, ThreadLocalStorage& tls) noexcept : stack_(stack), tls_(tls) {}
+    explicit ThreadRootSet(ThreadData& threadData) noexcept;
+
+    Iterator begin() noexcept { return Iterator(Iterator::begin, *this); }
+    Iterator end() noexcept { return Iterator(Iterator::end, *this); }
+
+private:
+    ShadowStack& stack_;
+    ThreadLocalStorage& tls_;
+};
+
+class GlobalRootSet {
+public:
+    class Iterator {
+    public:
+        struct begin_t {};
+        static constexpr inline begin_t begin = begin_t{};
+
+        struct end_t {};
+        static constexpr inline end_t end = end_t{};
+
+        Iterator(begin_t, GlobalRootSet& owner) noexcept;
+        Iterator(end_t, GlobalRootSet& owner) noexcept;
+
+        ObjHeader*& operator*() noexcept;
+
+        Iterator& operator++() noexcept;
+
+        bool operator==(const Iterator& rhs) const noexcept;
+        bool operator!=(const Iterator& rhs) const noexcept { return !(*this == rhs); }
+
+    private:
+        enum class Phase {
+            kGlobals,
+            kStableRefs,
+            kDone,
+        };
+
+        void Init() noexcept;
+
+        GlobalRootSet& owner_;
+        Phase phase_;
+        union {
+            GlobalsRegistry::Iterator globalsIterator_;
+            StableRefRegistry::Iterator stableRefsIterator_;
+        };
+    };
+
+    GlobalRootSet(GlobalsRegistry& globalsRegistry, StableRefRegistry& stableRefRegistry) noexcept :
+        globalsIterable_(globalsRegistry.Iter()), stableRefsIterable_(stableRefRegistry.Iter()) {}
+    GlobalRootSet() noexcept;
+
+    Iterator begin() noexcept { return Iterator(Iterator::begin, *this); }
+    Iterator end() noexcept { return Iterator(Iterator::end, *this); }
+
+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_;
+};
+
+} // namespace mm
+} // namespace kotlin
+
+#endif // RUNTIME_MM_ROOT_SET_H

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

@@ -0,0 +1,126 @@
+/*
+ * Copyright 2010-2021 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 "RootSet.hpp"
+
+#include "gmock/gmock.h"
+#include "gtest/gtest.h"
+
+#include "ShadowStack.hpp"
+
+using namespace kotlin;
+
+namespace {
+
+// TODO: All the test helpers to create the rootset should be abstracted out.
+
+template <size_t LocalsCount>
+class StackEntry : private Pinned {
+public:
+    static_assert(LocalsCount > 0, "Must have at least 1 object on stack");
+
+    explicit StackEntry(mm::ShadowStack& shadowStack) : shadowStack_(shadowStack), value_(make_unique<ObjHeader>()) {
+        // Fill `locals_` with some values.
+        for (size_t i = 0; i < LocalsCount; ++i) {
+            (*this)[i] = value_.get() + i;
+        }
+
+        shadowStack_.EnterFrame(data_.data(), 0, kTotalCount);
+    }
+
+    ~StackEntry() { shadowStack_.LeaveFrame(data_.data(), 0, kTotalCount); }
+
+    ObjHeader*& operator[](size_t index) { return data_[kFrameOverlayCount + index]; }
+
+private:
+    mm::ShadowStack& shadowStack_;
+    KStdUniquePtr<ObjHeader> value_;
+
+    // The following is what the compiler creates on the stack.
+    static inline constexpr int kFrameOverlayCount = sizeof(FrameOverlay) / sizeof(ObjHeader**);
+    static inline constexpr int kTotalCount = kFrameOverlayCount + LocalsCount;
+    std::array<ObjHeader*, kTotalCount> data_;
+};
+
+struct TLSKey {};
+
+} // namespace
+
+TEST(ThreadRootSetTest, Basic) {
+    mm::ShadowStack stack;
+    StackEntry<2> entry(stack);
+
+    TLSKey key;
+    mm::ThreadLocalStorage tls;
+    tls.AddRecord(&key, 3);
+    tls.Commit();
+
+    mm::ThreadRootSet iter(stack, tls);
+
+    KStdVector<ObjHeader*> actual;
+    for (auto& object : iter) {
+        actual.push_back(object);
+    }
+
+    EXPECT_THAT(actual, testing::ElementsAre(entry[0], entry[1], *tls.Lookup(&key, 0), *tls.Lookup(&key, 1), *tls.Lookup(&key, 2)));
+}
+
+TEST(ThreadRootSetTest, Empty) {
+    mm::ShadowStack stack;
+    mm::ThreadLocalStorage tls;
+
+    mm::ThreadRootSet iter(stack, tls);
+
+    KStdVector<ObjHeader*> actual;
+    for (auto& object : iter) {
+        actual.push_back(object);
+    }
+
+    EXPECT_THAT(actual, testing::IsEmpty());
+}
+
+TEST(GlobalRootSetTest, Basic) {
+    mm::GlobalsRegistry globals;
+    mm::GlobalsRegistry::ThreadQueue globalsProducer(globals);
+    ObjHeader* global1 = reinterpret_cast<ObjHeader*>(1);
+    ObjHeader* global2 = reinterpret_cast<ObjHeader*>(2);
+    globalsProducer.Insert(&global1);
+    globalsProducer.Insert(&global2);
+
+    mm::StableRefRegistry stableRefs;
+    mm::StableRefRegistry::ThreadQueue stableRefsProducer(stableRefs);
+    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);
+
+    globalsProducer.Publish();
+    stableRefsProducer.Publish();
+
+    mm::GlobalRootSet iter(globals, stableRefs);
+
+    KStdVector<ObjHeader*> actual;
+    for (auto& object : iter) {
+        actual.push_back(object);
+    }
+
+    EXPECT_THAT(actual, testing::ElementsAre(global1, global2, stableRef1, stableRef2, stableRef3));
+}
+
+TEST(GlobalRootSetTest, Empty) {
+    mm::GlobalsRegistry globals;
+    mm::StableRefRegistry stableRefs;
+
+    mm::GlobalRootSet iter(globals, stableRefs);
+
+    KStdVector<ObjHeader*> actual;
+    for (auto& object : iter) {
+        actual.push_back(object);
+    }
+
+    EXPECT_THAT(actual, testing::IsEmpty());
+}

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

@@ -26,6 +26,9 @@ public:
     using Iterator = MultiSourceQueue<ObjHeader*>::Iterator;
     using Node = MultiSourceQueue<ObjHeader*>::Node;
 
+    StableRefRegistry();
+    ~StableRefRegistry();
+
     static StableRefRegistry& Instance() noexcept;
 
     Node* RegisterStableRef(mm::ThreadData* threadData, ObjHeader* object) noexcept;
@@ -46,11 +49,6 @@ public:
     Iterable Iter() noexcept { return stableRefs_.Iter(); }
 
 private:
-    friend class GlobalData;
-
-    StableRefRegistry();
-    ~StableRefRegistry();
-
     // 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

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

@@ -59,6 +59,13 @@ public:
 
     GC::ThreadData& gc() noexcept { return gc_; }
 
+    void Publish() noexcept {
+        // TODO: These use separate locks, which is inefficient.
+        globalsThreadQueue_.Publish();
+        stableRefThreadQueue_.Publish();
+        objectFactoryThreadQueue_.Publish();
+    }
+
 private:
     const pthread_t threadId_;
     GlobalsRegistry::ThreadQueue globalsThreadQueue_;