Remake stream gate inlining, fix unit tests

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

@@ -107,8 +107,8 @@ class AquaCompilerSpec extends AnyFlatSpec with Matchers {
 
   private val init = LiteralModel.fromRaw(ValueRaw.InitPeerId)
 
-  private def join(vm: VarModel, idx: ValueModel) =
+  private def join(vm: VarModel, size: ValueModel) =
-    ResBuilder.join(vm, idx, init)
+    ResBuilder.join(vm, size, init)
 
   "aqua compiler" should "create right topology" in {
 
@@ -148,6 +148,7 @@ class AquaCompilerSpec extends AnyFlatSpec with Matchers {
     val canonResult = VarModel("-" + results.name + "-fix-0", CanonStreamType(resultsType.element))
     val flatResult = VarModel("-results-flat-0", ArrayType(ScalarType.string))
     val initPeer = LiteralModel.fromRaw(ValueRaw.InitPeerId)
+    val sizeVar = VarModel("results_size", LiteralType.unsigned)
     val retVar = VarModel("ret", ScalarType.string)
 
     val expected =
@@ -187,7 +188,13 @@ class AquaCompilerSpec extends AnyFlatSpec with Matchers {
                   )
                 )
               ),
-              join(results, LiteralModel.fromRaw(LiteralRaw.number(2))),
+              ResBuilder.add(
+                LiteralModel.number(2),
+                LiteralModel.number(1),
+                sizeVar,
+                initPeer
+              ),
+              join(results, sizeVar),
               CanonRes(results, init, CallModel.Export(canonResult.name, canonResult.`type`)).leaf,
               ApRes(
                 canonResult,

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

@@ -6,12 +6,12 @@ import aqua.model.*
 import aqua.model.inline.raw.{
   ApplyBinaryOpRawInliner,
   ApplyFunctorRawInliner,
-  ApplyGateRawInliner,
   ApplyPropertiesRawInliner,
   ApplyUnaryOpRawInliner,
   CallArrowRawInliner,
   CollectionRawInliner,
-  MakeAbilityRawInliner
+  MakeAbilityRawInliner,
+  StreamGateInliner
 }
 import aqua.raw.ops.*
 import aqua.raw.value.*
@@ -48,9 +48,6 @@ object RawValueInliner extends Logging {
       case alr: ApplyPropertyRaw =>
         ApplyPropertiesRawInliner(alr, propertiesAllowed)
 
-      case agr: ApplyGateRaw =>
-        ApplyGateRawInliner(agr, propertiesAllowed)
-
       case cr: CollectionRaw =>
         CollectionRawInliner(cr, propertiesAllowed)
 

model/inline/src/main/scala/aqua/model/inline/raw/ApplyGateRawInliner.scala

@@ -1,159 +0,0 @@
-package aqua.model.inline.raw
-
-import aqua.errors.Errors.internalError
-import aqua.model.*
-import aqua.model.inline.Inline
-import aqua.model.inline.state.{Arrows, Exports, Mangler}
-import aqua.raw.value.{ApplyGateRaw, LiteralRaw, VarRaw}
-import aqua.model.inline.RawValueInliner.unfold
-import aqua.types.{ArrayType, CanonStreamType, ScalarType, StreamType}
-
-import cats.data.State
-import cats.data.Chain
-import cats.syntax.monoid.*
-import cats.syntax.option.*
-import cats.syntax.applicative.*
-import scribe.Logging
-
-object ApplyGateRawInliner extends RawInliner[ApplyGateRaw] with Logging {
-
-  /**
-   * To wait for the element of a stream by the given index, the following model is generated:
-   * (seq
-   * (seq
-   *  (seq
-   *   (call <peer> ("math" "add") [0 1] stream_incr)
-   *   (fold $stream s
-   *    (seq
-   *     (seq
-   *      (ap s $stream_test)
-   *      (canon <peer> $stream_test  #stream_iter_canon)
-   *     )
-   *     (xor
-   *      (match #stream_iter_canon.length stream_incr
-   *       (null)
-   *      )
-   *      (next s)
-   *     )
-   *    )
-   *    (never)
-   *   )
-   *  )
-   *  (canon <peer> $stream_test  #stream_result_canon)
-   * )
-   * (ap #stream_result_canon stream_gate)
-   * )
-   */
-  def joinStreamOnIndexModel(
-    streamName: String,
-    streamType: StreamType,
-    idxModel: ValueModel,
-    idxIncrName: String,
-    testName: String,
-    iterName: String,
-    canonName: String,
-    iterCanonName: String,
-    resultName: String
-  ): OpModel.Tree = {
-    val varSTest = VarModel(testName, streamType)
-    val iter = VarModel(iterName, streamType.element)
-
-    val iterCanon = VarModel(iterCanonName, CanonStreamType(streamType.element))
-
-    val resultCanon =
-      VarModel(canonName, CanonStreamType(streamType.element))
-
-    val incrVar = VarModel(idxIncrName, ScalarType.u32)
-
-    RestrictionModel(varSTest.name, streamType).wrap(
-      increment(idxModel, incrVar),
-      ForModel(iter.name, VarModel(streamName, streamType), ForModel.Mode.Never.some).wrap(
-        PushToStreamModel(
-          iter,
-          CallModel.Export(varSTest.name, varSTest.`type`)
-        ).leaf,
-        CanonicalizeModel(
-          varSTest,
-          CallModel.Export(iterCanon.name, iterCanon.`type`)
-        ).leaf,
-        XorModel.wrap(
-          MatchMismatchModel(
-            iterCanon
-              .copy(properties = Chain.one(FunctorModel("length", ScalarType.`u32`))),
-            incrVar,
-            true
-          ).leaf,
-          NextModel(iter.name).leaf
-        )
-      ),
-      CanonicalizeModel(
-        varSTest,
-        CallModel.Export(resultCanon.name, CanonStreamType(streamType.element))
-      ).leaf,
-      FlattenModel(
-        resultCanon,
-        resultName
-      ).leaf
-    )
-  }
-
-  override def apply[S: Mangler: Exports: Arrows](
-    afr: ApplyGateRaw,
-    propertyAllowed: Boolean
-  ): State[S, (ValueModel, Inline)] =
-    for {
-      uniqueCanonName <- Mangler[S].findAndForbidName(afr.name + "_result_canon")
-      uniqueResultName <- Mangler[S].findAndForbidName(afr.name + "_gate")
-      uniqueTestName <- Mangler[S].findAndForbidName(afr.name + "_test")
-      uniqueIdxIncr <- Mangler[S].findAndForbidName(afr.name + "_incr")
-      uniqueIterCanon <- Mangler[S].findAndForbidName(afr.name + "_iter_canon")
-      uniqueIter <- Mangler[S].findAndForbidName(afr.name + "_fold_var")
-      idxFolded <- unfold(afr.idx)
-      (idxModel, idxInline) = idxFolded
-      idxFlattened <- idxModel match {
-        case vr: VarModel => ApplyPropertiesRawInliner.removeProperties(vr)
-        case _ => (idxModel, Inline.empty).pure[State[S, *]]
-      }
-      (idxFlatModel, idxFlatInline) = idxFlattened
-    } yield {
-      val gate = joinStreamOnIndexModel(
-        streamName = afr.name,
-        streamType = afr.streamType,
-        idxModel = idxModel,
-        idxIncrName = uniqueIdxIncr,
-        testName = uniqueTestName,
-        iterName = uniqueIter,
-        canonName = uniqueCanonName,
-        iterCanonName = uniqueIterCanon,
-        resultName = uniqueResultName
-      )
-
-      val tree = SeqModel.wrap(
-        idxInline.predo.toList ++
-          idxFlatInline.predo.toList :+
-          gate
-      )
-      val treeInline = Inline(predo = Chain.one(tree))
-      val idx = IntoIndexModel
-        .fromValueModel(idxFlatModel, afr.streamType.element)
-        .getOrElse(
-          internalError(s"Unexpected: cant convert ($idxFlatModel) to IntoIndexModel")
-        )
-      val value = VarModel(
-        uniqueResultName,
-        ArrayType(afr.streamType.element)
-      ).withProperty(idx)
-
-      (value, treeInline)
-    }
-
-  private def increment(v: ValueModel, result: VarModel) =
-    CallServiceModel(
-      LiteralModel("\"math\"", ScalarType.string),
-      "add",
-      CallModel(
-        v :: LiteralModel.fromRaw(LiteralRaw.number(1)) :: Nil,
-        CallModel.Export(result.name, result.`type`) :: Nil
-      )
-    ).leaf
-}

model/inline/src/main/scala/aqua/model/inline/raw/ApplyPropertiesRawInliner.scala

@@ -254,29 +254,47 @@ object ApplyPropertiesRawInliner extends RawInliner[ApplyPropertyRaw] with Loggi
     propertiesAllowed: Boolean
   ): State[S, (ValueModel, Inline)] = {
     ((raw, properties.uncons) match {
-      case (vr @ VarRaw(_, st @ StreamType(_)), Some(IntoIndexRaw(idx, _), otherProperties)) =>
+      case (
+            vr @ VarRaw(_, st @ StreamType(_)),
+            Some(IntoIndexRaw(idx, _), otherProperties)
+          ) =>
         unfold(vr).flatMap {
           case (VarModel(nameVM, _, _), inl) =>
-            val gateRaw = ApplyGateRaw(nameVM, st, idx)
+            for {
-            unfold(gateRaw).flatMap {
+              idxInlined <- unfold(idx)
-              case (gateResVal: VarModel, gateResInline) =>
+              (idxVM, idxInline) = idxInlined
-                unfoldProperties(
+              sizeName <- Mangler[S].findAndForbidName(s"${nameVM}_size")
-                  gateResInline,
+              sizeVar = VarModel(sizeName, idxVM.`type`)
-                  gateResVal,
+              sizeInline = CallServiceModel(
-                  otherProperties,
+                "math",
-                  propertiesAllowed
+                funcName = "add",
-                ).map { case (v, i) =>
+                args = List(idxVM, LiteralModel.number(1)),
-                  v -> Inline(
+                result = sizeVar
-                    inl.predo ++ i.predo,
+              ).leaf
-                    mergeMode = SeqMode
+              gateInlined <- StreamGateInliner(nameVM, st, sizeVar)
+              (gateVM, gateInline) = gateInlined
+              idxFlattened <- idxVM match {
+                case vr: VarModel => removeProperties(vr)
+                case _ => (idxVM, Inline.empty).pure[State[S, *]]
+              }
+              (idxFlat, idxFlatInline) = idxFlattened
+              gate = gateVM.withProperty(
+                IntoIndexModel
+                  .fromValueModel(idxFlat, st.element)
+                  .getOrElse(
+                    internalError(s"Unexpected: could not convert ($idxFlat) to IntoIndexModel")
                   )
-                }
+              )
-              case (v, i) =>
+              propsInlined <- unfoldProperties(
-                // what if pass nil as stream argument?
+                Inline(
-                internalError(
+                  (idxInline.predo :+ sizeInline) ++ gateInline.predo,
-                  s"Unfolded stream ($gateRaw) cannot be a literal"
+                  mergeMode = SeqMode
-                )
+                ),
-            }
+                gate,
+                otherProperties,
+                propertiesAllowed
+              )
+            } yield propsInlined
           case l =>
             internalError(
               s"Unfolded stream ($vr) cannot be a literal"
@@ -308,7 +326,7 @@ object ApplyPropertiesRawInliner extends RawInliner[ApplyPropertyRaw] with Loggi
   /**
    * Remove properties from the var and return a new var without them
    */
-  def removeProperties[S: Mangler](
+  private def removeProperties[S: Mangler](
     varModel: VarModel
   ): State[S, (VarModel, Inline)] =
     if (varModel.properties.isEmpty) (varModel, Inline.empty).pure

model/inline/src/main/scala/aqua/model/inline/raw/StreamGateInliner.scala

@@ -0,0 +1,126 @@
+package aqua.model.inline.raw
+
+import aqua.errors.Errors.internalError
+import aqua.model.*
+import aqua.model.inline.Inline
+import aqua.model.inline.state.{Arrows, Exports, Mangler}
+import aqua.raw.value.{LiteralRaw, VarRaw}
+import aqua.model.inline.RawValueInliner.unfold
+import aqua.types.{ArrayType, CanonStreamType, ScalarType, StreamType}
+
+import cats.data.State
+import cats.data.Chain
+import cats.syntax.monoid.*
+import cats.syntax.option.*
+import cats.syntax.applicative.*
+import scribe.Logging
+import cats.instances.stream
+
+object StreamGateInliner extends Logging {
+
+  /**
+   * To wait for size elements of a stream,
+   * the following model is generated:
+   * (seq
+   * (seq
+   *  (fold $stream s
+   *   (seq
+   *    (seq
+   *     (ap s $stream_test)
+   *     (canon <peer> $stream_test  #stream_iter_canon)
+   *    )
+   *    (xor
+   *     (match #stream_iter_canon.length size
+   *      (null)
+   *     )
+   *     (next s)
+   *    )
+   *   )
+   *   (never)
+   *  )
+   *  (canon <peer> $stream_test  #stream_result_canon)
+   * )
+   * (ap #stream_result_canon stream_gate)
+   * )
+   */
+  def joinStreamOnIndexModel(
+    streamName: String,
+    streamType: StreamType,
+    sizeModel: ValueModel,
+    testName: String,
+    iterName: String,
+    canonName: String,
+    iterCanonName: String,
+    resultName: String
+  ): OpModel.Tree = {
+    val varSTest = VarModel(testName, streamType)
+    val iter = VarModel(iterName, streamType.element)
+    val iterCanon = VarModel(iterCanonName, CanonStreamType(streamType.element))
+    val resultCanon = VarModel(canonName, CanonStreamType(streamType.element))
+
+    RestrictionModel(varSTest.name, streamType).wrap(
+      ForModel(iter.name, VarModel(streamName, streamType), ForModel.Mode.Never.some).wrap(
+        PushToStreamModel(
+          iter,
+          CallModel.Export(varSTest.name, varSTest.`type`)
+        ).leaf,
+        CanonicalizeModel(
+          varSTest,
+          CallModel.Export(iterCanon.name, iterCanon.`type`)
+        ).leaf,
+        XorModel.wrap(
+          MatchMismatchModel(
+            iterCanon
+              .withProperty(
+                FunctorModel("length", ScalarType.`u32`)
+              ),
+            sizeModel,
+            true
+          ).leaf,
+          NextModel(iter.name).leaf
+        )
+      ),
+      CanonicalizeModel(
+        varSTest,
+        CallModel.Export(resultCanon.name, CanonStreamType(streamType.element))
+      ).leaf,
+      FlattenModel(
+        resultCanon,
+        resultName
+      ).leaf
+    )
+  }
+
+  def apply[S: Mangler: Exports: Arrows](
+    streamName: String,
+    streamType: StreamType,
+    sizeModel: ValueModel
+  ): State[S, (VarModel, Inline)] =
+    for {
+      uniqueCanonName <- Mangler[S].findAndForbidName(streamName + "_result_canon")
+      uniqueResultName <- Mangler[S].findAndForbidName(streamName + "_gate")
+      uniqueTestName <- Mangler[S].findAndForbidName(streamName + "_test")
+      uniqueIdxIncr <- Mangler[S].findAndForbidName(streamName + "_incr")
+      uniqueIterCanon <- Mangler[S].findAndForbidName(streamName + "_iter_canon")
+      uniqueIter <- Mangler[S].findAndForbidName(streamName + "_fold_var")
+    } yield {
+      val gate = joinStreamOnIndexModel(
+        streamName = streamName,
+        streamType = streamType,
+        sizeModel = sizeModel,
+        testName = uniqueTestName,
+        iterName = uniqueIter,
+        canonName = uniqueCanonName,
+        iterCanonName = uniqueIterCanon,
+        resultName = uniqueResultName
+      )
+
+      val inline = Inline(predo = Chain.one(gate))
+      val value = VarModel(
+        uniqueResultName,
+        ArrayType(streamType.element)
+      )
+
+      (value, inline)
+    }
+}

model/raw/src/main/scala/aqua/raw/value/ValueRaw.scala

@@ -58,8 +58,8 @@ object ValueRaw {
     errorType
   )
 
-  type ApplyRaw = ApplyGateRaw | ApplyPropertyRaw | CallArrowRaw | CollectionRaw |
+  type ApplyRaw = ApplyPropertyRaw | CallArrowRaw | CollectionRaw | ApplyBinaryOpRaw |
-    ApplyBinaryOpRaw | ApplyUnaryOpRaw
+    ApplyUnaryOpRaw
 }
 
 case class ApplyPropertyRaw(value: ValueRaw, property: PropertyRaw) extends ValueRaw {
@@ -94,22 +94,6 @@ object ApplyPropertyRaw {
     }
 }
 
-case class ApplyGateRaw(name: String, streamType: StreamType, idx: ValueRaw) extends ValueRaw {
-  override def baseType: Type = streamType
-
-  override def `type`: Type = idx.`type`
-
-  override def renameVars(map: Map[String, String]): ValueRaw =
-    copy(name = map.getOrElse(name, name), idx = idx.renameVars(map))
-
-  override def map(f: ValueRaw => ValueRaw): ValueRaw =
-    f(copy(idx = f(idx)))
-
-  override def toString: String = s"gate $name.$idx"
-
-  override def varNames: Set[String] = Set(name) ++ idx.varNames
-}
-
 case class VarRaw(name: String, baseType: Type) extends ValueRaw {
 
   override def map(f: ValueRaw => ValueRaw): ValueRaw = f(this)

model/res/src/test/scala/aqua/res/ResBuilder.scala

@@ -9,32 +9,22 @@ import cats.syntax.option.*
 
 object ResBuilder {
 
-  def join(stream: VarModel, onIdx: ValueModel, peer: ValueModel) = {
+  def join(stream: VarModel, sizeModel: ValueModel, peer: ValueModel) = {
     val testVM = VarModel(stream.name + "_test", stream.`type`)
     val testStreamType = stream.`type`.asInstanceOf[StreamType] // Unsafe
     val iter = VarModel(stream.name + "_fold_var", ScalarType.string)
     val canon = VarModel(stream.name + "_iter_canon", CanonStreamType(ScalarType.string))
     val canonRes = VarModel(stream.name + "_result_canon", CanonStreamType(ScalarType.string))
     val arrayRes = VarModel(stream.name + "_gate", ArrayType(ScalarType.string))
-    val idx = VarModel(stream.name + "_incr", ScalarType.u32)
 
     RestrictionRes(testVM.name, testStreamType).wrap(
-      CallServiceRes(
-        LiteralModel("\"math\"", ScalarType.string),
-        "add",
-        CallRes(
-          onIdx :: LiteralModel.fromRaw(LiteralRaw.number(1)) :: Nil,
-          Some(CallModel.Export(idx.name, idx.`type`))
-        ),
-        peer
-      ).leaf,
       FoldRes(iter.name, stream, ForModel.Mode.Never.some).wrap(
         ApRes(iter, CallModel.Export(testVM.name, testVM.`type`)).leaf,
         CanonRes(testVM, peer, CallModel.Export(canon.name, canon.`type`)).leaf,
         XorRes.wrap(
           MatchMismatchRes(
             canon.copy(properties = Chain.one(FunctorModel("length", ScalarType.u32))),
-            idx,
+            sizeModel,
             true
           ).leaf,
           NextRes(iter.name).leaf
@@ -45,4 +35,20 @@ object ResBuilder {
     )
   }
 
+  def add(
+    a: ValueModel,
+    b: ValueModel,
+    res: VarModel,
+    peer: ValueModel
+  ): ResolvedOp.Tree =
+    CallServiceRes(
+      LiteralModel.quote("math"),
+      "add",
+      CallRes(
+        a :: b :: Nil,
+        Some(CallModel.Export(res.name, res.`type`))
+      ),
+      peer
+    ).leaf
+
 }

model/transform/src/test/scala/aqua/model/transform/ModelBuilder.scala

@@ -8,7 +8,7 @@ import aqua.res.{CallRes, CallServiceRes, MakeRes}
 import aqua.types.{ArrayType, LiteralType, ScalarType}
 import aqua.types.StreamType
 import aqua.model.IntoIndexModel
-import aqua.model.inline.raw.ApplyGateRawInliner
+import aqua.model.inline.raw.StreamGateInliner
 import aqua.model.OnModel
 import aqua.model.FailModel
 import aqua.res.ResolvedOp
@@ -142,21 +142,20 @@ object ModelBuilder {
 
   /**
    * @param stream stream [[VarModel]]
-   * @param idx id [[ValueModel]]
+   * @param size size [[ValueModel]]
-   * @return [[OpModel.Tree]] of join of `stream[idx]`
+   * @return [[OpModel.Tree]] of join of size elements of stream
    */
-  def join(stream: VarModel, idx: ValueModel): OpModel.Tree =
+  def join(stream: VarModel, size: ValueModel): OpModel.Tree =
     stream match {
       case VarModel(
             streamName,
             streamType: StreamType,
             Chain.`nil`
           ) =>
-        ApplyGateRawInliner.joinStreamOnIndexModel(
+        StreamGateInliner.joinStreamOnIndexModel(
           streamName = streamName,
           streamType = streamType,
-          idxModel = idx,
+          sizeModel = size,
-          idxIncrName = streamName + "_incr",
           testName = streamName + "_test",
           iterName = streamName + "_fold_var",
           canonName = streamName + "_result_canon",
@@ -165,4 +164,12 @@ object ModelBuilder {
         )
       case _ => ???
     }
+
+  def add(a: ValueModel, b: ValueModel, res: VarModel): OpModel.Tree =
+    CallServiceModel(
+      "math",
+      "add",
+      args = List(a, b),
+      result = res
+    ).leaf
 }

model/transform/src/test/scala/aqua/model/transform/topology/TopologySpec.scala

@@ -25,22 +25,45 @@ class TopologySpec extends AnyFlatSpec with Matchers {
   import ModelBuilder.{join as joinModel, *}
   import ResBuilder.join as joinRes
 
-  def joinModelRes(streamEl: ValueRaw | ValueModel): (OpModel.Tree, ResolvedOp.Tree) =
+  def joinModelRes(
+    streamEl: ValueRaw | ValueModel
+  ): (Chain[OpModel.Tree], Chain[ResolvedOp.Tree]) =
     streamEl match {
       case vm: ValueModel => vm
       case vr: ValueRaw => ValueModel.fromRaw(vr)
     } match {
       case stream @ VarModel(name, baseType, IntoIndexModel(idx, idxType) ==: Chain.`nil`) =>
         val idxModel =
-          if (idx.forall(Character.isDigit)) LiteralModel(idx, idxType)
+          if (idx.forall(Character.isDigit)) LiteralModel(idx, ScalarType.u32)
-          else VarModel(idx, idxType)
+          else VarModel(idx, ScalarType.u32)
 
         val streamWithoutIdx = stream.copy(properties = Chain.`nil`)
 
-        (
+        val sizeModel = VarModel(s"${name}_size", ScalarType.u32)
-          joinModel(streamWithoutIdx, idxModel),
+        val sizeTree = ModelBuilder.add(
-          joinRes(streamWithoutIdx, idxModel, ValueModel.fromRaw(initPeer))
+          idxModel,
+          LiteralModel.number(1),
+          sizeModel
         )
+
+        val model = Chain(
+          sizeTree,
+          joinModel(streamWithoutIdx, sizeModel)
+        )
+
+        val sizeTreeResolved = ResBuilder.add(
+          idxModel,
+          LiteralModel.number(1),
+          sizeModel,
+          ValueModel.fromRaw(initPeer)
+        )
+
+        val resolved = Chain(
+          sizeTreeResolved,
+          joinRes(streamWithoutIdx, sizeModel, ValueModel.fromRaw(initPeer))
+        )
+
+        (model, resolved)
       case _ => ???
     }
 
@@ -462,54 +485,58 @@ class TopologySpec extends AnyFlatSpec with Matchers {
 
     val (joinModel, joinRes) = joinModelRes(streamEl)
 
+    val foldModel = foldPar(
+      "i",
+      valueArray,
+      OnModel(iRelay, Chain.empty).wrap(
+        XorModel.wrap(
+          callModel(2, CallModel.Export(streamRaw.name, streamRaw.`type`) :: Nil),
+          OnModel(initPeer, Chain.one(relay)).wrap(
+            callModel(4, Nil, Nil)
+          )
+        )
+      )
+    )
     val init = SeqModel.wrap(
       DeclareStreamModel(stream).leaf,
       OnModel(initPeer, Chain.one(relay)).wrap(
-        foldPar(
+        foldModel +:
-          "i",
+          joinModel :+
-          valueArray,
+          callModel(3, Nil, streamRaw :: Nil)
-          OnModel(iRelay, Chain.empty).wrap(
-            XorModel.wrap(
-              callModel(2, CallModel.Export(streamRaw.name, streamRaw.`type`) :: Nil),
-              OnModel(initPeer, Chain.one(relay)).wrap(
-                callModel(4, Nil, Nil)
-              )
-            )
-          )
-        ),
-        joinModel,
-        callModel(3, Nil, streamRaw :: Nil)
       )
     )
 
     val proc = Topology.resolve(init).value
 
-    val expected = SeqRes.wrap(
+    val foldRes = ParRes.wrap(
-      through(relay),
+      FoldRes("i", valueArray, ForModel.Mode.Never.some).wrap(
-      ParRes.wrap(
+        ParRes.wrap(
-        FoldRes("i", valueArray, ForModel.Mode.Never.some).wrap(
+          // better if first relay will be outside `for`
-          ParRes.wrap(
+          SeqRes.wrap(
-            // better if first relay will be outside `for`
+            through(relay),
-            SeqRes.wrap(
+            XorRes.wrap(
-              through(relay),
+              SeqRes.wrap(
-              XorRes.wrap(
+                callRes(2, iRelay, Some(CallModel.Export(streamRaw.name, streamRaw.`type`))),
-                SeqRes.wrap(
+                through(relay),
-                  callRes(2, iRelay, Some(CallModel.Export(streamRaw.name, streamRaw.`type`))),
+                through(initPeer)
-                  through(relay),
+              ),
-                  through(initPeer)
+              SeqRes.wrap(
-                ),
+                through(relay),
-                SeqRes.wrap(
+                callRes(4, initPeer)
-                  through(relay),
-                  callRes(4, initPeer)
-                )
               )
-            ),
+            )
-            NextRes("i").leaf
+          ),
-          )
+          NextRes("i").leaf
         )
-      ),
+      )
-      joinRes,
+    )
-      callRes(3, initPeer, None, stream :: Nil)
+    val expected = SeqRes.wrap(
+      Chain(
+        through(relay),
+        foldRes
+      ) ++
+        joinRes :+
+        callRes(3, initPeer, None, stream :: Nil)
     )
 
     proc.equalsOrShowDiff(expected) should be(true)
@@ -543,42 +570,45 @@ class TopologySpec extends AnyFlatSpec with Matchers {
               )
             )
           )
-        ),
+        ) +:
-        joinModel,
+          joinModel :+
-        callModel(3, Nil, streamRaw :: Nil)
+          callModel(3, Nil, streamRaw :: Nil)
       )
     )
 
     val proc = Topology.resolve(init).value
 
-    val expected = SeqRes.wrap(
+    val fold = ParRes.wrap(
-      through(relay),
+      FoldRes("i", valueArray, ForModel.Mode.Never.some).wrap(
-      ParRes.wrap(
+        ParRes.wrap(
-        FoldRes("i", valueArray, ForModel.Mode.Never.some).wrap(
+          // better if first relay will be outside `for`
-          ParRes.wrap(
+          SeqRes.wrap(
-            // better if first relay will be outside `for`
+            through(relay),
-            SeqRes.wrap(
+            XorRes.wrap(
-              through(relay),
               XorRes.wrap(
-                XorRes.wrap(
-                  SeqRes.wrap(
-                    callRes(2, iRelay, Some(CallModel.Export(streamRaw.name, streamRaw.`type`))),
-                    through(relay),
-                    through(initPeer)
-                  )
-                ),
                 SeqRes.wrap(
+                  callRes(2, iRelay, Some(CallModel.Export(streamRaw.name, streamRaw.`type`))),
                   through(relay),
-                  callRes(4, initPeer)
+                  through(initPeer)
                 )
+              ),
+              SeqRes.wrap(
+                through(relay),
+                callRes(4, initPeer)
               )
-            ),
+            )
-            NextRes("i").leaf
+          ),
-          )
+          NextRes("i").leaf
         )
-      ),
+      )
-      joinRes,
+    )
-      callRes(3, initPeer, None, stream :: Nil)
+    val expected = SeqRes.wrap(
+      Chain(
+        through(relay),
+        fold
+      ) ++
+        joinRes :+
+        callRes(3, initPeer, None, stream :: Nil)
     )
 
     // println(Console.MAGENTA + init.show + Console.RESET)
@@ -804,9 +834,9 @@ class TopologySpec extends AnyFlatSpec with Matchers {
         OnModel(i, Chain.empty).wrap(
           callModel(1, CallModel.Export(used.name, used.`type`) :: Nil)
         )
-      ),
+      ) +:
-      joinModel,
+        joinModel :+
-      callModel(3, Nil, used :: Nil)
+        callModel(3, Nil, used :: Nil)
     )
 
     val proc = Topology.resolve(init).value
@@ -828,9 +858,9 @@ class TopologySpec extends AnyFlatSpec with Matchers {
             NextRes("i").leaf
           )
         )
-      ),
+      ) +:
-      joinRes,
+        joinRes :+
-      callRes(3, initPeer, None, ValueModel.fromRaw(used) :: Nil)
+        callRes(3, initPeer, None, ValueModel.fromRaw(used) :: Nil)
     )
 
     proc.equalsOrShowDiff(expected) should be(true)
@@ -844,46 +874,48 @@ class TopologySpec extends AnyFlatSpec with Matchers {
 
     val (joinModel, joinRes) = joinModelRes(usedWithIdx)
 
-    val init = OnModel(initPeer, Chain.one(relay)).wrap(
+    val foldModel = foldPar(
-      foldPar(
+      "i",
-        "i",
+      valueArray,
-        valueArray,
+      OnModel(i, Chain.empty).wrap(
-        OnModel(i, Chain.empty).wrap(
+        XorModel.wrap(
-          XorModel.wrap(
+          callModel(1, CallModel.Export(used.name, used.`type`) :: Nil)
-            callModel(1, CallModel.Export(used.name, used.`type`) :: Nil)
-          )
         )
-      ),
+      )
-      joinModel,
+    )
-      callModel(3, Nil, used :: Nil)
+    val init = OnModel(initPeer, Chain.one(relay)).wrap(
+      foldModel +:
+        joinModel :+
+        callModel(3, Nil, used :: Nil)
     )
 
     val proc = Topology.resolve(init).value
 
-    val expected = SeqRes.wrap(
+    val foldRes = ParRes.wrap(
-      ParRes.wrap(
+      FoldRes("i", ValueModel.fromRaw(valueArray), ForModel.Mode.Never.some).wrap(
-        FoldRes("i", ValueModel.fromRaw(valueArray), ForModel.Mode.Never.some).wrap(
+        ParRes.wrap(
-          ParRes.wrap(
+          SeqRes.wrap(
-            SeqRes.wrap(
+            through(relay),
-              through(relay),
+            XorRes.wrap(
-              XorRes.wrap(
+              SeqRes.wrap(
-                SeqRes.wrap(
+                callRes(
-                  callRes(
+                  1,
-                    1,
+                  ValueModel.fromRaw(i),
-                    ValueModel.fromRaw(i),
+                  Some(CallModel.Export(used.name, used.`type`))
-                    Some(CallModel.Export(used.name, used.`type`))
+                ),
-                  ),
+                through(relay),
-                  through(relay),
+                through(initPeer)
-                  through(initPeer)
-                )
               )
-            ),
+            )
-            NextRes("i").leaf
+          ),
-          )
+          NextRes("i").leaf
         )
-      ),
+      )
-      joinRes,
+    )
-      callRes(3, initPeer, None, ValueModel.fromRaw(used) :: Nil)
+    val expected = SeqRes.wrap(
+      foldRes +:
+        joinRes :+
+        callRes(3, initPeer, None, ValueModel.fromRaw(used) :: Nil)
     )
 
     proc.equalsOrShowDiff(expected) should be(true)