[Repl] Some attempts at fixing K2 Repl tests
diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirTypeLazyResolver.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirTypeLazyResolver.kt
index 7cfbe0e..b7b35eb 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirTypeLazyResolver.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirTypeLazyResolver.kt
@@ -90,6 +90,7 @@
             is FirFile,
             is FirTypeAlias,
             is FirScript,
+            is FirReplSnippet,
             is FirRegularClass,
             is FirAnonymousInitializer,
                 -> rawResolve(target)
@@ -115,6 +116,7 @@
             }
 
             is FirScript -> resolveScriptTypes(target)
+            is FirReplSnippet -> resolveReplTypes(target)
             is FirField if (target.origin == FirDeclarationOrigin.Synthetic.DelegateField) -> {
                 // delegated field should be resolved in the same context as super types
                 resolveOutsideClassBody(target, transformer::transformDelegateField)
@@ -163,6 +165,11 @@
         firScript.transformReceivers(transformer, null)
     }
 
+    private fun resolveReplTypes(firSnippet: FirReplSnippet) {
+        // Repl statements are lazy here, unlike Script. Problem?
+        TODO("Unclear if this is correct")
+    }
+
     private fun resolveClassTypes(firClass: FirRegularClass) {
         transformer.transformClassTypeParameters(firClass, null)
         transformer.withScopeCleanup {
diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/FirElementFinder.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/FirElementFinder.kt
index a292c19..190b9ef 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/FirElementFinder.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/FirElementFinder.kt
@@ -350,6 +350,7 @@
             is FirAnonymousInitializer -> SpecialNames.ANONYMOUS
             is FirTypeAlias -> declaration.name
             is FirCodeFragment, is FirDanglingModifierList -> SpecialNames.NO_NAME_PROVIDED
+            is FirReplSnippet -> declaration.name
 
             is FirFile,
             is FirAnonymousFunction,
@@ -357,7 +358,6 @@
             is FirPropertyAccessor,
             is FirAnonymousObject,
             is FirReceiverParameter,
-            is FirReplSnippet,
             is FirTypeParameter,
                 -> errorWithFirSpecificEntries("Unexpected declaration ${declaration::class.simpleName}", fir = declaration)
         }
diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/containingFileUtils.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/containingFileUtils.kt
index f0383fc..8da6473 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/containingFileUtils.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/containingFileUtils.kt
@@ -43,6 +43,14 @@
                 ?: error("Fir file for dandling modifier list cannot be null")
         }
         is FirReceiverParameter -> containingDeclarationSymbol.fir.getContainingFile()
+        is FirReplSnippet -> {
+            val ktFile = psi?.containingFile as? KtFile
+                ?: error("File for FirReplSnippet cannot be null")
+            val moduleComponents = llFirResolvableSession?.moduleComponents
+                ?: error("LLFirResolvableModuleSession for FirReplSnippet cannot be null")
+            moduleComponents.cache.getCachedFirFile(ktFile)
+                ?: error("Fir file for FirReplSnippet cannot be null")
+        }
         else -> errorWithFirSpecificEntries("Unsupported declaration ${this::class}", fir = this)
     }
 }
diff --git a/compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/PsiRawFirBuilder.kt b/compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/PsiRawFirBuilder.kt
index b66aa50..1cc768f 100644
--- a/compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/PsiRawFirBuilder.kt
+++ b/compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/PsiRawFirBuilder.kt
@@ -1427,10 +1427,12 @@
                 name = snippetName
                 symbol = snippetSymbol
 
+                var lastStatement: Any? = null // Unclear how to get this from the LazyFirL
                 body = buildOrLazyBlock {
                     withContainerSymbol(snippetSymbol, isLocal = true) {
                         buildBlock {
                             script.declarations.forEach { declaration ->
+                                ;
                                 when (declaration) {
                                     is KtScriptInitializer -> {
                                         val initializer = buildAnonymousInitializer(
@@ -1441,10 +1443,12 @@
                                         )
 
                                         statements.addAll(initializer.body!!.statements)
+                                        lastStatement = initializer.body!!.statements.lastOrNull()
                                     }
                                     is KtDestructuringDeclaration -> {
                                         val destructuringContainerVar = buildScriptDestructuringDeclaration(declaration)
                                         statements.add(destructuringContainerVar)
+                                        lastStatement = destructuringContainerVar
 
                                         addDestructuringVariables(
                                             statements,
@@ -1462,12 +1466,14 @@
                                         withForcedLocalContext {
                                             val firProperty = convertProperty(declaration, null, forceLocal = true)
                                             statements.add(firProperty)
+                                            lastStatement = firProperty
                                         }
                                     }
                                     else -> {
                                         val firStatement = declaration.toFirStatement()
                                         if (firStatement is FirDeclaration) {
                                             statements.add(firStatement)
+                                            lastStatement = firStatement
                                         } else {
                                             error("unexpected declaration type in script")
                                         }
@@ -1479,7 +1485,7 @@
                 }
                 // TODO: proper lazy support - see the script
                 resultTypeRef =
-                    if (body.statements.lastOrNull() is FirDeclaration) implicitUnitType else FirImplicitTypeRefImplWithoutSource
+                    if (lastStatement is FirDeclaration) implicitUnitType else FirImplicitTypeRefImplWithoutSource
                 setup()
             }
         }