feat(language-server): Resolve paths for imports (#1079)

aqua-src/antithesis.aqua

@@ -1,5 +1,7 @@
 aqua A
 
+import "aqua-src/gen/OneMore.aqua"
+
 export main
 
 alias SomeAlias: string

aqua-src/gen/OneMore.aqua

@@ -1,3 +1,5 @@
+aqua One
+
 service OneMore:
   more_call()
   consume(s: string)

build.sbt

@@ -82,7 +82,7 @@ lazy val `language-server-api` = crossProject(JSPlatform, JVMPlatform)
       "co.fs2"        %%% "fs2-io"      % fs2V
     )
   )
-  .dependsOn(compiler, io)
+  .dependsOn(compiler, io, compiler % "test->test")
 
 lazy val `language-server-apiJS` = `language-server-api`.js
   .settings(

compiler/src/main/scala/aqua/compiler/AquaCompiler.scala

@@ -3,17 +3,18 @@ package aqua.compiler
 import aqua.compiler.AquaError.*
 import aqua.linker.Linker
 import aqua.parser.{Ast, ParserError}
+import aqua.semantics.header.Picker.setImportPaths
 import aqua.semantics.header.{HeaderHandler, Picker}
-import aqua.semantics.{SemanticError, Semantics}
-
+import aqua.semantics.{FileId, SemanticError, Semantics}
 import cats.arrow.FunctionK
 import cats.data.*
 import cats.syntax.either.*
 import cats.syntax.functor.*
-import cats.{Comonad, Monad, Monoid, Order, ~>}
+import cats.syntax.show.*
+import cats.{~>, Comonad, Monad, Monoid, Order, Show}
 import scribe.Logging
 
-class AquaCompiler[F[_]: Monad, E, I: Order, S[_]: Comonad, C: Monoid: Picker](
+class AquaCompiler[F[_]: Monad, E, I: FileId, S[_]: Comonad, C: Monoid: Picker](
   headerHandler: HeaderHandler[S, C],
   semantics: Semantics[S, C]
 ) extends Logging {
@@ -27,7 +28,7 @@ class AquaCompiler[F[_]: Monad, E, I: Order, S[_]: Comonad, C: Monoid: Picker](
   // (Imports contexts => Compilation result)
   type TP = Map[String, C] => CompileRes[C]
 
-  private def transpile(body: Ast[S]): TP =
+  private def transpile(body: Ast[S], importPaths: Map[String, String]): TP =
     imports =>
       for {
         // Process header, get initial context
@@ -43,7 +44,7 @@ class AquaCompiler[F[_]: Monad, E, I: Order, S[_]: Comonad, C: Monoid: Picker](
         rc <- headerSem
           .finCtx(processed)
           .toCompileRes
-      } yield rc
+      } yield rc.setImportPaths(importPaths)
 
   def compileRaw(
     sources: AquaSources[F, E, I],
@@ -58,7 +59,10 @@ class AquaCompiler[F[_]: Monad, E, I: Order, S[_]: Comonad, C: Monoid: Picker](
         // Lift resolution to CompileRes
         modules <- resolution.toEitherT[CompileWarns]
         // Generate transpilation functions for each module
-        transpiled = modules.map(body => transpile(body))
+        transpiled = modules.map { m =>
+          val importPaths = m.imports.view.mapValues(_.show).toMap
+          m.copy(body = transpile(m.body, importPaths))
+        }
         // Link modules
         linked <- Linker.link(transpiled, CycleError.apply)
       } yield linked

compiler/src/main/scala/aqua/compiler/CompilerAPI.scala

@@ -5,21 +5,21 @@ import aqua.compiler.AquaError.*
 import aqua.model.AquaContext
 import aqua.parser.{Ast, ParserError}
 import aqua.raw.RawContext
-import aqua.semantics.RawSemantics
 import aqua.semantics.header.{HeaderHandler, HeaderSem}
-import aqua.semantics.rules.locations.{LocationsAlgebra, DummyLocationsInterpreter}
+import aqua.semantics.rules.locations.{DummyLocationsInterpreter, LocationsAlgebra}
+import aqua.semantics.{FileId, RawSemantics}
 
 import cats.data.*
 import cats.syntax.either.*
 import cats.syntax.flatMap.*
 import cats.syntax.functor.*
 import cats.syntax.traverse.*
-import cats.{Comonad, Monad, Monoid, Order}
+import cats.{Comonad, Monad, Monoid, Order, Show}
 import scribe.Logging
 
 object CompilerAPI extends Logging {
 
-  private def toAquaProcessed[I: Order, E, S[_]: Comonad](
+  private def toAquaProcessed[I: FileId, S[_]: Comonad](
     filesWithContext: Map[I, RawContext]
   ): Chain[AquaProcessed[I]] = {
     logger.trace("linking finished")
@@ -41,7 +41,7 @@ object CompilerAPI extends Logging {
       .value
   }
 
-  private def getAquaCompiler[F[_]: Monad, E, I: Order, S[_]: Comonad](
+  private def getAquaCompiler[F[_]: Monad, E, I: FileId, S[_]: Comonad](
     config: AquaCompilerConf
   ): AquaCompiler[F, E, I, S, RawContext] = {
     given Monoid[RawContext] = RawContext
@@ -66,7 +66,7 @@ object CompilerAPI extends Logging {
   }
 
   // Get result generated by backend
-  def compile[F[_]: Monad, E, I: Order, S[_]: Comonad](
+  def compile[F[_]: Monad, E, I: FileId, S[_]: Comonad](
     sources: AquaSources[F, E, I],
     parser: I => String => ValidatedNec[ParserError[S], Ast[S]],
     airValidator: AirValidator[F],
@@ -104,7 +104,7 @@ object CompilerAPI extends Logging {
     } yield result
   }
 
-  def compileToContext[F[_]: Monad, E, I: Order, S[_]: Comonad](
+  def compileToContext[F[_]: Monad, E, I: FileId, S[_]: Comonad](
     sources: AquaSources[F, E, I],
     parser: I => String => ValidatedNec[ParserError[S], Ast[S]],
     config: AquaCompilerConf

compiler/src/test/scala/aqua/compiler/AquaCompilerSpec.scala

@@ -1,5 +1,6 @@
 package aqua.compiler
 
+import aqua.compiler.FileIdString.given
 import aqua.model.AquaContext
 import aqua.model.CallServiceModel
 import aqua.model.FlattenModel
@@ -16,6 +17,7 @@ import aqua.raw.ConstantRaw
 import aqua.raw.value.{LiteralRaw, ValueRaw, VarRaw}
 import aqua.res.*
 import aqua.res.ResBuilder
+import aqua.semantics.FileId
 import aqua.types.{ArrayType, CanonStreamType, LiteralType, ScalarType, StreamType, Type}
 
 import cats.Id

compiler/src/test/scala/aqua/compiler/FileIdString.scala

@@ -0,0 +1,11 @@
+package aqua.compiler
+
+import aqua.semantics.FileId
+
+object FileIdString {
+  given FileId[String] with {
+    override def show(t: String): String = t
+
+    override def compare(x: String, y: String): Int = x.compare(y)
+  }
+}

io/src/main/scala/aqua/files/FileModuleId.scala

@@ -1,18 +1,20 @@
 package aqua.files
 
+import aqua.semantics.FileId
 import fs2.io.file.Path
-import cats.Order
+import cats.{Order, Show}
 
 case class FileModuleId private (file: Path) {
-  override def toString: String = s"${file}"
+  override def toString: String = s"$file"
 }
 
 object FileModuleId {
 
-  implicit object FileModuleIdOrder extends Order[FileModuleId] {
-
+  given FileId[FileModuleId] with {
     override def compare(x: FileModuleId, y: FileModuleId): Int =
       x.file.toString.compareTo(y.file.toString)
+
+    override def show(t: FileModuleId): String = t.toString
   }
 
   def apply(file: Path): FileModuleId =

language-server/language-server-api/.js/src/main/scala/aqua/lsp/ResultHelper.scala

@@ -107,11 +107,12 @@ object ResultHelper extends Logging {
       link.toList
     }.toJSArray
 
-  private def importsToTokenImport(imports: List[LiteralToken[FileSpan.F]]): js.Array[TokenImport] =
+  private def importsToTokenImport(imports: List[LiteralToken[FileSpan.F]], importPaths: Map[String, String]): js.Array[TokenImport] =
     imports.flatMap { lt =>
       val (span, str) = lt.valueToken
       val unquoted = str.substring(1, str.length - 1)
-      TokenLocation.fromSpan(span).map(l => TokenImport(l, unquoted))
+      val path = importPaths.getOrElse(unquoted, unquoted)
+      TokenLocation.fromSpan(span).map(l => TokenImport(l, path))
     }.toJSArray
 
   def lspToCompilationResult(lsp: LspContext[FileSpan.F]): CompilationResult = {
@@ -124,11 +125,13 @@ object ResultHelper extends Logging {
       case errs =>
         logger.debug("Errors: " + errs.mkString("\n"))
 
+    val importTokens = importsToTokenImport(lsp.importTokens, lsp.importPaths)
+
     CompilationResult(
       errors.toJSArray,
       warnings.toJSArray,
       locationsToJs(lsp.variables.flatMap(v => v.allLocations)),
-      importsToTokenImport(lsp.importTokens),
+      importTokens,
       tokensToJs(lsp.variables.map(_.definition))
     )
   }

language-server/language-server-api/src/main/scala/aqua/lsp/LSPCompiler.scala

@@ -5,18 +5,18 @@ import aqua.parser.{Ast, ParserError}
 import aqua.raw.RawContext
 import aqua.semantics.header.{HeaderHandler, HeaderSem}
 import aqua.semantics.rules.locations.LocationsAlgebra
-
+import aqua.semantics.FileId
 import cats.data.Validated.validNec
-import cats.data.{State, Chain, Validated, ValidatedNec}
+import cats.data.{Chain, State, Validated, ValidatedNec}
 import cats.syntax.either.*
 import cats.syntax.functor.*
 import cats.syntax.monoid.*
 import cats.syntax.semigroup.*
-import cats.{Comonad, Monad, Monoid, Order}
+import cats.{Comonad, Monad, Monoid, Order, Show}
 
 object LSPCompiler {
 
-  private def getLspAquaCompiler[F[_]: Monad, E, I: Order, S[_]: Comonad](
+  private def getLspAquaCompiler[F[_]: Monad, E, I: FileId, S[_]: Comonad](
     config: AquaCompilerConf
   ): AquaCompiler[F, E, I, S, LspContext[S]] = {
     given Monoid[LspContext[S]] = LspContext
@@ -57,7 +57,7 @@ object LSPCompiler {
     )
   }
 
-  def compileToLsp[F[_]: Monad, E, I: Order, S[_]: Comonad](
+  def compileToLsp[F[_]: Monad, E, I: FileId, S[_]: Comonad](
     sources: AquaSources[F, E, I],
     parser: I => String => ValidatedNec[ParserError[S], Ast[S]],
     config: AquaCompilerConf

language-server/language-server-api/src/main/scala/aqua/lsp/LspContext.scala

@@ -22,7 +22,8 @@ case class LspContext[S[_]](
   variables: List[VariableInfo[S]] = Nil,
   importTokens: List[LiteralToken[S]] = Nil,
   errors: List[SemanticError[S]] = Nil,
-  warnings: List[SemanticWarning[S]] = Nil
+  warnings: List[SemanticWarning[S]] = Nil,
+  importPaths: Map[String, String] = Map.empty
 ) {
   lazy val allLocations: List[TokenLocation[S]] = variables.flatMap(_.allLocations)
 }
@@ -41,7 +42,8 @@ object LspContext {
         importTokens = x.importTokens ++ y.importTokens,
         variables = x.variables ++ y.variables,
         errors = x.errors ++ y.errors,
-        warnings = x.warnings ++ y.warnings
+        warnings = x.warnings ++ y.warnings,
+        importPaths = x.importPaths ++ y.importPaths
       )
 
   trait Implicits[S[_]] {
@@ -101,6 +103,9 @@ object LspContext {
         )
       )
 
+    override def setImportPaths(ctx: LspContext[S], importPaths: Map[String, String]): LspContext[S] =
+      ctx.copy(importPaths = importPaths)
+
     override def setModule(
       ctx: LspContext[S],
       name: Option[String],

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

@@ -1,5 +1,6 @@
 package aqua.lsp
 
+import aqua.compiler.FileIdString.given_FileId_String
 import aqua.compiler.{AquaCompilerConf, AquaError, AquaSources}
 import aqua.parser.Parser
 import aqua.parser.lift.Span
@@ -10,7 +11,6 @@ import aqua.types.*
 
 import cats.Id
 import cats.data.*
-import cats.instances.string.*
 import org.scalatest.Inside
 import org.scalatest.flatspec.AnyFlatSpec
 import org.scalatest.matchers.should.Matchers
@@ -423,7 +423,8 @@ class AquaLSPSpec extends AnyFlatSpec with Matchers with Inside {
             ),
             ProductType(nestedType :: Nil)
           )
-        ), ("check2", ArrowType(NilType, ProductType(someStr :: Nil))),
+        ),
+        ("check2", ArrowType(NilType, ProductType(someStr :: Nil))),
         ("check3", ArrowType(NilType, ProductType(ScalarType.string :: Nil)))
       )
     )

linker/src/main/scala/aqua/linker/AquaModule.scala

@@ -1,10 +1,3 @@
 package aqua.linker
 
-case class AquaModule[I, E, T](id: I, imports: Map[String, I], dependsOn: Map[I, E], body: T) {
-  def map[TT](f: T => TT): AquaModule[I, E, TT] = copy(body = f(body))
-
-  def mapWithId[TT](f: (I, T) => TT): AquaModule[I, E, TT] = copy(body = f(id, body))
-
-  def mapErr[EE](f: E => EE): AquaModule[I, EE, T] =
-    copy(dependsOn = dependsOn.view.mapValues(f).toMap)
-}
+case class AquaModule[I, E, T](id: I, imports: Map[String, I], dependsOn: Map[I, E], body: T)

linker/src/main/scala/aqua/linker/Modules.scala

@@ -30,14 +30,8 @@ case class Modules[I, E, T](
 
   def isResolved: Boolean = dependsOn.isEmpty
 
-  def map[TT](f: T => TT): Modules[I, E, TT] =
-    copy(loaded = loaded.view.mapValues(_.map(f)).toMap)
-
-  def mapErr[EE](f: E => EE): Modules[I, EE, T] =
-    copy(
-      loaded = loaded.view.mapValues(_.mapErr(f)).toMap,
-      dependsOn = dependsOn.view.mapValues(_.map(f)).toMap
-    )
+  def map[TT](f: AquaModule[I, E, T] => AquaModule[I, E, TT]): Modules[I, E, TT] =
+    copy(loaded = loaded.view.mapValues(f).toMap)
 }
 
 object Modules {

semantics/src/main/scala/aqua/semantics/FileId.scala

@@ -0,0 +1,7 @@
+package aqua.semantics
+
+import cats.Show
+import cats.kernel.Order
+import cats.syntax.order.*
+
+trait FileId[I] extends Show[I] with Order[I]

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

@@ -22,6 +22,7 @@ trait Picker[A] {
   def funcAcceptAbility(ctx: A, name: String): Boolean
   def declares(ctx: A): Set[String]
   def setAbility(ctx: A, name: String, ctxAb: A): A
+  def setImportPaths(ctx: A, importPaths: Map[String, String]): A
   def setModule(ctx: A, name: Option[String], declares: Set[String]): A
   def setExports(ctx: A, exports: Map[String, Option[String]]): A
   def setInit(ctx: A, ctxInit: Option[A]): A
@@ -51,6 +52,7 @@ object Picker {
     def funcAcceptAbility(name: String): Boolean = Picker[A].funcAcceptAbility(p, name)
     def declares: Set[String] = Picker[A].declares(p)
     def setAbility(name: String, ctx: A): A = Picker[A].setAbility(p, name, ctx)
+    def setImportPaths(importPaths: Map[String, String]): A = Picker[A].setImportPaths(p, importPaths)
     def setInit(ctx: Option[A]): A = Picker[A].setInit(p, ctx)
     def addPart(part: (A, RawPart)): A = Picker[A].addPart(p, part)
 
@@ -117,6 +119,10 @@ object Picker {
     override def setAbility(ctx: RawContext, name: String, ctxAb: RawContext): RawContext =
       ctx.copy(abilities = Map(name -> ctxAb))
 
+    // dummy
+    override def setImportPaths(ctx: RawContext, importPaths: Map[String, String]): RawContext =
+      ctx
+
     override def setModule(
       ctx: RawContext,
       name: Option[String],