Implement freezing in the new MM

Currently, it uses suboptimal storage for the freeze bit. This'll be revised later when we finalize stack objects handling in the new MM.

kotlin-native/backend.native/tests/build.gradle

@@ -985,22 +985,20 @@ task freeze0(type: KonanLocalTest) {
 }
 
 task freeze1(type: KonanLocalTest) {
-    enabled = (project.testTarget != 'wasm32') && // No exceptions on WASM.
-        !isExperimentalMM  // Experimental MM does not support freezing yet.
+    enabled = (project.testTarget != 'wasm32') // No exceptions on WASM.
     goldValue = "OK, cannot mutate frozen\n"
     source = "runtime/workers/freeze1.kt"
 }
 
 task freeze_stress(type: KonanLocalTest) {
-    enabled = (project.testTarget != 'wasm32') && // No exceptions on WASM.
-        !isExperimentalMM  // Experimental MM does not support freezing yet.
+    enabled = (project.testTarget != 'wasm32') // No exceptions on WASM.
     goldValue = "OK\n"
     source = "runtime/workers/freeze_stress.kt"
 }
 
 task freeze2(type: KonanLocalTest) {
     enabled = (project.testTarget != 'wasm32') && // No exceptions on WASM.
-        !isExperimentalMM  // Experimental MM does not support freezing yet.
+        !isExperimentalMM  // Experimental MM doesn't support multiple mutators yet.
     goldValue =
             "Worker 1: Hello world\n" + "Worker2: 42\n" +
             "Worker3: 239.0\n" + "Worker4: a\n" + "OK\n"
@@ -1009,27 +1007,25 @@ task freeze2(type: KonanLocalTest) {
 
 task freeze3(type: KonanLocalTest) {
     enabled = (project.testTarget != 'wasm32') && // No exceptions on WASM.
-        !isExperimentalMM  // Experimental MM does not support freezing yet.
+        !isExperimentalMM  // Experimental MM doesn't support multiple mutators yet.
     goldValue = "OK\n"
     source = "runtime/workers/freeze3.kt"
 }
 
 task freeze4(type: KonanLocalTest) {
-    enabled = (project.testTarget != 'wasm32') && // No exceptions on WASM.
-        !isExperimentalMM  // Experimental MM does not support freezing yet.
+    enabled = (project.testTarget != 'wasm32') // No exceptions on WASM.
     goldValue = "OK\n"
     source = "runtime/workers/freeze4.kt"
 }
 
 task freeze5(type: KonanLocalTest) {
-    enabled = !isExperimentalMM  // Experimental MM does not support freezing yet.
     goldValue = "OK\n"
     source = "runtime/workers/freeze5.kt"
 }
 
 task freeze6(type: KonanLocalTest) {
     enabled = (project.testTarget != 'wasm32') && // No exceptions on WASM.
-        !isExperimentalMM  // Experimental MM does not support freezing yet.
+        !isExperimentalMM // Experimental MM does not have a GC yet.
     goldValue = "OK\nOK\n"
     source = "runtime/workers/freeze6.kt"
 }
@@ -1055,8 +1051,7 @@ task lazy0(type: KonanLocalTest) {
 }
 
 task lazy1(type: KonanLocalTest) {
-    enabled = (project.testTarget != 'wasm32') && // Need exceptions.
-        !isExperimentalMM  // Experimental MM does not support freezing yet.
+    enabled = (project.testTarget != 'wasm32') // Need exceptions.
     goldValue = "OK\n"
     source = "runtime/workers/lazy1.kt"
 }
@@ -1178,7 +1173,6 @@ task enum_nested(type: KonanLocalTest) {
 }
 
 task enum_isFrozen(type: KonanLocalTest) {
-    enabled = !isExperimentalMM  // Experimental MM does not support freezing yet.
     goldValue = "true\n"
     source = "codegen/enum/isFrozen.kt"
 }
@@ -1929,8 +1923,7 @@ task typed_array0(type: KonanLocalTest) {
 }
 
 task typed_array1(type: KonanLocalTest) {
-    enabled  = (project.testTarget != 'wasm32') && // No exceptions on WASM.
-        !isExperimentalMM  // Experimental MM does not support freezing yet.
+    enabled  = (project.testTarget != 'wasm32') // No exceptions on WASM.
     goldValue = "OK\n"
     source = "runtime/collections/typed_array1.kt"
 }
@@ -2416,8 +2409,7 @@ standaloneTest("kt-37572") {
 }
 
 standaloneTest("custom_hook") {
-    enabled = (project.testTarget != 'wasm32') && // Uses exceptions.
-        !isExperimentalMM  // Experimental MM does not support freezing yet.
+    enabled = (project.testTarget != 'wasm32') // Uses exceptions.
     goldValue = "value 42: Error\n"
     expectedExitStatus = 1
     source = "runtime/exceptions/custom_hook.kt"

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

@@ -467,10 +467,6 @@ ContainerHeader* containerFor(const ObjHeader* obj) {
   return (reinterpret_cast<MetaObjHeader*>(clearPointerBits(obj->typeInfoOrMeta_, OBJECT_TAG_MASK)))->container_;
 }
 
-ALWAYS_INLINE bool isFrozen(const ObjHeader* obj) {
-    return containerFor(obj)->frozen();
-}
-
 ALWAYS_INLINE bool isPermanentOrFrozen(const ObjHeader* obj) {
     auto* container = containerFor(obj);
     return container == nullptr || container->frozen();

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

@@ -29,7 +29,7 @@ namespace {
 
 ALWAYS_INLINE inline void mutabilityCheck(KConstRef thiz) {
   // TODO: optimize it!
-  if (!thiz->local() && isFrozen(thiz)) {
+  if (!thiz->local() && isPermanentOrFrozen(thiz)) {
       ThrowInvalidMutabilityException(thiz);
   }
 }

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

@@ -176,10 +176,6 @@ KNativePtr Kotlin_AtomicNativePtr_get(KRef thiz) {
 }
 
 void Kotlin_AtomicReference_checkIfFrozen(KRef value) {
-    if (CurrentMemoryModel == MemoryModel::kExperimental) {
-        // TODO: Remove when freezing is implemented.
-        return;
-    }
     if (value != nullptr && !isPermanentOrFrozen(value)) {
         ThrowInvalidMutabilityException(value);
     }

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

@@ -66,6 +66,8 @@ struct ObjHeader {
       return createMetaObject(this);
   }
 
+  MetaObjHeader* meta_object_or_null() const noexcept { return AsMetaObject(typeInfoOrMeta_); }
+
   ALWAYS_INLINE ObjHeader** GetWeakCounterLocation();
 
 #ifdef KONAN_OBJC_INTEROP
@@ -109,7 +111,6 @@ struct ArrayHeader {
   uint32_t count_;
 };
 
-ALWAYS_INLINE bool isFrozen(const ObjHeader* obj);
 ALWAYS_INLINE bool isPermanentOrFrozen(const ObjHeader* obj);
 ALWAYS_INLINE bool isShareable(const ObjHeader* obj);
 

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

@@ -6,6 +6,7 @@
 #ifndef RUNTIME_MM_EXTRA_OBJECT_DATA_H
 #define RUNTIME_MM_EXTRA_OBJECT_DATA_H
 
+#include <atomic>
 #include <cstddef>
 #include <cstdint>
 
@@ -20,9 +21,21 @@ namespace mm {
 // Optional data that's lazily allocated only for objects that need it.
 class ExtraObjectData : private Pinned, public KonanAllocatorAware {
 public:
+    enum Flags : uint32_t {
+        FLAGS_NONE = 0,
+        FLAGS_FROZEN = 1 << 0,
+        FLAGS_NEVER_FROZEN = 1 << 1,
+    };
+
     MetaObjHeader* AsMetaObjHeader() noexcept { return reinterpret_cast<MetaObjHeader*>(this); }
     static ExtraObjectData& FromMetaObjHeader(MetaObjHeader* header) noexcept { return *reinterpret_cast<ExtraObjectData*>(header); }
 
+    // Get installed `ExtraObjectData` or `nullptr`.
+    static ExtraObjectData* Get(const ObjHeader* object) noexcept {
+        return reinterpret_cast<ExtraObjectData*>(object->meta_object_or_null());
+    }
+    static ExtraObjectData& GetOrInstall(ObjHeader* object) noexcept { return FromMetaObjHeader(object->meta_object()); }
+
     static ExtraObjectData& Install(ObjHeader* object) noexcept;
     static void Uninstall(ObjHeader* object) noexcept;
 
@@ -32,6 +45,8 @@ public:
 
     ObjHeader** GetWeakCounterLocation() noexcept { return &weakReferenceCounter_; }
 
+    std::atomic<Flags>& flags() noexcept { return flags_; }
+
 private:
     explicit ExtraObjectData(const TypeInfo* typeInfo) noexcept : typeInfo_(typeInfo) {}
     ~ExtraObjectData();
@@ -39,6 +54,8 @@ private:
     // Must be first to match `TypeInfo` layout.
     const TypeInfo* typeInfo_;
 
+    std::atomic<Flags> flags_ = FLAGS_NONE;
+
 #ifdef KONAN_OBJC_INTEROP
     void* associatedObject_ = nullptr;
 #endif

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

@@ -0,0 +1,69 @@
+/*
+ * 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 "Freezing.hpp"
+
+#include "ExtraObjectData.hpp"
+#include "FreezeHooks.hpp"
+#include "Memory.h"
+#include "Natives.h"
+#include "ObjectTraversal.hpp"
+#include "Types.h"
+
+using namespace kotlin;
+
+bool mm::IsFrozen(const ObjHeader* object) noexcept {
+    if (object->permanent()) {
+        return true;
+    }
+
+    if (auto* extraObjectData = mm::ExtraObjectData::Get(object)) {
+        return (extraObjectData->flags() & mm::ExtraObjectData::FLAGS_FROZEN) != 0;
+    }
+    return false;
+}
+
+ObjHeader* mm::FreezeSubgraph(ObjHeader* root) noexcept {
+    if (IsFrozen(root)) return nullptr;
+
+    KStdVector<ObjHeader*> objects;
+    KStdVector<ObjHeader*> stack;
+    // TODO: This may be a suboptimal container for the job.
+    KStdUnorderedSet<ObjHeader*> visited;
+    stack.push_back(root);
+    while (!stack.empty()) {
+        ObjHeader* object = stack.back();
+        stack.pop_back();
+        auto visitedResult = visited.insert(object);
+        if (!visitedResult.second) continue;
+        objects.push_back(object);
+        RunFreezeHooks(object);
+        traverseReferredObjects(object, [&stack](ObjHeader* field) noexcept {
+            if (!IsFrozen(field)) {
+                stack.push_back(field);
+            }
+        });
+    }
+    for (auto* object : objects) {
+        if (auto* extraObjectData = mm::ExtraObjectData::Get(object)) {
+            if ((extraObjectData->flags() & mm::ExtraObjectData::FLAGS_NEVER_FROZEN) != 0) return object;
+        }
+    }
+    for (auto* object : objects) {
+        auto& flags = mm::ExtraObjectData::GetOrInstall(object).flags();
+        flags = static_cast<mm::ExtraObjectData::Flags>(flags | mm::ExtraObjectData::FLAGS_FROZEN);
+    }
+    return nullptr;
+}
+
+bool mm::EnsureNeverFrozen(ObjHeader* object) noexcept {
+    if (IsFrozen(object)) {
+        return false;
+    }
+
+    auto& flags = mm::ExtraObjectData::GetOrInstall(object).flags();
+    flags = static_cast<mm::ExtraObjectData::Flags>(flags | mm::ExtraObjectData::FLAGS_NEVER_FROZEN);
+    return true;
+}

kotlin-native/runtime/src/mm/cpp/Freezing.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_FREEZING_H
+#define RUNTIME_MM_FREEZING_H
+
+struct ObjHeader;
+
+namespace kotlin {
+namespace mm {
+
+bool IsFrozen(const ObjHeader* object) noexcept;
+
+// If some object in the `root` subgraph is marked with `EnsureNeverFrozen` only
+// freeze hooks will be executed, and the offending object will be return. Otherwise,
+// all objects in the subgraph will be frozen and `nullptr` will be returned.
+// Note: not thread safe.
+ObjHeader* FreezeSubgraph(ObjHeader* root) noexcept;
+
+// If `object` is already frozen returns false.
+// Note: not thread safe.
+bool EnsureNeverFrozen(ObjHeader* object) noexcept;
+
+} // namespace mm
+} // namespace kotlin
+
+#endif // RUNTIME_MM_FREEZING_H

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

@@ -0,0 +1,616 @@
+/*
+ * 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 "Freezing.hpp"
+
+#include "gmock/gmock.h"
+#include "gtest/gtest.h"
+
+#include "FreezeHooksTestSupport.hpp"
+#include "Memory.h"
+#include "ObjectTestSupport.hpp"
+#include "Utils.hpp"
+
+using namespace kotlin;
+
+using ::testing::_;
+
+namespace {
+
+struct WithFreezeHook {
+    static constexpr bool hasFreezeHook = true;
+};
+
+struct NoFreezeHook {
+    static constexpr bool hasFreezeHook = false;
+};
+
+struct EmptyPayload {
+    using Field = ObjHeader* EmptyPayload::*;
+    static constexpr std::array<Field, 0> kFields{};
+};
+
+struct Payload {
+    ObjHeader* field1;
+    ObjHeader* field2;
+    ObjHeader* field3;
+
+    static constexpr std::array kFields{
+            &Payload::field1,
+            &Payload::field2,
+            &Payload::field3,
+    };
+};
+
+template <typename Payload, typename Traits>
+class ObjectHolder : private Pinned {
+public:
+    static test_support::TypeInfoHolder::ObjectBuilder<Payload> TypeBuilder() {
+        auto builder = test_support::TypeInfoHolder::ObjectBuilder<Payload>();
+        if (Traits::hasFreezeHook) {
+            builder.addFlag(TF_HAS_FREEZE_HOOK);
+        }
+        return builder;
+    }
+
+    ObjectHolder() : type_(TypeBuilder()), object_(type_.typeInfo()) {}
+
+    ~ObjectHolder() {
+        if (header()->has_meta_object()) {
+            ObjHeader::destroyMetaObject(header());
+        }
+    }
+
+    ObjHeader* header() { return object_.header(); }
+
+    ObjHeader*& operator[](size_t field) { return object_.fields()[field]; }
+
+    void MakePermanent() { header()->typeInfoOrMeta_ = setPointerBits(header()->typeInfoOrMeta_, OBJECT_TAG_PERMANENT_CONTAINER); }
+
+private:
+    test_support::TypeInfoHolder type_;
+    test_support::Object<Payload> object_;
+};
+
+// Arrays types are predetermined, and none of them have freeze hooks.
+template <size_t Elements>
+class ArrayHolder : private Pinned {
+public:
+    ArrayHolder() {}
+
+    ~ArrayHolder() {
+        if (header()->has_meta_object()) {
+            ObjHeader::destroyMetaObject(header());
+        }
+    }
+
+    ObjHeader* header() { return array_.header(); }
+
+    ObjHeader*& operator[](size_t index) { return array_.elements()[index]; }
+
+    void MakePermanent() { header()->typeInfoOrMeta_ = setPointerBits(header()->typeInfoOrMeta_, OBJECT_TAG_PERMANENT_CONTAINER); }
+
+private:
+    test_support::ObjectArray<Elements> array_;
+};
+
+class FreezingTest : public testing::Test {
+public:
+    testing::MockFunction<void(ObjHeader*)>& freezeHook() { return freezeHooks_.freezeHook(); }
+
+private:
+    FreezeHooksTestSupport freezeHooks_;
+};
+
+class TypesNames {
+public:
+    template <typename T>
+    static std::string GetName(int i) {
+        switch (i) {
+            case 0:
+                return "object";
+            case 1:
+                return "array";
+            default:
+                return "unknown";
+        }
+    }
+};
+
+template <typename T>
+class FreezingEmptyNoHookTest : public FreezingTest {};
+using EmptyNoHookTypes = testing::Types<ObjectHolder<EmptyPayload, NoFreezeHook>, ArrayHolder<0>>;
+TYPED_TEST_SUITE(FreezingEmptyNoHookTest, EmptyNoHookTypes, TypesNames);
+
+template <typename T>
+class FreezingEmptyWithHookTest : public FreezingTest {};
+using EmptyWithHookTypes = testing::Types<ObjectHolder<EmptyPayload, WithFreezeHook>>;
+TYPED_TEST_SUITE(FreezingEmptyWithHookTest, EmptyWithHookTypes, TypesNames);
+
+template <typename T>
+class FreezingNoHookTest : public FreezingTest {};
+using NoHookTypes = testing::Types<ObjectHolder<Payload, NoFreezeHook>, ArrayHolder<3>>;
+TYPED_TEST_SUITE(FreezingNoHookTest, NoHookTypes, TypesNames);
+
+template <typename T>
+class FreezingWithHookTest : public FreezingTest {};
+using WithHookTypes = testing::Types<ObjectHolder<Payload, WithFreezeHook>>;
+TYPED_TEST_SUITE(FreezingWithHookTest, WithHookTypes, TypesNames);
+
+} // namespace
+
+TYPED_TEST(FreezingEmptyNoHookTest, UnfrozenByDefault) {
+    TypeParam object;
+    EXPECT_FALSE(mm::IsFrozen(object.header()));
+}
+
+TYPED_TEST(FreezingEmptyNoHookTest, PermanentIsFrozen) {
+    TypeParam object;
+    object.MakePermanent();
+    EXPECT_TRUE(mm::IsFrozen(object.header()));
+}
+
+TYPED_TEST(FreezingEmptyNoHookTest, FailToEnsureNeverFrozen) {
+    TypeParam object;
+    ASSERT_THAT(mm::FreezeSubgraph(object.header()), nullptr);
+    ASSERT_TRUE(mm::IsFrozen(object.header()));
+    EXPECT_FALSE(mm::EnsureNeverFrozen(object.header()));
+}
+
+TYPED_TEST(FreezingEmptyNoHookTest, FailToEnsureNeverFrozenPermanent) {
+    TypeParam object;
+    object.MakePermanent();
+    ASSERT_TRUE(mm::IsFrozen(object.header()));
+    EXPECT_FALSE(mm::EnsureNeverFrozen(object.header()));
+}
+
+TYPED_TEST(FreezingEmptyNoHookTest, Freeze) {
+    TypeParam object;
+    EXPECT_THAT(mm::FreezeSubgraph(object.header()), nullptr);
+    EXPECT_TRUE(mm::IsFrozen(object.header()));
+}
+
+TYPED_TEST(FreezingEmptyNoHookTest, FreezePermanent) {
+    TypeParam object;
+    object.MakePermanent();
+    EXPECT_THAT(mm::FreezeSubgraph(object.header()), nullptr);
+    EXPECT_TRUE(mm::IsFrozen(object.header()));
+}
+
+TYPED_TEST(FreezingEmptyNoHookTest, FreezeTwice) {
+    TypeParam object;
+    EXPECT_THAT(mm::FreezeSubgraph(object.header()), nullptr);
+    EXPECT_THAT(mm::FreezeSubgraph(object.header()), nullptr);
+    EXPECT_TRUE(mm::IsFrozen(object.header()));
+}
+
+TYPED_TEST(FreezingEmptyNoHookTest, FreezeForbidden) {
+    TypeParam object;
+    ASSERT_TRUE(mm::EnsureNeverFrozen(object.header()));
+    EXPECT_THAT(mm::FreezeSubgraph(object.header()), object.header());
+    EXPECT_FALSE(mm::IsFrozen(object.header()));
+}
+
+TYPED_TEST(FreezingEmptyWithHookTest, Freeze) {
+    TypeParam object;
+    EXPECT_CALL(this->freezeHook(), Call(object.header()));
+    EXPECT_THAT(mm::FreezeSubgraph(object.header()), nullptr);
+    EXPECT_TRUE(mm::IsFrozen(object.header()));
+}
+
+TYPED_TEST(FreezingEmptyWithHookTest, FreezePermanent) {
+    TypeParam object;
+    object.MakePermanent();
+    EXPECT_CALL(this->freezeHook(), Call(_)).Times(0);
+    EXPECT_THAT(mm::FreezeSubgraph(object.header()), nullptr);
+    EXPECT_TRUE(mm::IsFrozen(object.header()));
+}
+
+TYPED_TEST(FreezingEmptyWithHookTest, FreezeTwice) {
+    TypeParam object;
+    // Only called for the first freeze.
+    EXPECT_CALL(this->freezeHook(), Call(object.header()));
+    EXPECT_THAT(mm::FreezeSubgraph(object.header()), nullptr);
+    testing::Mock::VerifyAndClearExpectations(&this->freezeHook());
+    EXPECT_THAT(mm::FreezeSubgraph(object.header()), nullptr);
+    EXPECT_TRUE(mm::IsFrozen(object.header()));
+}
+
+TYPED_TEST(FreezingEmptyWithHookTest, FreezeForbidden) {
+    TypeParam object;
+    ASSERT_TRUE(mm::EnsureNeverFrozen(object.header()));
+    EXPECT_CALL(this->freezeHook(), Call(object.header()));
+    EXPECT_THAT(mm::FreezeSubgraph(object.header()), object.header());
+    EXPECT_FALSE(mm::IsFrozen(object.header()));
+}
+
+TYPED_TEST(FreezingNoHookTest, Freeze) {
+    TypeParam object;
+    EXPECT_THAT(mm::FreezeSubgraph(object.header()), nullptr);
+    EXPECT_TRUE(mm::IsFrozen(object.header()));
+}
+
+TYPED_TEST(FreezingNoHookTest, FreezePermanent) {
+    TypeParam object;
+    object.MakePermanent();
+    EXPECT_THAT(mm::FreezeSubgraph(object.header()), nullptr);
+    EXPECT_TRUE(mm::IsFrozen(object.header()));
+}
+
+TYPED_TEST(FreezingNoHookTest, FreezeTwice) {
+    TypeParam object;
+    EXPECT_THAT(mm::FreezeSubgraph(object.header()), nullptr);
+    EXPECT_THAT(mm::FreezeSubgraph(object.header()), nullptr);
+    EXPECT_TRUE(mm::IsFrozen(object.header()));
+}
+
+TYPED_TEST(FreezingNoHookTest, FreezeForbidden) {
+    TypeParam object;
+    ASSERT_TRUE(mm::EnsureNeverFrozen(object.header()));
+    EXPECT_THAT(mm::FreezeSubgraph(object.header()), object.header());
+    EXPECT_FALSE(mm::IsFrozen(object.header()));
+}
+
+TYPED_TEST(FreezingNoHookTest, FreezeTree) {
+    TypeParam object;
+    TypeParam field1;
+    TypeParam field2;
+    TypeParam field3;
+    object[0] = field1.header();
+    object[1] = field2.header();
+    object[2] = field3.header();
+    EXPECT_THAT(mm::FreezeSubgraph(object.header()), nullptr);
+    EXPECT_TRUE(mm::IsFrozen(object.header()));
+    EXPECT_TRUE(mm::IsFrozen(field1.header()));
+    EXPECT_TRUE(mm::IsFrozen(field2.header()));
+    EXPECT_TRUE(mm::IsFrozen(field3.header()));
+}
+
+TYPED_TEST(FreezingNoHookTest, FreezeTreePermanent) {
+    TypeParam object;
+    object.MakePermanent();
+    TypeParam field1;
+    field1.MakePermanent();
+    TypeParam field2;
+    field2.MakePermanent();
+    TypeParam field3;
+    field3.MakePermanent();
+    object[0] = field1.header();
+    object[1] = field2.header();
+    object[2] = field3.header();
+    EXPECT_THAT(mm::FreezeSubgraph(object.header()), nullptr);
+    EXPECT_TRUE(mm::IsFrozen(object.header()));
+    EXPECT_TRUE(mm::IsFrozen(field1.header()));
+    EXPECT_TRUE(mm::IsFrozen(field2.header()));
+    EXPECT_TRUE(mm::IsFrozen(field3.header()));
+}
+
+TYPED_TEST(FreezingNoHookTest, FreezeTreePermanentLeaf) {
+    TypeParam object;
+    TypeParam field1;
+    field1.MakePermanent();
+    TypeParam field2;
+    TypeParam field3;
+    object[0] = field1.header();
+    object[1] = field2.header();
+    object[2] = field3.header();
+    EXPECT_THAT(mm::FreezeSubgraph(object.header()), nullptr);
+    EXPECT_TRUE(mm::IsFrozen(object.header()));
+    EXPECT_TRUE(mm::IsFrozen(field1.header()));
+    EXPECT_TRUE(mm::IsFrozen(field2.header()));
+    EXPECT_TRUE(mm::IsFrozen(field3.header()));
+}
+
+TYPED_TEST(FreezingNoHookTest, FreezeTreeTwice) {
+    TypeParam object;
+    TypeParam field1;
+    TypeParam field2;
+    TypeParam field3;
+    object[0] = field1.header();
+    object[1] = field2.header();
+    object[2] = field3.header();
+    EXPECT_THAT(mm::FreezeSubgraph(object.header()), nullptr);
+    EXPECT_THAT(mm::FreezeSubgraph(object.header()), nullptr);
+    EXPECT_TRUE(mm::IsFrozen(object.header()));
+    EXPECT_TRUE(mm::IsFrozen(field1.header()));
+    EXPECT_TRUE(mm::IsFrozen(field2.header()));
+    EXPECT_TRUE(mm::IsFrozen(field3.header()));
+}
+
+TYPED_TEST(FreezingNoHookTest, FreezeTreeForbidden) {
+    TypeParam object;
+    TypeParam field1;
+    TypeParam field2;
+    TypeParam field3;
+    object[0] = field1.header();
+    object[1] = field2.header();
+    object[2] = field3.header();
+    ASSERT_TRUE(mm::EnsureNeverFrozen(object.header()));
+    EXPECT_THAT(mm::FreezeSubgraph(object.header()), object.header());
+    EXPECT_FALSE(mm::IsFrozen(object.header()));
+    EXPECT_FALSE(mm::IsFrozen(field1.header()));
+    EXPECT_FALSE(mm::IsFrozen(field2.header()));
+    EXPECT_FALSE(mm::IsFrozen(field3.header()));
+}
+
+TYPED_TEST(FreezingNoHookTest, FreezeTreeForbiddenByField) {
+    TypeParam object;
+    TypeParam field1;
+    TypeParam field2;
+    TypeParam field3;
+    object[0] = field1.header();
+    object[1] = field2.header();
+    object[2] = field3.header();
+    ASSERT_TRUE(mm::EnsureNeverFrozen(field2.header()));
+    EXPECT_THAT(mm::FreezeSubgraph(object.header()), field2.header());
+    EXPECT_FALSE(mm::IsFrozen(object.header()));
+    EXPECT_FALSE(mm::IsFrozen(field1.header()));
+    EXPECT_FALSE(mm::IsFrozen(field2.header()));
+    EXPECT_FALSE(mm::IsFrozen(field3.header()));
+}
+
+TYPED_TEST(FreezingNoHookTest, FreezeTreeRecursive) {
+    TypeParam object;
+    TypeParam inner1;
+    TypeParam inner2;
+    object[0] = inner1.header();
+    inner1[0] = inner2.header();
+    inner2[0] = object.header();
+    EXPECT_THAT(mm::FreezeSubgraph(object.header()), nullptr);
+    EXPECT_TRUE(mm::IsFrozen(object.header()));
+    EXPECT_TRUE(mm::IsFrozen(inner1.header()));
+    EXPECT_TRUE(mm::IsFrozen(inner2.header()));
+}
+
+TYPED_TEST(FreezingNoHookTest, FreezeTreeRecursivePermanent) {
+    TypeParam object;
+    object.MakePermanent();
+    TypeParam inner1;
+    inner1.MakePermanent();
+    TypeParam inner2;
+    inner2.MakePermanent();
+    object[0] = inner1.header();
+    inner1[0] = inner2.header();
+    inner2[0] = object.header();
+    EXPECT_THAT(mm::FreezeSubgraph(object.header()), nullptr);
+    EXPECT_TRUE(mm::IsFrozen(object.header()));
+    EXPECT_TRUE(mm::IsFrozen(inner1.header()));
+    EXPECT_TRUE(mm::IsFrozen(inner2.header()));
+}
+
+TYPED_TEST(FreezingWithHookTest, Freeze) {
+    TypeParam object;
+    EXPECT_CALL(this->freezeHook(), Call(object.header()));
+    EXPECT_THAT(mm::FreezeSubgraph(object.header()), nullptr);
+    EXPECT_TRUE(mm::IsFrozen(object.header()));
+}
+
+TYPED_TEST(FreezingWithHookTest, FreezePermanent) {
+    TypeParam object;
+    object.MakePermanent();
+    EXPECT_CALL(this->freezeHook(), Call(_)).Times(0);
+    EXPECT_THAT(mm::FreezeSubgraph(object.header()), nullptr);
+    EXPECT_TRUE(mm::IsFrozen(object.header()));
+}
+
+TYPED_TEST(FreezingWithHookTest, FreezeTwice) {
+    TypeParam object;
+    // Only called for the first freeze.
+    EXPECT_CALL(this->freezeHook(), Call(object.header()));
+    EXPECT_THAT(mm::FreezeSubgraph(object.header()), nullptr);
+    testing::Mock::VerifyAndClearExpectations(&this->freezeHook());
+    EXPECT_THAT(mm::FreezeSubgraph(object.header()), nullptr);
+    EXPECT_TRUE(mm::IsFrozen(object.header()));
+}
+
+TYPED_TEST(FreezingWithHookTest, FreezeForbidden) {
+    TypeParam object;
+    ASSERT_TRUE(mm::EnsureNeverFrozen(object.header()));
+    EXPECT_CALL(this->freezeHook(), Call(object.header()));
+    EXPECT_THAT(mm::FreezeSubgraph(object.header()), object.header());
+    EXPECT_FALSE(mm::IsFrozen(object.header()));
+}
+
+TYPED_TEST(FreezingWithHookTest, FreezeTree) {
+    TypeParam object;
+    TypeParam field1;
+    TypeParam field2;
+    TypeParam field3;
+    object[0] = field1.header();
+    object[1] = field2.header();
+    object[2] = field3.header();
+    EXPECT_CALL(this->freezeHook(), Call(object.header()));
+    EXPECT_CALL(this->freezeHook(), Call(field1.header()));
+    EXPECT_CALL(this->freezeHook(), Call(field2.header()));
+    EXPECT_CALL(this->freezeHook(), Call(field3.header()));
+    EXPECT_THAT(mm::FreezeSubgraph(object.header()), nullptr);
+    EXPECT_TRUE(mm::IsFrozen(object.header()));
+    EXPECT_TRUE(mm::IsFrozen(field1.header()));
+    EXPECT_TRUE(mm::IsFrozen(field2.header()));
+    EXPECT_TRUE(mm::IsFrozen(field3.header()));
+}
+
+TYPED_TEST(FreezingWithHookTest, FreezeTreePermanent) {
+    TypeParam object;
+    object.MakePermanent();
+    TypeParam field1;
+    field1.MakePermanent();
+    TypeParam field2;
+    field2.MakePermanent();
+    TypeParam field3;
+    field3.MakePermanent();
+    object[0] = field1.header();
+    object[1] = field2.header();
+    object[2] = field3.header();
+    EXPECT_CALL(this->freezeHook(), Call(_)).Times(0);
+    EXPECT_THAT(mm::FreezeSubgraph(object.header()), nullptr);
+    EXPECT_TRUE(mm::IsFrozen(object.header()));
+    EXPECT_TRUE(mm::IsFrozen(field1.header()));
+    EXPECT_TRUE(mm::IsFrozen(field2.header()));
+    EXPECT_TRUE(mm::IsFrozen(field3.header()));
+}
+
+TYPED_TEST(FreezingWithHookTest, FreezeTreePermanentLeaf) {
+    TypeParam object;
+    TypeParam field1;
+    field1.MakePermanent();
+    TypeParam field2;
+    TypeParam field3;
+    object[0] = field1.header();
+    object[1] = field2.header();
+    object[2] = field3.header();
+    EXPECT_CALL(this->freezeHook(), Call(object.header()));
+    EXPECT_CALL(this->freezeHook(), Call(field1.header())).Times(0);
+    EXPECT_CALL(this->freezeHook(), Call(field2.header()));
+    EXPECT_CALL(this->freezeHook(), Call(field3.header()));
+    EXPECT_THAT(mm::FreezeSubgraph(object.header()), nullptr);
+    EXPECT_TRUE(mm::IsFrozen(object.header()));
+    EXPECT_TRUE(mm::IsFrozen(field1.header()));
+    EXPECT_TRUE(mm::IsFrozen(field2.header()));
+    EXPECT_TRUE(mm::IsFrozen(field3.header()));
+}
+
+TYPED_TEST(FreezingWithHookTest, FreezeTreeTwice) {
+    TypeParam object;
+    TypeParam field1;
+    TypeParam field2;
+    TypeParam field3;
+    object[0] = field1.header();
+    object[1] = field2.header();
+    object[2] = field3.header();
+    // Only called for the first freeze.
+    EXPECT_CALL(this->freezeHook(), Call(object.header()));
+    EXPECT_CALL(this->freezeHook(), Call(field1.header()));
+    EXPECT_CALL(this->freezeHook(), Call(field2.header()));
+    EXPECT_CALL(this->freezeHook(), Call(field3.header()));
+    EXPECT_THAT(mm::FreezeSubgraph(object.header()), nullptr);
+    testing::Mock::VerifyAndClearExpectations(&this->freezeHook());
+    EXPECT_THAT(mm::FreezeSubgraph(object.header()), nullptr);
+    EXPECT_TRUE(mm::IsFrozen(object.header()));
+    EXPECT_TRUE(mm::IsFrozen(field1.header()));
+    EXPECT_TRUE(mm::IsFrozen(field2.header()));
+    EXPECT_TRUE(mm::IsFrozen(field3.header()));
+}
+
+TYPED_TEST(FreezingWithHookTest, FreezeTreeForbidden) {
+    TypeParam object;
+    TypeParam field1;
+    TypeParam field2;
+    TypeParam field3;
+    object[0] = field1.header();
+    object[1] = field2.header();
+    object[2] = field3.header();
+    ASSERT_TRUE(mm::EnsureNeverFrozen(object.header()));
+    EXPECT_CALL(this->freezeHook(), Call(object.header()));
+    EXPECT_CALL(this->freezeHook(), Call(field1.header()));
+    EXPECT_CALL(this->freezeHook(), Call(field2.header()));
+    EXPECT_CALL(this->freezeHook(), Call(field3.header()));
+    EXPECT_THAT(mm::FreezeSubgraph(object.header()), object.header());
+    EXPECT_FALSE(mm::IsFrozen(object.header()));
+    EXPECT_FALSE(mm::IsFrozen(field1.header()));
+    EXPECT_FALSE(mm::IsFrozen(field2.header()));
+    EXPECT_FALSE(mm::IsFrozen(field3.header()));
+}
+
+TYPED_TEST(FreezingWithHookTest, FreezeTreeForbiddenByField) {
+    TypeParam object;
+    TypeParam field1;
+    TypeParam field2;
+    TypeParam field3;
+    object[0] = field1.header();
+    object[1] = field2.header();
+    object[2] = field3.header();
+    ASSERT_TRUE(mm::EnsureNeverFrozen(field2.header()));
+    EXPECT_CALL(this->freezeHook(), Call(object.header()));
+    EXPECT_CALL(this->freezeHook(), Call(field1.header()));
+    EXPECT_CALL(this->freezeHook(), Call(field2.header()));
+    EXPECT_CALL(this->freezeHook(), Call(field3.header()));
+    EXPECT_THAT(mm::FreezeSubgraph(object.header()), field2.header());
+    EXPECT_FALSE(mm::IsFrozen(object.header()));
+    EXPECT_FALSE(mm::IsFrozen(field1.header()));
+    EXPECT_FALSE(mm::IsFrozen(field2.header()));
+    EXPECT_FALSE(mm::IsFrozen(field3.header()));
+}
+
+TYPED_TEST(FreezingWithHookTest, FreezeTreeRecursive) {
+    TypeParam object;
+    TypeParam inner1;
+    TypeParam inner2;
+    object[0] = inner1.header();
+    inner1[0] = inner2.header();
+    inner2[0] = object.header();
+    EXPECT_CALL(this->freezeHook(), Call(object.header()));
+    EXPECT_CALL(this->freezeHook(), Call(inner1.header()));
+    EXPECT_CALL(this->freezeHook(), Call(inner2.header()));
+    EXPECT_THAT(mm::FreezeSubgraph(object.header()), nullptr);
+    EXPECT_TRUE(mm::IsFrozen(object.header()));
+    EXPECT_TRUE(mm::IsFrozen(inner1.header()));
+    EXPECT_TRUE(mm::IsFrozen(inner2.header()));
+}
+
+TYPED_TEST(FreezingWithHookTest, FreezeTreeRecursivePermanent) {
+    TypeParam object;
+    object.MakePermanent();
+    TypeParam inner1;
+    inner1.MakePermanent();
+    TypeParam inner2;
+    inner2.MakePermanent();
+    object[0] = inner1.header();
+    inner1[0] = inner2.header();
+    inner2[0] = object.header();
+    EXPECT_CALL(this->freezeHook(), Call(_)).Times(0);
+    EXPECT_THAT(mm::FreezeSubgraph(object.header()), nullptr);
+    EXPECT_TRUE(mm::IsFrozen(object.header()));
+    EXPECT_TRUE(mm::IsFrozen(inner1.header()));
+    EXPECT_TRUE(mm::IsFrozen(inner2.header()));
+}
+
+TYPED_TEST(FreezingWithHookTest, FreezeTreeWithHookRewrite) {
+    TypeParam object;
+    TypeParam field1;
+    TypeParam field2;
+    TypeParam field3;
+    TypeParam oldInner;
+    TypeParam newInner;
+    object[0] = field1.header();
+    object[1] = field2.header();
+    object[2] = field3.header();
+    field2[0] = oldInner.header();
+    ON_CALL(this->freezeHook(), Call(field2.header())).WillByDefault([&field2, &newInner](ObjHeader* obj) {
+        field2[0] = newInner.header();
+    });
+    EXPECT_CALL(this->freezeHook(), Call(object.header()));
+    EXPECT_CALL(this->freezeHook(), Call(field1.header()));
+    EXPECT_CALL(this->freezeHook(), Call(field2.header()));
+    EXPECT_CALL(this->freezeHook(), Call(field3.header()));
+    EXPECT_CALL(this->freezeHook(), Call(newInner.header()));
+    EXPECT_CALL(this->freezeHook(), Call(oldInner.header())).Times(0);
+    EXPECT_THAT(mm::FreezeSubgraph(object.header()), nullptr);
+    EXPECT_TRUE(mm::IsFrozen(object.header()));
+    EXPECT_TRUE(mm::IsFrozen(field1.header()));
+    EXPECT_TRUE(mm::IsFrozen(field2.header()));
+    EXPECT_TRUE(mm::IsFrozen(field3.header()));
+    EXPECT_TRUE(mm::IsFrozen(newInner.header()));
+    EXPECT_FALSE(mm::IsFrozen(oldInner.header()));
+}
+
+TYPED_TEST(FreezingWithHookTest, FreezeTreeForbiddenByHook) {
+    TypeParam object;
+    TypeParam field1;
+    TypeParam field2;
+    TypeParam field3;
+    object[0] = field1.header();
+    object[1] = field2.header();
+    object[2] = field3.header();
+    ON_CALL(this->freezeHook(), Call(field2.header())).WillByDefault([](ObjHeader* obj) { EXPECT_TRUE(mm::EnsureNeverFrozen(obj)); });
+    EXPECT_CALL(this->freezeHook(), Call(object.header()));
+    EXPECT_CALL(this->freezeHook(), Call(field1.header()));
+    EXPECT_CALL(this->freezeHook(), Call(field2.header()));
+    EXPECT_CALL(this->freezeHook(), Call(field3.header()));
+    EXPECT_THAT(mm::FreezeSubgraph(object.header()), field2.header());
+    EXPECT_FALSE(mm::IsFrozen(object.header()));
+    EXPECT_FALSE(mm::IsFrozen(field1.header()));
+    EXPECT_FALSE(mm::IsFrozen(field2.header()));
+}

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

@@ -8,6 +8,7 @@
 
 #include "Exceptions.h"
 #include "ExtraObjectData.hpp"
+#include "Freezing.hpp"
 #include "GlobalsRegistry.hpp"
 #include "InitializationScheme.hpp"
 #include "KAssert.h"
@@ -78,7 +79,7 @@ void ObjHeader::destroyMetaObject(ObjHeader* object) {
 }
 
 ALWAYS_INLINE bool isPermanentOrFrozen(const ObjHeader* obj) {
-    return obj->permanent() || isFrozen(obj);
+    return mm::IsFrozen(obj);
 }
 
 ALWAYS_INLINE bool isShareable(const ObjHeader* obj) {
@@ -375,6 +376,13 @@ extern "C" RUNTIME_NOTHROW OBJ_GETTER(AdoptStablePointer, void* pointer) {
     return object;
 }
 
+extern "C" void MutationCheck(ObjHeader* obj) {
+    if (obj->local()) return;
+    if (!isPermanentOrFrozen(obj)) return;
+
+    ThrowInvalidMutabilityException(obj);
+}
+
 extern "C" RUNTIME_NOTHROW void CheckLifetimesConstraint(ObjHeader* obj, ObjHeader* pointee) {
     // TODO: Consider making it a `RuntimeCheck`. Probably all `RuntimeCheck`s and `RuntimeAssert`s should specify
     //       that their firing is a compiler bug and should be reported.
@@ -385,6 +393,18 @@ extern "C" RUNTIME_NOTHROW void CheckLifetimesConstraint(ObjHeader* obj, ObjHead
     }
 }
 
+extern "C" void FreezeSubgraph(ObjHeader* obj) {
+    if (auto* blocker = mm::FreezeSubgraph(obj)) {
+        ThrowFreezingException(obj, blocker);
+    }
+}
+
+extern "C" void EnsureNeverFrozen(ObjHeader* obj) {
+    if (!mm::EnsureNeverFrozen(obj)) {
+        ThrowFreezingException(obj, obj);
+    }
+}
+
 extern "C" ForeignRefContext InitForeignRef(ObjHeader* object) {
     auto* threadData = mm::ThreadRegistry::Instance().CurrentThreadData();
     auto* node = mm::StableRefRegistry::Instance().RegisterStableRef(threadData, object);

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

@@ -7,25 +7,8 @@
 
 #include "KAssert.h"
 
-ALWAYS_INLINE bool isFrozen(const ObjHeader* obj) {
-    // TODO: Unimplemented
-    return false;
-}
-
 extern "C" {
 
-void MutationCheck(ObjHeader* obj) {
-    // TODO: Unimplemented
-}
-
-void FreezeSubgraph(ObjHeader* obj) {
-    // TODO: Unimplemented
-}
-
-void EnsureNeverFrozen(ObjHeader* obj) {
-    TODO();
-}
-
 void Kotlin_native_internal_GC_suspend(ObjHeader*) {
     TODO();
 }