Take snapshot for LexicalWritableScope when subscopes created

compiler/frontend/src/org/jetbrains/kotlin/resolve/BindingContextUtils.kt

@@ -31,6 +31,7 @@ import org.jetbrains.kotlin.resolve.calls.context.ResolutionContext
 import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowInfo
 import org.jetbrains.kotlin.resolve.scopes.LexicalWritableScope
 import org.jetbrains.kotlin.resolve.scopes.utils.asJetScope
+import org.jetbrains.kotlin.resolve.scopes.utils.takeSnapshot
 import org.jetbrains.kotlin.types.expressions.typeInfoFactory.noTypeInfo
 
 public fun JetReturnExpression.getTargetFunctionDescriptor(context: BindingContext): FunctionDescriptor? {
@@ -57,7 +58,7 @@ public fun JetExpression.isUsedAsStatement(context: BindingContext): Boolean = !
 public fun <C : ResolutionContext<C>> ResolutionContext<C>.recordScopeAndDataFlowInfo(expression: JetExpression?) {
     if (expression == null) return
 
-    val scopeToRecord = if (scope is LexicalWritableScope) scope.takeSnapshot() else scope
+    val scopeToRecord = scope.takeSnapshot()
     trace.record(BindingContext.RESOLUTION_SCOPE, expression, scopeToRecord.asJetScope())
     trace.record(BindingContext.LEXICAL_SCOPE, expression, scopeToRecord)
 

compiler/frontend/src/org/jetbrains/kotlin/resolve/scopes/LexicalChainedScope.kt

@@ -20,18 +20,20 @@ import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
 import org.jetbrains.kotlin.descriptors.ReceiverParameterDescriptor
 import org.jetbrains.kotlin.incremental.components.LookupLocation
 import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.resolve.scopes.utils.takeSnapshot
 import org.jetbrains.kotlin.util.collectionUtils.getFirstMatch
 import org.jetbrains.kotlin.util.collectionUtils.getFromAllScopes
 import org.jetbrains.kotlin.utils.Printer
 
 public class LexicalChainedScope(
-        override val parent: LexicalScope,
+        parent: LexicalScope,
         override val ownerDescriptor: DeclarationDescriptor,
         override val isOwnerDescriptorAccessibleByLabel: Boolean,
         override val implicitReceiver: ReceiverParameterDescriptor?,
         private val debugName: String,
         vararg memberScopes: JetScope // todo JetScope -> MemberScope
 ): LexicalScope {
+    override val parent = parent.takeSnapshot()
     private val scopeChain = memberScopes.clone()
 
     override fun getDeclaredDescriptors() = getFromAllScopes(scopeChain) { it.getAllDescriptors() }

compiler/frontend/src/org/jetbrains/kotlin/resolve/scopes/LexicalScopeImpl.kt

@@ -20,10 +20,11 @@ import com.intellij.util.SmartList
 import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.incremental.components.LookupLocation
 import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.resolve.scopes.utils.takeSnapshot
 import org.jetbrains.kotlin.utils.Printer
 
 public class LexicalScopeImpl jvmOverloads constructor(
-        override val parent: LexicalScope,
+        parent: LexicalScope,
         override val ownerDescriptor: DeclarationDescriptor,
         override val isOwnerDescriptorAccessibleByLabel: Boolean,
         override val implicitReceiver: ReceiverParameterDescriptor?,
@@ -31,6 +32,7 @@ public class LexicalScopeImpl jvmOverloads constructor(
         redeclarationHandler: RedeclarationHandler = RedeclarationHandler.DO_NOTHING,
         initialize: LexicalScopeImpl.InitializeHandler.() -> Unit = {}
 ): LexicalScope, WritableScopeStorage {
+    override val parent = parent.takeSnapshot()
     override val addedDescriptors: MutableList<DeclarationDescriptor> = SmartList()
     override val redeclarationHandler: RedeclarationHandler
         get() = RedeclarationHandler.DO_NOTHING

compiler/frontend/src/org/jetbrains/kotlin/resolve/scopes/LexicalWritableScope.kt

@@ -20,16 +20,18 @@ import com.intellij.util.SmartList
 import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.incremental.components.LookupLocation
 import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.resolve.scopes.utils.takeSnapshot
 import org.jetbrains.kotlin.utils.Printer
 
 class LexicalWritableScope(
-        override val parent: LexicalScope,
+        parent: LexicalScope,
         override val ownerDescriptor: DeclarationDescriptor,
         override val isOwnerDescriptorAccessibleByLabel: Boolean,
         override val implicitReceiver: ReceiverParameterDescriptor?,
         override val redeclarationHandler: RedeclarationHandler,
         private val debugName: String
 ) : LexicalScope, WritableScopeStorage {
+    override val parent = parent.takeSnapshot()
     override val addedDescriptors: MutableList<DeclarationDescriptor> = SmartList()
 
     override var functionsByName: MutableMap<Name, WritableScopeStorage.IntList>? = null

compiler/frontend/src/org/jetbrains/kotlin/resolve/scopes/utils/ScopeUtils.kt

@@ -21,10 +21,7 @@ import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.incremental.components.LookupLocation
 import org.jetbrains.kotlin.incremental.components.NoLookupLocation
 import org.jetbrains.kotlin.name.Name
-import org.jetbrains.kotlin.resolve.scopes.DescriptorKindFilter
-import org.jetbrains.kotlin.resolve.scopes.FileScope
-import org.jetbrains.kotlin.resolve.scopes.JetScope
-import org.jetbrains.kotlin.resolve.scopes.LexicalScope
+import org.jetbrains.kotlin.resolve.scopes.*
 import org.jetbrains.kotlin.types.JetType
 import org.jetbrains.kotlin.util.collectionUtils.concat
 import org.jetbrains.kotlin.utils.Printer
@@ -109,6 +106,8 @@ public fun LexicalScope.getClassifier(name: Name, location: LookupLocation): Cla
     return null
 }
 
+public fun LexicalScope.takeSnapshot(): LexicalScope = if (this is LexicalWritableScope) takeSnapshot() else this
+
 public fun LexicalScope.asJetScope(): JetScope {
     if (this is JetScope) return this
     if (this is MemberScopeToFileScopeAdapter) return this.memberScope
@@ -126,7 +125,8 @@ public fun JetScope.asLexicalScope(): LexicalScope
             memberScopeAsFileScope()
         }
 
-private class LexicalToJetScopeAdapter(val lexicalScope: LexicalScope): JetScope {
+private class LexicalToJetScopeAdapter(lexicalScope: LexicalScope): JetScope {
+    val lexicalScope = lexicalScope.takeSnapshot()
 
     override fun getClassifier(name: Name, location: LookupLocation) = lexicalScope.getClassifier(name, location)