fix(compiler): Allow returning resolved service as ability [LNG-266]

build.sbt

@@ -7,6 +7,7 @@ val catsV = "2.10.0"
 val catsParseV = "0.3.10"
 val monocleV = "3.1.0"
 val scalaTestV = "3.2.17"
+val scalaTestScalaCheckV = "3.2.17.0"
 val sourcecodeV = "0.3.0"
 val fs2V = "3.9.3"
 val catsEffectV = "3.6-1f95fd7"
@@ -23,8 +24,9 @@ val commons = Seq(
   },
   scalaVersion := scalaV,
   libraryDependencies ++= Seq(
-    "com.outr"      %%% "scribe"    % scribeV,
-    "org.scalatest" %%% "scalatest" % scalaTestV % Test
+    "com.outr"          %%% "scribe"          % scribeV,
+    "org.scalatest"     %%% "scalatest"       % scalaTestV           % Test,
+    "org.scalatestplus" %%% "scalacheck-1-17" % scalaTestScalaCheckV % Test
   ),
   scalacOptions ++= {
     Seq(

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

@@ -1,30 +1,30 @@
 package aqua.compiler
 
-import aqua.compiler.AquaError.{ParserError as AquaParserError, *}
 import aqua.backend.Backend
+import aqua.compiler.AquaError.{ParserError as AquaParserError, *}
 import aqua.linker.{AquaModule, Linker, Modules}
 import aqua.model.AquaContext
 import aqua.parser.lift.{LiftParser, Span}
 import aqua.parser.{Ast, ParserError}
 import aqua.raw.RawPart.Parts
 import aqua.raw.{RawContext, RawPart}
 import aqua.res.AquaRes
-import aqua.semantics.{CompilerState, Semantics}
 import aqua.semantics.header.{HeaderHandler, HeaderSem, Picker}
+import aqua.semantics.{CompilerState, Semantics}
 import aqua.semantics.{SemanticError, SemanticWarning}
 
+import cats.arrow.FunctionK
 import cats.data.*
-import cats.data.Validated.{validNec, Invalid, Valid}
+import cats.data.Validated.{Invalid, Valid, validNec}
 import cats.parse.Parser0
 import cats.syntax.applicative.*
+import cats.syntax.either.*
 import cats.syntax.flatMap.*
 import cats.syntax.functor.*
 import cats.syntax.monoid.*
-import cats.syntax.traverse.*
 import cats.syntax.semigroup.*
-import cats.syntax.either.*
-import cats.{~>, Comonad, Functor, Monad, Monoid, Order}
-import cats.arrow.FunctionK
+import cats.syntax.traverse.*
+import cats.{Comonad, Functor, Monad, Monoid, Order, ~>}
 import scribe.Logging
 
 class AquaCompiler[F[_]: Monad, E, I: Order, S[_]: Comonad, C: Monoid: Picker](

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

@@ -120,7 +120,7 @@ class AquaCompilerSpec extends AnyFlatSpec with Matchers with Inside {
 
   def getDataSrv(name: String, varName: String, t: Type) = {
     CallServiceRes(
-      LiteralModel.fromRaw(LiteralRaw.quote("getDataSrv")),
+      LiteralModel.quote("getDataSrv"),
       name,
       CallRes(Nil, Some(CallModel.Export(varName, t))),
       LiteralModel.fromRaw(ValueRaw.InitPeerId)
@@ -364,7 +364,66 @@ class AquaCompilerSpec extends AnyFlatSpec with Matchers with Inside {
       errorCall(transformCfg, 0, initPeer)
     )
 
-    insideRes(src, transformCfg = transformCfg)("main") { case main :: _ =>
+    insideRes(src, transformCfg = transformCfg)("main") { case main :: Nil =>
+      main.body.equalsOrShowDiff(expected) should be(true)
+    }
+  }
+
+  it should "allow returning and passing services as abilities" in {
+    val src = Map(
+      "main.aqua" -> """
+                       |aqua Test
+                       |
+                       |export test
+                       |
+                       |ability Ab:
+                       |    log(log: string)
+                       |
+                       |service Srv("default-id"):
+                       |    log(log: string)
+                       |
+                       |func useAb{Ab}():
+                       |  Ab.log("test")
+                       |
+                       |func genDefault() -> Ab:
+                       |    <- Srv
+                       |
+                       |func genResolved() -> Ab:
+                       |    Srv "resolved-id"
+                       |    <- Srv
+                       |
+                       |func test():
+                       |  resolved <- genResolved()
+                       |  useAb{resolved}()
+                       |  default <- genDefault()
+                       |  useAb{default}()
+                       |""".stripMargin
+    )
+
+    val transformCfg = TransformConfig()
+
+    insideRes(src, transformCfg = transformCfg)("test") { case main :: Nil =>
+      def srvCall(id: String) =
+        CallServiceRes(
+          serviceId = LiteralModel.quote(id),
+          funcName = "log",
+          call = CallRes(
+            List(LiteralModel.quote("test")),
+            None
+          ),
+          initPeer
+        ).leaf
+
+      val expected = XorRes.wrap(
+        SeqRes.wrap(
+          getDataSrv("-relay-", "-relay-", ScalarType.string),
+          srvCall("resolved-id"),
+          srvCall("default-id"),
+          emptyRespCall(transformCfg, initPeer)
+        ),
+        errorCall(transformCfg, 0, initPeer)
+      )
+
       main.body.equalsOrShowDiff(expected) should be(true)
     }
   }

model/inline/src/main/scala/aqua/model/inline/ArrowInliner.scala

@@ -6,15 +6,17 @@ import aqua.model.*
 import aqua.model.inline.state.{Arrows, Exports, Mangler}
 import aqua.raw.ops.RawTag
 import aqua.raw.value.{ValueRaw, VarRaw}
-import aqua.types.{AbilityType, ArrowType, CollectionType, NamedType, StreamType, Type}
+import aqua.types.*
 
 import cats.data.StateT
 import cats.data.{Chain, IndexedStateT, State}
+import cats.kernel.Semigroup
 import cats.syntax.applicative.*
 import cats.syntax.bifunctor.*
 import cats.syntax.foldable.*
 import cats.syntax.functor.*
 import cats.syntax.option.*
+import cats.syntax.semigroup.*
 import cats.syntax.show.*
 import cats.syntax.traverse.*
 import cats.{Eval, Monoid}
@@ -105,7 +107,9 @@ object ArrowInliner extends Logging {
     exports <- Exports[S].exports
     arrows <- Arrows[S].arrows
     // gather all arrows and variables from abilities
-    returnedAbilities = rets.collect { case VarModel(name, at: AbilityType, _) => name -> at }
+    returnedAbilities = rets.collect { case ValueModel.Ability(vm, at) =>
+      vm.name -> at
+    }
     varsFromAbilities = returnedAbilities.flatMap { case (name, at) =>
       getAbilityVars(name, None, at, exports)
     }.toMap
@@ -138,9 +142,9 @@ object ArrowInliner extends Logging {
   private def getAbilityFields[T <: Type](
     name: String,
     newName: Option[String],
-    `type`: NamedType,
+    `type`: GeneralAbilityType,
     exports: Map[String, ValueModel]
-  )(fields: NamedType => Map[String, T]): Map[String, ValueModel] =
+  )(fields: GeneralAbilityType => Map[String, T]): Map[String, ValueModel] =
     fields(`type`).flatMap { case (fName, _) =>
       val fullName = AbilityType.fullName(name, fName)
       val newFullName = AbilityType.fullName(newName.getOrElse(name), fName)
@@ -162,7 +166,7 @@ object ArrowInliner extends Logging {
   private def getAbilityVars(
     abilityName: String,
     abilityNewName: Option[String],
-    abilityType: AbilityType,
+    abilityType: GeneralAbilityType,
     exports: Map[String, ValueModel]
   ): Map[String, ValueModel] = {
     val get = getAbilityFields(
@@ -173,7 +177,7 @@ object ArrowInliner extends Logging {
     )
 
     get(_.variables) ++ get(_.arrows).flatMap {
-      case arrow @ (_, vm: VarModel) =>
+      case arrow @ (_, vm @ ValueModel.Arrow(_, _)) =>
         arrow.some
       case (_, m) =>
         internalError(s"($m) cannot be an arrow")
@@ -193,7 +197,7 @@ object ArrowInliner extends Logging {
   private def getAbilityArrows(
     name: String,
     newName: Option[String],
-    `type`: NamedType,
+    `type`: GeneralAbilityType,
     exports: Map[String, ValueModel],
     arrows: Map[String, FuncArrow]
   ): Map[String, FuncArrow] = {
@@ -205,16 +209,16 @@ object ArrowInliner extends Logging {
     )
 
     get(_.arrows).flatMap {
-      case (_, VarModel(name, _, _)) =>
-        arrows.get(name).map(name -> _)
+      case (_, ValueModel.Arrow(vm, _)) =>
+        arrows.get(vm.name).map(vm.name -> _)
       case (_, m) =>
         internalError(s"($m) cannot be an arrow")
     }
   }
 
   private def getAbilityArrows[S: Arrows: Exports](
     name: String,
-    `type`: NamedType
+    `type`: GeneralAbilityType
   ): State[S, Map[String, FuncArrow]] = for {
     exports <- Exports[S].exports
     arrows <- Arrows[S].arrows
@@ -225,6 +229,16 @@ object ArrowInliner extends Logging {
     renamed: Map[String, T]
   )
 
+  given [T]: Monoid[Renamed[T]] with {
+    override def empty: Renamed[T] = Renamed(Map.empty, Map.empty)
+
+    override def combine(x: Renamed[T], y: Renamed[T]): Renamed[T] =
+      Renamed(
+        x.renames ++ y.renames,
+        x.renamed ++ y.renamed
+      )
+  }
+
   // TODO: Make this extension private somehow?
   extension [T](vals: Map[String, T]) {
 
@@ -250,6 +264,115 @@ object ArrowInliner extends Logging {
       )
     }
 
+  /**
+   * Correctly rename captured values and arrows of a function
+   *
+   * @param fn Function
+   * @param exports Exports state before calling/inlining
+   * @param arrows Arrows state before calling/inlining
+   * @return Renamed values and arrows
+   */
+  def renamedCaptured[S: Mangler](
+    fn: FuncArrow,
+    exports: Map[String, ValueModel],
+    arrows: Map[String, FuncArrow]
+  ): State[S, (Renamed[ValueModel], Renamed[FuncArrow])] = {
+    // Gather abilities related values
+    val abilitiesValues = fn.capturedValues.collect {
+      // Gather only top level abilities
+      case (name, ValueModel.Ability(vm, at)) if vm.properties.isEmpty =>
+        name -> (
+          at,
+          /**
+           * Gather all values related to `name`
+           * NOTE: It is important that `capturedValues` are
+           * populated by all values related to ability `name`
+           * on creation of `FuncArrow`.
+           */
+          Exports.gatherFrom(
+            name :: Nil,
+            fn.capturedValues
+          )
+        )
+    }
+    // Gather all abilities related names
+    val abilitiesValuesKeys = abilitiesValues.flatMap { case (_, (_, values)) =>
+      values.keySet
+    }
+
+    // Gather abilities related arrows
+    val abilitiesArrows = abilitiesValues.toList.foldMap { case (_, (_, values)) =>
+      Arrows.arrowsByValues(fn.capturedArrows, values).toList
+    }.toMap
+
+    // Gather all other values and arrows that are not related to abilities
+    val otherValues = fn.capturedValues -- abilitiesValuesKeys
+    val otherArrows = fn.capturedArrows -- abilitiesArrows.keySet
+
+    for {
+      // Calculate renaming based on abilities
+      valuesRenamed <- abilitiesValues.toList.traverse { case (name, (at, values)) =>
+        Mangler[S]
+          .findAndForbidName(name)
+          .map(rename =>
+            // Get renaming map for this ability
+            AbilityType
+              .renames(at)(name, rename)
+              // Add ability rename too
+              .updated(name, rename)
+          )
+          .map(renames =>
+            // This code is HACKERY!!!
+            val valuesRenamed = values.renamed(renames).map {
+              /**
+               * `VarModel` is sometimes used to point to an arrow.
+               * So if it is renamed, we should rename the `VarModel` too.
+               * Otherwise renamed value will be resolved
+               * to previous name when trying to resolve the arrow.
+               * But this should be done only if the name in model
+               * is the same as the name of the export,
+               * because export could point to another arrow.
+               */
+              case (name, ValueModel.Arrow(vm, _)) if renames.contains(vm.name) =>
+                name -> vm.copy(name = name)
+              /**
+               * `VarModel` is used to point to an ability.
+               * So if it is renamed, we should rename the `VarModel` too.
+               * Otherwise renamed value will be resolved
+               * to previous name when trying to resolve the ability.
+               */
+              case (name, ValueModel.Ability(vm, _)) =>
+                name -> vm.copy(name = name)
+              case v => v
+            }
+            Renamed(renames, valuesRenamed)
+          )
+      }.map(_.combineAll)
+
+      // Rename arrows according to values
+      arrowsRenamed = Renamed(
+        valuesRenamed.renames.filterKeys(abilitiesArrows.keySet).toMap,
+        abilitiesArrows.renamed(valuesRenamed.renames)
+      )
+
+      // Rename values and arrows unrelated to abilities
+      otherValuesRenamed <- findNewNames(otherValues)
+      otherArrowsValues = Arrows.arrowsByValues(
+        otherArrows,
+        otherValues
+      )
+      otherArrowsValuesRenamed = Renamed(
+        otherValuesRenamed.renames.filterKeys(otherArrowsValues.keySet).toMap,
+        otherArrowsValues.renamed(otherValuesRenamed.renames)
+      )
+
+      otherArrowsRenamed <- findNewNames(otherArrows -- otherArrowsValues.keySet)
+
+      values = valuesRenamed |+| otherValuesRenamed
+      arrows = arrowsRenamed |+| otherArrowsValuesRenamed |+| otherArrowsRenamed
+    } yield values -> arrows
+  }
+
   /**
    * Prepare the function and the context for inlining
    *
@@ -283,26 +406,8 @@ object ArrowInliner extends Logging {
     arrowRenames = args.arrowArgsRenames
     abRenames = args.abilityArgsRenames
 
-    /**
-     * Find new names for captured values and arrows
-     * to avoid collisions, then resolve them in context.
-     */
-    capturedValues <- findNewNames(fn.capturedValues)
-    /**
-     * If arrow correspond to a value,
-     * rename in accordingly to the value
-     */
-    capturedArrowValues = Arrows.arrowsByValues(
-      fn.capturedArrows,
-      fn.capturedValues
-    )
-    capturedArrowValuesRenamed = capturedArrowValues.renamed(
-      capturedValues.renames
-    )
-    /**
-     * Rename arrows that are not values
-     */
-    capturedArrows <- findNewNames(fn.capturedArrows -- capturedArrowValues.keySet)
+    captured <- renamedCaptured(fn, exports, arrows)
+    (capturedValues, capturedArrows) = captured
 
     /**
      * Function defines variables inside its body.
@@ -331,7 +436,7 @@ object ArrowInliner extends Logging {
      * It seems that resolving whole `exports`
      * and `arrows` is not necessary.
      */
-    arrowsResolved = arrows ++ capturedArrowValuesRenamed ++ capturedArrows.renamed
+    arrowsResolved = arrows ++ capturedArrows.renamed
     exportsResolved = exports ++ data.renamed ++ capturedValues.renamed
 
     tree = fn.body.rename(renaming)