From 1cfbc8e86d665518809b3692ae8f78219d1f62cd Mon Sep 17 00:00:00 2001
From: Ian Lynagh <igfoo@github.com>
Date: Thu, 22 Aug 2024 18:08:54 +0100
Subject: [PATCH] KE2: Handle multiple files

---
 .../src/main/kotlin/KotlinExtractor.kt        | 20 ++++++++++++-------
 1 file changed, 13 insertions(+), 7 deletions(-)

diff --git a/java/kotlin-extractor2/src/main/kotlin/KotlinExtractor.kt b/java/kotlin-extractor2/src/main/kotlin/KotlinExtractor.kt
index 870f470a61d..be7d3537a78 100644
--- a/java/kotlin-extractor2/src/main/kotlin/KotlinExtractor.kt
+++ b/java/kotlin-extractor2/src/main/kotlin/KotlinExtractor.kt
@@ -45,14 +45,20 @@ fun main(args : Array<String>) {
         }
     }
 
-   val ktFile = session.modulesWithFiles.getValue(sourceModule).single() as KtFile
-    analyze(ktFile) {
-        val c = ktFile.getDeclarations()[0] as KtClass
-        for (d: KtDeclaration in c.getDeclarations()) {
-            val f = d as KtFunction
-            if (f.name == "f") {
-                dumpFunction(f)
+    val psiFiles = session.modulesWithFiles.getValue(sourceModule)
+    for (psiFile in psiFiles) {
+        if (psiFile is KtFile) {
+            analyze(psiFile) {
+                val c = psiFile.getDeclarations()[0] as KtClass
+                for (d: KtDeclaration in c.getDeclarations()) {
+                    val f = d as KtFunction
+                    if (f.name == "f") {
+                        dumpFunction(f)
+                    }
+                }
             }
+        } else {
+            System.out.println("Warning: Not a KtFile")
         }
     }
 }