From 9976eff03ad86a23e98536d3f51ddc7d20b1ecc5 Mon Sep 17 00:00:00 2001
From: Simon Ogorodnik <simon.ogorodnik@jetbrains.com>
Date: Tue, 20 Mar 2018 01:08:56 +0300
Subject: [PATCH] FIR: Add components storage to FirSession

---
 .../org/jetbrains/kotlin/fir/FirSession.kt    | 14 +++++++++++++-
 .../jetbrains/kotlin/fir/FirSessionBase.kt    | 19 +++++++++++++++++++
 .../builder/AbstractRawFirBuilderTestCase.kt  |  4 ++--
 3 files changed, 34 insertions(+), 3 deletions(-)
 create mode 100644 compiler/fir/tree/src/org/jetbrains/kotlin/fir/FirSessionBase.kt

diff --git a/compiler/fir/tree/src/org/jetbrains/kotlin/fir/FirSession.kt b/compiler/fir/tree/src/org/jetbrains/kotlin/fir/FirSession.kt
index 0c40134b589..c5961d34e3a 100644
--- a/compiler/fir/tree/src/org/jetbrains/kotlin/fir/FirSession.kt
+++ b/compiler/fir/tree/src/org/jetbrains/kotlin/fir/FirSession.kt
@@ -5,4 +5,16 @@
 
 package org.jetbrains.kotlin.fir
 
-interface FirSession
\ No newline at end of file
+import kotlin.reflect.KClass
+
+interface FirSession {
+
+
+    val components: Map<KClass<*>, Any>
+
+    fun <T : Any> getService(kclass: KClass<T>): T =
+        components[kclass] as T
+}
+
+inline fun <reified T : Any> FirSession.service(): T =
+    getService(T::class)
\ No newline at end of file
diff --git a/compiler/fir/tree/src/org/jetbrains/kotlin/fir/FirSessionBase.kt b/compiler/fir/tree/src/org/jetbrains/kotlin/fir/FirSessionBase.kt
new file mode 100644
index 00000000000..bcc4c613fd1
--- /dev/null
+++ b/compiler/fir/tree/src/org/jetbrains/kotlin/fir/FirSessionBase.kt
@@ -0,0 +1,19 @@
+/*
+ * Copyright 2010-2018 JetBrains s.r.o. 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.fir
+
+import kotlin.reflect.KClass
+
+abstract class FirSessionBase : FirSession {
+
+    protected fun <T : Any> registerComponent(tClass: KClass<T>, t: T) {
+        assert(tClass !in components) { "Already registered component" }
+        components[tClass] = t
+    }
+
+    override val components: MutableMap<KClass<*>, Any> = mutableMapOf()
+
+}
\ No newline at end of file
diff --git a/compiler/tests-common/tests/org/jetbrains/kotlin/fir/builder/AbstractRawFirBuilderTestCase.kt b/compiler/tests-common/tests/org/jetbrains/kotlin/fir/builder/AbstractRawFirBuilderTestCase.kt
index fca25869bcb..033cd93f061 100644
--- a/compiler/tests-common/tests/org/jetbrains/kotlin/fir/builder/AbstractRawFirBuilderTestCase.kt
+++ b/compiler/tests-common/tests/org/jetbrains/kotlin/fir/builder/AbstractRawFirBuilderTestCase.kt
@@ -16,7 +16,7 @@ import com.intellij.util.PathUtil
 import org.jetbrains.kotlin.KtNodeTypes
 import org.jetbrains.kotlin.cli.common.script.CliScriptDefinitionProvider
 import org.jetbrains.kotlin.fir.FirRenderer
-import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.FirSessionBase
 import org.jetbrains.kotlin.parsing.KotlinParserDefinition
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.psi.KtPsiFactory
@@ -58,7 +58,7 @@ abstract class AbstractRawFirBuilderTestCase : KtParsingTestCase(
         val file = createFile(filePath, fileType) as KtFile
         myFile = file
 
-        val firFile = RawFirBuilder(object : FirSession {}).buildFirFile(file)
+        val firFile = RawFirBuilder(object : FirSessionBase() {}).buildFirFile(file)
         val firFileDump = StringBuilder().also { FirRenderer(it).visitFile(firFile) }.toString()
         val expectedPath = filePath.replace(".kt", ".txt")
         KotlinTestUtils.assertEqualsToFile(File(expectedPath), firFileDump)