Use FreezableAtomicReference in unhandled exception hook

kotlin-native/backend.native/tests/runtime/exceptions/custom_hook.kt

@@ -6,7 +6,7 @@ import kotlin.test.*
 
 import kotlin.native.concurrent.*
 
-fun main(args : Array<String>) {
+fun mainLegacyMM() {
     assertFailsWith<InvalidMutabilityException> {
         setUnhandledExceptionHook { _ -> println("wrong") }
     }
@@ -19,4 +19,26 @@ fun main(args : Array<String>) {
     assertNull(old)
 
     throw Error("an error")
-}
+}
+
+fun mainExperimentalMM() {
+    val unset = setUnhandledExceptionHook { _ -> println("ok") }
+    assertNull(unset)
+
+    val x = 42
+    val old = setUnhandledExceptionHook {
+        throwable: Throwable -> println("value $x: ${throwable::class.simpleName}")
+    }
+
+    assertNotNull(old)
+
+    throw Error("an error")
+}
+
+fun main() {
+    if (Platform.memoryModel == MemoryModel.EXPERIMENTAL) {
+        mainExperimentalMM()
+    } else {
+        mainLegacyMM()
+    }
+}

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

@@ -379,6 +379,19 @@ public class FreezableAtomicReference<T>(private var value_: T) {
     public override fun toString(): String =
             "${debugString(this)} -> ${debugString(value)}"
 
+    // TODO: Consider making this public.
+    internal fun swap(new: T): T {
+        while (true) {
+            val old = value
+            if (old === new) {
+                return old
+            }
+            if (compareAndSet(old, new)) {
+                return old
+            }
+        }
+    }
+
     // Implementation details.
     @GCUnsafeCall("Kotlin_AtomicReference_set")
     private external fun setImpl(new: Any?): Unit

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

@@ -7,7 +7,8 @@ package kotlin.native.internal
 
 import kotlin.internal.getProgressionLastElement
 import kotlin.reflect.KClass
-import kotlin.native.concurrent.AtomicReference
+import kotlin.native.concurrent.FreezableAtomicReference
+import kotlin.native.concurrent.freeze
 
 @ExportForCppRuntime
 fun ThrowNullPointerException(): Nothing {
@@ -123,7 +124,12 @@ internal external fun TerminateWithUnhandledException(throwable: Throwable)
 // in a normal global initialization flow. This is important if some global happens
 // to throw an exception during it's initialization before this hook would've been initialized.
 internal object UnhandledExceptionHookHolder {
-    internal val hook: AtomicReference<ReportUnhandledExceptionHook?> = AtomicReference(null)
+    internal val hook: FreezableAtomicReference<ReportUnhandledExceptionHook?> =
+        if (Platform.memoryModel == MemoryModel.EXPERIMENTAL) {
+            FreezableAtomicReference<ReportUnhandledExceptionHook?>(null)
+        } else {
+            FreezableAtomicReference<ReportUnhandledExceptionHook?>(null).freeze()
+        }
 }
 
 @PublishedApi