From 07bea1a909d9d66b6eb4a32cb99a990458c3f77e Mon Sep 17 00:00:00 2001
From: Dima <dmitry.shakhtarin@fluence.ai>
Date: Mon, 8 Apr 2024 15:43:30 +0300
Subject: [PATCH] fix(lsp): Fix go-to-definition for abilities as arguments in
 functions [LNG-343] (#1116)

---
 .../scala/aqua/lsp/LocationsInterpreter.scala | 11 ++++--
 .../src/test/scala/aqua/lsp/AquaLSPSpec.scala | 37 ++++++++++++++++++-
 .../aqua/semantics/expr/func/ArrowSem.scala   |  2 -
 3 files changed, 44 insertions(+), 6 deletions(-)

diff --git a/language-server/language-server-api/src/main/scala/aqua/lsp/LocationsInterpreter.scala b/language-server/language-server-api/src/main/scala/aqua/lsp/LocationsInterpreter.scala
index 084d55c7..2ddacff3 100644
--- a/language-server/language-server-api/src/main/scala/aqua/lsp/LocationsInterpreter.scala
+++ b/language-server/language-server-api/src/main/scala/aqua/lsp/LocationsInterpreter.scala
@@ -3,6 +3,7 @@ package aqua.lsp
 import aqua.parser.lexer.Token
 import aqua.semantics.rules.locations.{DefinitionInfo, LocationsAlgebra, LocationsState}
 import aqua.types.AbilityType
+
 import cats.data.State
 import monocle.Lens
 import scribe.Logging
@@ -13,9 +14,13 @@ class LocationsInterpreter[S[_], X](using
 
   type SX[A] = State[X, A]
 
-  override def addDefinition(definition: DefinitionInfo[S]): State[X, Unit] = modify { st =>
-    st.addDefinition(definition)
-  }
+  override def addDefinition(definition: DefinitionInfo[S]): State[X, Unit] =
+    definition.`type` match {
+      // case where ability is an {Argument} in a function
+      case t: AbilityType if definition.name == t.name =>
+        pointLocation(definition.name, definition.token)
+      case _ => modify { st => st.addDefinition(definition) }
+    }
 
   override def addDefinitionWithFields(
     definition: DefinitionInfo[S],
diff --git a/language-server/language-server-api/src/test/scala/aqua/lsp/AquaLSPSpec.scala b/language-server/language-server-api/src/test/scala/aqua/lsp/AquaLSPSpec.scala
index da3612a6..de349f23 100644
--- a/language-server/language-server-api/src/test/scala/aqua/lsp/AquaLSPSpec.scala
+++ b/language-server/language-server-api/src/test/scala/aqua/lsp/AquaLSPSpec.scala
@@ -448,7 +448,7 @@ class AquaLSPSpec extends AnyFlatSpec with Matchers with Inside {
     // from {SomeAbility} to 'ability SomeAbility'
     res.checkLocations("SomeAbility", 0, 1, main) shouldBe true
     // from 'SomeAbility.someStr' to {SomeAbility}
-    res.checkLocations("SomeAbility", 1, 2, main) shouldBe true
+    res.checkLocations("SomeAbility", 0, 2, main) shouldBe true
 
     res.checkTokenLoc(main, "Srv", 0, srvType) shouldBe true
     Range.inclusive(1, 3).foreach { n =>
@@ -503,4 +503,39 @@ class AquaLSPSpec extends AnyFlatSpec with Matchers with Inside {
     res.checkLocations("someString", 1, 2, firstImport, Some(main)) shouldBe true
     res.checkLocations("someString", 1, 3, firstImport, Some(main)) shouldBe true
   }
+
+  it should "return right tokens for multiple abilities" in {
+    val main =
+      """aqua Import declares *
+        |
+        |export main
+        |
+        |ability Abilyy:
+        |  field: string
+        |
+        |func firstFunc{Abilyy}() -> string:
+        |  <- "str"
+        |
+        |func secondFunc{Abilyy}() -> string:
+        |  <- "str"
+        |
+        |func main() -> string:
+        |  ab = Abilyy(field = "123")
+        |  res <- firstFunc{ab}()
+        |  secondFunc{ab}()
+        |  <- res
+        |""".stripMargin
+    val src = Map(
+      "index.aqua" -> main
+    )
+
+    val imports = Map.empty[String, String]
+
+    val res = compile(src, imports).toOption.get.values.head
+
+    res.errors shouldBe empty
+    res.checkLocations("Abilyy", 0, 1, main) shouldBe true
+    res.checkLocations("Abilyy", 0, 2, main) shouldBe true
+    res.checkLocations("Abilyy", 0, 3, main) shouldBe true
+  }
 }
diff --git a/semantics/src/main/scala/aqua/semantics/expr/func/ArrowSem.scala b/semantics/src/main/scala/aqua/semantics/expr/func/ArrowSem.scala
index 7dc2522b..f4ac5f57 100644
--- a/semantics/src/main/scala/aqua/semantics/expr/func/ArrowSem.scala
+++ b/semantics/src/main/scala/aqua/semantics/expr/func/ArrowSem.scala
@@ -33,8 +33,6 @@ class ArrowSem[S[_]](val expr: ArrowExpr[S]) extends AnyVal {
     L: LocationsAlgebra[S, Alg]
   ): Alg[ArrowType] = for {
     arrowType <- T.beginArrowScope(arrowTypeExpr)
-    // add locations before ability will be defined as new variable definition
-    _ <- L.pointLocations(arrowTypeExpr.abilities.map(n => n.value -> n))
     absAsArgs = arrowTypeExpr.abilities.map(_.asName)
     // Create local variables
     _ <- arrowTypeExpr.absWithArgs.flatMap { case (name, _) => name }