Skip to content

Commit 26094df

Browse files
committed
Async, Performance, and SQLite as sole backend
1 parent 8fcbf6a commit 26094df

22 files changed

Lines changed: 1872 additions & 533 deletions
Lines changed: 224 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,224 @@
1+
package io.github.amichne.kast.standalone
2+
3+
import io.github.amichne.kast.api.ModuleName
4+
import io.github.amichne.kast.api.NormalizedPath
5+
import io.github.amichne.kast.standalone.cache.SourceIndexCache
6+
import io.github.amichne.kast.standalone.cache.SqliteSourceIndexStore
7+
import io.github.amichne.kast.standalone.cache.SymbolReferenceRow
8+
import java.nio.file.Path
9+
import java.util.concurrent.CompletableFuture
10+
import java.util.concurrent.atomic.AtomicInteger
11+
import java.util.concurrent.atomic.AtomicReference
12+
import kotlin.concurrent.thread
13+
14+
/**
15+
* Manages eager background indexing in two phases:
16+
*
17+
* - **Phase 1 (identifier index)**: A fast text-only scan that builds
18+
* [MutableSourceIdentifierIndex] from source files. This runs immediately on
19+
* [startPhase1] and completes [identifierIndexReady].
20+
*
21+
* - **Phase 2 (symbol references)**: A deeper scan that resolves K2 symbol
22+
* references and populates the `symbol_references` table in SQLite. Triggered
23+
* via [startPhase2] after Phase 1 and the K2 session are ready. Completes
24+
* [referenceIndexReady].
25+
*
26+
* The indexer is designed to be cancelled cleanly: [close] interrupts in-flight
27+
* work and completes both futures so callers never hang.
28+
*/
29+
internal class BackgroundIndexer(
30+
private val sourceRoots: List<Path>,
31+
private val sourceIndexFileReader: (Path) -> String,
32+
private val sourceModuleNameResolver: (NormalizedPath) -> ModuleName?,
33+
private val sourceIndexCache: SourceIndexCache,
34+
private val store: SqliteSourceIndexStore,
35+
private val initialSourceIndexBuilder: (() -> Map<String, List<String>>)? = null,
36+
) : AutoCloseable {
37+
38+
val identifierIndexReady = CompletableFuture<Unit>()
39+
val referenceIndexReady = CompletableFuture<Unit>()
40+
41+
private val generation = AtomicInteger(0)
42+
private val indexRef = AtomicReference<MutableSourceIdentifierIndex?>(null)
43+
44+
@Volatile
45+
private var cancelled = false
46+
private var phase1Thread: Thread? = null
47+
private var phase2Thread: Thread? = null
48+
49+
/**
50+
* Starts Phase 1 (identifier index) on a daemon thread. Returns the
51+
* generation counter so the caller can detect stale results.
52+
*/
53+
fun startPhase1(): Int {
54+
val gen = generation.incrementAndGet()
55+
phase1Thread = thread(
56+
start = true,
57+
isDaemon = true,
58+
name = "kast-background-indexer-phase1",
59+
) {
60+
runCatching {
61+
if (cancelled) return@thread
62+
initialSourceIndexBuilder
63+
?.invoke()
64+
?.let(MutableSourceIdentifierIndex::fromCandidatePathsByIdentifier)
65+
?: loadOrBuildIndex()
66+
}.onSuccess { index ->
67+
if (cancelled || generation.get() != gen) return@onSuccess
68+
indexRef.set(index)
69+
runCatching { sourceIndexCache.save(index = index, sourceRoots = sourceRoots) }
70+
identifierIndexReady.complete(Unit)
71+
}.onFailure { error ->
72+
if (cancelled || generation.get() != gen) return@onFailure
73+
identifierIndexReady.completeExceptionally(error)
74+
}
75+
}
76+
return gen
77+
}
78+
79+
/**
80+
* Starts Phase 2 (symbol reference index) on a daemon thread. The
81+
* [referenceScanner] callback resolves references for a single file path
82+
* and returns a list of [SymbolReferenceRow]s. It is called inside the
83+
* caller-provided read-access context (e.g., K2 analysis session).
84+
*/
85+
fun startPhase2(referenceScanner: (String) -> List<SymbolReferenceRow>) {
86+
phase2Thread = thread(
87+
start = true,
88+
isDaemon = true,
89+
name = "kast-background-indexer-phase2",
90+
) {
91+
runCatching {
92+
if (cancelled) return@thread
93+
val allPaths = store.loadManifest()?.keys ?: return@thread
94+
generation.incrementAndGet()
95+
for (filePath in allPaths) {
96+
if (cancelled || Thread.currentThread().isInterrupted) break
97+
runCatching {
98+
store.clearReferencesFromFile(filePath)
99+
val refs = referenceScanner(filePath)
100+
refs.forEach { ref ->
101+
store.upsertSymbolReference(
102+
sourcePath = ref.sourcePath,
103+
sourceOffset = ref.sourceOffset,
104+
targetFqName = ref.targetFqName,
105+
targetPath = ref.targetPath,
106+
targetOffset = ref.targetOffset,
107+
)
108+
}
109+
}
110+
}
111+
if (!cancelled) {
112+
referenceIndexReady.complete(Unit)
113+
}
114+
}.onFailure { error ->
115+
if (cancelled) return@onFailure
116+
if (!referenceIndexReady.isDone) {
117+
referenceIndexReady.completeExceptionally(error)
118+
}
119+
}
120+
}
121+
}
122+
123+
/** Returns the current identifier index, or null if Phase 1 hasn't completed. */
124+
fun getIndex(): MutableSourceIdentifierIndex? = indexRef.get()
125+
126+
/** Returns the current generation counter. */
127+
fun currentGeneration(): Int = generation.get()
128+
129+
/**
130+
* Re-indexes a set of changed file paths incrementally. Skips files that
131+
* no longer exist on disk (deleted between discovery and read).
132+
*/
133+
fun reindexFiles(
134+
index: MutableSourceIdentifierIndex,
135+
paths: Set<NormalizedPath>,
136+
) {
137+
paths.forEach { normalizedPath ->
138+
val filePath = normalizedPath.toJavaPath()
139+
if (!java.nio.file.Files.isRegularFile(filePath)) {
140+
index.removeFile(normalizedPath.value)
141+
sourceIndexCache.saveRemovedFile(normalizedPath.value)
142+
return@forEach
143+
}
144+
runCatching {
145+
index.updateFile(
146+
normalizedPath = normalizedPath.value,
147+
newContent = sourceIndexFileReader(filePath),
148+
moduleName = sourceModuleNameResolver(normalizedPath),
149+
)
150+
sourceIndexCache.saveFileIndex(index, normalizedPath)
151+
}
152+
}
153+
}
154+
155+
override fun close() {
156+
cancelled = true
157+
phase1Thread?.interrupt()
158+
phase2Thread?.interrupt()
159+
if (!identifierIndexReady.isDone) {
160+
identifierIndexReady.complete(Unit)
161+
}
162+
if (!referenceIndexReady.isDone) {
163+
referenceIndexReady.complete(Unit)
164+
}
165+
}
166+
167+
// -------------------------------------------------------------------------
168+
// Phase 1 internals
169+
// -------------------------------------------------------------------------
170+
171+
private fun loadOrBuildIndex(): MutableSourceIdentifierIndex {
172+
val incrementalResult = runCatching {
173+
sourceIndexCache.load(sourceRoots)
174+
}.getOrNull()
175+
val index = incrementalResult?.index ?: return buildFullIndex()
176+
incrementalResult.deletedPaths.forEach(index::removeFile)
177+
(incrementalResult.newPaths + incrementalResult.modifiedPaths).forEach { pathString ->
178+
if (cancelled || Thread.currentThread().isInterrupted) return index
179+
refreshFileIndex(index, NormalizedPath.ofNormalized(pathString))
180+
}
181+
return index
182+
}
183+
184+
private fun buildFullIndex(): MutableSourceIdentifierIndex {
185+
val index = MutableSourceIdentifierIndex(
186+
pathsByIdentifier = java.util.concurrent.ConcurrentHashMap(),
187+
identifiersByPath = java.util.concurrent.ConcurrentHashMap(),
188+
)
189+
allTrackedKotlinSourcePaths().forEach { normalizedFilePath ->
190+
if (cancelled || Thread.currentThread().isInterrupted) return index
191+
val normalizedPath = NormalizedPath.ofNormalized(normalizedFilePath)
192+
runCatching {
193+
index.updateFile(
194+
normalizedPath = normalizedFilePath,
195+
newContent = sourceIndexFileReader(normalizedPath.toJavaPath()),
196+
moduleName = sourceModuleNameResolver(normalizedPath),
197+
)
198+
}
199+
// Skip files that fail to read (e.g., deleted between discovery and read)
200+
}
201+
return index
202+
}
203+
204+
private fun refreshFileIndex(
205+
index: MutableSourceIdentifierIndex,
206+
normalizedPath: NormalizedPath,
207+
) {
208+
val filePath = normalizedPath.toJavaPath()
209+
if (!java.nio.file.Files.isRegularFile(filePath)) {
210+
index.removeFile(normalizedPath.value)
211+
return
212+
}
213+
runCatching {
214+
index.updateFile(
215+
normalizedPath = normalizedPath.value,
216+
newContent = sourceIndexFileReader(filePath),
217+
moduleName = sourceModuleNameResolver(normalizedPath),
218+
)
219+
}
220+
}
221+
222+
private fun allTrackedKotlinSourcePaths(): Set<String> =
223+
io.github.amichne.kast.standalone.cache.scanTrackedKotlinFileTimestamps(sourceRoots).keys
224+
}

backend-standalone/src/main/kotlin/io/github/amichne/kast/standalone/MutableSourceIdentifierIndex.kt

Lines changed: 30 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,8 @@ import io.github.amichne.kast.api.KotlinIdentifier
55
import io.github.amichne.kast.api.ModuleName
66
import io.github.amichne.kast.api.NormalizedPath
77
import io.github.amichne.kast.api.PackageName
8+
import io.github.amichne.kast.standalone.cache.FileIndexUpdate
9+
import io.github.amichne.kast.standalone.cache.SqliteSourceIndexStore
810
import java.util.concurrent.ConcurrentHashMap
911

1012
internal class MutableSourceIdentifierIndex(
@@ -14,6 +16,7 @@ internal class MutableSourceIdentifierIndex(
1416
private val packageByPath: ConcurrentHashMap<NormalizedPath, PackageName> = ConcurrentHashMap(),
1517
private val importsByPath: ConcurrentHashMap<NormalizedPath, Set<FqName>> = ConcurrentHashMap(),
1618
private val wildcardImportPackagesByPath: ConcurrentHashMap<NormalizedPath, Set<PackageName>> = ConcurrentHashMap(),
19+
private val backingStore: SqliteSourceIndexStore? = null,
1720
) {
1821
fun candidatePathsFor(identifier: String): List<String> =
1922
pathsByIdentifier[KotlinIdentifier(identifier)]?.map { it.value }?.sorted().orEmpty()
@@ -94,11 +97,24 @@ internal class MutableSourceIdentifierIndex(
9497
moduleName: ModuleName? = null,
9598
) {
9699
val path = NormalizedPath.ofNormalized(normalizedPath)
97-
replaceIdentifiers(
98-
normalizedPath = path,
99-
identifiers = identifierRegex.findAll(newContent).map { match -> KotlinIdentifier(match.value) }.toSet(),
100-
)
100+
val identifiers = identifierRegex.findAll(newContent).map { match -> KotlinIdentifier(match.value) }.toSet()
101+
replaceIdentifiers(normalizedPath = path, identifiers = identifiers)
101102
extractFileMetadata(path, newContent, moduleName)
103+
104+
backingStore?.let { store ->
105+
runCatching {
106+
store.saveFileIndex(
107+
FileIndexUpdate(
108+
path = normalizedPath,
109+
identifiers = identifiers.mapTo(mutableSetOf()) { it.value },
110+
packageName = packageByPath[path]?.value,
111+
moduleName = moduleName?.value,
112+
imports = importsByPath[path]?.mapTo(mutableSetOf()) { it.value }.orEmpty(),
113+
wildcardImports = wildcardImportPackagesByPath[path]?.mapTo(mutableSetOf()) { it.value }.orEmpty(),
114+
),
115+
)
116+
}
117+
}
102118
}
103119

104120
fun removeFile(normalizedPath: String) {
@@ -108,6 +124,7 @@ internal class MutableSourceIdentifierIndex(
108124
packageByPath.remove(path)
109125
importsByPath.remove(path)
110126
wildcardImportPackagesByPath.remove(path)
127+
backingStore?.let { store -> runCatching { store.removeFile(normalizedPath) } }
111128
}
112129

113130
fun knownPaths(): Set<String> = identifiersByPath.keys.mapTo(mutableSetOf()) { it.value }
@@ -231,6 +248,7 @@ internal class MutableSourceIdentifierIndex(
231248
packageByPath: Map<String, String> = emptyMap(),
232249
importsByPath: Map<String, List<String>> = emptyMap(),
233250
wildcardImportPackagesByPath: Map<String, List<String>> = emptyMap(),
251+
backingStore: SqliteSourceIndexStore? = null,
234252
): MutableSourceIdentifierIndex {
235253
val typedPathsByIdentifier = ConcurrentHashMap<KotlinIdentifier, MutableSet<NormalizedPath>>()
236254
val typedIdentifiersByPath = ConcurrentHashMap<NormalizedPath, Set<KotlinIdentifier>>()
@@ -262,7 +280,15 @@ internal class MutableSourceIdentifierIndex(
262280
wildcardImportPackagesByPath = wildcardImportPackagesByPath.entries.associateTo(ConcurrentHashMap()) { (path, packages) ->
263281
NormalizedPath.ofNormalized(path) to packages.mapTo(mutableSetOf()) { PackageName(it) }
264282
},
283+
backingStore = backingStore,
265284
)
266285
}
286+
287+
/**
288+
* Loads the hot cache from [store] and wires write-through so every
289+
* subsequent [updateFile]/[removeFile] persists to SQLite immediately.
290+
*/
291+
fun fromSqliteStore(store: SqliteSourceIndexStore): MutableSourceIdentifierIndex =
292+
store.loadFullIndex(backingStore = store)
267293
}
268294
}

backend-standalone/src/main/kotlin/io/github/amichne/kast/standalone/StandaloneAnalysisBackend.kt

Lines changed: 43 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,15 +25,19 @@ import io.github.amichne.kast.api.RenameQuery
2525
import io.github.amichne.kast.api.RenameResult
2626
import io.github.amichne.kast.api.RuntimeState
2727
import io.github.amichne.kast.api.RuntimeStatusResponse
28+
import io.github.amichne.kast.api.SearchScope
29+
import io.github.amichne.kast.api.SearchScopeKind
2830
import io.github.amichne.kast.api.SemanticInsertionQuery
2931
import io.github.amichne.kast.api.SemanticInsertionResult
3032
import io.github.amichne.kast.api.ServerLimits
3133
import io.github.amichne.kast.api.SymbolQuery
3234
import io.github.amichne.kast.api.SymbolResult
35+
import io.github.amichne.kast.api.SymbolVisibility
3336
import io.github.amichne.kast.api.TextEdit
3437
import io.github.amichne.kast.api.TypeHierarchyQuery
3538
import io.github.amichne.kast.api.TypeHierarchyResult
3639
import io.github.amichne.kast.standalone.analysis.CandidateFileResolver
40+
import io.github.amichne.kast.standalone.analysis.CandidateSearchResult
3741
import io.github.amichne.kast.standalone.analysis.ImportAnalysis
3842
import io.github.amichne.kast.standalone.analysis.SemanticInsertionPointResolver
3943
import io.github.amichne.kast.standalone.analysis.callHierarchyDeclaration
@@ -42,6 +46,7 @@ import io.github.amichne.kast.standalone.analysis.referenceSearchIdentifier
4246
import io.github.amichne.kast.standalone.analysis.resolveTarget
4347
import io.github.amichne.kast.standalone.analysis.resolvedFilePath
4448
import io.github.amichne.kast.standalone.analysis.supertypeNames
49+
import io.github.amichne.kast.standalone.analysis.targetFqNameAndPackage
4550
import io.github.amichne.kast.standalone.analysis.toApiDiagnostics
4651
import io.github.amichne.kast.standalone.analysis.toKastLocation
4752
import io.github.amichne.kast.standalone.analysis.toSymbolModel
@@ -185,7 +190,7 @@ internal class StandaloneAnalysisBackend internal constructor(
185190
) { span ->
186191
val file = session.findKtFile(query.position.filePath)
187192
val target = resolveTarget(file, query.position.offset)
188-
val (candidateFiles, searchScope) = candidateFileResolver.resolve(target)
193+
val (candidateFiles, searchScope) = resolveCandidateFilesForReferences(target, span)
189194
span.setAttribute("kast.references.candidateFileCount", candidateFiles.size)
190195
span.setAttribute("kast.references.searchScope", searchScope.scope.name)
191196

@@ -344,6 +349,43 @@ internal class StandaloneAnalysisBackend internal constructor(
344349
}
345350
}
346351

352+
/**
353+
* Resolve candidate files for a reference search. When the cached symbol reference index is
354+
* complete, use it to narrow candidates; otherwise fall back to the standard resolver.
355+
*/
356+
private fun resolveCandidateFilesForReferences(
357+
target: PsiElement,
358+
span: StandaloneTelemetrySpan,
359+
): CandidateSearchResult {
360+
if (session.isReferenceIndexReady()) {
361+
val fqNameAndPkg = target.targetFqNameAndPackage()
362+
if (fqNameAndPkg != null) {
363+
val (fqName, _) = fqNameAndPkg
364+
val cachedRefs = session.sqliteStore.referencesToSymbol(fqName.value)
365+
if (cachedRefs.isNotEmpty()) {
366+
val cachedPaths = cachedRefs.mapTo(mutableSetOf()) { it.sourcePath }
367+
val ktFiles = cachedPaths.mapNotNull { path ->
368+
runCatching { session.findKtFile(path) }.getOrNull()
369+
}
370+
span.setAttribute("kast.references.cacheHit", true)
371+
span.setAttribute("kast.references.cachedPathCount", cachedPaths.size)
372+
return CandidateSearchResult(
373+
files = ktFiles,
374+
scope = SearchScope(
375+
visibility = io.github.amichne.kast.api.SymbolVisibility.PUBLIC,
376+
scope = SearchScopeKind.DEPENDENT_MODULES,
377+
exhaustive = true,
378+
candidateFileCount = ktFiles.size,
379+
searchedFileCount = ktFiles.size,
380+
),
381+
)
382+
}
383+
}
384+
}
385+
span.setAttribute("kast.references.cacheHit", false)
386+
return candidateFileResolver.resolve(target)
387+
}
388+
347389
private fun KtFile.findReferenceLocations(target: PsiElement): List<io.github.amichne.kast.api.Location> {
348390
val references = mutableListOf<io.github.amichne.kast.api.Location>()
349391

0 commit comments

Comments
 (0)