K2 scripting: add initial scripting support to K2 frontend

compiler/fir/providers/src/org/jetbrains/kotlin/fir/resolve/calls/FirReceivers.kt

@@ -24,6 +24,7 @@ import org.jetbrains.kotlin.fir.scopes.FirTypeScope
 import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirClassSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirScriptSymbol
 import org.jetbrains.kotlin.fir.types.ConeErrorType
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
 import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
@@ -221,6 +222,9 @@ sealed class ContextReceiverValue<S : FirBasedSymbol<*>>(
     boundSymbol, type, useSiteSession, scopeSession, mutable, contextReceiverNumber,
 ) {
     abstract override fun createSnapshot(): ContextReceiverValue<S>
+
+    override val isContextReceiver: Boolean
+        get() = true
 }
 
 class ContextReceiverValueForCallable(
@@ -236,9 +240,6 @@ class ContextReceiverValueForCallable(
 ) {
     override fun createSnapshot(): ContextReceiverValue<FirCallableSymbol<*>> =
         ContextReceiverValueForCallable(boundSymbol, type, labelName, useSiteSession, scopeSession, mutable = false, contextReceiverNumber)
-
-    override val isContextReceiver: Boolean
-        get() = true
 }
 
 class ContextReceiverValueForClass(
@@ -254,7 +255,20 @@ class ContextReceiverValueForClass(
 ) {
     override fun createSnapshot(): ContextReceiverValue<FirClassSymbol<*>> =
         ContextReceiverValueForClass(boundSymbol, type, labelName, useSiteSession, scopeSession, mutable = false, contextReceiverNumber)
-
-    override val isContextReceiver: Boolean
-        get() = true
 }
+
+class ImplicitReceiverValueForScript(
+    boundSymbol: FirScriptSymbol,
+    type: ConeKotlinType,
+    labelName: Name?,
+    useSiteSession: FirSession,
+    scopeSession: ScopeSession,
+    mutable: Boolean = true,
+    contextReceiverNumber: Int,
+) : ContextReceiverValue<FirScriptSymbol>(
+    boundSymbol, type, labelName, useSiteSession, scopeSession, mutable, contextReceiverNumber
+) {
+    override fun createSnapshot(): ContextReceiverValue<FirScriptSymbol> =
+        ImplicitReceiverValueForScript(boundSymbol, type, labelName, useSiteSession, scopeSession, mutable = false, contextReceiverNumber)
+}
+

compiler/fir/providers/src/org/jetbrains/kotlin/fir/scopes/impl/FirScriptDeclarationsScope.kt

@@ -0,0 +1,84 @@
+/*
+ * Copyright 2010-2022 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.fir.scopes.impl
+
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.declarations.utils.isSynthetic
+import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
+import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutorByMap
+import org.jetbrains.kotlin.fir.scopes.FirContainingNamesAwareScope
+import org.jetbrains.kotlin.fir.symbols.impl.*
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.name.SpecialNames
+
+class FirScriptDeclarationsScope(
+    val useSiteSession: FirSession,
+    val script: FirScript,
+) : FirContainingNamesAwareScope() {
+
+    private val callablesIndex: Map<Name, List<FirCallableSymbol<*>>> = run {
+        val result = mutableMapOf<Name, MutableList<FirCallableSymbol<*>>>()
+        loop@ for (statement in script.statements) {
+            if (statement is FirCallableDeclaration) {
+                val name = when (statement) {
+                    is FirVariable -> if (statement.isSynthetic) continue@loop else statement.name
+                    is FirSimpleFunction -> statement.name
+                    // TODO: destructuring decl
+                    else -> continue@loop
+                }
+                result.getOrPut(name) { mutableListOf() } += statement.symbol
+            }
+        }
+        result
+    }
+
+    private val classIndex: Map<Name, FirRegularClassSymbol> = run {
+        val result = mutableMapOf<Name, FirRegularClassSymbol>()
+        for (declaration in script.statements) {
+            if (declaration is FirRegularClass) {
+                result[declaration.name] = declaration.symbol
+            }
+        }
+        result
+    }
+
+    override fun processFunctionsByName(name: Name, processor: (FirNamedFunctionSymbol) -> Unit) {
+        if (name == SpecialNames.INIT) return
+        processCallables(name, processor)
+    }
+
+    override fun processPropertiesByName(name: Name, processor: (FirVariableSymbol<*>) -> Unit) {
+        processCallables(name, processor)
+    }
+
+    private inline fun <reified D : FirCallableSymbol<*>> processCallables(
+        name: Name,
+        processor: (D) -> Unit
+    ) {
+        val symbols = callablesIndex[name] ?: emptyList()
+        for (symbol in symbols) {
+            if (symbol is D) {
+                processor(symbol)
+            }
+        }
+    }
+
+    override fun getCallableNames(): Set<Name> {
+        return callablesIndex.keys
+    }
+
+    override fun processClassifiersByNameWithSubstitution(
+        name: Name,
+        processor: (FirClassifierSymbol<*>, ConeSubstitutor) -> Unit
+    ) {
+        val matchedClass = classIndex[name] ?: return
+        val substitution = matchedClass.typeParameterSymbols.associateWith { it.toConeType() }
+        processor(matchedClass, ConeSubstitutorByMap(substitution, useSiteSession))
+    }
+
+    override fun getClassifierNames(): Set<Name> = classIndex.keys
+}