fix: properly implement TSTree edition and query execution

.idea/deploymentTargetDropDown.xml

@@ -12,6 +12,6 @@
         </deviceKey>
       </Target>
     </runningDeviceTargetSelectedWithDropDown>
-    <timeTargetWasSelectedWithDropDown value="2023-10-15T12:47:54.859420162Z" />
+    <timeTargetWasSelectedWithDropDown value="2023-10-15T16:16:37.701274383Z" />
   </component>
 </project>

app/proguard-rules.pro

@@ -7,6 +7,7 @@
 -keepattributes Signature
 -keepattributes Exceptions
 -keepattributes *Annotation*
+-keepattributes SourceFile,LineNumberTable
 
 -keep class javax.** { *; }
 -keep class jdkx.** { *; }
@@ -89,6 +90,7 @@
 -keepclasseswithmembers class ** {
     native <methods>;
 }
+
 -keep class com.itsaky.androidide.treesitter.** { *; }
 
 # Retrofit 2

editor-api/build.gradle.kts

@@ -30,4 +30,5 @@ dependencies {
   api(projects.lsp.models)
   
   implementation(projects.common)
+  implementation(projects.logger)
 }

editor-api/src/main/java/com/itsaky/androidide/treesitter/api/tsUtils.kt

@@ -0,0 +1,163 @@
+/*
+ *  This file is part of AndroidIDE.
+ *
+ *  AndroidIDE is free software: you can redistribute it and/or modify
+ *  it under the terms of the GNU General Public License as published by
+ *  the Free Software Foundation, either version 3 of the License, or
+ *  (at your option) any later version.
+ *
+ *  AndroidIDE is distributed in the hope that it will be useful,
+ *  but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *  GNU General Public License for more details.
+ *
+ *  You should have received a copy of the GNU General Public License
+ *   along with AndroidIDE.  If not, see <https://www.gnu.org/licenses/>.
+ */
+
+package com.itsaky.androidide.treesitter.api
+
+import com.itsaky.androidide.treesitter.TSNode
+import com.itsaky.androidide.treesitter.TSQuery
+import com.itsaky.androidide.treesitter.TSQueryCursor
+import com.itsaky.androidide.treesitter.TSQueryMatch
+import com.itsaky.androidide.treesitter.TSTree
+import com.itsaky.androidide.utils.ILogger
+
+private val log = ILogger.newInstance("TsUtilsKt")
+
+/**
+ * Safely create a query cursor and execute the given [TSQuery]. The given [action] will be called for
+ * every [TSQueryMatch]. If the tree, cursor, tree or the tree's node is closed or edited while the cursor is
+ * querying the node, the [onClosedOrEdited] will be called.
+ *
+ * This method does not close the [TSQueryCursor] instance.
+ */
+fun <ResultT> TSQueryCursor.safeExecQueryCursor(
+  query: TSQuery,
+  tree: TSTree?,
+  recycleNodeAfterUse: Boolean = true,
+  matchCondition: ((TSQueryMatch?) -> Boolean)? = null,
+  whileTrue: ((TSQueryMatch?) -> Boolean)? = null,
+  onClosedOrEdited: (() -> Unit)? = null,
+  debugName: String = "",
+  action: (TSQueryMatch) -> ResultT
+) : ResultT? {
+
+  if (tree == null || !tree.canAccess()) {
+    log.debug("$debugName: Cannot execute query, tree is null or not accessible", "tree=$tree", "tree.canAccess=${tree?.canAccess()}")
+    return null
+  }
+
+  val rootNode = tree.rootNode
+  if (!rootNode.canAccess() || rootNode.hasChanges()) {
+    log.debug("$debugName, Cannot execute query, tree's root node is not accessible or has been edited", "rootNode=$rootNode", "rootNode.canAccess=${rootNode.canAccess()}", "rootNode.hasChanges=${rootNode.canAccess() && rootNode.hasChanges()}")
+    return null
+  }
+
+  return safeExecQueryCursor(
+    query = query,
+    node = rootNode,
+    recycleNodeAfterUse = recycleNodeAfterUse,
+    matchCondition = {
+      val result = tree.canAccess() && (matchCondition?.invoke(it) ?: true)
+      if (!result) {
+        log.debug("$debugName: tree.canAccess=${tree.canAccess()}")
+      }
+      result
+    },
+    whileTrue = whileTrue,
+    onClosedOrEdited = onClosedOrEdited,
+    debugName = debugName,
+    action = action
+  )
+}
+
+/**
+ * Safely create a query cursor and execute the given [TSQuery]. The given [action] will be called for
+ * every [TSQueryMatch]. If the tree, cursor or the node is closed or edited while the cursor is
+ * querying the node, the [onClosedOrEdited] will be called.
+ *
+ * This method does not close the [TSQueryCursor] instance.
+ */
+fun <ResultT> TSQueryCursor.safeExecQueryCursor(
+  query: TSQuery,
+  node: TSNode,
+  recycleNodeAfterUse: Boolean = true,
+  matchCondition: ((TSQueryMatch?) -> Boolean)? = null,
+  whileTrue: ((TSQueryMatch?) -> Boolean)? = null,
+  onClosedOrEdited: (() -> Unit)? = null,
+  debugName: String = "",
+  action: (TSQueryMatch) -> ResultT
+) : ResultT? {
+
+  return doSafeExecQueryCursor(
+    query = query,
+    node = node,
+    recycleNodeAfterUse = recycleNodeAfterUse,
+    matchCondition = { match ->
+      match != null && canAccess() && node.canAccess() && !node.hasChanges() && (matchCondition?.invoke(
+        match) ?: true)
+    },
+    whileTrue = whileTrue,
+    onClosedOrEdited = onClosedOrEdited,
+    debugName = debugName,
+    action = action)
+}
+
+private fun <ResultT> TSQueryCursor.doSafeExecQueryCursor(
+  query: TSQuery,
+  node: TSNode,
+  recycleNodeAfterUse: Boolean = true,
+  matchCondition: (TSQueryMatch?) -> Boolean,
+  whileTrue: ((TSQueryMatch?) -> Boolean)? = null,
+  onClosedOrEdited: (() -> Unit)? = null,
+  debugName: String = "",
+  action: (TSQueryMatch) -> ResultT
+) : ResultT? {
+
+  if (!query.canAccess()) {
+    log.debug("$debugName: Cannot execute query, query is not accessible")
+    return null
+  }
+
+  if (!node.canAccess() || node.hasChanges()) {
+    log.debug("$debugName: Cannot execute query, node is not accessible or has been edited", "node.canAccess=${node.canAccess()}", "node.hasChanges=${node.hasChanges()}")
+    return null
+  }
+
+  exec(query, node)
+  var match = nextMatch()
+  while (matchCondition(match) && (whileTrue?.invoke(match) != false)) {
+
+    val result = action(match)
+
+    if (!matchCondition(match)) {
+      log.debug(
+        "$debugName: Cannot proceed with query operation.",
+        "cursor.canAccess=${canAccess()}",
+        "query.canAccess=${query.canAccess()}",
+        "node.canAccess=${node.canAccess()}",
+        "node.hasChanges=${node.canAccess() && node.hasErrors()}"
+      )
+      onClosedOrEdited?.invoke()
+      break
+    }
+
+    (match as? TreeSitterQueryMatch?)?.recycle()
+
+    // if the action does not produce any output and simply returns Unit (void)
+    // then ignore the result and continue with the capture
+    if (result != Unit && result != null) {
+      return result
+    }
+
+    match = nextMatch()
+  }
+
+  if (recycleNodeAfterUse && node is TreeSitterNode && !node.isRecycled) {
+    node.recycle()
+  }
+
+  return null
+}

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

@@ -28,9 +28,9 @@ import com.itsaky.androidide.treesitter.TSInputEdit
 import com.itsaky.androidide.treesitter.TSQueryCapture
 import com.itsaky.androidide.treesitter.TSQueryCursor
 import com.itsaky.androidide.treesitter.TSTree
-import com.itsaky.androidide.treesitter.api.TreeSitterNode
 import com.itsaky.androidide.treesitter.api.TreeSitterQueryCapture
-import com.itsaky.androidide.treesitter.api.TreeSitterQueryMatch
+import com.itsaky.androidide.treesitter.api.safeExecQueryCursor
+import com.itsaky.androidide.utils.ILogger
 import io.github.rosemoe.sora.editor.ts.spans.TsSpanFactory
 import io.github.rosemoe.sora.lang.styling.Span
 import io.github.rosemoe.sora.lang.styling.Spans
@@ -59,6 +59,7 @@ class LineSpansGenerator(
   companion object {
 
     const val CACHE_THRESHOLD = 60
+    private val log = ILogger.newInstance("LineSpansGenerator")
   }
 
   private val caches = mutableListOf<SpanCache>()
@@ -94,26 +95,27 @@ class LineSpansGenerator(
   fun captureRegion(startIndex: Int, endIndex: Int): MutableList<Span> {
     val list = mutableListOf<Span>()
 
-    if (!isValid.get()) {
+    if (!isValid.get() || !tree.canAccess()) {
       list.add(emptySpan(0))
       return list
     }
 
     val captures = mutableListOf<TSQueryCapture>()
+
     TSQueryCursor.create().use { cursor ->
       cursor.setByteRange(startIndex * 2, endIndex * 2)
 
-      if (!tree.canAccess()) {
-        throw IllegalStateException("Cannot access tree")
-      }
-
-      // TSTree.getRootNode() always returns a new TSNode instance
-      // make sure to recycle the instance
-      val rootNode = tree.rootNode
-      cursor.exec(languageSpec.tsQuery, rootNode)
-
-      var match = cursor.nextMatch()
-      while (match != null && isValid.get() && tree.canAccess()) {
+      cursor.safeExecQueryCursor(
+        query = languageSpec.tsQuery,
+        tree = tree,
+        recycleNodeAfterUse = true,
+        matchCondition = { isValid.get() },
+        onClosedOrEdited = {
+          captures.clear()
+          log.debug("Tree closed or edited while in captureRegion")
+                           },
+        debugName = "LineSpansGenerator.captureRegion"
+      ) { match ->
         if (languageSpec.queryPredicator.doPredicate(
             languageSpec.predicates,
             content,
@@ -122,35 +124,12 @@ class LineSpansGenerator(
         ) {
           captures.addAll(match.captures)
         }
-
-        (match as? TreeSitterQueryMatch?)?.recycle()
-
-        if (!isValid.get()) {
-          captures.clear()
-          break
-        }
-
-        if (!tree.canAccess() || rootNode.hasChanges()) {
-          val hasChanges = rootNode.hasChanges()
-          (rootNode as? TreeSitterNode?)?.recycle()
-          throw IllegalStateException(
-            "Tree closed while querying, rootNode.hasChanges=$hasChanges, start=$startIndex, end=$endIndex")
-        }
-
-        match = cursor.nextMatch()
       }
 
-      (rootNode as? TreeSitterNode?)?.recycle()
-
       captures.sortBy { it.node.startByte }
       var lastIndex = 0
 
       for (capture in captures) {
-        if (!isValid.get()) {
-          list.clear()
-          break
-        }
-
         val startByte = capture.node.startByte
         val endByte = capture.node.endByte
         val start = (startByte / 2 - startIndex).coerceAtLeast(0)

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

@@ -35,6 +35,7 @@ import com.itsaky.androidide.treesitter.api.TreeSitterInputEdit
 import com.itsaky.androidide.treesitter.api.TreeSitterNode
 import com.itsaky.androidide.treesitter.api.TreeSitterQueryCapture
 import com.itsaky.androidide.treesitter.api.TreeSitterQueryMatch
+import com.itsaky.androidide.treesitter.api.safeExecQueryCursor
 import com.itsaky.androidide.treesitter.string.UTF16String
 import com.itsaky.androidide.treesitter.string.UTF16StringFactory
 import com.itsaky.androidide.utils.ILogger
@@ -230,60 +231,53 @@ open class TsAnalyzeManager(val languageSpec: TsLanguageSpec, var theme: TsTheme
       }
 
       val blocks = mutableListOf<CodeBlock>()
-      TSQueryCursor.create().use {
+      TSQueryCursor.create().use { cursor ->
 
-        val rootNode = tree!!.rootNode
-        it.exec(languageSpec.blocksQuery, rootNode)
-        (rootNode as? TreeSitterNode?)?.recycle()
-
-        var match = it.nextMatch()
-        while (match != null) {
-          if (languageSpec.blocksPredicator.doPredicate(
+        cursor.safeExecQueryCursor(
+          query = languageSpec.blocksQuery,
+          tree = tree,
+          recycleNodeAfterUse = true,
+          matchCondition = null,
+          whileTrue = null,
+          onClosedOrEdited = { blocks.clear() },
+          debugName = "TsAnalyzeManager.updateCodeBlocks()"
+        ) { match ->
+          if (!languageSpec.blocksPredicator.doPredicate(
               languageSpec.predicates,
               localText,
               match
             )
           ) {
-            match.captures.forEach { capture ->
-              val block = ObjectAllocator.obtainBlockLine().also { block ->
-                var node = capture.node
-                val start = node.startPoint
-                block.startLine = start.row
-                block.startColumn = start.column / 2
-                val end = if (languageSpec.blocksQuery.getCaptureNameForId(capture.index)
-                    .endsWith(".marked")
-                ) {
-                  // Goto last terminal element
-                  while (node.childCount > 0) {
-                    node = node.getChild(node.childCount - 1)
-                  }
-                  node.startPoint
-                } else {
-                  node.endPoint
-                }
-                block.endLine = end.row
-                block.endColumn = end.column / 2
-              }
-              if (block.endLine - block.startLine > 1) {
-                blocks.add(block)
-              }
+            return@safeExecQueryCursor
+          }
 
-              (capture as? TreeSitterQueryCapture?)?.recycle()
+          match.captures.forEach { capture ->
+            val block = ObjectAllocator.obtainBlockLine()
+            var node = capture.node
+            val start = node.startPoint
+
+            block.startLine = start.row
+            block.startColumn = start.column / 2
+
+            val end = if (languageSpec.blocksQuery.getCaptureNameForId(capture.index)
+                .endsWith(".marked")
+            ) {
+              // Goto last terminal element
+              while (node.childCount > 0) {
+                node = node.getChild(node.childCount - 1)
+              }
+              node.startPoint
+            } else {
+              node.endPoint
             }
+            block.endLine = end.row
+            block.endColumn = end.column / 2
+            if (block.endLine - block.startLine > 1) {
+              blocks.add(block)
+            }
+
+            (capture as? TreeSitterQueryCapture?)?.recycle()
           }
-
-          (match as? TreeSitterQueryMatch?)?.recycle()
-
-          if (tree?.canAccess() != true || rootNode.hasChanges()) {
-            val hasChanges = rootNode.hasChanges()
-            (rootNode as? TreeSitterNode?)?.recycle()
-            blocks.clear()
-            log.info(
-              "TsLooperThread.updateCodeBlock: Tree closed or edited while querying for code blocks. rootNode.hasChanges=$hasChanges")
-            break
-          }
-
-          match = it.nextMatch()
         }
       }
       val distinct = blocks.distinct().toMutableList()

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

@@ -27,9 +27,8 @@ package io.github.rosemoe.sora.editor.ts
 import com.itsaky.androidide.treesitter.TSNode
 import com.itsaky.androidide.treesitter.TSQueryCursor
 import com.itsaky.androidide.treesitter.TSTree
-import com.itsaky.androidide.treesitter.api.TreeSitterNode
 import com.itsaky.androidide.treesitter.api.TreeSitterQueryCapture
-import com.itsaky.androidide.treesitter.api.TreeSitterQueryMatch
+import com.itsaky.androidide.treesitter.api.safeExecQueryCursor
 import com.itsaky.androidide.utils.ILogger
 import io.github.rosemoe.sora.lang.brackets.BracketsProvider
 import io.github.rosemoe.sora.lang.brackets.PairedBracket
@@ -40,6 +39,7 @@ class TsBracketPairs(private val tree: TSTree, private val languageSpec: TsLangu
   BracketsProvider {
 
   companion object {
+
     private val log = ILogger.newInstance("TsBracketPairs")
 
     val OPEN_NAME = "editor.brackets.open"
@@ -55,52 +55,53 @@ class TsBracketPairs(private val tree: TSTree, private val languageSpec: TsLangu
       return null
     }
 
-    TSQueryCursor.create().use { cursor ->
+    return TSQueryCursor.create().use { cursor ->
       cursor.setByteRange(max(0, index - 1) * 2, index * 2 + 1)
-      val rootNode = tree.rootNode
-      cursor.exec(languageSpec.bracketsQuery, rootNode)
-      (rootNode as? TreeSitterNode?)?.recycle()
-      var match = cursor.nextMatch()
+
       var matched = false
-      while (match != null && !matched) {
-        if (languageSpec.bracketsPredicator.doPredicate(languageSpec.predicates, text, match)) {
-          var startNode: TSNode? = null
-          var endNode: TSNode? = null
 
-          for (capture in match.captures) {
-            val captureName = languageSpec.bracketsQuery.getCaptureNameForId(capture.index)
-            if (captureName == OPEN_NAME || captureName == CLOSE_NAME) {
-              val node = capture.node
-              if (index >= node.startByte / 2 && index <= node.endByte / 2) {
-                matched = true
-              }
-              if (captureName == OPEN_NAME) {
-                startNode = node
-              } else {
-                endNode = node
-              }
+      return@use cursor.safeExecQueryCursor(
+        query = languageSpec.bracketsQuery,
+        tree = tree,
+        recycleNodeAfterUse = true,
+        matchCondition = null,
+        whileTrue = { !matched },
+        onClosedOrEdited = null,
+        debugName = "TsBracketPairs.getPairedBracketAt()"
+      ) { match ->
+        if (!languageSpec.bracketsPredicator.doPredicate(languageSpec.predicates, text, match)) {
+          return@safeExecQueryCursor null
+        }
+
+        var startNode: TSNode? = null
+        var endNode: TSNode? = null
+
+        for (capture in match.captures) {
+          val captureName = languageSpec.bracketsQuery.getCaptureNameForId(capture.index)
+          if (captureName == OPEN_NAME || captureName == CLOSE_NAME) {
+            val node = capture.node
+            if (index >= node.startByte / 2 && index <= node.endByte / 2) {
+              matched = true
+            }
+            if (captureName == OPEN_NAME) {
+              startNode = node
+            } else {
+              endNode = node
             }
-
-            (capture as? TreeSitterQueryCapture?)?.recycle()
           }
-          if (matched && startNode != null && endNode != null) {
-            return PairedBracket(startNode.startByte / 2,
-              (startNode.endByte - startNode.startByte) / 2, endNode.startByte / 2,
-              (endNode.endByte - endNode.startByte) / 2)
-          }
-        }
-        (match as? TreeSitterQueryMatch?)?.recycle()
 
-        if (!tree.canAccess() || rootNode.hasChanges()) {
-          val hasChanges = rootNode.hasChanges()
-          (rootNode as? TreeSitterNode?)?.recycle()
-          log.info("Tree closed or edited while querying, rootNode.hasChanges=$hasChanges")
-          return null
+          (capture as? TreeSitterQueryCapture?)?.recycle()
         }
-        match = cursor.nextMatch()
+
+        if (matched && startNode != null && endNode != null) {
+          return@safeExecQueryCursor PairedBracket(startNode.startByte / 2,
+            (startNode.endByte - startNode.startByte) / 2, endNode.startByte / 2,
+            (endNode.endByte - endNode.startByte) / 2)
+        }
+
+        return@safeExecQueryCursor null
       }
     }
-    return null
   }
 
 }

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

@@ -30,7 +30,7 @@ import com.itsaky.androidide.treesitter.TSQueryCursor
 import com.itsaky.androidide.treesitter.TSTree
 import com.itsaky.androidide.treesitter.api.TreeSitterNode
 import com.itsaky.androidide.treesitter.api.TreeSitterQueryCapture
-import com.itsaky.androidide.treesitter.api.TreeSitterQueryMatch
+import com.itsaky.androidide.treesitter.api.safeExecQueryCursor
 import com.itsaky.androidide.treesitter.string.UTF16String
 import com.itsaky.androidide.utils.ILogger
 import java.util.Stack
@@ -49,37 +49,38 @@ class TsScopedVariables(tree: TSTree, text: UTF16String, val spec: TsLanguageSpe
   private val rootScope: Scope
 
   companion object {
+
     private val log = ILogger.newInstance("TsScopedVariables")
   }
 
   init {
-    if (!tree.canAccess()) {
-      throw IllegalStateException("Cannot access tree")
+    val rootNode = tree.rootNode
+    var needsWalk = true
+    rootScope = if (rootNode.canAccess()) {
+      Scope(0, rootNode.endByte / 2)
+    } else {
+      needsWalk = false
+      Scope(0, 0)
     }
 
-    val rootNode = tree.rootNode
-    rootScope = Scope(0, rootNode.endByte / 2)
-    if (spec.localsDefinitionIndices.isNotEmpty()) {
+    if (needsWalk && spec.localsDefinitionIndices.isNotEmpty()) {
       TSQueryCursor.create().use { cursor ->
-        cursor.exec(spec.tsQuery, rootNode)
-        var match = cursor.nextMatch()
+
         val captures = mutableListOf<TSQueryCapture>()
-        while (match != null) {
+        cursor.safeExecQueryCursor(
+          query = spec.tsQuery,
+          tree = tree,
+          recycleNodeAfterUse = true,
+          matchCondition = null,
+          onClosedOrEdited = { captures.clear() },
+          debugName = "TsScopedVariables.init()"
+        ) { match ->
+
           if (spec.queryPredicator.doPredicate(spec.predicates, text, match)) {
             captures.addAll(match.captures)
           }
-          (match as? TreeSitterQueryMatch?)?.recycle()
-
-          if (!tree.canAccess()) {
-            val hasChanges = rootNode.hasChanges()
-            (rootNode as? TreeSitterNode?)?.recycle()
-            captures.clear()
-            log.info("Tree editor or closed while querying, rootNode.hasChanges=$hasChanges")
-            break
-          }
-
-          match = cursor.nextMatch()
         }
+
         captures.sortBy { it.node.startByte }
         val scopeStack = Stack<Scope>()
         var lastAddedVariableNode: TSNode? = null