[K/N][New MM] Add thread state asserts to reference updates

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

@@ -12,6 +12,8 @@
 using namespace kotlin;
 
 OBJ_GETTER(mm::InitThreadLocalSingleton, ThreadData* threadData, ObjHeader** location, const TypeInfo* typeInfo, void (*ctor)(ObjHeader*)) {
+    // TODO: Is it possible that threadData != CurrentThreadData?
+    AssertThreadState(threadData, ThreadState::kRunnable);
     if (auto* value = *location) {
         // Initialized by someone else.
         RETURN_OBJ(value);
@@ -33,6 +35,8 @@ OBJ_GETTER(mm::InitThreadLocalSingleton, ThreadData* threadData, ObjHeader** loc
 }
 
 OBJ_GETTER(mm::InitSingleton, ThreadData* threadData, ObjHeader** location, const TypeInfo* typeInfo, void (*ctor)(ObjHeader*)) {
+    // TODO: Is it possible that threadData != CurrentThreadData?
+    AssertThreadState(threadData, ThreadState::kRunnable);
     auto& initializingSingletons = threadData->initializingSingletons();
 
     // Search from the top of the stack.

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

@@ -135,6 +135,7 @@ extern "C" ALWAYS_INLINE OBJ_GETTER(InitSingleton, ObjHeader** location, const T
 
 extern "C" RUNTIME_NOTHROW void InitAndRegisterGlobal(ObjHeader** location, const ObjHeader* initialValue) {
     auto* threadData = mm::ThreadRegistry::Instance().CurrentThreadData();
+    AssertThreadState(threadData, ThreadState::kRunnable);
     mm::GlobalsRegistry::Instance().RegisterStorageForGlobal(threadData, location);
     // Null `initialValue` means that the appropriate value was already set by static initialization.
     if (initialValue != nullptr) {
@@ -211,11 +212,13 @@ extern "C" OBJ_GETTER(ReadHeapRefNoLock, ObjHeader* object, int32_t index) {
 
 extern "C" RUNTIME_NOTHROW void EnterFrame(ObjHeader** start, int parameters, int count) {
     auto* threadData = mm::ThreadRegistry::Instance().CurrentThreadData();
+    AssertThreadState(threadData, ThreadState::kRunnable);
     threadData->shadowStack().EnterFrame(start, parameters, count);
 }
 
 extern "C" RUNTIME_NOTHROW void LeaveFrame(ObjHeader** start, int parameters, int count) {
     auto* threadData = mm::ThreadRegistry::Instance().CurrentThreadData();
+    AssertThreadState(threadData, ThreadState::kRunnable);
     threadData->shadowStack().LeaveFrame(start, parameters, count);
 }
 

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

@@ -13,10 +13,12 @@ using namespace kotlin;
 // TODO: Memory barriers.
 
 ALWAYS_INLINE void mm::SetStackRef(ObjHeader** location, ObjHeader* value) noexcept {
+    AssertThreadState(ThreadState::kRunnable);
     *location = value;
 }
 
 ALWAYS_INLINE void mm::SetHeapRef(ObjHeader** location, ObjHeader* value) noexcept {
+    AssertThreadState(ThreadState::kRunnable);
     *location = value;
 }
 
@@ -27,16 +29,19 @@ ALWAYS_INLINE void mm::SetHeapRef(ObjHeader** location, ObjHeader* value) noexce
 #pragma clang diagnostic ignored "-Watomic-alignment"
 
 ALWAYS_INLINE void mm::SetHeapRefAtomic(ObjHeader** location, ObjHeader* value) noexcept {
+    AssertThreadState(ThreadState::kRunnable);
     __atomic_store_n(location, value, __ATOMIC_RELEASE);
 }
 
 ALWAYS_INLINE OBJ_GETTER(mm::ReadHeapRefAtomic, ObjHeader** location) noexcept {
+    AssertThreadState(ThreadState::kRunnable);
     // TODO: Make this work with GCs that can stop thread at any point.
     auto result = __atomic_load_n(location, __ATOMIC_ACQUIRE);
     RETURN_OBJ(result);
 }
 
 ALWAYS_INLINE OBJ_GETTER(mm::CompareAndSwapHeapRef, ObjHeader** location, ObjHeader* expected, ObjHeader* value) noexcept {
+    AssertThreadState(ThreadState::kRunnable);
     // TODO: Make this work with GCs that can stop thread at any point.
     ObjHeader* actual = expected;
     // TODO: Do we need this strong memory model? Do we need to use strong CAS?
@@ -50,12 +55,14 @@ ALWAYS_INLINE OBJ_GETTER(mm::CompareAndSwapHeapRef, ObjHeader** location, ObjHea
 #pragma clang diagnostic pop
 
 OBJ_GETTER(mm::AllocateObject, ThreadData* threadData, const TypeInfo* typeInfo) noexcept {
+    AssertThreadState(ThreadState::kRunnable);
     // TODO: Make this work with GCs that can stop thread at any point.
     auto* object = threadData->objectFactoryThreadQueue().CreateObject(typeInfo);
     RETURN_OBJ(object);
 }
 
 OBJ_GETTER(mm::AllocateArray, ThreadData* threadData, const TypeInfo* typeInfo, uint32_t elements) noexcept {
+    AssertThreadState(ThreadState::kRunnable);
     // TODO: Make this work with GCs that can stop thread at any point.
     auto* array = threadData->objectFactoryThreadQueue().CreateArray(typeInfo, static_cast<uint32_t>(elements));
     // `ArrayHeader` and `ObjHeader` are expected to be compatible.

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

@@ -24,7 +24,7 @@ const char* stateToString(ThreadState state) noexcept {
 
 } // namespace
 
-// Switches the state of the current thread to `newState` and returns the previous state.
+// Switches the state of the given thread to `newState` and returns the previous state.
 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.

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

@@ -14,6 +14,7 @@ namespace kotlin {
 // 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;
 
+// Asserts that the given thread is in the given state.
 ALWAYS_INLINE void AssertThreadState(mm::ThreadData* threadData, ThreadState expected) noexcept;
 
 } // namespace kotlin