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

aqua-src/antithesis.aqua

@@ -1,41 +1,14 @@
-aqua StreamMapAbilities
+aqua StreamMapTest declares *
 
-export streamMapAbilityTest
+export main
 
-ability Streams:
+func foo() -> %string, u64:
-  stream: *string
   map: %string
+  <- map, 42
 
-ability Adds:
+func get() -> string, string:
-  addToStream(s: string)
+  <- "123", "44"
-  addToMap(k: string, v: string)
 
-func addToStreamClosure(str: *string) -> string -> ():
+func main() -> string:
-  cl = func (s: string):
+  m <- foo()
-    str <<- s
+  <- m.get("key")!
-  <- cl
-
-func addToMapClosure(str: %string) -> string, string -> ():
-  cl = func (k: string, v: string):
-    str <<- k, v
-  <- cl
-
-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()

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

@@ -85,7 +85,8 @@ import {
   testContainsFuncCall,
   testForFuncCall,
   testForTupleFuncCall,
-  testParSeqMapCall
+  testParSeqMapCall,
+  testInsertMapFromFuncCall
 } from "../examples/streamMapCall.js";
 import {
   topologyBug205Call,
@@ -965,6 +966,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

@@ -16,7 +16,7 @@
 
 import {
   testGetFunc, testGetStreamFunc, testKeysFunc, testKeysStreamFunc, testContainsFunc,
-  testForFunc, testParSeqMap, testForTupleFunc
+  testForFunc, testParSeqMap, testForTupleFunc, testInsertMapFromFunc
 } from "../compiled/examples/streamMap.js";
 import { config } from "../config.js";
 
@@ -54,3 +54,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

@@ -70,22 +70,35 @@ object ArrowInliner extends Logging {
     RawValueInliner
       .valueListToModel(results)
       .map(resolvedResults =>
-        // Fix the return values
+        (exportTo, resolvedResults) match {
-        (exportTo zip resolvedResults).map {
+          // if export is a stream map and there is two results from function and first one is not a stream map
-          case (
+          // TODO: can it be better?
+          case ((cexp @ CallModel.Export(n, st@StreamMapType(_))) :: Nil, (res1, desugar1) :: (res2, desugar2) :: Nil) if !Type.isStreamMapType(res1.`type`)  =>
+            (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) =>
+                ) if !outsideStreamNames.contains(n) =>
-            resDesugar.toList -> res
+                resDesugar.toList -> res
-          case (
+              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
+                // pass nested function results to a stream
-            (resDesugar.toList :+ PushToStreamModel(res, cexp).leaf) -> cexp.asVar
+                (resDesugar.toList :+ PushToStreamModel(res, cexp).leaf) -> cexp.asVar
-          case (_, (res, resDesugar)) =>
+              case (_, (res, resDesugar)) =>
-            resDesugar.toList -> res
+                resDesugar.toList -> res
-        }.unzip.leftMap(_.flatten)
+            }.unzip.leftMap(_.flatten)
+        }
+
       )
 
   /**

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

@@ -377,14 +377,10 @@ object TagInliner extends Logging {
           }
         } yield model.fold(TagInlined.Empty())(m => TagInlined.Single(model = m))
 
-      case DeclareStreamTag(value) =>
+      case DeclareStreamTag(name, t) =>
-        value match
+        for {
-          case VarRaw(name, t: MutableStreamType) =>
+          _ <- Exports[S].resolved(name, VarModel(name, t))
-            for {
+        } yield TagInlined.Empty()
-              _ <- Exports[S].resolved(name, VarModel(name, t))
-            } yield TagInlined.Empty()
-          case _ =>
-            internalError(s"Cannot declare $value as stream, because it is not a stream type")
 
       case ServiceIdTag(id, serviceType, name) =>
         for {

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

@@ -151,7 +151,7 @@ class ArrowInlinerSpec extends AnyFlatSpec with Matchers with Inside {
         FuncArrow(
           "stream-callback",
           SeqTag.wrap(
-            DeclareStreamTag(streamVar).leaf,
+            DeclareStreamTag(streamName, streamType).leaf,
             CallArrowRawTag.func("cb", Call(streamVar :: Nil, Nil)).leaf
           ),
           ArrowType(
@@ -247,7 +247,7 @@ class ArrowInlinerSpec extends AnyFlatSpec with Matchers with Inside {
         FuncArrow(
           "call",
           SeqTag.wrap(
-            DeclareStreamTag(streamVar).leaf,
+            DeclareStreamTag(streamName, streamType).leaf,
             CallArrowRawTag.func(useArrow.funcName, Call(streamVar :: Nil, Nil)).leaf
           ),
           ArrowType(
@@ -308,7 +308,7 @@ class ArrowInlinerSpec extends AnyFlatSpec with Matchers with Inside {
     val returnNil = FuncArrow(
       "returnNil",
       SeqTag.wrap(
-        DeclareStreamTag(someStr).leaf,
+        DeclareStreamTag(someStr.name, streamType).leaf,
         ReturnTag(
           NonEmptyList.one(someStr)
         ).leaf
@@ -409,7 +409,7 @@ class ArrowInlinerSpec extends AnyFlatSpec with Matchers with Inside {
     val returnStream = FuncArrow(
       "returnStream",
       SeqTag.wrap(
-        DeclareStreamTag(streamVar).leaf,
+        DeclareStreamTag(streamVar.name, streamType).leaf,
         PushToStreamTag(
           LiteralRaw.quote("one"),
           Call.Export(streamVar.name, streamVar.`type`)
@@ -595,7 +595,7 @@ class ArrowInlinerSpec extends AnyFlatSpec with Matchers with Inside {
     val returnFunc = FuncArrow(
       "return",
       SeqTag.wrap(
-        DeclareStreamTag(streamVar).leaf,
+        DeclareStreamTag(streamVar.name, streamType).leaf,
         PushToStreamTag(
           LiteralRaw.quote("one"),
           Call.Export(streamVar.name, streamVar.`type`)
@@ -893,7 +893,7 @@ class ArrowInlinerSpec extends AnyFlatSpec with Matchers with Inside {
         FuncArrow(
           "stream-callback",
           SeqTag.wrap(
-            DeclareStreamTag(streamVar).leaf,
+            DeclareStreamTag(streamVar.name, streamType).leaf,
             CallArrowRawTag.func("cb", Call(streamVarLambda :: Nil, Nil)).leaf
           ),
           ArrowType(
@@ -976,7 +976,7 @@ class ArrowInlinerSpec extends AnyFlatSpec with Matchers with Inside {
         FuncArrow(
           "outer",
           SeqTag.wrap(
-            DeclareStreamTag(recordsVar).leaf,
+            DeclareStreamTag(recordsVar.name, streamType).leaf,
             CallArrowRawTag.func(innerName, Call(recordsVar :: Nil, Nil)).leaf,
             CallArrowRawTag
               .service(
@@ -2311,7 +2311,8 @@ class ArrowInlinerSpec extends AnyFlatSpec with Matchers with Inside {
   it should "generate result in right order" in {
     val innerName = "inner"
     val results = VarRaw("results", ScalarType.string)
-    val resultsOut = VarRaw("results", StreamType(ScalarType.string))
+    val streamType = StreamType(ScalarType.string)
+    val resultsOut = VarRaw("results", streamType)
 
     val inner = FuncArrow(
       innerName,
@@ -2339,7 +2340,7 @@ class ArrowInlinerSpec extends AnyFlatSpec with Matchers with Inside {
     val outer = FuncArrow(
       "outer",
       SeqTag.wrap(
-        DeclareStreamTag(resultsOut).leaf,
+        DeclareStreamTag(resultsOut.name, streamType).leaf,
         CallArrowRawTag
           .func(innerName, Call(Nil, Call.Export(resultsOut.name, resultsOut.baseType) :: Nil))
           .leaf

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

@@ -16,6 +16,7 @@
 
 package aqua.raw.ops
 
+import aqua.errors.Errors.internalError
 import aqua.raw.arrow.FuncRaw
 import aqua.raw.value.{CallArrowRaw, CallServiceRaw, ValueRaw, VarRaw}
 import aqua.tree.{TreeNode, TreeNodeCompanion}
@@ -178,10 +179,16 @@ case class NextTag(item: String) extends RawTag {
 case class ForKeyValue(key: String, value: String) {
   def toSet: Set[String] = Set(key, value)
 
-  def rename(map: Map[String, String]): ForKeyValue = copy(key = map.getOrElse(key, key), value = map.getOrElse(value, value))
+  def rename(map: Map[String, String]): ForKeyValue =
+    copy(key = map.getOrElse(key, key), value = map.getOrElse(value, value))
 }
 
-case class ForTag(item: String, iterable: ValueRaw, mode: ForTag.Mode, keyValue: Option[ForKeyValue] = None) extends SeqGroupTag {
+case class ForTag(
+  item: String,
+  iterable: ValueRaw,
+  mode: ForTag.Mode,
+  keyValue: Option[ForKeyValue] = None
+) extends SeqGroupTag {
 
   override def restrictsVarNames: Set[String] = Set(item) ++ keyValue.toSet.flatMap(_.toSet)
 
@@ -283,15 +290,18 @@ object CallArrowRawTag {
 }
 
 case class DeclareStreamTag(
-  // TODO: Why is it ValueRaw and
+  name: String,
-  // not just (stream name, stream type)?
+  `type`: MutableStreamType
-  value: ValueRaw
 ) extends RawTag {
 
-  override def exportsVarNames: Set[String] = value.varNames
+  override def exportsVarNames: Set[String] = Set(name)
 
   override def mapValues(f: ValueRaw => ValueRaw): RawTag =
-    DeclareStreamTag(value.map(f))
+    f(VarRaw(name, `type`)) match {
+      case VarRaw(name, t: MutableStreamType) => copy(name, t)
+      case v =>
+        internalError(s"DeclareStreamTag can be only VarRaw with stream type, currently: '$v' ")
+    }
 }
 
 case class AssignmentTag(

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

@@ -23,7 +23,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.*
@@ -39,13 +39,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) =>
+    variables.traverse(v => N.read(v, mustBeDefined = false).map(v -> _)).flatMap {
-      N.read(v, mustBeDefined = false).flatMap {
+      case (v, Some(map @ StreamMapType(_))) :: Nil =>
-        case Some(stream @ StreamType(st)) =>
+        T.ensureTypeMatches(v, map.elementProduct, codomain)
-          T.ensureTypeMatches(v, st, t).as(Call.Export(v.value, stream, isExistingStream = true))
+          .as(Call.Export(v.value, map, isExistingStream = true) :: Nil)
-        case _ =>
+      case vars =>
-          N.define(v, t).as(Call.Export(v.value, t))
+        (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
@@ -56,8 +63,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)) <*
+      getExports(at.codomain).flatMap(exports =>
-        T.checkArrowCallResults(callArrow, at, variables)
+        T.checkArrowCallResults(callArrow, at, variables, exports).as(CallArrowRawTag(exports, raw))
+      )
     }
   } yield tag.map(_.funcOpLeaf)
 

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

@@ -41,8 +41,7 @@ class DeclareStreamSem[S[_]](val expr: DeclareStreamExpr[S]) {
       _ <- OptionT.withFilterF(
         N.define(expr.name, streamType)
       )
-      valueModel = VarRaw(expr.name.value, streamType)
+    } yield DeclareStreamTag(expr.name.value, streamType).funcOpLeaf: Raw
-    } yield DeclareStreamTag(valueModel).funcOpLeaf: Raw
 
     sem.getOrElse(Raw.error(s"Name `${expr.name.value}` not defined"))
   }

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

@@ -17,13 +17,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.{NonEmptyList, NonEmptyMap, OptionT}
-import cats.data.NonEmptyMap
-import cats.data.OptionT
 
 trait TypesAlgebra[S[_], Alg[_]] {
 
@@ -35,7 +34,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],
@@ -159,7 +161,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

@@ -18,6 +18,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
@@ -611,7 +612,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)
@@ -636,9 +638,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,

semantics/src/test/scala/aqua/semantics/SemanticsSpec.scala

@@ -158,7 +158,7 @@ class SemanticsSpec extends AnyFlatSpec with Matchers with Inside {
     val stream = VarRaw(name, streamType)
 
     SeqTag.wrap(
-      DeclareStreamTag(stream).leaf,
+      DeclareStreamTag(stream.name, streamType).leaf,
       PushToStreamTag(
         LiteralRaw.quote(value),
         Call.Export(name, streamType)

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

@@ -414,6 +414,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 elementProduct: ProductType = ProductType(ScalarType.string :: element :: Nil)
 }
 
 object StreamMapType {