FIR IDE: introduce WeakFirRef for low level

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/annotations/annotations.kt

@@ -0,0 +1,10 @@
+/*
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.fir.low.level.api.annotations
+
+@Target(AnnotationTarget.PROPERTY, AnnotationTarget.FUNCTION)
+@RequiresOptIn
+annotation class PrivateForInline

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/api/WeakFirRef.kt

@@ -0,0 +1,28 @@
+/*
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.fir.low.level.api.api
+
+import org.jetbrains.kotlin.fir.declarations.FirDeclaration
+import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
+import org.jetbrains.kotlin.idea.fir.low.level.api.annotations.PrivateForInline
+import java.lang.ref.WeakReference
+
+class WeakFirRef<D : FirDeclaration>(fir: D, resolveState: FirModuleResolveState) {
+    @PrivateForInline
+    val firWeakRef = WeakReference(fir)
+
+    @PrivateForInline
+    val resolveStateWeakRef = WeakReference(resolveState)
+
+    @OptIn(PrivateForInline::class)
+    inline fun <R> withFir(phase: FirResolvePhase = FirResolvePhase.RAW_FIR, action: (fir: D) -> R): R {
+        val fir = firWeakRef.get()
+            ?: error("FirElement was garbage collected")
+        val resolveState = resolveStateWeakRef.get()
+            ?: error("FirModuleResolveState was garbage collected")
+        return action(LowLevelFirApiFacade.resolvedFirToPhase(fir, phase, resolveState))
+    }
+}

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/symbols/KtFirSymbol.kt

@@ -8,13 +8,13 @@ package org.jetbrains.kotlin.idea.frontend.api.fir.symbols
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirDeclarationOrigin
 import org.jetbrains.kotlin.idea.frontend.api.ValidityTokenOwner
-import org.jetbrains.kotlin.idea.frontend.api.fir.utils.FirRef
+import org.jetbrains.kotlin.idea.frontend.api.fir.utils.FirRefWithValidityCheck
 import org.jetbrains.kotlin.idea.frontend.api.fir.utils.overriddenDeclaration
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtSymbol
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtSymbolOrigin
 
 internal interface KtFirSymbol<F : FirDeclaration> : KtSymbol, ValidityTokenOwner {
-    val firRef: FirRef<F>
+    val firRef: FirRefWithValidityCheck<F>
 
     override val origin: KtSymbolOrigin get() = firRef.withFir { it.ktSymbolOrigin() }
 }

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/utils/FirRefWithValidityCheck.kt

@@ -14,7 +14,7 @@ import org.jetbrains.kotlin.idea.frontend.api.ValidityTokenOwner
 import org.jetbrains.kotlin.idea.frontend.api.assertIsValid
 import java.lang.ref.WeakReference
 
-internal class FirRef<D : FirDeclaration>(fir: D, resolveState: FirModuleResolveState, val token: ValidityToken) {
+internal class FirRefWithValidityCheck<D : FirDeclaration>(fir: D, resolveState: FirModuleResolveState, val token: ValidityToken) {
     private val firWeakRef = WeakReference(fir)
     private val resolveStateWeakRef = WeakReference(resolveState)
 
@@ -34,4 +34,4 @@ internal class FirRef<D : FirDeclaration>(fir: D, resolveState: FirModuleResolve
 
 @Suppress("NOTHING_TO_INLINE")
 internal inline fun <D : FirDeclaration> ValidityTokenOwner.firRef(fir: D, resolveState: FirModuleResolveState) =
-    FirRef(fir, resolveState, token)
+    FirRefWithValidityCheck(fir, resolveState, token)