Deal with FrameOverlay in the new mm (#4646)

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

@@ -142,6 +142,16 @@ extern "C" RUNTIME_NOTHROW void InitAndRegisterGlobal(ObjHeader** location, cons
 
 extern "C" const MemoryModel CurrentMemoryModel = MemoryModel::kExperimental;
 
+extern "C" RUNTIME_NOTHROW void EnterFrame(ObjHeader** start, int parameters, int count) {
+    auto* threadData = mm::ThreadRegistry::Instance().CurrentThreadData();
+    threadData->shadowStack().EnterFrame(start, parameters, count);
+}
+
+extern "C" RUNTIME_NOTHROW void LeaveFrame(ObjHeader** start, int parameters, int count) {
+    auto* threadData = mm::ThreadRegistry::Instance().CurrentThreadData();
+    threadData->shadowStack().LeaveFrame(start, parameters, count);
+}
+
 extern "C" RUNTIME_NOTHROW void AddTLSRecord(MemoryState* memory, void** key, int size) {
     GetThreadData(memory)->tls().AddRecord(key, size);
 }

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

@@ -0,0 +1,37 @@
+/*
+ * 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 "ShadowStack.hpp"
+
+using namespace kotlin;
+
+mm::ShadowStack::Iterator& mm::ShadowStack::Iterator::operator++() noexcept {
+    ++object_;
+    Init();
+    return *this;
+}
+
+void mm::ShadowStack::Iterator::Init() noexcept {
+    while (frame_) {
+        if (object_ < end_) return;
+        frame_ = frame_->previous;
+        object_ = begin();
+        end_ = end();
+    }
+}
+
+void mm::ShadowStack::EnterFrame(ObjHeader** start, int parameters, int count) noexcept {
+    FrameOverlay* frame = reinterpret_cast<FrameOverlay*>(start);
+    frame->previous = currentFrame_;
+    currentFrame_ = frame;
+    // TODO: maybe compress in single value somehow.
+    frame->parameters = parameters;
+    frame->count = count;
+}
+
+void mm::ShadowStack::LeaveFrame(ObjHeader** start, int parameters, int count) noexcept {
+    FrameOverlay* frame = reinterpret_cast<FrameOverlay*>(start);
+    currentFrame_ = frame->previous;
+}

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

@@ -0,0 +1,67 @@
+/*
+ * 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_SHADOW_STACK
+#define RUNTIME_MM_SHADOW_STACK
+
+#include "Memory.h"
+#include "Utils.hpp"
+
+struct FrameOverlay;
+struct ObjHeader;
+
+namespace kotlin {
+namespace mm {
+
+// Accessing current stack as provided by K/N compiler. The compiler calls `EnterFrame` when
+// it has allocated and zeroed stack space in the function prologue. And it calls `LeaveFrame` in
+// the function epilogue (both for regular return and for exception unwinding).
+//
+// Stack scanning does not lock anything and so must either be done while the mutator is stopped (or is
+// running code outside Kotlin), or by the mutator itself. So, in concurrent collection case, make sure
+// to do as little as possible while scanning the stack to free the mutator as soon as possible.
+//
+// TODO: This is currently incompatible with stack-allocated objects. Fix it.
+class ShadowStack : private Pinned {
+public:
+    class Iterator {
+    public:
+        explicit Iterator(FrameOverlay* frame) noexcept : frame_(frame), object_(begin()), end_(end()) { Init(); }
+
+        ObjHeader*& operator*() noexcept { return *object_; }
+        Iterator& operator++() noexcept;
+
+        bool operator==(const Iterator& rhs) const noexcept { return frame_ == rhs.frame_ && object_ == rhs.object_; }
+        bool operator!=(const Iterator& rhs) const noexcept { return !(*this == rhs); }
+
+    private:
+        void Init() noexcept;
+
+        // TODO: This copies the approach in the old MM. Do we need to also traverse function parameters in the new MM?
+        ObjHeader** begin() noexcept { return frame_ ? reinterpret_cast<ObjHeader**>(frame_ + 1) + frame_->parameters : nullptr; }
+        ObjHeader** end() noexcept {
+            constexpr int kFrameOverlaySlots = sizeof(FrameOverlay) / sizeof(ObjHeader**);
+            return frame_ ? begin() + frame_->count - kFrameOverlaySlots - frame_->parameters : nullptr;
+        }
+
+        FrameOverlay* frame_;
+        ObjHeader** object_ = nullptr;
+        ObjHeader** end_ = nullptr;
+    };
+
+    void EnterFrame(ObjHeader** start, int parameters, int count) noexcept;
+    void LeaveFrame(ObjHeader** start, int parameters, int count) noexcept;
+
+    Iterator begin() noexcept { return Iterator(currentFrame_); }
+    Iterator end() noexcept { return Iterator(nullptr); }
+
+private:
+    FrameOverlay* currentFrame_ = nullptr;
+};
+
+} // namespace mm
+} // namespace kotlin
+
+#endif // RUNTIME_MM_SHADOW_STACK

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

@@ -0,0 +1,121 @@
+/*
+ * 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 "ShadowStack.hpp"
+
+#include "gmock/gmock.h"
+#include "gtest/gtest.h"
+
+#include "Memory.h"
+#include "Types.h"
+#include "Utils.hpp"
+
+using namespace kotlin;
+
+namespace {
+
+template <size_t ParametersCount, size_t LocalsCount>
+class StackEntry : private Pinned {
+public:
+    static_assert(ParametersCount + 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(), ParametersCount, kTotalCount);
+    }
+
+    ~StackEntry() { shadowStack_.LeaveFrame(data_.data(), ParametersCount, kTotalCount); }
+
+    ObjHeader*& operator[](size_t index) { return data_[kFrameOverlayCount + ParametersCount + 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 + ParametersCount + LocalsCount;
+    std::array<ObjHeader*, kTotalCount> data_;
+};
+
+KStdVector<ObjHeader*> Collect(mm::ShadowStack& shadowStack) {
+    KStdVector<ObjHeader*> result;
+    for (ObjHeader* local : shadowStack) {
+        result.push_back(local);
+    }
+    return result;
+}
+
+} // namespace
+
+TEST(ShadowStackTest, Empty) {
+    mm::ShadowStack shadowStack;
+
+    auto actual = Collect(shadowStack);
+
+    EXPECT_THAT(actual, testing::IsEmpty());
+}
+
+TEST(ShadowStackTest, OneLocal) {
+    mm::ShadowStack shadowStack;
+    StackEntry<0, 1> frame1(shadowStack);
+
+    auto actual = Collect(shadowStack);
+
+    EXPECT_THAT(actual, testing::ElementsAre(frame1[0]));
+}
+
+TEST(ShadowStackTest, ThreeLocals) {
+    mm::ShadowStack shadowStack;
+    StackEntry<0, 3> frame1(shadowStack);
+
+    auto actual = Collect(shadowStack);
+
+    EXPECT_THAT(actual, testing::ElementsAre(frame1[0], frame1[1], frame1[2]));
+}
+
+TEST(ShadowStackTest, OneParameter) {
+    mm::ShadowStack shadowStack;
+    StackEntry<1, 0> frame1(shadowStack);
+
+    auto actual = Collect(shadowStack);
+
+    EXPECT_THAT(actual, testing::IsEmpty());
+}
+
+TEST(ShadowStackTest, ThreeLocalsAndOneParameter) {
+    mm::ShadowStack shadowStack;
+    StackEntry<1, 3> frame1(shadowStack);
+
+    auto actual = Collect(shadowStack);
+
+    EXPECT_THAT(actual, testing::ElementsAre(frame1[0], frame1[1], frame1[2]));
+}
+
+TEST(ShadowStackTest, TwoStackFrames) {
+    mm::ShadowStack shadowStack;
+    StackEntry<1, 3> frame1(shadowStack);
+    StackEntry<1, 3> frame2(shadowStack);
+
+    auto actual = Collect(shadowStack);
+
+    EXPECT_THAT(actual, testing::ElementsAre(frame2[0], frame2[1], frame2[2], frame1[0], frame1[1], frame1[2]));
+}
+
+TEST(ShadowStackTest, ManyStackFrames) {
+    mm::ShadowStack shadowStack;
+    StackEntry<0, 3> frame1(shadowStack);
+    StackEntry<1, 0> frame2(shadowStack);
+    StackEntry<3, 1> frame3(shadowStack);
+    StackEntry<3, 3> frame4(shadowStack);
+
+    auto actual = Collect(shadowStack);
+
+    EXPECT_THAT(actual, testing::ElementsAre(frame4[0], frame4[1], frame4[2], frame3[0], frame1[0], frame1[1], frame1[2]));
+}

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

@@ -70,14 +70,6 @@ RUNTIME_NOTHROW OBJ_GETTER(ReadHeapRefLocked, ObjHeader** location, int32_t* spi
     TODO();
 }
 
-RUNTIME_NOTHROW void EnterFrame(ObjHeader** start, int parameters, int count) {
-    TODO();
-}
-
-RUNTIME_NOTHROW void LeaveFrame(ObjHeader** start, int parameters, int count) {
-    TODO();
-}
-
 void MutationCheck(ObjHeader* obj) {
     TODO();
 }

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

@@ -11,6 +11,7 @@
 
 #include "ObjectFactory.hpp"
 #include "GlobalsRegistry.hpp"
+#include "ShadowStack.hpp"
 #include "StableRefRegistry.hpp"
 #include "ThreadLocalStorage.hpp"
 #include "Utils.hpp"
@@ -46,6 +47,8 @@ public:
 
     ObjectFactory::ThreadQueue& objectFactoryThreadQueue() noexcept { return objectFactoryThreadQueue_; }
 
+    ShadowStack& shadowStack() noexcept { return shadowStack_; }
+
 private:
     const pthread_t threadId_;
     GlobalsRegistry::ThreadQueue globalsThreadQueue_;
@@ -53,6 +56,7 @@ private:
     StableRefRegistry::ThreadQueue stableRefThreadQueue_;
     std::atomic<ThreadState> state_;
     ObjectFactory::ThreadQueue objectFactoryThreadQueue_;
+    ShadowStack shadowStack_;
 };
 
 } // namespace mm