Phase 5: add completion and symbol navigation

AGENTS.md

@@ -174,10 +174,13 @@ Tier 3: parsed 714 files, 714 clean (100.0%), 0 with errors, 0 PsiErrorElements
   `JaiReferenceHostMixin` from the BNF. Do not remove the mixin and expect a
   `PsiReferenceContributor` to be queried through the standard reference service.
 
-### Not done — pick up here
+### In progress — pick up here
 
-1. **Phase 5** — completion (keywords, directives, module names), rename, and
+1. **Phase 5 completion and navigation** — basic completion for keywords,
-   find-usages; see `docs/BUILD_PLAN.md`.
+   compiler directives, `#import`/`#load` module paths is implemented in
+   `dev.hgh.jai.completion.JaiCompletionContributor`. Identifier references now
+   resolve procedures/types in the same file and transitively loaded/imported
+   files, with headless fixture coverage. Rename and find-usages remain.
 2. **Phases 6+** — formatter, inspections, compiler integration, and other
    optional work remain. `#asm` bodies are intentionally consumed opaquely, so
    nothing inside them has PSI yet.

src/main/gen/dev/hgh/jai/psi/impl/JaiRefExprImpl.java

@@ -8,15 +8,15 @@ import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiElementVisitor;
 import com.intellij.psi.util.PsiTreeUtil;
 import static dev.hgh.jai.psi.JaiTypes.*;
+import dev.hgh.jai.psi.mixin.JaiReferenceHostMixin;
 import dev.hgh.jai.psi.*;
 
-public class JaiRefExprImpl extends JaiExprImpl implements JaiRefExpr {
+public class JaiRefExprImpl extends JaiReferenceHostMixin implements JaiRefExpr {
 
   public JaiRefExprImpl(@NotNull ASTNode node) {
     super(node);
   }
 
-  @Override
   public void accept(@NotNull JaiVisitor visitor) {
     visitor.visitRefExpr(this);
   }

src/main/grammar/Jai.bnf

@@ -361,6 +361,7 @@ primaryExpr ::= procLiteralExpr
               | uninitializedExpr
 
 refExpr ::= IDENT
+            { mixin="dev.hgh.jai.psi.mixin.JaiReferenceHostMixin" }
 literalExpr ::= NUMBER | STRING | HERE_STRING | 'true' | 'false' | 'null' | 'context'
                 { mixin="dev.hgh.jai.psi.mixin.JaiReferenceHostMixin" }
 uninitializedExpr ::= '---' | '--'

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

@@ -0,0 +1,370 @@
+package dev.hgh.jai.completion
+
+import com.intellij.codeInsight.completion.CompletionContributor
+import com.intellij.codeInsight.completion.CompletionParameters
+import com.intellij.codeInsight.completion.CompletionProvider
+import com.intellij.codeInsight.completion.CompletionResultSet
+import com.intellij.codeInsight.completion.CompletionType
+import com.intellij.codeInsight.lookup.LookupElementBuilder
+import com.intellij.openapi.project.DumbAware
+import com.intellij.openapi.vfs.LocalFileSystem
+import com.intellij.openapi.vfs.VirtualFile
+import com.intellij.patterns.PlatformPatterns
+import com.intellij.psi.tree.IElementType
+import com.intellij.psi.util.PsiTreeUtil
+import com.intellij.util.ProcessingContext
+import dev.hgh.jai.lexer.JaiLexer
+import dev.hgh.jai.lexer.JaiTokenTypes
+import dev.hgh.jai.psi.JaiDirectiveExpr
+
+/** Directives and their adjacent comma flags before an import/load string. */
+private val DIRECTIVE_WITH_FLAGS =
+    Regex("""#(import|load)((?:\s*,\s*[A-Za-z_][A-Za-z0-9_]*)*)\s*$""")
+
+/** Basic completion for Jai's context-free language and module names. */
+class JaiCompletionContributor :
+    CompletionContributor(),
+    DumbAware {
+    init {
+        extend(
+            CompletionType.BASIC,
+            PlatformPatterns.psiElement(),
+            JaiCompletionProvider(),
+        )
+    }
+}
+
+private class JaiCompletionProvider : CompletionProvider<CompletionParameters>() {
+    override fun addCompletions(
+        parameters: CompletionParameters,
+        context: ProcessingContext,
+        resultSet: CompletionResultSet,
+    ) {
+        val completionContext = JaiCompletionContext(parameters)
+        val moduleDirective = completionContext.moduleDirective
+        when {
+            moduleDirective != null -> {
+                addModuleNames(parameters, moduleDirective, completionContext.prefix, resultSet)
+            }
+
+            completionContext.isCommentOrStringOrHereString -> {
+                return
+            }
+
+            completionContext.isDirectivePrefix -> {
+                addDirectives(completionContext.prefix, resultSet)
+            }
+
+            else -> {
+                addKeywords(completionContext.prefix, resultSet)
+            }
+        }
+    }
+
+    private fun addKeywords(
+        prefix: String,
+        resultSet: CompletionResultSet,
+    ) {
+        JaiTokenTypes.KEYWORD_MAP.keys
+            .filter { it.startsWith(prefix) }
+            .sorted()
+            .forEach { keyword ->
+                resultSet.addElement(
+                    LookupElementBuilder
+                        .create(keyword)
+                        .withTypeText("keyword"),
+                )
+            }
+    }
+
+    private fun addDirectives(
+        prefix: String,
+        resultSet: CompletionResultSet,
+    ) {
+        JaiCompletionCatalog.DIRECTIVES
+            .map { "#$it" }
+            .filter { it.startsWith(prefix) }
+            .forEach { directive ->
+                val name = directive.removePrefix("#")
+                resultSet.addElement(
+                    LookupElementBuilder
+                        // Keep the insertion text free of `#`: the prefix starts after
+                        // `#` for a directive token, so inserting `#import` would yield
+                        // `##import`. The alternate lookup string still matches `#im`.
+                        .create(name)
+                        .withPresentableText(directive)
+                        .withLookupString(directive)
+                        .withTypeText("directive"),
+                )
+            }
+    }
+
+    private fun addModuleNames(
+        parameters: CompletionParameters,
+        directive: String,
+        prefix: String,
+        resultSet: CompletionResultSet,
+    ) {
+        val sourceFile = parameters.originalFile.virtualFile
+        val variants =
+            if (directive == "#import") {
+                JaiModuleCompletion.moduleNames(parameters.originalFile.project, sourceFile)
+            } else {
+                JaiModuleCompletion.loadPaths(sourceFile)
+            }
+        variants
+            .filter { it.startsWith(prefix) }
+            .forEach { path ->
+                resultSet.addElement(
+                    LookupElementBuilder
+                        .create(path)
+                        .withTypeText(if (directive == "#import") "module" else "file"),
+                )
+            }
+    }
+}
+
+private class JaiCompletionContext(
+    private val parameters: CompletionParameters,
+) {
+    private val text =
+        parameters.editor.document.charsSequence
+            .toString()
+    private val offset = parameters.offset.coerceIn(0, text.length)
+    private val token = tokenAtCaret(text, offset)
+
+    val isCommentOrStringOrHereString: Boolean
+        get() =
+            token?.type in
+                setOf(
+                    JaiTokenTypes.LINE_COMMENT,
+                    JaiTokenTypes.BLOCK_COMMENT,
+                    JaiTokenTypes.STRING,
+                    JaiTokenTypes.HERE_STRING,
+                    JaiTokenTypes.NOTE,
+                )
+
+    val isDirectivePrefix: Boolean
+        get() = token?.type === JaiTokenTypes.DIRECTIVE || token?.type === JaiTokenTypes.HASH
+
+    val prefix: String
+        get() {
+            val current = token ?: return ""
+            val start =
+                when {
+                    moduleDirective != null && current.type === JaiTokenTypes.STRING -> current.start + 1
+                    isDirectivePrefix -> current.start
+                    current.type === JaiTokenTypes.IDENT || current.type in JaiTokenTypes.KEYWORDS -> current.start
+                    else -> return ""
+                }
+            return if (start <= offset) text.substring(start, offset) else ""
+        }
+
+    val moduleDirective: String?
+        get() {
+            if (token?.type !== JaiTokenTypes.STRING || token.end < offset) return null
+            if (isAfterClosedString()) return null
+
+            val beforeString = text.substring(0, token.start)
+            val directiveMatch = DIRECTIVE_WITH_FLAGS.find(beforeString)
+            val directive =
+                PsiTreeUtil
+                    .getParentOfType(parameters.position, JaiDirectiveExpr::class.java)
+                    ?.node
+                    ?.findChildByType(JaiTokenTypes.DIRECTIVE)
+                    ?.text
+                    ?: directiveMatch?.let { "#${it.groupValues[1]}" }
+            if (directive != "#import" && directive != "#load") return null
+            if (directive == "#import" && hasStringFlag(directiveMatch?.groupValues?.get(2).orEmpty())) {
+                return null
+            }
+            return directive
+        }
+
+    private fun hasStringFlag(flags: String): Boolean =
+        Regex("""[A-Za-z_][A-Za-z0-9_]*""")
+            .findAll(flags)
+            .any { it.value == "string" }
+
+    private fun isAfterClosedString(): Boolean {
+        if (token == null || token.end != offset || token.start >= offset) return false
+        val tokenText = text.substring(token.start, offset)
+        return tokenText.length > 1 && tokenText.last() == '"' && isUnescapedQuote(tokenText.lastIndex)
+    }
+
+    private fun isUnescapedQuote(index: Int): Boolean {
+        var backslashes = 0
+        var position = index - 1
+        while (position >= 0 && text[token!!.start + position] == '\\') {
+            backslashes++
+            position--
+        }
+        return backslashes % 2 == 0
+    }
+
+    private data class LexedToken(
+        val type: IElementType,
+        val start: Int,
+        val end: Int,
+    )
+
+    private fun tokenAtCaret(
+        text: String,
+        offset: Int,
+    ): LexedToken? {
+        val lexer = JaiLexer()
+        lexer.start(text, 0, offset, 0)
+        var last: LexedToken? = null
+        while (lexer.tokenType != null) {
+            val current =
+                LexedToken(
+                    lexer.tokenType!!,
+                    lexer.tokenStart,
+                    lexer.tokenEnd,
+                )
+            last = current
+            if (current.end >= offset) return current
+            lexer.advance()
+        }
+        return last
+    }
+}
+
+private object JaiCompletionCatalog {
+    /** Core/compiler directives observed in the local Jai distribution. */
+    val DIRECTIVES: List<String> =
+        setOf(
+            "add_context",
+            "align",
+            "as",
+            "asm",
+            "assert",
+            "bake",
+            "bake_arguments",
+            "bake_constants",
+            "body_text",
+            "bytes",
+            "c_call",
+            "caller_location",
+            "char",
+            "code",
+            "complete",
+            "compiler",
+            "cpp_method",
+            "cpp_return_type_is_non_pod",
+            "define",
+            "deprecated",
+            "discard",
+            "dump",
+            "else",
+            "elsewhere",
+            "expand",
+            "file",
+            "filepath",
+            "foreign",
+            "if",
+            "ifdef",
+            "import",
+            "include",
+            "insert",
+            "intrinsic",
+            "library",
+            "load",
+            "location",
+            "module_parameters",
+            "modify",
+            "no_abc",
+            "no_alias",
+            "no_aoc",
+            "no_context",
+            "no_debug",
+            "no_padding",
+            "no_reset",
+            "overlay",
+            "place",
+            "placeholder",
+            "placeholders",
+            "poke_name",
+            "pragma",
+            "procedure_of_call",
+            "program_export",
+            "run",
+            "run_and_insert",
+            "scope_export",
+            "scope_file",
+            "scope_module",
+            "specified",
+            "string",
+            "symmetric",
+            "this",
+            "through",
+            "type",
+            "type_info_none",
+            "type_info_no_size_complaint",
+            "type_info_procedures_are_void_pointers",
+            "undef",
+            "version",
+        ).sorted()
+}
+
+private object JaiModuleCompletion {
+    private const val JAI_EXTENSION = "jai"
+    private val fileSystem: LocalFileSystem
+        get() = LocalFileSystem.getInstance()
+
+    fun moduleNames(
+        project: com.intellij.openapi.project.Project,
+        sourceFile: VirtualFile?,
+    ): List<String> =
+        collectCandidates(candidateRoots(project, sourceFile, moduleMode = true)) { child ->
+            when {
+                child.isDirectory && child.findChild("module.jai") != null -> child.name
+                !child.isDirectory && child.extension == JAI_EXTENSION -> child.nameWithoutExtension
+                else -> null
+            }
+        }
+
+    fun loadPaths(sourceFile: VirtualFile?): List<String> =
+        collectCandidates(
+            listOfNotNull(sourceFile?.parent),
+        ) { child ->
+            if (!child.isDirectory && child.extension == JAI_EXTENSION) child.name else null
+        }
+
+    private fun collectCandidates(
+        roots: List<VirtualFile>,
+        nameOf: (VirtualFile) -> String?,
+    ): List<String> {
+        val names = linkedSetOf<String>()
+        roots.forEach { root ->
+            root.children
+                .asSequence()
+                .mapNotNull(nameOf)
+                .forEach(names::add)
+        }
+        return names.sorted()
+    }
+
+    private fun candidateRoots(
+        project: com.intellij.openapi.project.Project,
+        sourceFile: VirtualFile?,
+        moduleMode: Boolean,
+    ): List<VirtualFile> {
+        val roots = linkedMapOf<String, VirtualFile>()
+
+        fun add(root: VirtualFile?) {
+            if (root != null && root.isValid && root.isDirectory) roots.putIfAbsent(root.path, root)
+        }
+
+        if (moduleMode) {
+            add(sourceFile?.parent?.findChild("modules"))
+            val projectRoot = project.basePath?.let(fileSystem::findFileByPath)
+            add(projectRoot?.findChild("modules"))
+            add(projectRoot)
+            add(fileSystem.findFileByPath("${System.getProperty("user.home")}/.local/jai/modules"))
+        } else {
+            add(sourceFile?.parent)
+        }
+        return roots.values.toList()
+    }
+}

src/main/kotlin/dev/hgh/jai/reference/JaiReferenceContributor.kt

@@ -6,24 +6,35 @@ import com.intellij.openapi.vfs.LocalFileSystem
 import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.patterns.PlatformPatterns
 import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiFile
 import com.intellij.psi.PsiManager
 import com.intellij.psi.PsiReference
 import com.intellij.psi.PsiReferenceBase
 import com.intellij.psi.PsiReferenceContributor
 import com.intellij.psi.PsiReferenceProvider
 import com.intellij.psi.PsiReferenceRegistrar
+import com.intellij.psi.PsiReferenceService
+import com.intellij.psi.util.PsiTreeUtil
 import com.intellij.util.ProcessingContext
 import dev.hgh.jai.lexer.JaiTokenTypes
+import dev.hgh.jai.psi.JaiBlock
+import dev.hgh.jai.psi.JaiDeclName
+import dev.hgh.jai.psi.JaiDeclaration
 import dev.hgh.jai.psi.JaiDirectiveExpr
 import dev.hgh.jai.psi.JaiLiteralExpr
+import dev.hgh.jai.psi.JaiRefExpr
 
-/** Adds file references to string operands of Jai's import and load directives. */
+/** Adds file references for directives and symbol references for Jai identifiers. */
 class JaiReferenceContributor : PsiReferenceContributor() {
     override fun registerReferenceProviders(registrar: PsiReferenceRegistrar) {
         registrar.registerReferenceProvider(
             PlatformPatterns.psiElement(JaiLiteralExpr::class.java),
             JaiModuleReferenceProvider(),
         )
+        registrar.registerReferenceProvider(
+            PlatformPatterns.psiElement(JaiRefExpr::class.java),
+            JaiSymbolReferenceProvider(),
+        )
     }
 }
 
@@ -78,6 +89,129 @@ private class JaiModuleReferenceProvider : PsiReferenceProvider() {
         }
 }
 
+/** Adds references from identifiers to declarations in the same or imported Jai files. */
+private class JaiSymbolReferenceProvider : PsiReferenceProvider() {
+    override fun getReferencesByElement(
+        element: PsiElement,
+        context: ProcessingContext,
+    ): Array<PsiReference> {
+        val reference = element as? JaiRefExpr ?: return PsiReference.EMPTY_ARRAY
+        if (reference.text.isEmpty()) return PsiReference.EMPTY_ARRAY
+        return arrayOf(JaiSymbolReference(reference))
+    }
+}
+
+private class JaiSymbolReference(
+    private val sourceElement: JaiRefExpr,
+) : PsiReferenceBase<JaiRefExpr>(sourceElement, TextRange(0, sourceElement.textLength), true) {
+    override fun resolve(): PsiElement? = JaiSymbolResolver.resolve(sourceElement)
+
+    override fun getVariants(): Array<Any> = emptyArray()
+}
+
+private object JaiSymbolResolver {
+    private data class Candidate(
+        val name: JaiDeclName,
+        val scope: JaiBlock?,
+    )
+
+    fun resolve(reference: JaiRefExpr): JaiDeclName? {
+        val file = reference.containingFile ?: return null
+        resolveInFile(file, reference)?.let { return it }
+        return resolveImported(file, reference.text, linkedSetOf())
+    }
+
+    private fun resolveInFile(
+        file: PsiFile,
+        reference: JaiRefExpr,
+    ): JaiDeclName? {
+        val candidates =
+            PsiTreeUtil
+                .findChildrenOfType(file, JaiDeclaration::class.java)
+                .flatMap { declaration ->
+                    declaration.declNames.declNameList.map { name ->
+                        Candidate(name, PsiTreeUtil.getParentOfType(name, JaiBlock::class.java))
+                    }
+                }.filter { candidate ->
+                    symbolName(candidate.name) == reference.text && isVisible(candidate, reference)
+                }.sortedWith(
+                    compareByDescending<Candidate> { scopeDepth(it.scope) }
+                        .thenByDescending {
+                            it.name.textRange.startOffset <= reference.textRange.startOffset
+                        }.thenBy {
+                            kotlin.math.abs(it.name.textRange.startOffset - reference.textRange.startOffset)
+                        },
+                )
+        return candidates.firstOrNull()?.name
+    }
+
+    private fun resolveImported(
+        file: PsiFile,
+        referenceName: String,
+        visited: MutableSet<String>,
+    ): JaiDeclName? {
+        val fileKey = file.virtualFile?.path ?: file.name
+        if (!visited.add(fileKey)) return null
+
+        for (importedFile in importedFiles(file)) {
+            topLevelDeclaration(importedFile, referenceName)?.let { return it }
+            resolveImported(importedFile, referenceName, visited)?.let { return it }
+        }
+        return null
+    }
+
+    private fun topLevelDeclaration(
+        file: PsiFile,
+        referenceName: String,
+    ): JaiDeclName? =
+        PsiTreeUtil
+            .findChildrenOfType(file, JaiDeclaration::class.java)
+            .asSequence()
+            .flatMap { declaration -> declaration.declNames.declNameList.asSequence() }
+            .firstOrNull {
+                PsiTreeUtil.getParentOfType(it, JaiBlock::class.java) == null &&
+                    symbolName(it) == referenceName
+            }
+
+    private fun importedFiles(file: PsiFile): List<PsiFile> =
+        PsiTreeUtil
+            .findChildrenOfType(file, JaiLiteralExpr::class.java)
+            .mapNotNull { literal ->
+                val directive = literal.parent as? JaiDirectiveExpr ?: return@mapNotNull null
+                val name = directive.directiveName() ?: return@mapNotNull null
+                if (name != "#import" && name != "#load") return@mapNotNull null
+                PsiReferenceService
+                    .getService()
+                    .getContributedReferences(literal)
+                    .firstOrNull()
+                    ?.resolve() as? PsiFile
+            }.distinctBy { it.virtualFile?.path ?: it.name }
+
+    private fun isVisible(
+        candidate: Candidate,
+        reference: JaiRefExpr,
+    ): Boolean = candidate.scope == null || PsiTreeUtil.isAncestor(candidate.scope, reference, false)
+
+    private fun scopeDepth(scope: JaiBlock?): Int {
+        var current: PsiElement? = scope
+        var depth = 0
+        while (current != null) {
+            if (current is JaiBlock) depth++
+            current = current.parent
+        }
+        return depth
+    }
+
+    private fun symbolName(name: JaiDeclName): String {
+        val text = name.text
+        return if (text.length >= 2 && text.first() == '`' && text.last() == '`') {
+            text.substring(1, text.length - 1)
+        } else {
+            text
+        }
+    }
+}
+
 private data class JaiModuleTarget(
     val directive: String,
     val path: String,

src/main/resources/META-INF/plugin.xml

@@ -38,6 +38,9 @@
         <psi.referenceContributor
             language="Jai"
             implementation="dev.hgh.jai.reference.JaiReferenceContributor"/>
+        <completion.contributor
+            language="Jai"
+            implementationClass="dev.hgh.jai.completion.JaiCompletionContributor"/>
     </extensions>
 
 </idea-plugin>

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

@@ -0,0 +1,61 @@
+package dev.hgh.jai.completion
+
+import com.intellij.codeInsight.completion.CompletionType
+import com.intellij.testFramework.fixtures.BasePlatformTestCase
+
+class JaiCompletionTest : BasePlatformTestCase() {
+    fun testCompletesKeywords() {
+        myFixture.configureByText("keyword.jai", "ret<caret>")
+
+        myFixture.completeBasic()
+        myFixture.checkResult("return")
+    }
+
+    fun testCompletesDirectivesWithHash() {
+        myFixture.configureByText("directive.jai", "#imp<caret>")
+
+        myFixture.completeBasic()
+        myFixture.checkResult("#import")
+    }
+
+    fun testDoesNotCompleteInsideStringLiterals() {
+        myFixture.configureByText("string.jai", "#char \"ret<caret>\";")
+
+        myFixture.complete(CompletionType.BASIC, 1)
+
+        assertTrue("strings should not offer keyword completion", myFixture.lookupElementStrings.isNullOrEmpty())
+    }
+
+    fun testDoesNotOfferModulesForStringImport() {
+        myFixture.configureByText("string-import.jai", "#import,string \"Bas<caret>\";")
+
+        myFixture.complete(CompletionType.BASIC, 1)
+
+        assertTrue("#import,string should not offer file modules", myFixture.lookupElementStrings.isNullOrEmpty())
+    }
+
+    fun testCompletesImportModulesFromLocalJaiModulesDirectory() {
+        myFixture.configureByText("main.jai", "#import \"Bas<caret>\";")
+
+        val strings = completeStrings()
+
+        assertNotNull("module completion should be available", strings)
+        assertTrue("expected Basic in $strings", strings!!.contains("Basic"))
+        assertTrue("expected Base64 in $strings", strings.contains("Base64"))
+        assertFalse("prefix should filter Bucket_Array from $strings", strings.contains("Bucket_Array"))
+        assertFalse("prefix should filter Bit_Array from $strings", strings.contains("Bit_Array"))
+    }
+
+    fun testCompletesLoadFilesRelativeToSource() {
+        myFixture.addFileToProject("other.jai", "value := 1;")
+        myFixture.configureByText("main.jai", "#load \"oth<caret>\";")
+
+        myFixture.completeBasic()
+        myFixture.checkResult("#load \"other.jai\";")
+    }
+
+    private fun completeStrings(): List<String>? {
+        myFixture.complete(CompletionType.BASIC, 1)
+        return myFixture.lookupElementStrings
+    }
+}

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

@@ -2,6 +2,7 @@ package dev.hgh.jai.reference
 
 import com.intellij.openapi.util.TextRange
 import com.intellij.openapi.vfs.LocalFileSystem
+import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiFile
 import com.intellij.psi.PsiManager
 import com.intellij.psi.PsiReference
@@ -9,6 +10,7 @@ import com.intellij.psi.PsiReferenceService
 import com.intellij.psi.util.PsiTreeUtil
 import com.intellij.testFramework.fixtures.BasePlatformTestCase
 import dev.hgh.jai.psi.JaiLiteralExpr
+import dev.hgh.jai.psi.JaiRefExpr
 
 class JaiReferenceTest : BasePlatformTestCase() {
     fun testLoadResolvesAFileRelativeToTheSource() {
@@ -24,6 +26,117 @@ class JaiReferenceTest : BasePlatformTestCase() {
         assertEquals(TextRange(1, literal.textLength - 1), reference.rangeInElement)
     }
 
+    fun testProcedureReferenceResolvesToSameFileDeclaration() {
+        val source =
+            myFixture.addFileToProject(
+                "procedures.jai",
+                """
+                helper :: () {}
+                main :: () { helper(); }
+                """.trimIndent(),
+            )
+        val referenceElement =
+            PsiTreeUtil
+                .findChildrenOfType(source, JaiRefExpr::class.java)
+                .single { it.text == "helper" }
+
+        val reference = onlyReference(referenceElement)
+
+        assertEquals("helper", reference.resolve()?.text)
+        assertEquals(TextRange(0, referenceElement.textLength), reference.rangeInElement)
+    }
+
+    fun testTypeReferenceResolvesToSameFileDeclaration() {
+        val source =
+            myFixture.addFileToProject(
+                "types.jai",
+                """
+                Point :: struct { x: int; }
+                use_point :: (point: Point) -> int { return point.x; }
+                """.trimIndent(),
+            )
+        val referenceElement =
+            PsiTreeUtil
+                .findChildrenOfType(source, JaiRefExpr::class.java)
+                .single { it.text == "Point" }
+
+        val reference = onlyReference(referenceElement)
+
+        assertEquals("Point", reference.resolve()?.text)
+    }
+
+    fun testSymbolsResolveThroughLoadedFile() {
+        val target =
+            myFixture.addFileToProject(
+                "library.jai",
+                """
+                Helper :: () {}
+                Point :: struct { x: int; }
+                """.trimIndent(),
+            )
+        val source =
+            myFixture.addFileToProject(
+                "consumer.jai",
+                """
+                #load "library.jai";
+                main :: (point: Point) { Helper(); }
+                """.trimIndent(),
+            )
+        val references = PsiTreeUtil.findChildrenOfType(source, JaiRefExpr::class.java)
+        val helperReference = references.single { it.text == "Helper" }
+        val pointReference = references.single { it.text == "Point" }
+
+        assertEquals(
+            target.virtualFile.path,
+            onlyReference(helperReference)
+                .resolve()
+                ?.containingFile
+                ?.virtualFile
+                ?.path,
+        )
+        assertEquals(
+            target.virtualFile.path,
+            onlyReference(pointReference)
+                .resolve()
+                ?.containingFile
+                ?.virtualFile
+                ?.path,
+        )
+    }
+
+    fun testImportedProcedureReferenceResolvesFromJaiModule() {
+        val source =
+            myFixture.addFileToProject(
+                "module-consumer.jai",
+                """
+                #import "Basic";
+                main :: () { alloc(1); }
+                """.trimIndent(),
+            )
+        val referenceElement =
+            PsiTreeUtil
+                .findChildrenOfType(source, JaiRefExpr::class.java)
+                .single { it.text == "alloc" }
+
+        val resolved = onlyReference(referenceElement).resolve()
+
+        assertEquals("alloc", resolved?.text)
+        assertEquals(
+            "${System.getProperty("user.home")}/.local/jai/modules/Basic/module.jai",
+            resolved?.containingFile?.virtualFile?.path,
+        )
+    }
+
+    fun testUnresolvedSymbolReferenceHasNoTarget() {
+        val source = myFixture.addFileToProject("unresolved.jai", "main :: () { missing(); }")
+        val referenceElement =
+            PsiTreeUtil
+                .findChildrenOfType(source, JaiRefExpr::class.java)
+                .single { it.text == "missing" }
+
+        assertNull(onlyReference(referenceElement).resolve())
+    }
+
     fun testImportResolvesAJaiModuleDirectory() {
         val sourceVirtualFile =
             LocalFileSystem
@@ -90,12 +203,12 @@ class JaiReferenceTest : BasePlatformTestCase() {
         return literal!!
     }
 
-    private fun onlyReference(literal: JaiLiteralExpr): PsiReference {
+    private fun onlyReference(element: PsiElement): PsiReference {
-        val references = references(literal)
+        val references = references(element)
-        assertEquals("expected exactly one path reference", 1, references.size)
+        assertEquals("expected exactly one reference", 1, references.size)
         return references.single()
     }
 
-    private fun references(literal: JaiLiteralExpr): List<PsiReference> =
+    private fun references(element: PsiElement): List<PsiReference> =
-        PsiReferenceService.getService().getContributedReferences(literal).toList()
+        PsiReferenceService.getService().getContributedReferences(element).toList()
 }