feat(compiler): Add feature to insert map from a function result with correct type [LNG-367]

aqua-src/antithesis.aqua

@@ -1,41 +1,20 @@
-aqua StreamMapAbilities
+aqua StreamMapTest declares *
 
-export streamMapAbilityTest
+export main
 
-ability Streams:
-  stream: *string
-  map: %string
+func foo() -> string, u64:
+  <- "123", 42
 
-ability Adds:
-  addToStream(s: string)
-  addToMap(k: string, v: string)
+func zzz() -> string, string:
+  <- "123", "str"
 
-func addToStreamClosure(str: *string) -> string -> ():
-  cl = func (s: string):
-    str <<- s
-  <- cl
+func create() -> %u64:
+  map: %u64
+  map <- foo()
+  <- map
 
-func addToMapClosure(str: %string) -> string, string -> ():
-  cl = func (k: string, v: string):
-    str <<- k, v
-  <- cl
+func main():
+  map <- create()
+  a <- zzz()
 
-func addTo{Streams}() -> Adds:
-  addStream = addToStreamClosure(Streams.stream)
-  addMap = addToMapClosure(Streams.map)
-  adds = Adds(addToStream = addStream, addToMap = addMap)
-  <- adds
-
-func add{Adds}(s: string, k: string):
-  Adds.addToStream(s)
-  Adds.addToMap(k, k)
-
-func streamMapAbilityTest() -> []string, []string:
-  stream: *string
-  map: %string
-  ab = Streams(stream = stream, map = map)
-  adds <- addTo{ab}()
-  add{adds}("one", "1")
-  add{adds}("two", "2")
-  add{adds}("three", "3")
-  <- stream, map.keys()
+  map <- foo()

integration-tests/aqua/examples/closureReturnRename.aqua

@@ -30,9 +30,12 @@ func addToStreamClosure(str: *string) -> string -> ():
     str <<- s
   <- cl
 
+func toMap(k: string, v: string) -> string, string:
+  <- k, v
+
 func addToMapClosure(str: %string) -> string, string -> ():
   cl = func (k: string, v: string):
-    str <<- k, v
+    str <- toMap(k, v)
   <- cl
 
 func addTo{Streams}() -> Adds:

integration-tests/aqua/examples/streamMap.aqua

@@ -1,7 +1,7 @@
 aqua StreamMapTest declares *
 
 export testGetFunc, testGetStreamFunc, testKeysFunc, testKeysStreamFunc
-export testContainsFunc, testForFunc, testParSeqMap, testForTupleFunc
+export testContainsFunc, testForFunc, testParSeqMap, testForTupleFunc, testInsertMapFromFunc
 
 import "@fluencelabs/aqua-lib/builtin.aqua"
 
@@ -127,4 +127,18 @@ func testForTupleFunc() -> []string, []string, []string:
   for k, v <- streamMap:
     streamThird <<- streamMap.get(k)!
 
-  <- streamFirst, streamSecond, streamThird
+  <- streamFirst, streamSecond, streamThird
+
+func foo() -> string, u64:
+  <- "123", 42
+
+func create() -> %u64:
+  map: %u64
+  map <- foo()
+  <- map
+
+func testInsertMapFromFunc() -> []string:
+  map <- create()
+
+  map <- foo()
+  <- map.keys()

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

@@ -69,7 +69,8 @@ import {
   testContainsFuncCall,
   testForFuncCall,
   testForTupleFuncCall,
-  testParSeqMapCall
+  testParSeqMapCall,
+  testInsertMapFromFuncCall
 } from "../examples/streamMapCall.js";
 import {
   topologyBug205Call,
@@ -949,6 +950,11 @@ describe("Testing examples", () => {
     expect(res).toEqual("ok");
   });
 
+  it("streamMap.aqua insert from func", async () => {
+    const res = await testInsertMapFromFuncCall();
+    expect(res).toEqual(["123", "123"]);
+  });
+
   it("stream.aqua", async () => {
     const streamResult = await streamCall();
     expect(streamResult).toEqual([

integration-tests/src/examples/streamMapCall.ts

@@ -1,6 +1,6 @@
 import {
   testGetFunc, testGetStreamFunc, testKeysFunc, testKeysStreamFunc, testContainsFunc,
-  testForFunc, testParSeqMap, testForTupleFunc
+  testForFunc, testParSeqMap, testForTupleFunc, testInsertMapFromFunc
 } from "../compiled/examples/streamMap.js";
 import { config } from "../config.js";
 
@@ -38,3 +38,7 @@ export async function testParSeqMapCall() {
     return testParSeqMap(relays[3].peerId, relays[4].peerId, relays[5].peerId)
 }
 
+export async function testInsertMapFromFuncCall() {
+    return testInsertMapFromFunc()
+}
+

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

@@ -54,22 +54,34 @@ object ArrowInliner extends Logging {
     RawValueInliner
       .valueListToModel(results)
       .map(resolvedResults =>
-        // Fix the return values
-        (exportTo zip resolvedResults).map {
-          case (
+        (exportTo, resolvedResults) match {
+          // if export is a stream map and there is two results from function
+          case ((cexp @ CallModel.Export(n, st@StreamMapType(_))) :: Nil, (res1, desugar1) :: (res2, desugar2) :: Nil)  =>
+            (desugar1.toList ++ desugar2.toList :+ InsertKeyValueModel(res1, res2, n, st).leaf) -> (cexp.asVar :: Nil)
+          case _ =>
+            // Fix the return values
+            (exportTo zip resolvedResults).map {
+              case (
                 CallModel.Export(n, StreamType(_)),
                 (res @ VarModel(_, StreamType(_), _), resDesugar)
-              ) if !outsideStreamNames.contains(n) =>
-            resDesugar.toList -> res
-          case (
+                ) if !outsideStreamNames.contains(n) =>
+                resDesugar.toList -> res
+              case (
+                CallModel.Export(n, StreamMapType(_)),
+                (res @ VarModel(_, StreamMapType(_), _), resDesugar)
+                ) if !outsideStreamNames.contains(n) =>
+                resDesugar.toList -> res
+              case (
                 cexp @ CallModel.Export(_, StreamType(_)),
                 (res, resDesugar)
-              ) =>
-            // pass nested function results to a stream
-            (resDesugar.toList :+ PushToStreamModel(res, cexp).leaf) -> cexp.asVar
-          case (_, (res, resDesugar)) =>
-            resDesugar.toList -> res
-        }.unzip.leftMap(_.flatten)
+                ) =>
+                // pass nested function results to a stream
+                (resDesugar.toList :+ PushToStreamModel(res, cexp).leaf) -> cexp.asVar
+              case (_, (res, resDesugar)) =>
+                resDesugar.toList -> res
+            }.unzip.leftMap(_.flatten)
+        }
+
       )
 
   /**

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

@@ -7,7 +7,7 @@ import aqua.semantics.Prog
 import aqua.semantics.rules.ValuesAlgebra
 import aqua.semantics.rules.names.NamesAlgebra
 import aqua.semantics.rules.types.TypesAlgebra
-import aqua.types.{ProductType, StreamType, Type}
+import aqua.types.{ProductType, StreamMapType, StreamType, Type}
 
 import cats.Monad
 import cats.syntax.apply.*
@@ -23,13 +23,20 @@ class CallArrowSem[S[_]](val expr: CallArrowExpr[S]) extends AnyVal {
     N: NamesAlgebra[S, Alg],
     T: TypesAlgebra[S, Alg]
   ): Alg[List[Call.Export]] =
-    (variables zip codomain.toList).traverse { case (v, t) =>
-      N.read(v, mustBeDefined = false).flatMap {
-        case Some(stream @ StreamType(st)) =>
-          T.ensureTypeMatches(v, st, t).as(Call.Export(v.value, stream, isExistingStream = true))
-        case _ =>
-          N.define(v, t).as(Call.Export(v.value, t))
-      }
+    variables.traverse(v => N.read(v, mustBeDefined = false).map(v -> _)).flatMap {
+      case (v, Some(map @ StreamMapType(_))) :: Nil =>
+        T.ensureTypeMatches(v, map.toProduct, codomain)
+          .as(Call.Export(v.value, map, isExistingStream = true) :: Nil)
+      case vars =>
+        (vars zip codomain.toList).traverse { case ((v, vType), t) =>
+          vType match {
+            case Some(stream @ StreamType(st)) =>
+              T.ensureTypeMatches(v, st, t)
+                .as(Call.Export(v.value, stream, isExistingStream = true))
+            case _ =>
+              N.define(v, t).as(Call.Export(v.value, t))
+          }
+        }
     }
 
   private def toModel[Alg[_]: Monad](using
@@ -40,8 +47,9 @@ class CallArrowSem[S[_]](val expr: CallArrowExpr[S]) extends AnyVal {
     // TODO: Accept other expressions
     callArrowRaw <- V.valueToCall(expr.callArrow)
     tag <- callArrowRaw.traverse { case (raw, at) =>
-      getExports(at.codomain).map(CallArrowRawTag(_, raw)) <*
-        T.checkArrowCallResults(callArrow, at, variables)
+      getExports(at.codomain).flatMap(exports =>
+        T.checkArrowCallResults(callArrow, at, variables, exports).as(CallArrowRawTag(exports, raw))
+      )
     }
   } yield tag.map(_.funcOpLeaf)
 

semantics/src/main/scala/aqua/semantics/rules/types/TypesAlgebra.scala

@@ -1,13 +1,12 @@
 package aqua.semantics.rules.types
 
 import aqua.parser.lexer.*
+import aqua.raw.ops.Call
 import aqua.raw.value.{PropertyRaw, ValueRaw}
 import aqua.types.*
 import aqua.types.Type.*
 
-import cats.data.NonEmptyList
-import cats.data.NonEmptyMap
-import cats.data.OptionT
+import cats.data.{NonEmptyList, NonEmptyMap, OptionT}
 
 trait TypesAlgebra[S[_], Alg[_]] {
 
@@ -19,7 +18,10 @@ trait TypesAlgebra[S[_], Alg[_]] {
 
   def resolveArrowDef(arrowDef: ArrowTypeToken[S]): Alg[Option[ArrowType]]
 
-  def resolveServiceType(name: NamedTypeToken[S], mustBeDefined: Boolean = true): Alg[Option[ServiceType]]
+  def resolveServiceType(
+    name: NamedTypeToken[S],
+    mustBeDefined: Boolean = true
+  ): Alg[Option[ServiceType]]
 
   def defineAbilityType(
     name: NamedTypeToken[S],
@@ -143,7 +145,8 @@ trait TypesAlgebra[S[_], Alg[_]] {
   def checkArrowCallResults(
     token: Token[S],
     arrowType: ArrowType,
-    results: List[Name[S]]
+    results: List[Name[S]],
+    exports: List[Call.Export]
   ): Alg[Unit]
 
   def checkArgumentsNumber(token: Token[S], expected: Int, givenNum: Int): Alg[Boolean]

semantics/src/main/scala/aqua/semantics/rules/types/TypesInterpreter.scala

@@ -2,6 +2,7 @@ package aqua.semantics.rules.types
 
 import aqua.errors.Errors.internalError
 import aqua.parser.lexer.*
+import aqua.raw.ops.Call
 import aqua.raw.value.*
 import aqua.semantics.Levenshtein
 import aqua.semantics.rules.StackInterpreter
@@ -595,7 +596,8 @@ class TypesInterpreter[S[_], X](using
   override def checkArrowCallResults(
     token: Token[S],
     arrowType: ArrowType,
-    results: List[Name[S]]
+    results: List[Name[S]],
+    exports: List[Call.Export]
   ): State[X, Unit] = for {
     _ <- results
       .drop(arrowType.codomain.length)
@@ -620,9 +622,22 @@ class TypesInterpreter[S[_], X](using
           case i => s"only $i are"
         }} used"
       )
-      .whenA(arrowType.codomain.length > results.length)
+      .whenA(checkNumberOfResults(arrowType, results, exports))
   } yield ()
 
+  private def checkNumberOfResults(
+    at: ArrowType,
+    results: List[Name[S]],
+    exports: List[Call.Export]
+  ): Boolean = {
+    val checkLength = at.codomain.length > results.length
+    val isOneStreamMapInExport =
+      exports.headOption.exists(e => isStreamMapType(e.`type`)) && exports.length == 1
+    val twoResultsToStreamMap = isOneStreamMapInExport && at.codomain.length == 2
+
+    !twoResultsToStreamMap && checkLength
+  }
+
   override def checkArgumentsNumber(
     token: Token[S],
     expected: Int,

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

@@ -398,6 +398,8 @@ case class StreamMapType(override val element: DataType) extends MutableStreamTy
     StructType(name, NonEmptyMap.of("key" -> ScalarType.string, "value" -> element))
 
   def toCanon: ImmutableCollectionType = CanonStreamMapType(element)
+
+  def toProduct: ProductType = ProductType(ScalarType.string :: element :: Nil)
 }
 
 object StreamMapType {