Extract FreezeHooks and make them testable (#4654)

kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/KonanFqNames.kt

@@ -31,4 +31,5 @@ object KonanFqNames {
     val typedIntrinsic = FqName("kotlin.native.internal.TypedIntrinsic")
     val objCMethod = FqName("kotlinx.cinterop.ObjCMethod")
     val hasFinalizer = FqName("kotlin.native.internal.HasFinalizer")
+    val hasFreezeHook = FqName("kotlin.native.internal.HasFreezeHook")
 }

kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/llvm/RTTIGenerator.kt

@@ -75,6 +75,10 @@ internal class RTTIGenerator(override val context: Context) : ContextUtils {
             result = result or TF_HAS_FINALIZER
         }
 
+        if (irClass.hasAnnotation(KonanFqNames.hasFreezeHook)) {
+            result = result or TF_HAS_FREEZE_HOOK
+        }
+
         return result
     }
 
@@ -638,4 +642,5 @@ private const val TF_OBJC_DYNAMIC = 8
 private const val TF_LEAK_DETECTOR_CANDIDATE = 16
 private const val TF_SUSPEND_FUNCTION = 32
 private const val TF_HAS_FINALIZER = 64
+private const val TF_HAS_FREEZE_HOOK = 128
 

kotlin-native/build-tools/src/main/kotlin/org/jetbrains/kotlin/bitcode/CompileToBitcode.kt

@@ -33,7 +33,9 @@ open class CompileToBitcode @Inject constructor(
     val linkerArgs = mutableListOf<String>()
     var excludeFiles: List<String> = listOf(
             "**/*Test.cpp",
+            "**/*TestSupport.cpp",
             "**/*Test.mm",
+            "**/*TestSupport.mm",
     )
     var includeFiles: List<String> = listOf(
             "**/*.cpp",

kotlin-native/build-tools/src/main/kotlin/org/jetbrains/kotlin/testing/native/NativeTest.kt

@@ -206,7 +206,7 @@ private fun createTestTask(
                     ).apply {
                 this.sanitizer = sanitizer
                 excludeFiles = emptyList()
-                includeFiles = listOf("**/*Test.cpp", "**/*Test.mm")
+                includeFiles = listOf("**/*Test.cpp", "**/*TestSupport.cpp", "**/*Test.mm", "**/*TestSupport.mm")
                 dependsOn(it)
                 dependsOn("downloadGoogleTest")
                 compilerArgs.addAll(it.compilerArgs)

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

@@ -39,6 +39,7 @@
 #include "CyclicCollector.h"
 #endif  // USE_CYCLIC_GC
 #include "Exceptions.h"
+#include "FreezeHooks.hpp"
 #include "KString.h"
 #include "Memory.h"
 #include "MemoryPrivate.hpp"
@@ -2861,13 +2862,6 @@ void freezeCyclic(ObjHeader* root,
   }
 }
 
-// These hooks are only allowed to modify `obj` subgraph.
-void runFreezeHooks(ObjHeader* obj) {
-  if (obj->type_info() == theWorkerBoundReferenceTypeInfo) {
-    WorkerBoundReferenceFreezeHook(obj);
-  }
-}
-
 void runFreezeHooksRecursive(ObjHeader* root) {
   KStdUnorderedSet<KRef> seen;
   KStdVector<KRef> toVisit;
@@ -2877,7 +2871,7 @@ void runFreezeHooksRecursive(ObjHeader* root) {
     KRef obj = toVisit.back();
     toVisit.pop_back();
 
-    runFreezeHooks(obj);
+    kotlin::RunFreezeHooks(obj);
 
     traverseReferredObjects(obj, [&seen, &toVisit](ObjHeader* field) {
       auto wasNotSeenYet = seen.insert(field).second;

kotlin-native/runtime/src/main/cpp/FreezeHooks.cpp

@@ -0,0 +1,42 @@
+/*
+ * 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 "FreezeHooks.hpp"
+
+#include "Memory.h"
+#include "Types.h"
+#include "WorkerBoundReference.h"
+
+using namespace kotlin;
+
+namespace {
+
+void (*g_hookOverrideForTesting)(ObjHeader*) = nullptr;
+
+NO_INLINE void RunFreezeHooksImpl(ObjHeader* object, const TypeInfo* type) noexcept {
+    if (g_hookOverrideForTesting != nullptr) {
+        g_hookOverrideForTesting(object);
+        return;
+    }
+    // TODO: Consider some global registration.
+    if (type == theWorkerBoundReferenceTypeInfo) {
+        WorkerBoundReferenceFreezeHook(object);
+    }
+}
+
+} // namespace
+
+void kotlin::RunFreezeHooks(ObjHeader* object) noexcept {
+    auto* type = object->type_info();
+    if ((type->flags_ & TF_HAS_FREEZE_HOOK) == 0) {
+        return;
+    }
+    // This is a cold path.
+    RunFreezeHooksImpl(object, type);
+}
+
+void kotlin::SetFreezeHookForTesting(void (*hook)(ObjHeader*)) noexcept {
+    g_hookOverrideForTesting = hook;
+}

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

@@ -0,0 +1,21 @@
+/*
+ * 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_FREEZE_HOOKS_H
+#define RUNTIME_MM_FREEZE_HOOKS_H
+
+struct ObjHeader;
+struct TypeInfo;
+
+namespace kotlin {
+
+// These hooks are only allowed to modify `object` subgraph.
+void RunFreezeHooks(ObjHeader* object) noexcept;
+
+void SetFreezeHookForTesting(void (*hook)(ObjHeader*)) noexcept;
+
+} // namespace kotlin
+
+#endif // RUNTIME_MM_FREEZE_HOOKS_H

kotlin-native/runtime/src/main/cpp/FreezeHooksTest.cpp

@@ -0,0 +1,48 @@
+/*
+ * 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 "FreezeHooks.hpp"
+
+#include "gtest/gtest.h"
+#include "gmock/gmock.h"
+
+#include "FreezeHooksTestSupport.hpp"
+#include "Memory.h"
+
+using namespace kotlin;
+
+using ::testing::_;
+
+namespace {
+
+class FreezeHooksTest : public testing::Test {
+public:
+    testing::MockFunction<void(ObjHeader*)>& freezeHook() { return freezeHooks_.freezeHook(); }
+
+private:
+    FreezeHooksTestSupport freezeHooks_;
+};
+
+} // namespace
+
+TEST_F(FreezeHooksTest, TypeWithFreezeHook) {
+    TypeInfo type;
+    type.typeInfo_ = &type;
+    type.flags_ |= TF_HAS_FREEZE_HOOK;
+    ObjHeader obj = {&type};
+    EXPECT_CALL(freezeHook(), Call(&obj));
+    RunFreezeHooks(&obj);
+    testing::Mock::VerifyAndClearExpectations(&freezeHook());
+}
+
+TEST_F(FreezeHooksTest, TypeWithoutFreezeHook) {
+    TypeInfo type;
+    type.typeInfo_ = &type;
+    type.flags_ &= ~TF_HAS_FREEZE_HOOK;
+    ObjHeader obj = {&type};
+    EXPECT_CALL(freezeHook(), Call(_)).Times(0);
+    RunFreezeHooks(&obj);
+    testing::Mock::VerifyAndClearExpectations(&freezeHook());
+}

kotlin-native/runtime/src/main/cpp/FreezeHooksTestSupport.cpp

@@ -0,0 +1,30 @@
+/*
+ * 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 "FreezeHooksTestSupport.hpp"
+
+#include "FreezeHooks.hpp"
+
+using namespace kotlin;
+
+namespace {
+
+testing::MockFunction<void(ObjHeader*)>* g_freezeHook = nullptr;
+
+void freezeHook(ObjHeader* object) {
+    g_freezeHook->Call(object);
+}
+
+} // namespace
+
+kotlin::FreezeHooksTestSupport::FreezeHooksTestSupport() {
+    g_freezeHook = &freezeHook_;
+    SetFreezeHookForTesting(&::freezeHook);
+}
+
+kotlin::FreezeHooksTestSupport::~FreezeHooksTestSupport() {
+    SetFreezeHookForTesting(nullptr);
+    g_freezeHook = nullptr;
+}

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

@@ -0,0 +1,29 @@
+/*
+ * 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_FREEZE_HOOKS_TEST_SUPPORT_H
+#define RUNTIME_MM_FREEZE_HOOKS_TEST_SUPPORT_H
+
+#include "gtest/gtest.h"
+#include "gmock/gmock.h"
+
+struct ObjHeader;
+
+namespace kotlin {
+
+class FreezeHooksTestSupport {
+public:
+    FreezeHooksTestSupport();
+    ~FreezeHooksTestSupport();
+
+    testing::MockFunction<void(ObjHeader*)>& freezeHook() { return freezeHook_; }
+
+private:
+    testing::StrictMock<testing::MockFunction<void(ObjHeader*)>> freezeHook_;
+};
+
+} // namespace kotlin
+
+#endif // RUNTIME_MM_FREEZE_HOOKS_TEST_SUPPORT_H

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

@@ -63,6 +63,7 @@ enum Konan_TypeFlags {
   TF_LEAK_DETECTOR_CANDIDATE = 1 << 4,
   TF_SUSPEND_FUNCTION = 1 << 5,
   TF_HAS_FINALIZER = 1 << 6,
+  TF_HAS_FREEZE_HOOK = 1 << 7,
 };
 
 // Flags per object instance.

kotlin-native/runtime/src/main/kotlin/kotlin/native/concurrent/WorkerBoundReference.kt

@@ -29,6 +29,7 @@ external private fun describeWorkerBoundReference(ref: NativePtr): String
 @NoReorderFields
 @ExportTypeInfo("theWorkerBoundReferenceTypeInfo")
 @HasFinalizer
+@HasFreezeHook
 public class WorkerBoundReference<out T : Any>(value: T) {
 
     private var ptr = NativePtr.NULL

kotlin-native/runtime/src/main/kotlin/kotlin/native/internal/Annotations.kt

@@ -145,3 +145,9 @@ internal annotation class HasFinalizer
 @RequiresOptIn(level = RequiresOptIn.Level.ERROR)
 @Retention(value = AnnotationRetention.BINARY)
 internal annotation class InternalForKotlinNative
+
+/**
+ * Marks a class that has a freeze hook.
+ */
+@Target(AnnotationTarget.CLASS)
+internal annotation class HasFreezeHook