ADFA-3099 | Prevent UI stalls by generating TreeSitter line spans asynchronously

editor-treesitter/src/main/java/io/github/rosemoe/sora/editor/ts/LineSpansGenerator.kt

@@ -41,6 +41,10 @@
 
 package io.github.rosemoe.sora.editor.ts
 
+import android.os.Handler
+import android.os.Looper
+import android.util.Log
+import android.util.LruCache
 import com.itsaky.androidide.treesitter.TSInputEdit
 import com.itsaky.androidide.treesitter.TSQueryCapture
 import com.itsaky.androidide.treesitter.TSQueryCursor
@@ -55,6 +59,15 @@ import io.github.rosemoe.sora.lang.styling.TextStyle
 import io.github.rosemoe.sora.text.CharPosition
 import io.github.rosemoe.sora.text.Content
 import io.github.rosemoe.sora.widget.schemes.EditorColorScheme
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.SupervisorJob
+import kotlinx.coroutines.asCoroutineDispatcher
+import kotlinx.coroutines.cancel
+import kotlinx.coroutines.launch
+import java.util.concurrent.ConcurrentHashMap
+import java.util.concurrent.Executors
+import java.util.concurrent.atomic.AtomicBoolean
+import java.util.concurrent.atomic.AtomicInteger
 
 /**
  * Spans generator for tree-sitter. Results are cached.
@@ -66,42 +79,67 @@ import io.github.rosemoe.sora.widget.schemes.EditorColorScheme
 class LineSpansGenerator(internal var tree: TSTree, internal var lineCount: Int,
   private val content: Content, internal var theme: TsTheme,
   private val languageSpec: TsLanguageSpec, var scopedVariables: TsScopedVariables,
-  private val spanFactory: TsSpanFactory) : Spans {
+  private val spanFactory: TsSpanFactory, private val requestRedraw: () -> Unit) : Spans {
 
   companion object {
 
     const val CACHE_THRESHOLD = 60
+    const val TAG = "LineSpansGenerator"
+    /**
+     * Delay in milliseconds to batch UI redraws, preventing frame drops
+     * when rapidly calculating multiple lines.
+     */
+    const val REDRAW_DEBOUNCE_DELAY_MS = 32L
   }
 
-  private val caches = mutableListOf<SpanCache>()
+  /**
+   * Thread-safe cache for calculated line spans.
+   * Automatically evicts the least recently used lines.
+   */
+  private val caches = LruCache<Int, MutableList<Span>>(CACHE_THRESHOLD)
+  private val calculatingLines = ConcurrentHashMap.newKeySet<Int>()
 
-  fun edit(edit: TSInputEdit) {
-    tree.edit(edit)
+  private val tsExecutor = Executors.newSingleThreadExecutor { r ->
+    Thread(r, "TreeSitterWorker")
   }
+  private val tsDispatcher = tsExecutor.asCoroutineDispatcher()
+  private val scope = CoroutineScope(SupervisorJob() + tsDispatcher)
 
-  fun queryCache(line: Int): MutableList<Span>? {
-    for (i in 0 until caches.size) {
-      val cache = caches[i]
-      if (cache.line == line) {
-        caches.removeAt(i)
-        caches.add(0, cache)
-        return cache.spans
-      }
+  /**
+   * Tracks content changes so the worker can instantly abort
+   * outdated calculations when the user types.
+   */
+  private val contentVersion = AtomicInteger(0)
+  private val mainHandler = Handler(Looper.getMainLooper())
+  private var isRefreshScheduled = AtomicBoolean(false)
+
+  fun edit(edit: TSInputEdit) {
+    contentVersion.incrementAndGet()
+    scope.launch {
+      tree.edit(edit)
+      calculatingLines.clear()
     }
-    return null
   }
 
-  fun pushCache(line: Int, spans: MutableList<Span>) {
-    while (caches.size >= CACHE_THRESHOLD) {
-      caches.removeAt(caches.size - 1)
-    }
-    caches.add(0, SpanCache(spans, line))
+  /**
+   * Queues the native tree destruction in the background
+   * so it doesn't close while a query is running.
+   */
+  fun destroy() {
+    scope.cancel()
+    caches.evictAll()
+    calculatingLines.clear()
+
+    mainHandler.removeCallbacksAndMessages(null)
+
+    tsExecutor.execute { runCatching { tree.close() } }
+    tsExecutor.shutdown()
   }
 
   fun captureRegion(startIndex: Int, endIndex: Int): MutableList<Span> {
     val list = mutableListOf<Span>()
 
-    if (!tree.canAccess()) {
+    if (!tree.canAccess() || tree.rootNode.hasChanges()) {
       list.add(emptySpan(0))
       return list
     }
@@ -199,58 +237,133 @@ class LineSpansGenerator(internal var tree: TSTree, internal var lineCount: Int,
   }
 
   override fun adjustOnInsert(start: CharPosition, end: CharPosition) {
+    val lineDiff = end.line - start.line
 
+    if (lineDiff == 0) {
+      val colDiff = end.column - start.column
+      shiftSpansOnLine(start.line, start.column, colDiff)
+      return
+    }
+
+    rebuildCache { line, spans, cache ->
+      when {
+        line < start.line -> cache.put(line, spans)
+        line == start.line -> {
+          cache.put(line, spans)
+          cache.put(line + lineDiff, spans)
+        }
+        else -> cache.put(line + lineDiff, spans)
+      }
+    }
   }
 
   override fun adjustOnDelete(start: CharPosition, end: CharPosition) {
+    val lineDiff = end.line - start.line
+
+    if (lineDiff == 0) {
+      val colDiff = start.column - end.column
+      shiftSpansOnLine(start.line, end.column, colDiff)
+      return
+    }
+
+    rebuildCache { line, spans, cache ->
+      when {
+        line < start.line -> cache.put(line, spans)
+        line == start.line -> cache.put(line, spans)
+        line > end.line -> cache.put(line - lineDiff, spans)
+      }
+    }
+  }
+
+  /**
+   * Shifts span columns horizontally to prevent visual flickering during inline edits.
+   *
+   * @param line Line index of the modification.
+   * @param startColumn Column index where the shift begins.
+   * @param colDiff Number of columns to shift.
+   */
+  private fun shiftSpansOnLine(line: Int, startColumn: Int, colDiff: Int) {
+    caches.get(line)?.forEach { span ->
+      if (span.column >= startColumn) {
+        span.column += colDiff
+      }
+    }
+  }
 
+  /**
+   * Rebuilds the line cache for vertical text shifts (line additions or deletions).
+   *
+   * @param action Logic to determine how each cached line is re-inserted.
+   */
+  private inline fun rebuildCache(action: (line: Int, spans: MutableList<Span>, cache: LruCache<Int, MutableList<Span>>) -> Unit) {
+    val snapshot = caches.snapshot()
+    caches.evictAll()
+
+    for ((line, spans) in snapshot) {
+      action(line, spans, caches)
+    }
   }
 
   override fun read() = object : Spans.Reader {
 
     private var spans = mutableListOf<Span>()
 
     override fun moveToLine(line: Int) {
-      try {
-        if (line < 0 || line >= lineCount) {
-          spans = mutableListOf()
-          return
-        }
-        val cached = queryCache(line)
-        if (cached != null) {
-          spans = cached
-          return
-        }
-        val start = content.indexer.getCharPosition(line, 0).index
-        val end = start + content.getColumnCount(line)
-        spans = captureRegion(start, end)
-        pushCache(line, spans)
-      } catch (err: Throwable) {
-        err.printStackTrace()
-      }
+      spans = getSpansForLine(line)
     }
 
     override fun getSpanCount() = spans.size
 
     override fun getSpanAt(index: Int) = spans[index]
+    override fun getSpansOnLine(line: Int): MutableList<Span> = getSpansForLine(line)
+
+    private fun getSpansForLine(line: Int): MutableList<Span> {
+      if (line !in 0..<lineCount) return mutableListOf()
+
+      caches.get(line)?.let { return it }
+
+      if (!calculatingLines.add(line)) return mutableListOf(emptySpan(0))
 
-    override fun getSpansOnLine(line: Int): MutableList<Span> {
-      try {
-        val cached = queryCache(line)
-        if (cached != null) {
-          return ArrayList(cached)
+      val requestedVersion = contentVersion.get()
+
+      scope.launch {
+        try {
+          if (requestedVersion != contentVersion.get()) return@launch
+
+          val start = content.indexer.getCharPosition(line, 0).index
+          val end = start + content.getColumnCount(line)
+
+          val resultSpans = captureRegion(start, end)
+
+          if (requestedVersion == contentVersion.get()) {
+            caches.put(line, resultSpans)
+            scheduleRefresh()
+          }
+        } catch (e: Exception) {
+          Log.e(TAG, "Error processing spans for line $line", e)
+          e.printStackTrace()
+        } finally {
+          calculatingLines.remove(line)
         }
-        val start = content.indexer.getCharPosition(line, 0).index
-        val end = start + content.getColumnCount(line)
-        return captureRegion(start, end)
-      } catch (err: Throwable) {
-        err.printStackTrace()
-        throw err
-      }
+			}
+			return mutableListOf(emptySpan(0))
     }
 
   }
 
+  /**
+   * Groups redraw requests together to avoid overloading the UI thread.
+   */
+  private fun scheduleRefresh() {
+    if (!isRefreshScheduled.compareAndSet(false, true)) return
+
+    mainHandler.postDelayed({
+      isRefreshScheduled.set(false)
+      requestRedraw()
+      Log.d(TAG, "Refreshing UI with newly processed lines")
+    }, REDRAW_DEBOUNCE_DELAY_MS)
+  }
+
   override fun supportsModify() = false
 
   override fun modify(): Spans.Modifier {
@@ -259,5 +372,3 @@ class LineSpansGenerator(internal var tree: TSTree, internal var lineCount: Int,
 
   override fun getLineCount() = lineCount
 }
-
-data class SpanCache(val spans: MutableList<Span>, val line: Int)

editor-treesitter/src/main/java/io/github/rosemoe/sora/editor/ts/TsAnalyzeManager.kt

@@ -37,6 +37,7 @@ open class TsAnalyzeManager(val languageSpec: TsLanguageSpec, var theme: TsTheme
   var spanFactory: TsSpanFactory = DefaultSpanFactory()
 
   open var styles = Styles()
+  internal var currentBracketPairs: TsBracketPairs? = null
 
   private var _analyzeWorker: TsAnalyzeWorker? = null
   val analyzeWorker: TsAnalyzeWorker?
@@ -131,7 +132,7 @@ open class TsAnalyzeManager(val languageSpec: TsLanguageSpec, var theme: TsTheme
     _analyzeWorker?.stop()
     _analyzeWorker = null
 
-    (styles.spans as LineSpansGenerator?)?.tree?.close()
+    (styles.spans as LineSpansGenerator?)?.destroy()
     styles.spans = null
     styles = Styles()
 
@@ -149,9 +150,11 @@ open class TsAnalyzeManager(val languageSpec: TsLanguageSpec, var theme: TsTheme
     _analyzeWorker?.stop()
     _analyzeWorker = null
 
-    (styles.spans as LineSpansGenerator?)?.tree?.close()
+    (styles.spans as LineSpansGenerator?)?.destroy()
     styles.spans = null
 
+    currentBracketPairs?.close()
+    currentBracketPairs = null
     spanFactory.close()
   }
 }

editor-treesitter/src/main/java/io/github/rosemoe/sora/editor/ts/TsAnalyzeWorker.kt

@@ -17,6 +17,8 @@
 
 package io.github.rosemoe.sora.editor.ts
 
+import android.os.Handler
+import android.os.Looper
 import com.itsaky.androidide.syntax.colorschemes.SchemeAndroidIDE
 import com.itsaky.androidide.treesitter.TSInputEdit
 import com.itsaky.androidide.treesitter.TSQueryCursor
@@ -270,28 +272,35 @@ class TsAnalyzeWorker(
 
     val tree = tree!!
     val scopedVariables = TsScopedVariables(tree, text, languageSpec)
-    val oldTree = (styles.spans as? LineSpansGenerator?)?.tree
-    val copied = tree.copy()
+    val oldSpans = (styles.spans as? LineSpansGenerator?)
+    val oldBrackets = analyzer.currentBracketPairs
 
+    oldSpans?.destroy()
+
+    // Use separate tree copies for the background worker and the UI thread
+    // to prevent concurrent access crashes.
     styles.spans = LineSpansGenerator(
-      copied,
+      tree.copy(),
       reference.lineCount,
       reference.reference,
       theme,
       languageSpec,
       scopedVariables,
-      spanFactory
+      spanFactory,
+      requestRedraw = { stylesReceiver?.setStyles(analyzer, styles) }
     )
 
+    val newBrackets = TsBracketPairs(tree.copy(), languageSpec)
+    analyzer.currentBracketPairs = newBrackets
+
     val oldBlocks = styles.blocks
     updateCodeBlocks()
     oldBlocks?.also { ObjectAllocator.recycleBlockLines(it) }
 
-    stylesReceiver?.setStyles(analyzer, styles) {
-      oldTree?.close()
-    }
+    stylesReceiver?.setStyles(analyzer, styles)
+    stylesReceiver?.updateBracketProvider(analyzer, newBrackets)
 
-    stylesReceiver?.updateBracketProvider(analyzer, TsBracketPairs(copied, languageSpec))
+    oldBrackets?.let { Handler(Looper.getMainLooper()).post { it.close() } }
   }
 
   private fun updateCodeBlocks() {

editor-treesitter/src/main/java/io/github/rosemoe/sora/editor/ts/TsBracketPairs.kt

@@ -43,6 +43,15 @@ class TsBracketPairs(private val tree: TSTree, private val languageSpec: TsLangu
 
   }
 
+  /**
+   * Frees the native TreeSitter object. Must run on the UI thread
+   * to avoid crashing if a query is active.
+   */
+  fun close() {
+    if (!tree.canAccess()) return
+    tree.close()
+  }
+
   override fun getPairedBracketAt(text: Content, index: Int): PairedBracket? {
     if (!languageSpec.bracketsQuery.canAccess() || languageSpec.bracketsQuery.patternCount <= 0 || !tree.canAccess()) {
       return null