From abc13589c891e2b7417a9269577ee800cb57b0ca Mon Sep 17 00:00:00 2001
From: Ilya Chernikov <ilya.chernikov@jetbrains.com>
Date: Tue, 11 Dec 2018 11:19:49 +0100
Subject: [PATCH] Fix locking in script definition provider

fixes possible parallel modifications on definitions loading, e.g.
as in https://github.com/shyiko/ktlint/issues/287
---
 .../script/KotlinScriptDefinitionProvider.kt      | 15 +++++++++++----
 1 file changed, 11 insertions(+), 4 deletions(-)

diff --git a/compiler/psi/src/org/jetbrains/kotlin/script/KotlinScriptDefinitionProvider.kt b/compiler/psi/src/org/jetbrains/kotlin/script/KotlinScriptDefinitionProvider.kt
index 0a729ed1a9b..9078d1a6a8a 100644
--- a/compiler/psi/src/org/jetbrains/kotlin/script/KotlinScriptDefinitionProvider.kt
+++ b/compiler/psi/src/org/jetbrains/kotlin/script/KotlinScriptDefinitionProvider.kt
@@ -121,11 +121,18 @@ private class CachingSequence<T>(from: Sequence<T>) : Sequence<T> {
 
         private var cacheCursor = 0
 
-        override fun hasNext(): Boolean = lock.read { cacheCursor < cache.size || sequenceIterator.hasNext() }
+        override fun hasNext(): Boolean =
+            lock.read { cacheCursor < cache.size } || lock.write { cacheCursor < cache.size || sequenceIterator.hasNext() }
 
-        override fun next(): T = lock.write {
-            if (cacheCursor < cache.size) cache[cacheCursor++]
-            else sequenceIterator.next().also { cache.add(it) }
+        override fun next(): T {
+            lock.read {
+                if (cacheCursor < cache.size) return cache[cacheCursor++]
+            }
+            // lock.write is not an upgrade but retake, therefore - one more check needed
+            lock.write {
+                return if (cacheCursor < cache.size) cache[cacheCursor++]
+                else sequenceIterator.next().also { cache.add(it) }
+            }
         }
     }