fix(compiler): Fix arrows capture in closures [fixes LNG-242] (#903)

* Fix arrows capture

* Add comment

* Add test

* Add integration test

integration-tests/aqua/examples/closureArrowCapture.aqua

@@ -0,0 +1,36 @@
+aqua Test
+
+export test, TestService
+
+service TestService:
+    call(s: string) -> string
+
+ability TestAbility:
+    arrow(s: string) -> string
+
+func returnCapture() -> string -> string:
+    TestService "test-service"
+
+    closure = (s: string) -> string:
+        <- TestService.call(s)
+    
+    closure1 = closure
+    closure2 = closure1
+    closure3 = closure2
+
+    Ab = TestAbility(
+        arrow = closure
+    )
+
+    capture = (s: string) -> string:
+        s1 <- closure(s) -- capture closure
+        s2 <- closure3(s1) -- capture renamed closure
+        s3 <- Ab.arrow(s2) -- capture ability
+        s4 <- TestService.call(s3) -- capture service
+        <- s4
+    
+    <- capture
+
+func test(s: string) -> string:
+    capture <- returnCapture()
+    <- capture(s)

integration-tests/src/__test__/examples.spec.ts

@@ -100,6 +100,7 @@ import { declareCall } from "../examples/declareCall.js";
 import { genOptions, genOptionsEmptyString } from "../examples/optionsCall.js";
 import { lng193BugCall } from "../examples/closureReturnRename.js";
 import { closuresCall } from "../examples/closures.js";
+import { closureArrowCaptureCall } from "../examples/closureArrowCapture.js";
 import {
   bugLNG63_2Call,
   bugLNG63_3Call,
@@ -842,6 +843,11 @@ describe("Testing examples", () => {
     expect(closuresResult).toEqual(["in", res1, res1, res2]);
   }, 20000);
 
+  it("closureArrowCapture.aqua", async () => {
+    let result = await closureArrowCaptureCall("input");
+    expect(result).toEqual("call: ".repeat(4) + "input");
+  });
+
   it("tryOtherwise.aqua", async () => {
     let tryOtherwiseResult = await tryOtherwiseCall(relayPeerId1);
     expect(tryOtherwiseResult).toBe("error");

integration-tests/src/examples/closureArrowCapture.ts

@@ -0,0 +1,14 @@
+import {
+  test,
+  registerTestService,
+} from "../compiled/examples/closureArrowCapture.js";
+
+export async function closureArrowCaptureCall(s: string) {
+  registerTestService("test-service", {
+    call: (s: string) => {
+      return "call: " + s;
+    },
+  });
+
+  return await test(s);
+}

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

@@ -223,19 +223,28 @@ object ArrowInliner extends Logging {
     renamed: Map[String, T]
   )
 
+  // TODO: Make this extension private somehow?
+  extension [T](vals: Map[String, T]) {
+
+    def renamed(renames: Map[String, String]): Map[String, T] =
+      vals.map { case (name, value) =>
+        renames.getOrElse(name, name) -> value
+      }
+  }
+
   /**
    * Rename values and forbid new names
    *
    * @param values Mapping name -> value
    * @return Renamed values and renames
    */
-  private def findNewNames[S: Mangler, T](values: Map[String, T]): State[S, Renamed[T]] =
+  private def findNewNames[S: Mangler, T](
+    values: Map[String, T]
+  ): State[S, Renamed[T]] =
     Mangler[S].findAndForbidNames(values.keySet).map { renames =>
       Renamed(
         renames,
-        values.map { case (name, value) =>
-          renames.getOrElse(name, name) -> value
-        }
+        values.renamed(renames)
       )
     }
 
@@ -281,18 +290,17 @@ object ArrowInliner extends Logging {
      * If arrow correspond to a value,
      * rename in accordingly to the value
      */
-    capturedArrowValues = fn.capturedArrows.flatMap { case (arrowName, arrow) =>
+    capturedArrowValues = Arrows.arrowsByValues(
+      fn.capturedArrows,
+      fn.capturedValues
+    )
+    capturedArrowValuesRenamed = capturedArrowValues.renamed(
       capturedValues.renames
-        .get(arrowName)
-        .orElse(fn.capturedValues.get(arrowName).as(arrowName))
-        .map(_ -> arrow)
-    }
+    )
     /**
      * Rename arrows that are not values
      */
-    capturedArrows <- findNewNames(fn.capturedArrows.filterNot { case (arrowName, _) =>
-      capturedArrowValues.contains(arrowName)
-    })
+    capturedArrows <- findNewNames(fn.capturedArrows -- capturedArrowValues.keySet)
 
     /**
      * Function defines variables inside its body.
@@ -322,7 +330,7 @@ object ArrowInliner extends Logging {
      * It seems that resolving whole `exports`
      * and `arrows` is not necessary.
      */
-    arrowsResolved = arrows ++ capturedArrowValues ++ capturedArrows.renamed
+    arrowsResolved = arrows ++ capturedArrowValuesRenamed ++ capturedArrows.renamed
     exportsResolved = exports ++ data.renamed ++ capturedValues.renamed
 
     tree = fn.body.rename(renaming)

model/inline/src/main/scala/aqua/model/inline/state/Arrows.scala

@@ -2,6 +2,7 @@ package aqua.model.inline.state
 
 import aqua.model.{ArgsCall, FuncArrow}
 import aqua.raw.arrow.FuncRaw
+import aqua.model.ValueModel
 
 import cats.data.State
 import cats.instances.list.*
@@ -32,9 +33,10 @@ trait Arrows[S] extends Scoped[S] {
     for {
       exps <- Exports[S].exports
       arrs <- arrows
-      captuedVars = exps.filterKeys(arrow.capturedVars).toMap
-      capturedArrows = arrs.filterKeys(arrow.capturedVars).toMap
-      funcArrow = FuncArrow.fromRaw(arrow, capturedArrows, captuedVars, topology)
+      capturedVars = exps.filterKeys(arrow.capturedVars).toMap
+      capturedArrows = arrs.filterKeys(arrow.capturedVars).toMap ++
+        Arrows.arrowsByValues(arrs, capturedVars)
+      funcArrow = FuncArrow.fromRaw(arrow, capturedArrows, capturedVars, topology)
       _ <- save(arrow.name, funcArrow)
     } yield ()
 
@@ -97,6 +99,25 @@ trait Arrows[S] extends Scoped[S] {
 }
 
 object Arrows {
+
+  /**
+   * Retrieve all arrows that correspond to values
+   */
+  def arrowsByValues(
+    arrows: Map[String, FuncArrow],
+    values: Map[String, ValueModel]
+  ): Map[String, FuncArrow] = {
+    val arrowKeys = arrows.keySet ++ arrows.values.map(_.funcName)
+    val varsKeys = values.keySet ++ values.values.collect { case ValueModel.Arrow(name, _) =>
+      name
+    }
+    val keys = arrowKeys.intersect(varsKeys)
+
+    arrows.filter { case (arrowName, arrow) =>
+      keys.contains(arrowName) || keys.contains(arrow.funcName)
+    }
+  }
+
   def apply[S](implicit arrows: Arrows[S]): Arrows[S] = arrows
 
   // Default implementation with the most straightforward state – just a Map

model/inline/src/test/scala/aqua/model/inline/ArrowInlinerSpec.scala

@@ -2334,4 +2334,212 @@ class ArrowInlinerSpec extends AnyFlatSpec with Matchers with Inside {
 
     model.equalsOrShowDiff(expected) shouldEqual true
   }
+
+  it should "handle captured arrows" in {
+    val sArg = VarRaw("s", ScalarType.string)
+    val ret = VarRaw("ret", ScalarType.string)
+    val captureType = ArrowType(
+      ProductType.labelled(sArg.name -> sArg.`type` :: Nil),
+      ProductType(ScalarType.string :: Nil)
+    )
+    val captureTypeUnlabelled = captureType.copy(
+      domain = ProductType(sArg.`type` :: Nil)
+    )
+    val captureVar = VarRaw("capture", captureTypeUnlabelled)
+    val returnCaptureName = "returnCapture"
+
+    /**
+     * func returnCapture() -> string -> string:
+     *   <captureGen>
+     *   capture = (s: string) -> string:
+     *     ret <- <captureName>(s)
+     *     <- ret
+     *   <- capture
+     *
+     * func main(s: string) -> string:
+     *   capture <- returnCapture()
+     *   ret <- capture(s)
+     *   <- ret
+     *
+     * -- inlining:
+     * main("test")
+     */
+    def test(
+      capturedGen: List[RawTag.Tree],
+      capturedName: String
+    ) = {
+      val mainBody = SeqTag.wrap(
+        CallArrowRawTag
+          .func(
+            "returnCapture",
+            Call(Nil, Call.Export(captureVar.name, captureType) :: Nil)
+          )
+          .leaf,
+        CallArrowRawTag
+          .func(
+            captureVar.name,
+            Call(sArg :: Nil, Call.Export(ret.name, ret.`type`) :: Nil)
+          )
+          .leaf,
+        ReturnTag(NonEmptyList.one(ret)).leaf
+      )
+
+      val captureBody = SeqTag.wrap(
+        CallArrowRawTag
+          .func(
+            capturedName,
+            Call(sArg :: Nil, Call.Export(ret.name, ret.`type`) :: Nil)
+          )
+          .leaf,
+        ReturnTag(NonEmptyList.one(ret)).leaf
+      )
+
+      val returnCaptureBody = SeqTag.wrap(
+        capturedGen ++ (ClosureTag(
+          FuncRaw(
+            captureVar.name,
+            ArrowRaw(
+              captureType,
+              ret :: Nil,
+              captureBody
+            )
+          ),
+          false
+        ).leaf :: ReturnTag(
+          NonEmptyList.one(captureVar)
+        ).leaf :: Nil)
+      )
+
+      val returnCapture = FuncArrow(
+        returnCaptureName,
+        returnCaptureBody,
+        ArrowType(
+          ProductType(Nil),
+          ProductType(captureTypeUnlabelled :: Nil)
+        ),
+        captureVar :: Nil,
+        Map.empty,
+        Map.empty,
+        None
+      )
+
+      val main = FuncArrow(
+        "main",
+        mainBody,
+        captureType,
+        ret :: Nil,
+        Map(returnCaptureName -> returnCapture),
+        Map.empty,
+        None
+      )
+
+      val model = ArrowInliner
+        .callArrow[InliningState](
+          FuncArrow(
+            "wrapper",
+            CallArrowRawTag
+              .func(
+                "main",
+                Call(LiteralRaw.quote("test") :: Nil, Nil)
+              )
+              .leaf,
+            ArrowType(
+              ProductType(Nil),
+              ProductType(Nil)
+            ),
+            Nil,
+            Map("main" -> main),
+            Map.empty,
+            None
+          ),
+          CallModel(Nil, Nil)
+        )
+        .runA(InliningState())
+        .value
+
+      // TODO: Don't know for what to test here
+      // inliner will just log an error in case of failure
+      model.head should not equal EmptyModel
+    }
+
+    /**
+     * closure = (s: string) -> string:
+     *    ret <- s
+     *    <-ret
+     * closure1 = closure
+     * closure2 = closure1
+     * closure3 = closure2
+     *
+     * -- captureName = closure3
+     */
+    val closureRename = List(
+      ClosureTag(
+        FuncRaw(
+          "closure",
+          ArrowRaw(
+            captureType,
+            ret :: Nil,
+            ReturnTag(NonEmptyList.one(ret)).leaf
+          )
+        ),
+        false
+      ).leaf,
+      AssignmentTag(
+        VarRaw("closure", captureType),
+        "closure1"
+      ).leaf,
+      AssignmentTag(
+        VarRaw("closure1", captureType),
+        "closure2"
+      ).leaf,
+      AssignmentTag(
+        VarRaw("closure2", captureType),
+        "closure3"
+      ).leaf
+    )
+
+    test(closureRename, "closure3")
+
+    /**
+     * closure = (s: string) -> string:
+     *    ret <- s
+     *    <-ret
+     * Ab = TestAbility(
+     *   arrow = closure
+     * )
+     *
+     * -- captureName = Ab.arrow
+     */
+    val makeAbility = List(
+      ClosureTag(
+        FuncRaw(
+          "closure",
+          ArrowRaw(
+            captureType,
+            ret :: Nil,
+            ReturnTag(NonEmptyList.one(ret)).leaf
+          )
+        ),
+        false
+      ).leaf,
+      AssignmentTag(
+        AbilityRaw(
+          fieldsAndArrows = NonEmptyMap.one(
+            "arrow",
+            VarRaw("closure", captureType)
+          ),
+          AbilityType(
+            "TestAbility",
+            NonEmptyMap.one(
+              "arrow",
+              captureType
+            )
+          )
+        ),
+        "Ab"
+      ).leaf
+    )
+
+    test(makeAbility, "Ab.arrow")
+  }
 }

model/src/main/scala/aqua/model/ValueModel.scala

@@ -53,7 +53,7 @@ object ValueModel {
 
   object Arrow {
 
-    def unapply(vm: VarModel): Option[(String, ArrowType)] =
+    def unapply(vm: ValueModel): Option[(String, ArrowType)] =
       vm match {
         case VarModel(name, t: ArrowType, _) =>
           (name, t).some