Merge pull request #309 from alephium/scope_walker

Scope Walker

compiler-access/src/main/scala/org/alephium/ralph/lsp/access/compiler/ast/AstExtra.scala

@@ -17,6 +17,7 @@
 package org.alephium.ralph.lsp.access.compiler.ast
 
 import org.alephium.protocol.vm.StatelessContext
+import org.alephium.ralph.lsp.access.compiler.message.SourceIndexExtra
 import org.alephium.ralph.{Type, Ast}
 
 object AstExtra {
@@ -47,6 +48,36 @@ object AstExtra {
           false
       }
 
+  /**
+   * Checks if the `current` AST's position is before the `anchor` AST's position.
+   *
+   * @param current The AST whose position is being tested.
+   * @param anchor  The AST with which the position of `current` is compared.
+   * @return `true` if `current`'s position is before `anchor`'s position, `false` otherwise.
+   */
+  def isBehind(
+      current: Ast.Positioned,
+      anchor: Ast.Positioned): Boolean =
+    SourceIndexExtra.isBehind(
+      current = current.sourceIndex,
+      anchor = anchor.sourceIndex
+    )
+
+  /**
+   * Checks if the `current` AST's position is after the `anchor` AST's position.
+   *
+   * @param current The AST whose position is being tested.
+   * @param anchor  The AST with which the position of `current` is compared.
+   * @return `true` if `current`'s position is after `anchor`'s position, `false` otherwise.
+   */
+  def isAhead(
+      current: Ast.Positioned,
+      anchor: Ast.Positioned): Boolean =
+    SourceIndexExtra.isAhead(
+      current = current.sourceIndex,
+      anchor = anchor.sourceIndex
+    )
+
   /**
    * Fetches the type identifier for a given type.
    *

compiler-access/src/main/scala/org/alephium/ralph/lsp/access/compiler/message/SourceIndexExtra.scala

@@ -71,6 +71,51 @@ object SourceIndexExtra {
       fileURI = fileURI
     )
 
+  def contains(
+      parent: Option[SourceIndex],
+      child: Option[SourceIndex]): Boolean =
+    (parent, child) match {
+      case (Some(parent), Some(child)) =>
+        parent contains child
+
+      case (_, _) =>
+        false
+    }
+
+  /**
+   * Checks if the `current` position is before the `anchor`'s position.
+   *
+   * @param current The SourceIndex whose position is being tested.
+   * @param anchor  The SourceIndex with which the position of `current` is compared.
+   * @return `true` if `current`'s position is before `anchor`'s position, `false` otherwise.
+   */
+  def isBehind(
+      current: Option[SourceIndex],
+      anchor: Option[SourceIndex]): Boolean =
+    current
+      .zip(anchor)
+      .exists {
+        case (current, anchor) =>
+          current isBehind anchor
+      }
+
+  /**
+   * Checks if the `current` SourceIndex's position is after the `anchor` SourceIndex's position.
+   *
+   * @param current The SourceIndex whose position is being tested.
+   * @param anchor  The SourceIndex with which the position of `current` is compared.
+   * @return `true` if `current`'s position is after `anchor`'s position, `false` otherwise.
+   */
+  def isAhead(
+      current: Option[SourceIndex],
+      anchor: Option[SourceIndex]): Boolean =
+    current
+      .zip(anchor)
+      .exists {
+        case (current, anchor) =>
+          current isAhead anchor
+      }
+
   implicit class SourceIndexExtension(val sourceIndex: SourceIndex) extends AnyVal {
 
     def from: Int =
@@ -86,6 +131,27 @@ object SourceIndexExtra {
     def contains(index: Int): Boolean =
       index >= from && index <= to
 
+    def contains(child: SourceIndex): Boolean =
+      (sourceIndex.fileURI, child.fileURI) match {
+        case (Some(parentURI), Some(childURI)) if parentURI == childURI =>
+          sourceIndex contains child.from
+
+        case (_, _) =>
+          false
+      }
+
+    def isBehind(that: SourceIndex): Boolean =
+      isBehind(that.from)
+
+    def isBehind(index: Int): Boolean =
+      sourceIndex.from < index
+
+    def isAhead(that: SourceIndex): Boolean =
+      isAhead(that.from)
+
+    def isAhead(index: Int): Boolean =
+      sourceIndex.from > index
+
     /** Offset this SourceIndex */
     def +(right: Int): SourceIndex =
       sourceIndex.copy(index = from + right)

presentation-compiler/src/main/scala/org/alephium/ralph/lsp/pc/search/gotodef/GoToDefIdent.scala

@@ -331,12 +331,13 @@ private object GoToDefIdent extends StrictImplicitLogging {
       .iterator
       .flatMap {
         block =>
-          block
-            .walkDown
-            .collect {
-              case Node(varDef: Ast.VarDef[_], _) if AstExtra.containsNamedVar(varDef, identNode.data) =>
-                SourceLocation.Node(varDef, sourceCode)
-            }
+          ScopeWalker.walk(
+            from = block,
+            anchor = identNode.data
+          ) {
+            case Node(varDef: Ast.VarDef[_], _) if AstExtra.containsNamedVar(varDef, identNode.data) =>
+              SourceLocation.Node(varDef, sourceCode)
+          }
       }
 
   /**

presentation-compiler/src/main/scala/org/alephium/ralph/lsp/pc/search/gotodef/ScopeWalker.scala

@@ -0,0 +1,74 @@
+// Copyright 2024 The Alephium Authors
+// This file is part of the alephium project.
+//
+// The library is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Lesser General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// The library 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 Lesser General Public License for more details.
+//
+// You should have received a copy of the GNU Lesser General Public License
+// along with the library. If not, see http://www.gnu.org/licenses/.
+
+package org.alephium.ralph.lsp.pc.search.gotodef
+
+import org.alephium.ralph.Ast
+import org.alephium.ralph.lsp.access.compiler.ast.AstExtra
+import org.alephium.ralph.lsp.access.compiler.ast.node.Node
+import org.alephium.ralph.lsp.access.compiler.message.SourceIndexExtra
+
+import scala.collection.mutable.ListBuffer
+
+private object ScopeWalker {
+
+  /**
+   * Navigates the nodes within the scope of the `anchor` node, starting from the `from` node.
+   *
+   * @param from   The node where the search starts.
+   * @param anchor The node which is being scoped and where the search ends.
+   *               If the collected result is empty, nodes after the `anchor`'s position
+   *               are processed until at least one item is collected.
+   * @param pf     Only the Nodes defined by this partial function are collected.
+   * @return Nodes within the scope of the anchor AST.
+   */
+  def walk[T](
+      from: Node[Ast.Positioned, Ast.Positioned],
+      anchor: Ast.Positioned
+    )(pf: PartialFunction[Node[Ast.Positioned, Ast.Positioned], T]): Iterable[T] = {
+    val found  = ListBuffer.empty[T]
+    var walker = from.walkDown
+
+    while (walker.hasNext)
+      walker.next() match {
+        // Check: Is this a scoped node that does not contain the anchor node within its scope? If yes, drop all its child nodes.
+        // format: off
+        case block @ Node(_: Ast.While[_] | _: Ast.ForLoop[_] | _: Ast.IfBranch[_] | _: Ast.ElseBranch[_], _) if !SourceIndexExtra.contains(block.data.sourceIndex, anchor.sourceIndex) =>
+        // format: on
+          walker = walker dropWhile { // drop all child nodes
+            next =>
+              SourceIndexExtra.contains(
+                parent = block.data.sourceIndex,
+                child = next.data.sourceIndex
+              )
+          }
+
+        // Check:
+        // - Is this node (i.e., within the scope) defined by the partial-function?
+        // - And is it before the anchor node?
+        // - If it's defined after the anchor node (node in scope), then only add it if currently collected items are empty.
+        case node @ Node(ast, _) if pf.isDefinedAt(node) && (AstExtra.isBehind(ast, anchor) || found.isEmpty) =>
+          found addOne pf(node)
+
+        case _ =>
+        // ignore the rest
+      }
+
+    found
+
+  }
+
+}