fix(lsp): Go-to-definition for used types (LNG-345) (#1128)

aqua-src/antithesis.aqua

@@ -1,29 +1,13 @@
 aqua Job declares *
 
-export aaa, bbb
+use "declare.aqua"
 
-data Peer:
-  id: string
+export timeout
 
-func aaa() -> Peer:
-  peer1 = Peer(id = "123")
-  peer2 = Peer(id = peer1.id)
-  <- peer2
+data Worker:
+  field: string
 
-data BrokenStruct:
-  fff: UnknownType
-
-alias BrokenAlias: str3ng
-
-ability BrokenAbility:
-  fff: str4ng
-
-const BROKEN_CONST = UNKNOWN_CONST
-
-func bbb() -> string:
-  <- 323
-
-func ccc() -> Peer:
-  peer1 = Peer(id = "123")
-  peer2 = Peer(id = peer1.id)
-  <- peer2
+func timeout() -> Worker:
+  w <- AquaName.getWorker()
+  a = w.host_id
+  <- w

aqua-src/declare.aqua

@@ -1,15 +1,8 @@
-aqua DeclareModule declares decl_foo, decl_bar, SuperFoo, DECLARE_CONST, DECLARE_CONST2
-export SuperFoo
+aqua AquaName declares getWorker, Worker
 
-const DECLARE_CONST = "declare_const"
-const DECLARE_CONST2 = "declare_const2"
+data Worker:
+  host_id: string
 
-service SuperFoo("super_foo"):
-    small_foo() -> string
+func getWorker() -> Worker:
+  <- Worker(host_id = "")
 
-func decl_foo() -> string:
-    res1 <- SuperFoo.small_foo()
-    <- res1
-
-func decl_bar() -> string:
-    <- DECLARE_CONST

language-server/language-server-api/src/test/scala/aqua/lsp/AquaLSPSpec.scala

@@ -23,6 +23,12 @@ class AquaLSPSpec extends AnyFlatSpec with Matchers with Inside {
     str.r.findAllMatchIn(code).toList.lift(position).map(r => (r.start, r.end))
   }
 
+  extension [T](o: Option[T]) {
+
+    def tapNone(f: => Unit): Option[T] =
+      o.orElse { f; None }
+  }
+
   extension (c: LspContext[Span.S]) {
 
     def checkLocations(
@@ -34,12 +40,14 @@ class AquaLSPSpec extends AnyFlatSpec with Matchers with Inside {
       fieldOrSynonym: Option[String] = None
     ): Boolean = {
       (for {
-        defPos <- getByPosition(defCode, name, defPosition)
+        defPos <- getByPosition(defCode, name, defPosition).tapNone(
+          fail(s"Didn't find definition of '$name'")
+        )
         usePos <- getByPosition(
           useCode.getOrElse(defCode),
           fieldOrSynonym.getOrElse(name),
           usePosition
-        )
+        ).tapNone(fail(s"Didn't find usage of '$name'"))
       } yield {
         val (defStart, defEnd) = defPos
         val (useStart, useEnd) = usePos
@@ -520,6 +528,56 @@ class AquaLSPSpec extends AnyFlatSpec with Matchers with Inside {
     res.checkLocations("someString", 1, 3, firstImport, Some(main)) shouldBe true
   }
 
+  it should "return right tokens in 'use'd structures" in {
+    val main =
+      """aqua Job declares *
+        |
+        |use "declare.aqua"
+        |
+        |export timeout
+        |
+        |func timeout() -> string:
+        |  w <- AquaName.getWorker()
+        |  a = w.host_id
+        |  ab = AquaName.SomeAbility(getWrk = AquaName.getWorker, someField = "123")
+        |  c = ab.getWrk()
+        |  d = ab.someField
+        |  <- a
+        |""".stripMargin
+    val src = Map(
+      "index.aqua" -> main
+    )
+
+    val firstImport =
+      """aqua AquaName declares getWorker, Worker, SomeAbility
+        |
+        |data Worker:
+        |  host_id: string
+        |
+        |ability SomeAbility:
+        |  getWrk() -> Worker
+        |  someField: string
+        |
+        |func getWorker() -> Worker:
+        |  <- Worker(host_id = "")
+        |
+        |""".stripMargin
+
+    val imports = Map(
+      "declare.aqua" ->
+        firstImport
+    )
+
+    val res = compile(src, imports).toOption.get.values.head
+
+    res.errors shouldBe empty
+    res.checkLocations("host_id", 0, 0, firstImport, Some(main)) shouldBe true
+    res.checkLocations("getWorker", 1, 0, firstImport, Some(main)) shouldBe true
+    res.checkLocations("getWorker", 1, 0, firstImport) shouldBe true
+    res.checkLocations("getWrk", 0, 1, firstImport, Some(main)) shouldBe true
+    res.checkLocations("someField", 0, 1, firstImport, Some(main)) shouldBe true
+  }
+
   it should "return right tokens for multiple abilities" in {
     val main =
       """aqua Import declares *

model/raw/src/main/scala/aqua/raw/ConstantRaw.scala

@@ -7,6 +7,8 @@ case class ConstantRaw(name: String, value: ValueRaw, allowOverrides: Boolean) e
   override def rename(s: String): RawPart = copy(name = s)
 
   override def rawPartType: Type = value.`type`
+
+  def addAbilityName(s: String): RawPart = this
 }
 
 object ConstantRaw {

model/raw/src/main/scala/aqua/raw/ErroredPart.scala

@@ -7,4 +7,6 @@ case class ErroredPart(name: String) extends RawPart {
   override def rawPartType: Type = BottomType
 
   override def rename(s: String): RawPart = copy(name = s)
+
+  def addAbilityName(s: String): RawPart = this
 }

model/raw/src/main/scala/aqua/raw/RawPart.scala

@@ -11,6 +11,8 @@ trait RawPart extends Raw {
   def rawPartType: Type
 
   def rename(s: String): RawPart
+
+  def addAbilityName(s: String): RawPart
 }
 
 object RawPart {

model/raw/src/main/scala/aqua/raw/ServiceRaw.scala

@@ -1,6 +1,6 @@
 package aqua.raw
 
-import aqua.types.ServiceType
+import aqua.types.{ServiceType, Type}
 import aqua.raw.value.ValueRaw
 
 case class ServiceRaw(
@@ -12,4 +12,5 @@ case class ServiceRaw(
 
   override def rename(s: String): RawPart = copy(name = s)
 
+  def addAbilityName(s: String): RawPart = copy(`type` = Type.addAbilityNameService(s, `type`))
 }

model/raw/src/main/scala/aqua/raw/TypeRaw.scala

@@ -6,4 +6,6 @@ case class TypeRaw(name: String, `type`: Type) extends RawPart {
   override def rename(s: String): RawPart = copy(name = s)
 
   override def rawPartType: Type = `type`
+
+  def addAbilityName(s: String): RawPart = copy(`type` = Type.addAbilityName(s, `type`))
 }

model/raw/src/main/scala/aqua/raw/arrow/FuncRaw.scala

@@ -9,6 +9,8 @@ case class FuncRaw(
 ) extends RawPart {
   override def rename(s: String): RawPart = copy(name = s)
 
+  def addAbilityName(s: String): RawPart = copy(arrow = arrow.copy(`type` = Type.addAbilityNameArrow(s, arrow.`type`)))
+
   override def rawPartType: Type = arrow.`type`
 
   // vars that we capture from external space (outer functions, etc)

parser/src/main/scala/aqua/parser/lexer/ValueToken.scala

@@ -133,6 +133,7 @@ case class PropertyToken[F[_]: Comonad](
           case IntoField(name) => name.extract.some
           case _ => none
         }.map { props =>
+          // TODO: this loses the correct token location
           val typeName = name
             .rename(
               (name.value +: props).mkString(".")

semantics/src/main/scala/aqua/semantics/header/Picker.scala

@@ -130,7 +130,10 @@ object Picker {
     override def allNames(ctx: RawContext): Set[String] = ctx.allNames
 
     override def setAbility(ctx: RawContext, name: String, ctxAb: RawContext): RawContext =
-      ctx.copy(abilities = Map(name -> ctxAb))
+      ctx.copy(abilities = Map(name -> ctxAb.copy(parts = ctxAb.parts.map {
+        case (partContext, part) =>
+          (partContext, part.addAbilityName(name))
+      })))
 
     // dummy
     override def setImportPaths(ctx: RawContext, importPaths: Map[String, String]): RawContext =

semantics/src/main/scala/aqua/semantics/rules/ValuesAlgebra.scala

@@ -18,7 +18,6 @@ import cats.data.{NonEmptyList, OptionT}
 import cats.instances.list.*
 import cats.syntax.applicative.*
 import cats.syntax.apply.*
-import cats.syntax.bifunctor.*
 import cats.syntax.flatMap.*
 import cats.syntax.foldable.*
 import cats.syntax.functor.*

types/src/main/scala/aqua/types/Type.scala

@@ -57,6 +57,8 @@ sealed trait ProductType extends Type {
 
   def length: Int
 
+  def map(f: Type => Type): ProductType
+
   def uncons: Option[(Type, ProductType)] = this match {
     case ConsType(t, pt) => Some(t -> pt)
     case _ => None
@@ -149,14 +151,20 @@ object ConsType {
 }
 
 case class LabeledConsType(label: String, `type`: Type, tail: ProductType) extends ConsType {
+  def map(f: Type => Type): ProductType = copy(`type` = f(`type`), tail = tail.map(f))
+
   override def toString: String = s"($label: " + `type` + s") :: $tail"
 }
 
 case class UnlabeledConsType(`type`: Type, tail: ProductType) extends ConsType {
+  def map(f: Type => Type): ProductType = copy(`type` = f(`type`), tail = tail.map(f))
+
   override def toString: String = `type`.toString + s" :: $tail"
 }
 
 object NilType extends ProductType {
+  def map(f: Type => Type): ProductType = this
+
   override def toString: String = "∅"
 
   override def isInhabited: Boolean = false
@@ -546,6 +554,56 @@ object Type {
       t.toString
   }
 
+  def addAbilityNameProduct(abName: String, t: ProductType): ProductType =
+    t.map(t => addAbilityName(abName, t))
+
+  def addAbilityNameArrow(abName: String, t: ArrowType): ArrowType = {
+    t.copy(
+      domain = addAbilityNameProduct(abName, t.domain),
+      codomain = addAbilityNameProduct(abName, t.codomain)
+    )
+  }
+
+  def addAbilityNameData(abName: String, dt: DataType): DataType =
+    dt match {
+      case st @ StructType(name, fields) =>
+        st.copy(
+          name = AbilityType.fullName(abName, name),
+          fields = fields.map(Type.addAbilityName(abName, _))
+        )
+      case ot @ OptionType(el) => ot.copy(element = addAbilityNameData(abName, el))
+      case at @ ArrayType(el) => at.copy(element = addAbilityNameData(abName, el))
+      case t => t
+    }
+
+  def addAbilityNameService(abName: String, t: ServiceType): ServiceType = {
+    t.copy(
+      name = AbilityType.fullName(abName, t.name),
+      fields = t.fields.map(Type.addAbilityNameArrow(abName, _))
+    )
+  }
+
+  // Add ability name to type names
+  def addAbilityName(abName: String, t: Type): Type = {
+    t match {
+      case at @ AbilityType(name, fields) =>
+        at.copy(
+          name = AbilityType.fullName(abName, name),
+          fields = fields.map(Type.addAbilityName(abName, _))
+        )
+      case st: ServiceType =>
+        addAbilityNameService(abName, st)
+
+      case at: ArrowType =>
+        addAbilityNameArrow(abName, at)
+      case st @ CanonStreamType(el) => st.copy(element = addAbilityNameData(abName, el))
+      case st @ StreamType(el) => st.copy(element = addAbilityNameData(abName, el))
+      case smt @ StreamMapType(el) => smt.copy(element = addAbilityNameData(abName, el))
+      case pt: ProductType => addAbilityNameProduct(abName, pt)
+      case t: DataType => addAbilityNameData(abName, t)
+    }
+  }
+
   // pretty print for Type
   given Show[Type] = {
     case ArrayType(el) =>