From e7f9d60be886e73f451a8f7c56554e9059c91d47 Mon Sep 17 00:00:00 2001
From: hgranthorner <37941012+hgranthorner@users.noreply.github.com>
Date: Wed, 5 Aug 2026 16:44:16 -0400
Subject: [PATCH] Complete symbols loaded by imported modules

---
 .../hgh/jai/reference/JaiSymbolResolver.kt    | 44 +++++++++++++++----
 .../hgh/jai/completion/JaiCompletionTest.kt   | 16 +++++++
 2 files changed, 52 insertions(+), 8 deletions(-)

diff --git a/src/main/kotlin/dev/hgh/jai/reference/JaiSymbolResolver.kt b/src/main/kotlin/dev/hgh/jai/reference/JaiSymbolResolver.kt
index 18c51c7..1beab1e 100644
--- a/src/main/kotlin/dev/hgh/jai/reference/JaiSymbolResolver.kt
+++ b/src/main/kotlin/dev/hgh/jai/reference/JaiSymbolResolver.kt
@@ -50,6 +50,7 @@ internal object JaiSymbolResolver {
 
     private data class ImportBinding(
         val file: PsiFile,
+        val directive: String,
         val scope: PsiElement?,
         val visibleFrom: Int,
         val boundName: String?,
@@ -244,6 +245,7 @@ internal object JaiSymbolResolver {
 
                     ImportBinding(
                         file = target,
+                        directive = name,
                         scope = PsiTreeUtil.getParentOfType(literal, JaiBlock::class.java),
                         visibleFrom = literal.textOffset,
                         boundName = boundName,
@@ -288,6 +290,7 @@ internal object JaiSymbolResolver {
                 val targetPsi = PsiManager.getInstance(file.project).findFile(target) ?: return@mapNotNull null
                 ImportBinding(
                     file = targetPsi,
+                    directive = directive,
                     scope = null,
                     visibleFrom = match.range.first,
                     boundName = null,
@@ -331,14 +334,39 @@ internal object JaiSymbolResolver {
             .toSet()
     }
 
-    private fun topLevelDeclarations(file: PsiFile): List<JaiDeclName> =
-        PsiTreeUtil
-            .findChildrenOfType(file, JaiDeclaration::class.java)
-            .asSequence()
-            .flatMap { declaration -> declaration.declNames.declNameList.asSequence() }
-            .filter { PsiTreeUtil.getParentOfType(it, JaiBlock::class.java) == null }
-            .sortedBy { it.textOffset }
-            .toList()
+    /**
+     * Returns declarations exported by a file, including files it textually loads. A module's
+     * public surface is often assembled by loading several sibling files (for example, Basic
+     * loads Print.jai). Do not follow #import here: an imported module's private dependencies
+     * must not become unqualified candidates in the importing file.
+     */
+    private fun topLevelDeclarations(file: PsiFile): List<JaiDeclName> = topLevelDeclarations(file, linkedSetOf())
+
+    private fun topLevelDeclarations(
+        file: PsiFile,
+        visited: MutableSet<String>,
+    ): List<JaiDeclName> {
+        val fileKey = file.virtualFile?.path ?: file.name
+        if (!visited.add(fileKey)) return emptyList()
+
+        val directDeclarations =
+            PsiTreeUtil
+                .findChildrenOfType(file, JaiDeclaration::class.java)
+                .asSequence()
+                .flatMap { declaration -> declaration.declNames.declNameList.asSequence() }
+                .filter { PsiTreeUtil.getParentOfType(it, JaiBlock::class.java) == null }
+                .sortedBy { it.textOffset }
+                .toList()
+
+        val loadedDeclarations =
+            importBindings(file)
+                .asSequence()
+                .filter { it.directive == "#load" }
+                .flatMap { topLevelDeclarations(it.file, visited).asSequence() }
+                .toList()
+
+        return directDeclarations + loadedDeclarations
+    }
 
     private fun isVisible(
         symbol: ScopedSymbol,
diff --git a/src/test/kotlin/dev/hgh/jai/completion/JaiCompletionTest.kt b/src/test/kotlin/dev/hgh/jai/completion/JaiCompletionTest.kt
index df088cf..09abdc7 100644
--- a/src/test/kotlin/dev/hgh/jai/completion/JaiCompletionTest.kt
+++ b/src/test/kotlin/dev/hgh/jai/completion/JaiCompletionTest.kt
@@ -111,6 +111,22 @@ class JaiCompletionTest : BasePlatformTestCase() {
         assertTrue("expected alloc in $strings", strings!!.contains("alloc"))
     }
 
+    fun testCompletesSymbolsLoadedByImportedModule() {
+        myFixture.configureByText(
+            "nested-module-consumer.jai",
+            """
+            #import "Basic";
+            main :: () { pri<caret>("hello"); }
+            """.trimIndent(),
+        )
+
+        myFixture.complete(CompletionType.BASIC, 1)
+        assertTrue(
+            "expected print from a file loaded by the imported Basic module: ${myFixture.lookupElementStrings}",
+            myFixture.lookupElementStrings.orEmpty().contains("print"),
+        )
+    }
+
     fun testCompletesSymbolsFromLoadedFile() {
         myFixture.addFileToProject("library.jai", "Helper :: () {}")
         myFixture.configureByText(