Preserve navigation targets for module completions

src/main/kotlin/dev/hgh/jai/completion/JaiCompletionContributor.kt

@@ -128,8 +128,7 @@ private class JaiCompletionProvider : CompletionProvider<CompletionParameters>()
             .forEach { candidate ->
                 resultSet.addElement(
                     LookupElementBuilder
-                        .create(candidate.name)
-                        .withPsiElement(candidate.element)
+                        .createWithSmartPointer(candidate.name, candidate.element)
                         .withTypeText(candidate.kind),
                 )
             }
@@ -152,8 +151,7 @@ private class JaiCompletionProvider : CompletionProvider<CompletionParameters>()
             .forEach { candidate ->
                 resultSet.addElement(
                     LookupElementBuilder
-                        .create(candidate.name)
-                        .withPsiElement(candidate.element)
+                        .createWithSmartPointer(candidate.name, candidate.element)
                         .withTypeText(candidate.kind),
                 )
             }

src/test/kotlin/dev/hgh/jai/completion/JaiCompletionTest.kt

@@ -120,11 +120,21 @@ class JaiCompletionTest : BasePlatformTestCase() {
             """.trimIndent(),
         )
 
-        myFixture.complete(CompletionType.BASIC, 1)
+        val lookupElements = myFixture.complete(CompletionType.BASIC, 1)
         assertTrue(
             "expected print from a file loaded by the imported Basic module: ${myFixture.lookupElementStrings}",
             myFixture.lookupElementStrings.orEmpty().contains("print"),
         )
+        val printLookup = lookupElements.firstOrNull { it.lookupString == "print" }
+        assertNotNull("print completion should retain its declaration target", printLookup)
+        assertEquals(
+            "${System.getProperty("user.home")}/.local/jai/modules/Basic/Print.jai",
+            printLookup!!
+                .psiElement
+                ?.containingFile
+                ?.virtualFile
+                ?.path,
+        )
     }
 
     fun testCompletesSymbolsFromLoadedFile() {

src/test/kotlin/dev/hgh/jai/reference/JaiReferenceTest.kt

@@ -1,5 +1,6 @@
 package dev.hgh.jai.reference
 
+import com.intellij.codeInsight.TargetElementUtil
 import com.intellij.openapi.util.TextRange
 import com.intellij.openapi.vfs.LocalFileSystem
 import com.intellij.psi.PsiElement
@@ -127,6 +128,38 @@ class JaiReferenceTest : BasePlatformTestCase() {
         )
     }
 
+    fun testGotoDeclarationResolvesProcedureFromFileLoadedByImportedModule() {
+        myFixture.configureByText(
+            "nested-module-consumer.jai",
+            """
+            #import "Basic";
+            main :: () { print<caret>("hello"); }
+            """.trimIndent(),
+        )
+        val source = myFixture.file
+        val referenceElement =
+            PsiTreeUtil
+                .findChildrenOfType(source, JaiRefExpr::class.java)
+                .single { it.text == "print" }
+
+        val resolved = onlyReference(referenceElement).resolve()
+
+        assertEquals("print", resolved?.text)
+        val expectedPath = "${System.getProperty("user.home")}/.local/jai/modules/Basic/Print.jai"
+        assertEquals(expectedPath, resolved?.containingFile?.virtualFile?.path)
+
+        val targetElementUtil = TargetElementUtil.getInstance()
+        val gotoTarget =
+            targetElementUtil.findTargetElement(
+                myFixture.editor,
+                targetElementUtil.referenceSearchFlags,
+                myFixture.caretOffset,
+            )
+        assertNotNull("go to declaration should find the loaded procedure", gotoTarget)
+        assertEquals("print", gotoTarget!!.text)
+        assertEquals(expectedPath, gotoTarget.containingFile?.virtualFile?.path)
+    }
+
     fun testUnresolvedSymbolReferenceHasNoTarget() {
         val source = myFixture.addFileToProject("unresolved.jai", "main :: () { missing(); }")
         val referenceElement =