Improve declaring type information for Kotlin top-level functions

rewrite-kotlin/src/main/kotlin/org/openrewrite/kotlin/KotlinTypeMapping.kt

@@ -39,9 +39,11 @@ import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
 import org.jetbrains.kotlin.fir.references.FirSuperReference
 import org.jetbrains.kotlin.fir.references.toResolvedBaseSymbol
 import org.jetbrains.kotlin.fir.resolve.calls.FirSyntheticFunctionSymbol
+import org.jetbrains.kotlin.fir.resolve.providers.firProvider
 import org.jetbrains.kotlin.fir.resolve.providers.toSymbol
 import org.jetbrains.kotlin.fir.resolve.toFirRegularClass
 import org.jetbrains.kotlin.fir.resolve.toSymbol
+import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.SymbolInternals
 import org.jetbrains.kotlin.fir.symbols.impl.*
 import org.jetbrains.kotlin.fir.types.*
@@ -59,11 +61,11 @@ import org.openrewrite.java.JavaTypeMapping
 import org.openrewrite.java.internal.JavaTypeCache
 import org.openrewrite.java.tree.JavaType
 import org.openrewrite.java.tree.JavaType.*
+import org.openrewrite.java.tree.JavaType.Array
 import org.openrewrite.java.tree.TypeUtils
 import org.openrewrite.kotlin.KotlinTypeSignatureBuilder.Companion.convertClassIdToFqn
 import org.openrewrite.kotlin.KotlinTypeSignatureBuilder.Companion.methodName
 import org.openrewrite.kotlin.KotlinTypeSignatureBuilder.Companion.variableName
-import kotlin.collections.ArrayList
 
 @Suppress("DuplicatedCode")
 class KotlinTypeMapping(
@@ -149,7 +151,7 @@ class KotlinTypeMapping(
             }
 
             is FirFile -> {
-                fileType(signature)
+                fileType(type, signature)
             }
 
             is FirFunction -> {
@@ -241,8 +243,18 @@ class KotlinTypeMapping(
         return jt
     }
 
-    private fun fileType(signature: String): JavaType {
+    private fun fileType(file: FirFile, signature: String): JavaType {
+        val functions = buildList {
+            file.declarations.forEach {
+                when (it) {
+                    is FirSimpleFunction -> add(it)
+                    is FirScript -> it.statements.filterIsInstance<FirSimpleFunction>().forEach(::add)
+                    else -> {}
+                }
+            }
+        }
         val fileType = ShallowClass.build(signature)
+            .withMethods(functions.map { methodDeclarationType(it, null) })
         typeCache.put(signature, fileType)
         return fileType
     }
@@ -777,6 +789,8 @@ class KotlinTypeMapping(
                     is Parameterized -> type.type
                     else -> Unknown.getInstance()
                 }
+            } else {
+                declaringType = TypeUtils.asFullyQualified(type(resolvedSymbol.getContainingFile()))
             }
         } else {
             declaringType = TypeUtils.asFullyQualified(type(function.typeRef))
@@ -1376,3 +1390,10 @@ class KotlinTypeMapping(
         }
     }
 }
+
+internal fun FirBasedSymbol<*>.getContainingFile() =
+    when (this) {
+        is FirCallableSymbol<*> -> moduleData.session.firProvider.getFirCallableContainerFile(this)
+        is FirClassLikeSymbol<*> -> moduleData.session.firProvider.getFirClassifierContainerFileIfAny(this)
+        else -> null
+    }

rewrite-kotlin/src/main/kotlin/org/openrewrite/kotlin/KotlinTypeSignatureBuilder.kt

@@ -45,7 +45,6 @@ import org.jetbrains.kotlin.types.Variance
 import org.openrewrite.java.JavaTypeSignatureBuilder
 import org.openrewrite.java.tree.JavaType
 import java.util.*
-import kotlin.collections.HashMap
 
 @Suppress("DuplicatedCode")
 class KotlinTypeSignatureBuilder(private val firSession: FirSession, private val firFile: FirFile) :
@@ -381,6 +380,8 @@ class KotlinTypeSignatureBuilder(private val firSession: FirSession, private val
                 ) {
                     declaringSig = "kotlin.Library"
                 }
+            } else {
+                declaringSig = signature(resolvedSymbol.getContainingFile())
             }
         } else if (sym is FirFunctionSymbol<*>) {
             declaringSig = signature(function.typeRef)

rewrite-kotlin/src/test/java/org/openrewrite/kotlin/KotlinTypeMappingTest.java

@@ -631,19 +631,21 @@ public J.Binary visitBinary(J.Binary binary, AtomicBoolean b) {
 
         @CsvSource(value = {
           // Method type on overload with no named arguments.
-          "foo(\"\", 1, true)~openRewriteFile0Kt{name=foo,return=kotlin.Unit,parameters=[kotlin.String,kotlin.Int,kotlin.Boolean]}",
+          "foo(\"\", 1, true)~org.example.openRewriteFile0Kt{name=foo,return=kotlin.Unit,parameters=[kotlin.String,kotlin.Int,kotlin.Boolean]}",
           // Method type on overload with named arguments.
-          "foo(b = 1)~openRewriteFile0Kt{name=foo,return=kotlin.Unit,parameters=[kotlin.Int,kotlin.Boolean]}",
+          "foo(b = 1)~org.example.openRewriteFile0Kt{name=foo,return=kotlin.Unit,parameters=[kotlin.Int,kotlin.Boolean]}",
           // Method type when named arguments are declared out of order.
-          "foo(trailingLambda = {}, noDefault = true, c = true, b = 1)~openRewriteFile0Kt{name=foo,return=kotlin.Unit,parameters=[kotlin.String,kotlin.Int,kotlin.Boolean,kotlin.Boolean,kotlin.Function0<kotlin.Unit>]}",
+          "foo(trailingLambda = {}, noDefault = true, c = true, b = 1)~org.example.openRewriteFile0Kt{name=foo,return=kotlin.Unit,parameters=[kotlin.String,kotlin.Int,kotlin.Boolean,kotlin.Boolean,kotlin.Function0<kotlin.Unit>]}",
           // Method type with trailing lambda
-          "foo(b = 1, noDefault = true) {}~openRewriteFile0Kt{name=foo,return=kotlin.Unit,parameters=[kotlin.String,kotlin.Int,kotlin.Boolean,kotlin.Boolean,kotlin.Function0<kotlin.Unit>]}"
+          "foo(b = 1, noDefault = true) {}~org.example.openRewriteFile0Kt{name=foo,return=kotlin.Unit,parameters=[kotlin.String,kotlin.Int,kotlin.Boolean,kotlin.Boolean,kotlin.Function0<kotlin.Unit>]}"
         }, delimiter = '~')
         @ParameterizedTest
         void methodInvocationWithDefaults(String invocation, String methodType) {
             rewriteRun(
               kotlin(
                 """
+                  package org.example
+
                   fun <T> foo(b: T, c: Boolean = true) {
                       foo("", b, true, c) {}
                   }
@@ -656,7 +658,7 @@ fun m() {
                       %s
                   }
                   """.formatted(invocation), spec -> spec.afterRecipe(cu -> {
-                    MethodMatcher matcher = new MethodMatcher("* foo(..)");
+                    MethodMatcher matcher = new MethodMatcher("*..* foo(..)");
                     AtomicBoolean methodFound = new AtomicBoolean(false);
                     new KotlinIsoVisitor<AtomicBoolean>() {
                         @Override
@@ -670,7 +672,14 @@ public J.MethodDeclaration visitMethodDeclaration(J.MethodDeclaration method, At
                         @Override
                         public J.MethodInvocation visitMethodInvocation(J.MethodInvocation method, AtomicBoolean found) {
                             if (matcher.matches(method)) {
-                                assertThat(method.getMethodType().toString()).isEqualTo(methodType);
+                                assertThat(method.getMethodType()).satisfies(m -> {
+                                    assertThat(m.toString()).isEqualTo(methodType);
+                                    assertThat(m.getDeclaringType())
+                                      .satisfies(it -> {
+                                          assertThat(it.getFullyQualifiedName()).isEqualTo("org.example.openRewriteFile0Kt");
+                                          assertThat(it.getMethods()).extracting(JavaType.Method::getName).containsExactlyInAnyOrder("foo", "foo", "foo", "m");
+                                      });
+                                });
                                 found.set(true);
                             }
                             return super.visitMethodInvocation(method, found);