Files
intellijai/src/main/kotlin/dev/hgh/jai/reference/JaiSymbolResolver.kt
2026-08-05 16:44:16 -04:00

439 lines
18 KiB
Kotlin

package dev.hgh.jai.reference
import com.intellij.openapi.diagnostic.Logger
import com.intellij.psi.PsiElement
import com.intellij.psi.PsiFile
import com.intellij.psi.PsiManager
import com.intellij.psi.PsiReferenceService
import com.intellij.psi.util.PsiTreeUtil
import dev.hgh.jai.lexer.JaiTokenTypes
import dev.hgh.jai.module.JaiModuleResolver
import dev.hgh.jai.module.JaiModuleTarget
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.JaiForName
import dev.hgh.jai.psi.JaiForStatement
import dev.hgh.jai.psi.JaiInitializer
import dev.hgh.jai.psi.JaiLiteralExpr
import dev.hgh.jai.psi.JaiParamName
import dev.hgh.jai.psi.JaiProcLiteralExpr
import dev.hgh.jai.psi.JaiRefExpr
import dev.hgh.jai.settings.JaiProjectSettings
private const val COMPLETION_DEBUG_PROPERTY = "jai.completion.debug"
private val LOG = Logger.getInstance("dev.hgh.jai.reference.JaiSymbolResolver")
private val RECOVERABLE_IMPORT =
Regex("""(?m)(#(?:import|load))((?:\s*,\s*[A-Za-z_][A-Za-z0-9_]*)*)\s*"((?:\\.|[^"\\])*)"""")
private val RECOVERABLE_PROCEDURE =
Regex("""(?s)(?:^|[;}\n])\s*[A-Za-z_][A-Za-z0-9_]*\s*::\s*\([^{}]*\)\s*(?:->[^\{;}]*)?\{""")
private val RECOVERABLE_LOCAL_DECLARATION =
Regex("""(?m)(?:^|[;{}])\s*([A-Za-z_][A-Za-z0-9_]*)\s*(?::=|:(?!=))""")
/** A declaration-like PSI element that can be inserted by identifier completion. */
internal data class JaiSymbolCandidate(
val element: PsiElement,
val name: String,
val kind: String = "symbol",
)
/** Shared symbol lookup for references and symbol-aware completion. */
internal object JaiSymbolResolver {
private data class ScopedSymbol(
val element: PsiElement,
val name: String,
val scope: PsiElement?,
val fallbackScope: PsiElement?,
val visibleFrom: Int,
)
private data class ImportBinding(
val file: PsiFile,
val directive: String,
val scope: PsiElement?,
val visibleFrom: Int,
val boundName: String?,
val importsIntoScope: Boolean,
val excludedNames: Set<String>,
)
fun resolve(reference: JaiRefExpr): PsiElement? {
val file = reference.containingFile ?: return null
return completionCandidates(file, reference)
.firstOrNull { it.name == reference.text }
?.element
}
/** Returns visible local declarations followed by visible unqualified imports. */
fun completionCandidates(
file: PsiFile,
position: PsiElement,
): List<JaiSymbolCandidate> {
val candidates = linkedMapOf<String, JaiSymbolCandidate>()
val localSymbols = localSymbols(file)
val recoveredLocals = recoverLocalSymbols(file, position)
localSymbols
.filter { isVisible(it, position) && !isInside(it.element, position) }
.sortedWith(
compareByDescending<ScopedSymbol> { scopeDepth(it.scope ?: it.fallbackScope) }
.thenBy { it.visibleFrom }
.thenBy { it.element.textOffset },
).forEach { symbol ->
candidates.putIfAbsent(
symbol.name,
JaiSymbolCandidate(symbol.element, symbol.name),
)
}
recoveredLocals.forEach { symbol ->
candidates.putIfAbsent(symbol.name, symbol)
}
importedCandidates(file, position).forEach { symbol ->
candidates.putIfAbsent(symbol.name, symbol)
}
val result = candidates.values.toList()
debug(
"file=${file.virtualFile?.path ?: file.name} roots=${JaiProjectSettings.getInstance(file.project).rootPaths()} " +
"locals=${localSymbols.joinToString {
"${it.name}@${it.element.textOffset}/scope=${it.scope?.textRange}/fallback=${it.fallbackScope?.textRange}"
}} recovered=${recoveredLocals.joinToString { it.name }} " +
"imports=${importBindings(file).joinToString { it.file.virtualFile?.path ?: it.file.name }} " +
"candidates=${result.take(40).joinToString { it.name }}",
)
return result
}
/** Returns members of a module bound to [qualifier] in the visible lexical scope. */
fun memberCompletionCandidates(
file: PsiFile,
position: PsiElement,
qualifier: String,
): List<JaiSymbolCandidate> {
val binding =
importBindings(file)
.filter { it.boundName == normalizeName(qualifier) && isVisible(it, position) }
.sortedWith(
compareByDescending<ImportBinding> { scopeDepth(it.scope) }
.thenByDescending { it.visibleFrom },
).firstOrNull()
?: return emptyList()
return topLevelDeclarations(binding.file)
.asSequence()
.filter { symbolName(it) !in binding.excludedNames }
.mapNotNull { name ->
symbolName(name)?.let { JaiSymbolCandidate(name, it, "member") }
}.distinctBy { it.name }
.toList()
}
private fun localSymbols(file: PsiFile): List<ScopedSymbol> {
val declarations =
PsiTreeUtil
.findChildrenOfType(file, JaiDeclaration::class.java)
.flatMap { declaration ->
declaration.declNames.declNameList.mapNotNull { name ->
symbolName(name)?.let {
ScopedSymbol(
element = name,
name = it,
scope = PsiTreeUtil.getParentOfType(name, JaiBlock::class.java),
fallbackScope = PsiTreeUtil.getParentOfType(name, JaiProcLiteralExpr::class.java),
visibleFrom = name.textOffset,
)
}
}
}
val parameters =
PsiTreeUtil
.findChildrenOfType(file, JaiParamName::class.java)
.mapNotNull { name ->
val symbol = symbolName(name) ?: return@mapNotNull null
val procedure = PsiTreeUtil.getParentOfType(name, JaiProcLiteralExpr::class.java)
ScopedSymbol(
element = name,
name = symbol,
// Keep the procedure itself as the parameter scope. The body block is a
// nested scope, so a local declaration correctly shadows a parameter.
scope = procedure,
fallbackScope = null,
visibleFrom = procedure?.textOffset ?: name.textOffset,
)
}
val forNames =
PsiTreeUtil
.findChildrenOfType(file, JaiForName::class.java)
.mapNotNull { name ->
val symbol = symbolName(name) ?: return@mapNotNull null
ScopedSymbol(
element = name,
name = symbol,
scope = PsiTreeUtil.getParentOfType(name, JaiForStatement::class.java),
fallbackScope = PsiTreeUtil.getParentOfType(name, JaiProcLiteralExpr::class.java),
visibleFrom = name.textOffset,
)
}
return declarations + parameters + forNames
}
private fun importedCandidates(
file: PsiFile,
position: PsiElement,
): List<JaiSymbolCandidate> =
importBindings(file)
.asSequence()
.filter { (it.boundName == null || it.importsIntoScope) && isVisible(it, position) }
.flatMap { binding ->
topLevelDeclarations(binding.file)
.asSequence()
.filter { symbolName(it) !in binding.excludedNames }
.mapNotNull { name ->
symbolName(name)?.let { JaiSymbolCandidate(name, it, "import") }
}
}.distinctBy { it.name }
.toList()
/** Recovers simple locals when an unfinished statement prevents the procedure PSI from closing. */
private fun recoverLocalSymbols(
file: PsiFile,
position: PsiElement,
): List<JaiSymbolCandidate> {
val caretOffset = position.textOffset.coerceIn(0, file.textLength)
val beforeCaret = file.text.substring(0, caretOffset)
val procedure = RECOVERABLE_PROCEDURE.findAll(beforeCaret).lastOrNull() ?: return emptyList()
val bodyStart = procedure.range.last + 1
if (bodyStart >= beforeCaret.length) return emptyList()
return RECOVERABLE_LOCAL_DECLARATION
.findAll(beforeCaret.substring(bodyStart))
.map { match -> JaiSymbolCandidate(position, match.groupValues[1], "local") }
.distinctBy { it.name }
.toList()
}
private fun importBindings(file: PsiFile): List<ImportBinding> {
val parsed =
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
val declaration = boundDeclaration(literal)
val boundName =
declaration
?.declNames
?.declNameList
?.firstOrNull()
?.let(::symbolName)
val declarationPrefix = declaration?.let { prefixBeforeNames(it) }.orEmpty()
val importsIntoScope = declaration == null || declarationPrefix.trimStart().startsWith("using")
val excludedNames = excludedNames(declarationPrefix)
val target =
PsiReferenceService
.getService()
.getContributedReferences(literal)
.firstOrNull()
?.resolve() as? PsiFile
?: return@mapNotNull null
ImportBinding(
file = target,
directive = name,
scope = PsiTreeUtil.getParentOfType(literal, JaiBlock::class.java),
visibleFrom = literal.textOffset,
boundName = boundName,
importsIntoScope = importsIntoScope,
excludedNames = excludedNames,
)
}
// An unfinished statement immediately before the caret can make Grammar-Kit retain the
// file's top-level declarations but omit later directive PSI. Recover bare imports from
// the source text so completion still works while the user is typing.
val recovered = if (parsed.isEmpty()) recoverImportBindings(file) else emptyList()
return (parsed + recovered).distinctBy { binding ->
listOf(
binding.file.virtualFile?.path ?: binding.file.name,
binding.scope?.textRange?.startOffset,
binding.boundName,
binding.visibleFrom,
)
}
}
private fun recoverImportBindings(file: PsiFile): List<ImportBinding> {
val sourceFile = file.virtualFile ?: return emptyList()
return RECOVERABLE_IMPORT
.findAll(file.text)
.mapNotNull { match ->
val directive = match.groupValues[1]
val flags =
Regex("""[A-Za-z_][A-Za-z0-9_]*""")
.findAll(match.groupValues[2])
.map { it.value }
.toSet()
if ("string" in flags) return@mapNotNull null
val path = unescapeImportPath(match.groupValues[3])
val target =
JaiModuleResolver.resolve(
file.project,
sourceFile,
JaiModuleTarget(directive, path, flags),
) ?: return@mapNotNull null
val targetPsi = PsiManager.getInstance(file.project).findFile(target) ?: return@mapNotNull null
ImportBinding(
file = targetPsi,
directive = directive,
scope = null,
visibleFrom = match.range.first,
boundName = null,
importsIntoScope = true,
excludedNames = emptySet(),
)
}.toList()
}
private fun unescapeImportPath(text: String): String =
buildString(text.length) {
var index = 0
while (index < text.length) {
if (text[index] == '\\' && index + 1 < text.length) {
append(text[index + 1])
index += 2
} else {
append(text[index++])
}
}
}
/** The declaration that binds an import, not an enclosing procedure declaration. */
private fun boundDeclaration(literal: JaiLiteralExpr): JaiDeclaration? {
val initializer = PsiTreeUtil.getParentOfType(literal, JaiInitializer::class.java) ?: return null
return PsiTreeUtil.getParentOfType(initializer, JaiDeclaration::class.java)
}
private fun prefixBeforeNames(declaration: JaiDeclaration): String {
val start = declaration.textRange.startOffset
val end = declaration.declNames.textRange.startOffset
return declaration.containingFile.text.substring(start, end)
}
private fun excludedNames(prefix: String): Set<String> {
val body = Regex("""except\s*\(([^)]*)\)""").find(prefix)?.groupValues?.get(1) ?: return emptySet()
return body
.split(',')
.map { normalizeName(it.trim()) }
.filter(String::isNotEmpty)
.toSet()
}
/**
* 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,
position: PsiElement,
): Boolean {
val inScope =
when {
symbol.scope == null && symbol.fallbackScope == null -> true
symbol.scope?.textRange?.contains(position.textOffset) == true -> true
symbol.fallbackScope?.textRange?.contains(position.textOffset) == true -> true
else -> false
}
if (!inScope) return false
// Block-local declarations and loop names are not visible before their declaration. Keep
// top-level declarations available for forward references, which Jai permits.
if (symbol.scope != null && symbol.visibleFrom > position.textOffset) return false
return true
}
private fun isVisible(
binding: ImportBinding,
position: PsiElement,
): Boolean =
// File-scope imports are declarations for the whole file even when the directive is
// written below its first use (a common Jai layout). Block-local imports remain ordered.
(binding.scope == null || binding.scope.textRange.contains(position.textOffset)) &&
(binding.scope == null || binding.visibleFrom <= position.textOffset)
private fun isInside(
element: PsiElement,
position: PsiElement,
): Boolean = element == position || element.textRange.contains(position.textOffset)
private fun scopeDepth(scope: PsiElement?): Int {
var current = scope
var depth = 0
while (current != null) {
if (current is JaiBlock || current is JaiForStatement || current is JaiProcLiteralExpr) depth++
current = current.parent
}
return depth
}
private fun debug(message: String) {
if (java.lang.Boolean.getBoolean(COMPLETION_DEBUG_PROPERTY)) LOG.info(message)
}
private fun symbolName(element: PsiElement): String? {
val text = element.text.trim()
if (text.isEmpty()) return null
val name =
when (element) {
is JaiParamName -> text.removePrefix("$$").removePrefix("$")
else -> text
}
return normalizeName(name).takeIf(String::isNotEmpty)
}
private fun normalizeName(name: String): String {
val trimmed = name.trim()
return if (trimmed.length >= 2 && trimmed.first() == '`' && trimmed.last() == '`') {
trimmed.substring(1, trimmed.length - 1)
} else {
trimmed
}
}
}
private fun JaiDirectiveExpr.directiveName(): String? = node.findChildByType(JaiTokenTypes.DIRECTIVE)?.text