fix(compiler): Multiple closures in one function [LNG-262] (#941)

aqua-src/antithesis.aqua

@@ -1,34 +1,35 @@
 aqua A
 
-export test
+export bugLNG260
 
-ability InnerAb:
+func create(a: i8) -> -> i8:
-    arrow() -> i8, i8
+    closureArrow = () -> i8:
+        <- a
+    <- closureArrow
 
-ability TestAb:
+func test() -> i8, i8:
-    inner: InnerAb
+    arr1 <- create(1)
+    arr2 <- create(2)
+    <- arr1(), arr2()
 
--- func create(a: i8, b: i8) -> TestAb:
+func cmp(a: i32, b: i32, pred: i8 -> bool) -> bool:
---     arrow = () -> i8, i8:
+  result: ?bool
---         <- a, b
---     <- TestAb(inner = InnerAb(arrow = arrow))
---
--- func test() -> i8, i8, i8, i8, i8, i8:
---     Ab <- create(1, 2)
---     ab <- create(3, 4)
---     AB <- create(5, 6)
---     res1, res2 <- ab.inner.arrow()
---     res3, res4 <- Ab.inner.arrow()
---     res5, res6 <- AB.inner.arrow()
---     <- res1, res2, res3, res4, res5, res6
 
-func create(a: i8, b: i8) -> i8 -> i8, i8:
+  if a < b:
-    arrow = (c: i8) -> i8, i8:
+    result <- pred(-1)
-        <- a, b
+  else:
-    <- arrow
+    if a == b:
+      result <- pred(0)
+    else:
+      result <- pred(1)
 
-func test() -> i8, i8, i8, i8, i8, i8:
+  <- result!
-    Ab <- create(1, 2)
+
-    ab <- create(3, 4)
+func gt(a: i32, b: i32) -> bool:
-    AB <- create(5, 6)
+  pred = (ord: i8) -> bool:
-    <- res1, res2, res3, res4, res5, res6
+        <- ord > 0
+
+  <- cmp(a, b, pred)
+
+func bugLNG260(a: i32, b: i32) -> bool:
+  <- gt(a, b)

integration-tests/aqua/examples/abilities.aqua

@@ -2,7 +2,7 @@ aqua Main
 
 use DECLARE_CONST, decl_bar from "imports_exports/declare.aqua" as Declare
 
-export handleAb, SomeService, bug214, checkAbCalls, bugLNG258_1, bugLNG258_2, bugLNG258_3
+export handleAb, SomeService, bug214, checkAbCalls, bugLNG258_1, bugLNG258_2, bugLNG258_3, multipleAbilityWithClosure
 
 service SomeService("wed"):
   getStr(s: string) -> string
@@ -114,4 +114,18 @@ func bugLNG258_3() -> i8, i8:
     res1, res2 <- aB.inner.arrow()
     <- res1, res2
 
+ability TestAbWithClosure:
+    arrow: -> i8
+
+func createAb(a: i8) -> TestAbWithClosure:
+    closureArrow = () -> i8:
+        <- a
+    ab = TestAbWithClosure(arrow = closureArrow)
+    <- ab
+
+func multipleAbilityWithClosure() -> i8, i8:
+    ab <- createAb(1)
+    ab2 <- createAb(2)
+    <- ab.arrow(), ab2.arrow()
+
 

integration-tests/aqua/examples/closures.aqua

@@ -2,7 +2,7 @@ module Closure declares *
 
 import "@fluencelabs/aqua-lib/builtin.aqua"
 
-export LocalSrv, closureIn, closureOut, closureBig, closureOut2, lng58Bug
+export LocalSrv, closureIn, closureOut, closureBig, closureOut2, lng58Bug, multipleClosuresBugLNG262
 
 service MyOp("op"):
   identity(s: string) -> string
@@ -71,3 +71,13 @@ func lng58Bug() -> string:
   waiting()
 
   <- status!
+
+func create(a: i8) -> -> i8:
+    closureArrow = () -> i8:
+        <- a
+    <- closureArrow
+
+func multipleClosuresBugLNG262() -> i8, i8:
+    arr1 <- create(1)
+    arr2 <- create(2)
+    <- arr1(), arr2()

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

@@ -33,7 +33,7 @@ import {
 import {
   abilityCall,
   complexAbilityCall,
-  checkAbCallsCall, bugLNG258Call1, bugLNG258Call2, bugLNG258Call3,
+  checkAbCallsCall, bugLNG258Call1, bugLNG258Call2, bugLNG258Call3, multipleAbilityWithClosureCall,
 } from "../examples/abilityCall.js";
 import {
   nilLengthCall,
@@ -104,7 +104,7 @@ import { multiReturnCall } from "../examples/multiReturnCall.js";
 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 {closuresCall, multipleClosuresLNG262BugCall} from "../examples/closures.js";
 import { closureArrowCaptureCall } from "../examples/closureArrowCapture.js";
 import {
   bugLNG63_2Call,
@@ -533,7 +533,7 @@ describe("Testing examples", () => {
     });
   });
 
-  it("ability.aqua", async () => {
+  it("abilities.aqua", async () => {
     let result = await abilityCall();
     expect(result).toStrictEqual([
       "declare_const123",
@@ -543,17 +543,17 @@ describe("Testing examples", () => {
     ]);
   });
 
-  it("ability.aqua complex", async () => {
+  it("abilities.aqua complex", async () => {
     let result = await complexAbilityCall();
     expect(result).toStrictEqual([false, true]);
   });
 
-  it("ability.aqua ability calls", async () => {
+  it("abilities.aqua ability calls", async () => {
     let result = await checkAbCallsCall();
     expect(result).toStrictEqual([true, false, true]);
   });
 
-  it("ability.aqua bug LNG-258", async () => {
+  it("abilities.aqua bug LNG-258", async () => {
     let result1 = await bugLNG258Call1();
     expect(result1).toStrictEqual([1, 2]);
 
@@ -564,6 +564,11 @@ describe("Testing examples", () => {
     expect(result3).toStrictEqual([5, 6]);
   });
 
+  it("abilities.aqua multiple abilities with closures", async () => {
+    let result1 = await multipleAbilityWithClosureCall();
+    expect(result1).toStrictEqual([1, 2]);
+  });
+
   it("functors.aqua LNG-119 bug", async () => {
     let result = await bugLng119Call();
     expect(result).toEqual([1]);
@@ -949,6 +954,11 @@ describe("Testing examples", () => {
     expect(closuresResult).toEqual(["in", res1, res1, res2]);
   }, 20000);
 
+  it("closures.aqua bug LNG-262", async () => {
+    let result = await multipleClosuresLNG262BugCall();
+    expect(result).toEqual([1, 2]);
+  });
+
   it("closureArrowCapture.aqua", async () => {
     let result = await closureArrowCaptureCall("input");
     expect(result).toEqual("call: ".repeat(4) + "input");

integration-tests/src/examples/abilityCall.ts

@@ -6,6 +6,7 @@ import {
   bugLNG258_1,
   bugLNG258_2,
   bugLNG258_3,
+  multipleAbilityWithClosure
 } from "../compiled/examples/abilities";
 
 export async function abilityCall(): Promise<[string, string, string, number]> {
@@ -37,3 +38,7 @@ export async function bugLNG258Call2(): Promise<[number, number]> {
 export async function bugLNG258Call3(): Promise<[number, number]> {
   return await bugLNG258_3();
 }
+
+export async function multipleAbilityWithClosureCall(): Promise<[number, number]> {
+  return await multipleAbilityWithClosure()
+}

integration-tests/src/examples/closures.ts

@@ -5,6 +5,7 @@ import {
   registerLocalSrv,
   closureOut2,
   lng58Bug,
+  multipleClosuresBugLNG262
 } from "../compiled/examples/closures.js";
 import { config } from "../config.js";
 
@@ -32,3 +33,7 @@ export async function closuresCall(): Promise<
 export async function lng58CBugCall(): Promise<string> {
   return lng58Bug();
 }
+
+export async function multipleClosuresLNG262BugCall(): Promise<[number, number]> {
+  return multipleClosuresBugLNG262();
+}

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

@@ -60,7 +60,7 @@ object ArrowInliner extends Logging {
               ) if !outsideStreamNames.contains(n) =>
             resDesugar.toList -> res
           case (
-                cexp @ CallModel.Export(exp, st @ StreamType(_)),
+                cexp @ CallModel.Export(_, StreamType(_)),
                 (res, resDesugar)
               ) =>
             // pass nested function results to a stream

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

@@ -1,13 +1,13 @@
 package aqua.model.inline
 
 import aqua.model.*
+import aqua.model.MetaModel.CallArrowModel
 import aqua.model.inline.state.InliningState
 import aqua.raw.ops.*
 import aqua.raw.value.*
 import aqua.types.*
 import aqua.raw.value.{CallArrowRaw, ValueRaw}
 import aqua.raw.arrow.{ArrowRaw, FuncRaw}
-
 import cats.Eval
 import cats.syntax.show.*
 import cats.syntax.option.*
@@ -520,6 +520,137 @@ class ArrowInlinerSpec extends AnyFlatSpec with Matchers with Inside {
     }
   }
 
+  /**
+   * func return(a: i8) -> -> i8:
+   *    closure = () -> i8:
+   *        <- a
+   *    <- closure
+   *
+   * func test() -> i8, i8:
+   *    closure <- return(1)
+   *    closure2 <- return(2)
+   *    <- closure(), closure2()
+   */
+  it should "correct renaming on multiple closures from same function" in {
+    val resType = ScalarType.i8
+    val resVar = VarRaw("a", resType)
+
+    val closureType = ArrowType(
+      ProductType(Nil),
+      ProductType(resType :: Nil)
+    )
+
+    val innerClosure = VarRaw("closureArrow", closureType)
+    val closureFunc = FuncRaw(
+      innerClosure.name,
+      ArrowRaw(
+        closureType,
+        List(resVar),
+        ReturnTag(NonEmptyList.one(resVar)).leaf
+      )
+    )
+
+    val returnFunc = FuncArrow(
+      "return",
+      SeqTag.wrap(
+        ClosureTag(
+          closureFunc,
+          detach = false
+        ).leaf,
+        ReturnTag(
+          NonEmptyList.one(innerClosure)
+        ).leaf
+      ),
+      ArrowType(
+        ProductType.labelled((resVar.name, ScalarType.i8) :: Nil),
+        ProductType(closureType :: Nil)
+      ),
+      List(innerClosure),
+      Map.empty,
+      Map.empty,
+      None
+    )
+
+    val closureVar = VarRaw("closure", closureType)
+    val closureVar2 = VarRaw("closure2", closureType)
+
+    val res1 = VarRaw("res1", ScalarType.i8)
+    val res2 = VarRaw("res2", ScalarType.i8)
+
+    val testFunc = FuncArrow(
+      "test",
+      SeqTag.wrap(
+        CallArrowRawTag
+          .func(
+            returnFunc.funcName,
+            Call(LiteralRaw.number(1) :: Nil, Call.Export(closureVar.name, closureType) :: Nil)
+          )
+          .leaf,
+        CallArrowRawTag
+          .func(
+            returnFunc.funcName,
+            Call(LiteralRaw.number(2) :: Nil, Call.Export(closureVar2.name, closureType) :: Nil)
+          )
+          .leaf,
+        CallArrowRawTag
+          .func(
+            closureVar.name,
+            Call(Nil, Call.Export(res1.name, res1.baseType) :: Nil)
+          )
+          .leaf,
+        CallArrowRawTag
+          .func(
+            closureVar2.name,
+            Call(Nil, Call.Export(res2.name, res2.baseType) :: Nil)
+          )
+          .leaf,
+        CallArrowRawTag
+          .service(LiteralRaw.quote("Srv"), "callSrv", Call(res1 :: res2 :: Nil, Nil))
+          .leaf
+      ),
+      ArrowType(
+        ProductType(Nil),
+        ProductType(Nil)
+      ),
+      Nil,
+      Map(returnFunc.funcName -> returnFunc),
+      Map.empty,
+      None
+    )
+
+    println("testFunc: ")
+    println(testFunc.body.show)
+
+    val model = ArrowInliner
+      .callArrow[InliningState](testFunc, CallModel(Nil, Nil))
+      .runA(InliningState())
+      .value
+
+    model.tailForced
+
+    model.equalsOrShowDiff(
+      SeqModel.wrap(
+        CallArrowModel("return").wrap(
+          CaptureTopologyModel("closureArrow").leaf
+        ),
+        CallArrowModel("return").wrap(
+          CaptureTopologyModel("closureArrow-0").leaf
+        ),
+        CallArrowModel("closureArrow").wrap(
+          ApplyTopologyModel("closureArrow").wrap(EmptyModel.leaf)
+        ),
+        CallArrowModel("closureArrow-0").wrap(
+          ApplyTopologyModel("closureArrow-0").wrap(EmptyModel.leaf)
+        ),
+        CallServiceModel(
+          LiteralModel.quote("Srv"),
+          "callSrv",
+          CallModel(LiteralModel.number(1) :: LiteralModel.number(2) :: Nil, Nil)
+        ).leaf
+      )
+    ) should be(true)
+  }
+
   /*
     func stream-callback(cb: string -> ()):
       records: *string
@@ -1652,7 +1783,7 @@ class ArrowInlinerSpec extends AnyFlatSpec with Matchers with Inside {
     val mainFunc = FuncArrow(
       funcName = "main",
       body = mainBody,
-      arrowType = ArrowType(ProductType(Nil), ProductType(Nil)),
+      arrowType = mainType,
       ret = Nil,
       capturedArrows = Map(testName -> testFunc),
       capturedValues = Map.empty,

model/raw/src/main/scala/aqua/raw/ops/RawTag.scala

@@ -290,13 +290,17 @@ case class ClosureTag(
   override def usesVarNames: Set[String] = Set.empty
 
   override def renameExports(map: Map[String, String]): RawTag =
-    copy(func = func.copy(name = map.getOrElse(func.name, func.name)))
+    copy(func =
+      func.copy(
+        name = map.getOrElse(func.name, func.name)
+      )
+    )
 
   override def mapValues(f: ValueRaw => ValueRaw): RawTag =
     copy(
       func.copy(arrow =
         func.arrow.copy(
-          ret = func.arrow.ret.map(_.mapValues(f)),
+          ret = func.arrow.ret.map(_.map(f)),
           body = func.arrow.body.map(_.mapValues(f))
         )
       )

semantics/src/main/scala/aqua/semantics/ExprSem.scala

@@ -10,8 +10,10 @@ import aqua.semantics.rules.ValuesAlgebra
 import aqua.semantics.rules.abilities.AbilitiesAlgebra
 import aqua.semantics.rules.definitions.DefinitionsAlgebra
 import aqua.semantics.rules.locations.LocationsAlgebra
+import aqua.semantics.rules.mangler.ManglerAlgebra
 import aqua.semantics.rules.names.NamesAlgebra
 import aqua.semantics.rules.types.TypesAlgebra
+
 import cats.Monad
 
 object ExprSem {
@@ -24,7 +26,8 @@ object ExprSem {
     T: TypesAlgebra[S, G],
     V: ValuesAlgebra[S, G],
     D: DefinitionsAlgebra[S, G],
-    L: LocationsAlgebra[S, G]
+    L: LocationsAlgebra[S, G],
+    M: ManglerAlgebra[G]
   ): Prog[G, Raw] =
     expr match {
       case expr: ServiceIdExpr[S] => new ServiceIdSem(expr).program[G]

semantics/src/main/scala/aqua/semantics/RawSemantics.scala

@@ -276,7 +276,8 @@ object RawSemantics extends Logging {
     T: TypesAlgebra[S, G],
     D: DefinitionsAlgebra[S, G],
     L: LocationsAlgebra[S, G],
-    E: ReportAlgebra[S, G]
+    E: ReportAlgebra[S, G],
+    M: ManglerAlgebra[G]
   ): (Expr[S], Chain[G[RawWithToken[S]]]) => Eval[G[RawWithToken[S]]] = (expr, inners) =>
     Eval later ExprSem
       .getProg[S, G](expr)

semantics/src/main/scala/aqua/semantics/expr/func/ArrowSem.scala

@@ -1,33 +1,27 @@
 package aqua.semantics.expr.func
 
-import aqua.parser.expr.FuncExpr
 import aqua.parser.expr.func.ArrowExpr
-import aqua.parser.lexer.{Arg, DataTypeToken}
 import aqua.raw.Raw
 import aqua.raw.arrow.ArrowRaw
-import aqua.raw.ops.{SeqTag, *}
+import aqua.raw.ops.*
 import aqua.raw.value.*
 import aqua.semantics.Prog
-import aqua.semantics.rules.ValuesAlgebra
 import aqua.semantics.rules.abilities.AbilitiesAlgebra
 import aqua.semantics.rules.locations.LocationsAlgebra
+import aqua.semantics.rules.mangler.ManglerAlgebra
 import aqua.semantics.rules.names.NamesAlgebra
 import aqua.semantics.rules.types.TypesAlgebra
 import aqua.types.*
 
-import cats.Eval
+import cats.Monad
-import cats.data.{Chain, NonEmptyList}
+import cats.data.Chain
-import cats.free.{Cofree, Free}
+import cats.free.Cofree
-import cats.data.OptionT
-import cats.syntax.show.*
 import cats.syntax.applicative.*
-import cats.syntax.apply.*
-import cats.syntax.foldable.*
 import cats.syntax.bifunctor.*
 import cats.syntax.flatMap.*
+import cats.syntax.foldable.*
 import cats.syntax.functor.*
 import cats.syntax.traverse.*
-import cats.{Applicative, Monad}
 
 class ArrowSem[S[_]](val expr: ArrowExpr[S]) extends AnyVal {
 
@@ -35,9 +29,7 @@ class ArrowSem[S[_]](val expr: ArrowExpr[S]) extends AnyVal {
 
   def before[Alg[_]: Monad](implicit
     T: TypesAlgebra[S, Alg],
-    N: NamesAlgebra[S, Alg],
+    N: NamesAlgebra[S, Alg]
-    A: AbilitiesAlgebra[S, Alg],
-    L: LocationsAlgebra[S, Alg]
   ): Alg[ArrowType] = for {
     arrowType <- T.beginArrowScope(arrowTypeExpr)
     // Create local variables
@@ -57,13 +49,11 @@ class ArrowSem[S[_]](val expr: ArrowExpr[S]) extends AnyVal {
   )(using
     T: TypesAlgebra[S, Alg],
     N: NamesAlgebra[S, Alg],
-    A: AbilitiesAlgebra[S, Alg],
+    M: ManglerAlgebra[Alg]
-    L: LocationsAlgebra[S, Alg]
   ): Alg[Raw] = for {
     streamsInScope <- N.streamsDefinedWithinScope()
     retValues <- T.endArrowScope(expr.arrowTypeExpr)
-    retValuesDerivedFrom <- N.getDerivedFrom(retValues.map(_.varNames))
+    res <- bodyGen match {
-    res = bodyGen match {
       case FuncOp(bodyModel) =>
         // TODO: wrap with local on...via...
         val retsAndArgs = retValues zip funcArrow.codomain.toList
@@ -77,13 +67,17 @@ class ArrowSem[S[_]](val expr: ArrowExpr[S]) extends AnyVal {
         val localStreams = streamsInScope -- dataArgsNames -- streamsThatReturnAsStreams
 
         // process stream that returns as not streams and all Apply*Raw
-        val (bodyRets, retVals) = retsAndArgs.mapWithIndex {
+        retsAndArgs.traverse {
-          case ((v @ VarRaw(_, StreamType(_)), StreamType(_)), _) =>
+          case (v @ VarRaw(_, StreamType(_)), StreamType(_)) =>
-            (Chain.empty, v)
+            (Chain.empty, v).pure[Alg]
           // canonicalize and change return value
-          case ((VarRaw(streamName, streamType @ StreamType(streamElement)), _), idx) =>
+          case (VarRaw(streamName, streamType @ StreamType(streamElement)), _) =>
-            val canonReturnVar = VarRaw(s"-$streamName-fix-$idx", CanonStreamType(streamElement))
+            for {
-            val returnVar = VarRaw(s"-$streamName-flat-$idx", ArrayType(streamElement))
+              canonName <- M.rename(s"-$streamName-fix")
+              returnVarName <- M.rename(s"-$streamName-flat")
+            } yield {
+              val canonReturnVar = VarRaw(canonName, CanonStreamType(streamElement))
+              val returnVar = VarRaw(returnVarName, ArrayType(streamElement))
               val body = Chain(
                 CanonicalizeTag(
                   VarRaw(streamName, streamType),
@@ -96,9 +90,15 @@ class ArrowSem[S[_]](val expr: ArrowExpr[S]) extends AnyVal {
               )
 
               (body, returnVar)
+            }
+
           // assign and change return value for all `Apply*Raw`
-          case ((v: ValueRaw.ApplyRaw, _), idx) =>
+          case (v: ValueRaw.ApplyRaw, _) =>
-            val assignedReturnVar = VarRaw(s"-return-fix-$idx", v.`type`)
+            for {
+              assignedReturnName <- M.rename(s"-return-fix")
+            } yield {
+              val assignedReturnVar = VarRaw(assignedReturnName, v.`type`)
+
               val body = Chain.one(
                 AssignmentTag(
                   v,
@@ -107,21 +107,23 @@ class ArrowSem[S[_]](val expr: ArrowExpr[S]) extends AnyVal {
               )
 
               (body, assignedReturnVar)
-          case ((v, _), _) => (Chain.empty, v)
+            }
-        }.unzip.leftMap(_.combineAll)
 
+          case (v, _) => (Chain.empty, v).pure[Alg]
+        }.map(_.unzip.leftMap(_.combineAll)).map { case (bodyRets, retVals) =>
           val bodyModified = SeqTag.wrap(
             bodyModel +: bodyRets
           )
 
           // wrap streams with restrictions
-        val bodyWithRestrictions = localStreams.foldLeft(bodyModified) {
+          val bodyWithRestrictions =
-          case (bm, (streamName, streamType)) =>
+            localStreams.foldLeft(bodyModified) { case (bm, (streamName, streamType)) =>
               RestrictionTag(streamName, streamType).wrap(bm)
             }
-
           ArrowRaw(funcArrow, retVals, bodyWithRestrictions)
-      case _ => Raw.error("Invalid arrow body")
+        }
+
+      case _ => Raw.error("Invalid arrow body").pure[Alg]
     }
   } yield res
 
@@ -129,7 +131,8 @@ class ArrowSem[S[_]](val expr: ArrowExpr[S]) extends AnyVal {
     T: TypesAlgebra[S, Alg],
     N: NamesAlgebra[S, Alg],
     A: AbilitiesAlgebra[S, Alg],
-    L: LocationsAlgebra[S, Alg]
+    L: LocationsAlgebra[S, Alg],
+    M: ManglerAlgebra[Alg]
   ): Prog[Alg, Raw] =
     Prog
       .around(