[K/N][New MM] Move ThreadStateGuard to the main Memory.h

Integrating the stdlib with the thread states machinery requires
accessing to some parts of the new MM. This patch provides this
access by moving these parts to the main Memory.h header.

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

@@ -3719,3 +3719,16 @@ ALWAYS_INLINE ObjHeader* ExceptionObjHolder::GetExceptionObject() noexcept {
     return static_cast<ExceptionObjHolderImpl*>(this)->obj();
 }
 #endif
+
+ALWAYS_INLINE kotlin::ThreadState kotlin::SwitchThreadState(MemoryState* thread, ThreadState newState) noexcept {
+    // no-op, used by the new MM only.
+    return ThreadState::kRunnable;
+}
+
+ALWAYS_INLINE void kotlin::AssertThreadState(MemoryState* thread, ThreadState expected) noexcept {
+    // no-op, used by the new MM only.
+}
+
+MemoryState* kotlin::mm::GetMemoryState() {
+    return ::memoryState;
+}

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

@@ -1,14 +0,0 @@
-/*
- * 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 "MemoryPrivate.hpp"
-#include "TestSupport.hpp"
-
-
-// No-op, required for the new MM only.
-MemoryState* kotlin::InitMemoryForTests() { return nullptr; }
-
-// No-op, required for the new MM only.
-void kotlin::DeinitMemoryForTests(MemoryState* state) {}

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

@@ -17,6 +17,8 @@
 #ifndef RUNTIME_MEMORY_H
 #define RUNTIME_MEMORY_H
 
+#include <utility>
+
 #include "KAssert.h"
 #include "Common.h"
 #include "TypeInfo.h"
@@ -372,4 +374,66 @@ public:
     virtual ~ExceptionObjHolder() = default;
 };
 
+namespace kotlin {
+namespace mm {
+
+// Returns the MemoryState for the current thread. The runtime must be initialized.
+// Try not to use it very often, as (1) thread local access can be slow on some platforms,
+// (2) TLS gets deallocated before our thread destruction hooks run.
+MemoryState* GetMemoryState();
+
+} // namespace mm
+
+enum class ThreadState {
+    kRunnable, kNative
+};
+
+// Switches the state of the given thread to `newState` and returns the previous thread state.
+ALWAYS_INLINE ThreadState SwitchThreadState(MemoryState* thread, ThreadState newState) noexcept;
+
+// Asserts that the given thread is in the given state.
+ALWAYS_INLINE void AssertThreadState(MemoryState* thread, ThreadState expected) noexcept;
+
+// Asserts that the current thread is in the the given state.
+ALWAYS_INLINE inline void AssertThreadState(ThreadState expected) noexcept {
+    AssertThreadState(mm::GetMemoryState(), expected);
+}
+
+// Scopely sets the given thread state for the given thread.
+class ThreadStateGuard final : private Pinned {
+public:
+    // Set the state for the given thread.
+    ThreadStateGuard(MemoryState* thread, ThreadState state) noexcept : thread_(thread) {
+        oldState_ = SwitchThreadState(thread_, state);
+    }
+
+    // Sets the state for the current thread.
+    explicit ThreadStateGuard(ThreadState state) noexcept
+        : ThreadStateGuard(mm::GetMemoryState(), state) {};
+
+    ~ThreadStateGuard() noexcept {
+        SwitchThreadState(thread_, oldState_);
+    }
+private:
+    MemoryState* thread_;
+    ThreadState oldState_;
+};
+
+// Calls the given function in the `Runnable` thread state.
+template <typename R, typename... Args>
+ALWAYS_INLINE inline R CallKotlin(R(*kotlinFunction)(Args...), Args... args) {
+    ThreadStateGuard guard(ThreadState::kRunnable);
+    return kotlinFunction(std::forward<Args>(args)...);
+}
+
+// Calls the given function in the `Runnable` thread state. The function must be marked as RUNTIME_NORETURN.
+// If the function returns, behaviour is undefined.
+template <typename... Args>
+ALWAYS_INLINE RUNTIME_NORETURN inline void CallKotlinNoReturn(void(*noreturnKotlinFunction)(Args...), Args... args) {
+    CallKotlin(noreturnKotlinFunction, std::forward<Args>(args)...);
+    RuntimeFail("The function must not return");
+}
+
+} // namespace kotlin
+
 #endif // RUNTIME_MEMORY_H

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

@@ -52,7 +52,6 @@ bool Kotlin_cleanersLeakCheckerEnabled();
 bool Kotlin_forceCheckedShutdown();
 
 #ifdef __cplusplus
-}
+} // extern "C"
 #endif
-
 #endif // RUNTIME_RUNTIME_H

kotlin-native/runtime/src/main/cpp/TestSupport.hpp

@@ -7,6 +7,7 @@
 #include <thread>
 
 #include "Memory.h"
+#include "Runtime.h"
 
 namespace kotlin {
 
@@ -20,24 +21,13 @@ constexpr int kDefaultThreadCount = 10;
 constexpr int kDefaultThreadCount = 100;
 #endif
 
-// Performs minimal initialization of the current thread memory subsystem.
-// This initilization is enough for running tests for the C++ part for the stdlib.
-//  - For the new MM: registeres the current thread in the thread registry.
-//  - For the legacy MM: does nothing, returns nullptr.
-MemoryState* InitMemoryForTests();
-
-// Deinitiliazes memory subsystem used for C++ stdlib tests.
-//  - For the new MM: unregisteres the current thread, nullify current thread data.
-//  - For the legacy MM: does nothing.
-void DeinitMemoryForTests(MemoryState* state);
-
-// Scopely initializes the memory subsystem using the functions above.
-class ScopedMemoryInit : private kotlin::Pinned {
+// Scopely initializes the memory subsystem of the current thread for tests.
+class ScopedRuntimeInit : private kotlin::Pinned {
 public:
-    ScopedMemoryInit() : memoryState_(InitMemoryForTests()) {}
-    ~ScopedMemoryInit() {
+    ScopedRuntimeInit() : memoryState_(InitMemory(false)) {}
+    ~ScopedRuntimeInit() {
         ClearMemoryForTests(memoryState());
-        DeinitMemoryForTests(memoryState_);
+        DeinitMemory(memoryState_, false);
     }
 
     MemoryState* memoryState() { return memoryState_; }
@@ -45,15 +35,15 @@ private:
     MemoryState* memoryState_;
 };
 
-// Runs the given function in a separate thread with minimally initialized memory subsystem.
+// Runs the given function in a separate thread with minimally initialized runtime.
 inline void RunInNewThread(std::function<void(MemoryState*)> f) {
     std::thread([&f]() {
-        ScopedMemoryInit init;
+        ScopedRuntimeInit init;
         f(init.memoryState());
     }).join();
 }
 
-// Runs the given function in a separate thread with minimally initialized memory subsystem.
+// Runs the given function in a separate thread with minimally initialized runtime.
 inline void RunInNewThread(std::function<void()> f) {
     RunInNewThread([&f](MemoryState* unused) {
         f();

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

@@ -427,3 +427,15 @@ extern "C" RUNTIME_NOTHROW void Kotlin_mm_safePointExceptionUnwind() {
     auto* threadData = mm::ThreadRegistry::Instance().CurrentThreadData();
     threadData->gc().SafePointExceptionUnwind();
 }
+
+extern "C" ALWAYS_INLINE RUNTIME_NOTHROW void Kotlin_mm_switchThreadStateNative() {
+    SwitchThreadState(mm::ThreadRegistry::Instance().CurrentThreadData(), ThreadState::kNative);
+}
+
+extern "C" ALWAYS_INLINE RUNTIME_NOTHROW void Kotlin_mm_switchThreadStateRunnable() {
+    SwitchThreadState(mm::ThreadRegistry::Instance().CurrentThreadData(), ThreadState::kRunnable);
+}
+
+MemoryState* kotlin::mm::GetMemoryState() {
+    return ToMemoryState(ThreadRegistry::Instance().CurrentThreadDataNode());
+}

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

@@ -1,20 +0,0 @@
-/*
- * 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 "MemoryPrivate.hpp"
-#include "TestSupport.hpp"
-#include "ThreadRegistry.hpp"
-
-MemoryState* kotlin::InitMemoryForTests() {
-    auto threadDataNode = mm::ThreadRegistry::Instance().RegisterCurrentThread();
-    return mm::ToMemoryState(threadDataNode);
-}
-
-void kotlin::DeinitMemoryForTests(MemoryState* state) {
-    auto threadDataNode = mm::FromMemoryState(state);
-    mm::ThreadRegistry::Instance().Unregister(threadDataNode);
-    // Nullify current thread data. The thread is still alive, so this is safe.
-    mm::ThreadRegistry::TestSupport::SetCurrentThreadData(nullptr);
-}

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

@@ -17,9 +17,9 @@ mm::ThreadRegistry& mm::ThreadRegistry::Instance() noexcept {
 
 mm::ThreadRegistry::Node* mm::ThreadRegistry::RegisterCurrentThread() noexcept {
     auto* threadDataNode = list_.Emplace(pthread_self());
-    ThreadData*& currentData = currentThreadData_;
-    RuntimeAssert(currentData == nullptr, "This thread already had some data assigned to it.");
-    currentData = threadDataNode->Get();
+    Node*& currentDataNode = currentThreadDataNode_;
+    RuntimeAssert(currentDataNode == nullptr, "This thread already had some data assigned to it.");
+    currentDataNode = threadDataNode;
     return threadDataNode;
 }
 
@@ -32,8 +32,12 @@ mm::ThreadRegistry::Iterable mm::ThreadRegistry::Iter() noexcept {
     return list_.Iter();
 }
 
+ALWAYS_INLINE mm::ThreadData* mm::ThreadRegistry::CurrentThreadData() const noexcept {
+    return CurrentThreadDataNode()->Get();
+}
+
 mm::ThreadRegistry::ThreadRegistry() = default;
 mm::ThreadRegistry::~ThreadRegistry() = default;
 
 // static
-thread_local mm::ThreadData* mm::ThreadRegistry::currentThreadData_ = nullptr;
+thread_local mm::ThreadRegistry::Node* mm::ThreadRegistry::currentThreadDataNode_ = nullptr;

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

@@ -31,26 +31,19 @@ public:
     // Locks `ThreadRegistry` for safe iteration.
     Iterable Iter() noexcept;
 
-    // Try not to use it very often, as (1) thread local access can be slow on some platforms,
+    // Try not to use these methods very often, as (1) thread local access can be slow on some platforms,
     // (2) TLS gets deallocated before our thread destruction hooks run.
     // Using this after `Unregister` for the thread has been called is undefined behaviour.
-    ThreadData* CurrentThreadData() const noexcept { return currentThreadData_; }
-
-    class TestSupport {
-    public:
-        static void SetCurrentThreadData(ThreadData* newThreadData) {
-            ThreadRegistry::currentThreadData_ = newThreadData;
-        }
-    };
+    ALWAYS_INLINE ThreadData* CurrentThreadData() const noexcept;
+    Node* CurrentThreadDataNode() const noexcept { return currentThreadDataNode_; }
 
 private:
     friend class GlobalData;
-    friend class ThreadRegistry::TestSupport;
 
     ThreadRegistry();
     ~ThreadRegistry();
 
-    static thread_local ThreadData* currentThreadData_;
+    static thread_local Node* currentThreadDataNode_;
 
     SingleLockList<ThreadData> list_;
 };

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

@@ -3,22 +3,21 @@
  * that can be found in the LICENSE file.
  */
 
+#include "MemoryPrivate.hpp"
 #include "ThreadData.hpp"
 #include "ThreadState.hpp"
 
-using namespace kotlin;
-
 namespace {
 
-ALWAYS_INLINE bool isStateSwitchAllowed(mm::ThreadState oldState, mm::ThreadState newState) noexcept {
+ALWAYS_INLINE bool isStateSwitchAllowed(ThreadState oldState, ThreadState newState) noexcept {
     return oldState != newState;
 }
 
-const char* stateToString(mm::ThreadState state) noexcept {
+const char* stateToString(ThreadState state) noexcept {
     switch (state) {
-        case mm::ThreadState::kRunnable:
+        case ThreadState::kRunnable:
             return "RUNNABLE";
-        case mm::ThreadState::kNative:
+        case ThreadState::kNative:
             return "NATIVE";
     }
 }
@@ -26,7 +25,7 @@ const char* stateToString(mm::ThreadState state) noexcept {
 } // namespace
 
 // Switches the state of the current thread to `newState` and returns the previous state.
-ALWAYS_INLINE mm::ThreadState mm::SwitchThreadState(ThreadData* threadData, ThreadState newState) noexcept {
+ALWAYS_INLINE ThreadState kotlin::SwitchThreadState(mm::ThreadData* threadData, ThreadState newState) noexcept {
     auto oldState = threadData->setState(newState);
     // TODO(perf): Mesaure the impact of this assert in debug and opt modes.
     RuntimeAssert(isStateSwitchAllowed(oldState, newState),
@@ -35,26 +34,17 @@ ALWAYS_INLINE mm::ThreadState mm::SwitchThreadState(ThreadData* threadData, Thre
     return oldState;
 }
 
-ALWAYS_INLINE void mm::AssertThreadState(ThreadData* threadData, ThreadState expected) noexcept {
+ALWAYS_INLINE ThreadState kotlin::SwitchThreadState(MemoryState* thread, ThreadState newState) noexcept {
+    return SwitchThreadState(thread->GetThreadData(), newState);
+}
+
+ALWAYS_INLINE void kotlin::AssertThreadState(mm::ThreadData* threadData, ThreadState expected) noexcept {
     auto actual = threadData->state();
     RuntimeAssert(actual == expected,
                   "Unexpected thread state. Expected: %s. Actual: %s.",
                   stateToString(expected), stateToString(actual));
 }
 
-mm::ThreadStateGuard::ThreadStateGuard(ThreadData* threadData, ThreadState state) noexcept : threadData_(threadData) {
-    oldState_ = SwitchThreadState(threadData, state);
+ALWAYS_INLINE void kotlin::AssertThreadState(MemoryState* thread, ThreadState expected) noexcept {
+    AssertThreadState(thread->GetThreadData(), expected);
 }
-
-mm::ThreadStateGuard::~ThreadStateGuard() noexcept {
-    SwitchThreadState(threadData_, oldState_);
-}
-
-extern "C" ALWAYS_INLINE RUNTIME_NOTHROW void Kotlin_mm_switchThreadStateNative() {
-    mm::SwitchThreadState(mm::ThreadRegistry::Instance().CurrentThreadData(), mm::ThreadState::kNative);
-}
-
-extern "C" ALWAYS_INLINE RUNTIME_NOTHROW void Kotlin_mm_switchThreadStateRunnable() {
-    mm::SwitchThreadState(mm::ThreadRegistry::Instance().CurrentThreadData(), mm::ThreadState::kRunnable);
-}
-

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

@@ -10,27 +10,12 @@
 #include <Utils.hpp>
 
 namespace kotlin {
-namespace mm {
 
-enum class ThreadState {
-    kRunnable, kNative
-};
+// Switches the state of the given thread to `newState` and returns the previous thread state.
+ALWAYS_INLINE ThreadState SwitchThreadState(mm::ThreadData* threadData, ThreadState newState) noexcept;
 
-// Switches the state of the current thread to `newState` and returns the previous thread state.
-ALWAYS_INLINE ThreadState SwitchThreadState(ThreadData* threadData, ThreadState newState) noexcept;
+ALWAYS_INLINE void AssertThreadState(mm::ThreadData* threadData, ThreadState expected) noexcept;
 
-ALWAYS_INLINE void AssertThreadState(ThreadData* threadData, ThreadState expected) noexcept;
-
-class ThreadStateGuard final : private Pinned {
-public:
-    explicit ThreadStateGuard(ThreadData* threadData, ThreadState state) noexcept;
-    ~ThreadStateGuard() noexcept;
-private:
-    ThreadData* threadData_;
-    ThreadState oldState_;
-};
-
-} // namespace mm
 } // namespace kotlin
 
 #endif // RUNTIME_MM_THREAD_STATE_H

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

@@ -6,58 +6,165 @@
 #include <thread>
 
 #include "gtest/gtest.h"
+#include "gmock/gmock.h"
 
+#include "MemoryPrivate.hpp"
 #include "TestSupport.hpp"
 #include "ThreadData.hpp"
 #include "ThreadState.hpp"
 
 using namespace kotlin;
 
-TEST(ThreadStateTest, StateSwitch) {
+namespace {
+
+class ThreadStateTest : public testing::Test {
+public:
+    ThreadStateTest() {
+        globalKotlinFunctionMock = &kotlinFunctionMock_;
+    }
+
+    ~ThreadStateTest() {
+        globalKotlinFunctionMock = nullptr;
+    }
+
+    testing::MockFunction<int32_t(int32_t)>& kotlinFunctionMock() { return kotlinFunctionMock_; }
+
+    static int32_t kotlinFunction(int32_t arg) {
+        return globalKotlinFunctionMock->Call(arg);
+    }
+
+    static RUNTIME_NORETURN void noReturnKotlinFunciton(int32_t arg) {
+        globalKotlinFunctionMock->Call(arg);
+        throw std::exception();
+    }
+private:
+    testing::MockFunction<int32_t(int32_t)> kotlinFunctionMock_;
+    static testing::MockFunction<int32_t(int32_t)>* globalKotlinFunctionMock;
+};
+
+//static
+testing::MockFunction<int32_t(int32_t)>* ThreadStateTest::globalKotlinFunctionMock = nullptr;
+
+} // namespace
+
+TEST_F(ThreadStateTest, StateSwitchWithThreadData) {
     RunInNewThread([](mm::ThreadData& threadData) {
         auto initialState = threadData.state();
-        EXPECT_EQ(mm::ThreadState::kRunnable, initialState);
+        EXPECT_EQ(ThreadState::kRunnable, initialState);
 
-        mm::ThreadState oldState = mm::SwitchThreadState(&threadData, mm::ThreadState::kNative);
+        ThreadState oldState = SwitchThreadState(&threadData, ThreadState::kNative);
         EXPECT_EQ(initialState, oldState);
-        EXPECT_EQ(mm::ThreadState::kNative, threadData.state());
-
-        // Check functions exported for the compiler too.
-        Kotlin_mm_switchThreadStateRunnable();
-        EXPECT_EQ(mm::ThreadState::kRunnable, threadData.state());
-
-        Kotlin_mm_switchThreadStateNative();
-        EXPECT_EQ(mm::ThreadState::kNative, threadData.state());
+        EXPECT_EQ(ThreadState::kNative, threadData.state());
     });
 }
 
-TEST(ThreadStateTest, StateGuard) {
+TEST_F(ThreadStateTest, StateSwitchWithMemoryState) {
+    RunInNewThread([](MemoryState* memoryState) {
+        auto threadData = memoryState->GetThreadData();
+        auto initialState = threadData->state();
+        EXPECT_EQ(ThreadState::kRunnable, initialState);
+
+        ThreadState oldState = SwitchThreadState(memoryState, ThreadState::kNative);
+        EXPECT_EQ(initialState, oldState);
+        EXPECT_EQ(ThreadState::kNative, threadData->state());
+    });
+}
+
+TEST_F(ThreadStateTest, StateSwitchExported) {
     RunInNewThread([](mm::ThreadData& threadData) {
+        // Check functions exported for the compiler.
+        EXPECT_EQ(ThreadState::kRunnable, threadData.state());
+
+        Kotlin_mm_switchThreadStateNative();
+        EXPECT_EQ(ThreadState::kNative, threadData.state());
+
+        Kotlin_mm_switchThreadStateRunnable();
+        EXPECT_EQ(ThreadState::kRunnable, threadData.state());
+    });
+}
+
+TEST_F(ThreadStateTest, StateGuard) {
+    RunInNewThread([](MemoryState* memoryState) {
+        mm::ThreadData& threadData = *memoryState->GetThreadData();
         auto initialState = threadData.state();
-        EXPECT_EQ(mm::ThreadState::kRunnable, initialState);
+        EXPECT_EQ(ThreadState::kRunnable, initialState);
         {
-            mm::ThreadStateGuard guard(&threadData, mm::ThreadState::kNative);
-            EXPECT_EQ(mm::ThreadState::kNative, threadData.state());
+            ThreadStateGuard guard(memoryState, ThreadState::kNative);
+            EXPECT_EQ(ThreadState::kNative, threadData.state());
         }
         EXPECT_EQ(initialState, threadData.state());
     });
 }
 
+TEST_F(ThreadStateTest, StateGuardForCurrentThread) {
+    RunInNewThread([]() {
+        auto* memoryState = mm::GetMemoryState();
+        auto initialState = memoryState->GetThreadData()->state();
+        EXPECT_EQ(ThreadState::kRunnable, initialState);
+        {
+            ThreadStateGuard guard(memoryState, ThreadState::kNative);
+            EXPECT_EQ(ThreadState::kNative, memoryState->GetThreadData()->state());
+        }
+        EXPECT_EQ(initialState, memoryState->GetThreadData()->state());
+    });
+}
+
+TEST_F(ThreadStateTest, CallKotlin) {
+    RunInNewThread([this](mm::ThreadData& threadData) {
+        SwitchThreadState(&threadData, ThreadState::kNative);
+        ASSERT_THAT(threadData.state(), ThreadState::kNative);
+
+        EXPECT_CALL(kotlinFunctionMock(), Call(42))
+            .WillOnce([&threadData](int32_t arg) {
+                EXPECT_THAT(threadData.state(), ThreadState::kRunnable);
+                return 24;
+            });
+        int32_t result = CallKotlin(kotlinFunction, 42);
+        EXPECT_THAT(threadData.state(), ThreadState::kNative);
+        EXPECT_THAT(result, 24);
+    });
+}
+
+TEST_F(ThreadStateTest, CallKotlinNoReturn) {
+    RunInNewThread([this](mm::ThreadData& threadData) {
+        SwitchThreadState(&threadData, ThreadState::kNative);
+        ASSERT_THAT(threadData.state(), ThreadState::kNative);
+
+        EXPECT_CALL(kotlinFunctionMock(), Call(42))
+            .WillOnce([&threadData](int32_t arg){
+                EXPECT_THAT(threadData.state(), ThreadState::kRunnable);
+                return 24;
+            });
+
+        EXPECT_THROW(CallKotlinNoReturn(noReturnKotlinFunciton, 42), std::exception);
+        EXPECT_THAT(threadData.state(), ThreadState::kNative);
+    });
+}
+
 TEST(ThreadStateDeathTest, StateAsserts) {
-    RunInNewThread([](mm::ThreadData& threadData) {
-        EXPECT_DEATH(mm::AssertThreadState(&threadData, mm::ThreadState::kNative),
+    RunInNewThread([](MemoryState* memoryState) {
+        mm::ThreadData* threadData = memoryState->GetThreadData();
+        EXPECT_DEATH(AssertThreadState(memoryState, ThreadState::kNative),
+                     "runtime assert: Unexpected thread state. Expected: NATIVE. Actual: RUNNABLE");
+        EXPECT_DEATH(AssertThreadState(threadData, ThreadState::kNative),
+                     "runtime assert: Unexpected thread state. Expected: NATIVE. Actual: RUNNABLE");
+        EXPECT_DEATH(AssertThreadState(ThreadState::kNative),
                      "runtime assert: Unexpected thread state. Expected: NATIVE. Actual: RUNNABLE");
     });
 }
 
 TEST(ThreadStateDeathTest, IncorrectStateSwitch) {
-    RunInNewThread([](mm::ThreadData& threadData) {
-        EXPECT_DEATH(mm::SwitchThreadState(&threadData, kotlin::mm::ThreadState::kRunnable),
+    RunInNewThread([](MemoryState* memoryState) {
+        auto* threadData = memoryState->GetThreadData();
+        EXPECT_DEATH(SwitchThreadState(memoryState, ThreadState::kRunnable),
                      "runtime assert: Illegal thread state switch. Old state: RUNNABLE. New state: RUNNABLE");
+        EXPECT_DEATH(SwitchThreadState(threadData, ThreadState::kRunnable),
+                     "runtime assert: Illegal thread state switch. Old state: RUNNABLE. New state: RUNNABLE");
+
         EXPECT_DEATH(Kotlin_mm_switchThreadStateRunnable(),
                      "runtime assert: Illegal thread state switch. Old state: RUNNABLE. New state: RUNNABLE");
 
-        mm::SwitchThreadState(&threadData, kotlin::mm::ThreadState::kNative);
+        SwitchThreadState(threadData, kotlin::ThreadState::kNative);
         EXPECT_DEATH(Kotlin_mm_switchThreadStateNative(),
                      "runtime assert: Illegal thread state switch. Old state: NATIVE. New state: NATIVE");
     });