From e1b6f8198a768f361240620fa66700e77509dc25 Mon Sep 17 00:00:00 2001
From: Abduqodiri Qurbonzoda <abduqodiri.qurbonzoda@jetbrains.com>
Date: Tue, 30 May 2023 00:16:33 +0300
Subject: [PATCH] [K/N] Make Sequence<T>.constrainOnce() multi-thread friendly

---
 .../runtime/src/main/kotlin/kotlin/sequences/Sequences.kt  | 7 ++-----
 1 file changed, 2 insertions(+), 5 deletions(-)

diff --git a/kotlin-native/runtime/src/main/kotlin/kotlin/sequences/Sequences.kt b/kotlin-native/runtime/src/main/kotlin/kotlin/sequences/Sequences.kt
index 66b59747230..0612d88252c 100644
--- a/kotlin-native/runtime/src/main/kotlin/kotlin/sequences/Sequences.kt
+++ b/kotlin-native/runtime/src/main/kotlin/kotlin/sequences/Sequences.kt
@@ -8,13 +8,10 @@ package kotlin.sequences
 import kotlin.comparisons.*
 
 internal actual class ConstrainedOnceSequence<T> actual constructor(sequence: Sequence<T>) : Sequence<T> {
-    // TODO: not MT friendly.
-    private var sequenceRef : Sequence<T>? = sequence
+    private val sequenceRef = kotlin.concurrent.AtomicReference<Sequence<T>?>(sequence)
 
     override actual fun iterator(): Iterator<T> {
-        val sequence = sequenceRef
-        if (sequence == null) throw IllegalStateException("This sequence can be consumed only once.")
-        sequenceRef = null
+        val sequence = sequenceRef.getAndSet(null) ?: throw IllegalStateException("This sequence can be consumed only once.")
         return sequence.iterator()
     }
 }