67 plain air gen (#78)

* Refactoring ForClient to make its parts reusable/recomposable

* Func transformation decomposed into parts

* Improves AIR compilation target

aqua-src/demo.aqua

@@ -0,0 +1,4 @@
+service Demo("demo"):
+	get4() -> u64
+	get5(arg: u32)
+	get6(a: string) -> bool

aqua-src/error.aqua

@@ -1,7 +1,4 @@
-service Demo("demo"):
-	get4() -> u64
-	get5(arg: u32)
-	get6(a: string) -> bool
+import "demo.aqua"
 
 func one() -> u64:
     variable <- Demo.get4()

backend/air/src/main/scala/aqua/backend/air/FuncAirGen.scala

@@ -1,15 +1,17 @@
 package aqua.backend.air
 
 import aqua.model.func.FuncCallable
-import aqua.model.transform.{BodyConfig, ForClient}
+import aqua.model.transform.{BodyConfig, Transform}
 
 case class FuncAirGen(func: FuncCallable) {
 
   /**
-   * Generates AIR from the function body as it is, with no modifications and optimizations
+   * Generates AIR from the function body
    */
-  def generateAir: Air =
-    AirGen(func.body.tree).generate
+  def generateAir(conf: BodyConfig = BodyConfig()): Air =
+    AirGen(
+      Transform.forClient(func, conf)
+    ).generate
 
   /**
    * Generates AIR from the optimized function body, assuming client is behind a relay
@@ -17,6 +19,6 @@ case class FuncAirGen(func: FuncCallable) {
    */
   def generateClientAir(conf: BodyConfig = BodyConfig()): Air =
     AirGen(
-      ForClient.resolve(func, conf)
+      Transform.forClient(func, conf)
     ).generate
 }

build.sbt

@@ -17,7 +17,7 @@ val declineV = "2.0.0-RC1"
 name := "aqua-hll"
 
 val commons = Seq(
-  baseAquaVersion := "0.1.1",
+  baseAquaVersion                        := "0.1.1",
   version                                := baseAquaVersion.value + "-" + sys.env.getOrElse("BUILD_NUMBER", "SNAPSHOT"),
   scalaVersion                           := dottyVersion,
   libraryDependencies += "org.scalatest" %% "scalatest" % scalaTestV % Test,

cli/src/main/scala/aqua/Aqua.scala

@@ -1,13 +1,8 @@
 package aqua
 
-import aqua.backend.air.FuncAirGen
-import aqua.backend.ts.TypescriptFile
-import aqua.model.{Model, ScriptModel}
 import aqua.parser.Ast
 import cats.data.ValidatedNec
 import aqua.parser.lift.{FileSpan, LiftParser, Span}
-import aqua.semantics.Semantics
-import cats.syntax.show._
 
 object Aqua {
 
@@ -21,29 +16,4 @@ object Aqua {
       .leftMap(_.map(pe => SyntaxError(pe.failedAtOffset, pe.expected)))
   }
 
-  // Will fail if imports are used
-  def generateModel(input: String): ValidatedNec[AquaError, Model] =
-    parseString(input).andThen(ast =>
-      Semantics.generateModel(ast).leftMap(_.map(ts => CompilerError(ts._1.unit._1, ts._2)))
-    )
-
-  def generate(input: String, air: Boolean): ValidatedNec[AquaError, String] =
-    generateModel(input).map {
-      case m: ScriptModel if air =>
-        // TODO it's meaningless to compile all functions to air, as resulting .air file is incorrect; only one function should be taken
-        m.resolveFunctions
-          .map(FuncAirGen)
-          .map(g =>
-            // add function name before body
-            s";; function name: ${g.func.funcName}\n\n" + g.generateAir.show
-          )
-          .toList
-          .mkString("\n\n\n")
-
-      case m: ScriptModel =>
-        TypescriptFile(m).generateTS()
-
-      case _ => "//No input given"
-    }
-
 }

cli/src/main/scala/aqua/AquaCli.scala

@@ -1,5 +1,6 @@
 package aqua
 
+import aqua.model.transform.BodyConfig
 import cats.data.Validated
 import cats.effect.{ExitCode, IO, IOApp}
 import com.monovore.decline.Opts
@@ -35,7 +36,8 @@ object AquaCli extends IOApp {
           input,
           imports,
           output,
-          if (toAir) AquaCompiler.AirTarget else AquaCompiler.TypescriptTarget
+          if (toAir) AquaCompiler.AirTarget else AquaCompiler.TypescriptTarget,
+          BodyConfig()
         )
         .map {
           case Validated.Invalid(errs) =>

cli/src/main/scala/aqua/AquaCompiler.scala

@@ -5,6 +5,7 @@ import aqua.backend.ts.TypescriptFile
 import aqua.io.{AquaFileError, AquaFiles, FileModuleId, Unresolvable}
 import aqua.linker.Linker
 import aqua.model.ScriptModel
+import aqua.model.transform.BodyConfig
 import aqua.parser.lexer.Token
 import aqua.parser.lift.FileSpan
 import aqua.semantics.{CompilerState, Semantics}
@@ -25,7 +26,12 @@ object AquaCompiler {
   case object TypescriptTarget extends CompileTarget
   case object AirTarget extends CompileTarget
 
-  case class Prepared(target: String => Path, model: ScriptModel)
+  case class Prepared(target: String => Path, model: ScriptModel) {
+
+    def hasOutput(target: CompileTarget): Boolean = target match {
+      case _ => model.funcs.nonEmpty
+    }
+  }
 
   def prepareFiles[F[_]: Files: Concurrent](
     srcPath: Path,
@@ -98,50 +104,50 @@ object AquaCompiler {
     srcPath: Path,
     imports: LazyList[Path],
     targetPath: Path,
-    compileTo: CompileTarget
+    compileTo: CompileTarget,
+    bodyConfig: BodyConfig
   ): F[ValidatedNec[String, Chain[String]]] =
-    prepareFiles(srcPath, imports, targetPath).flatMap[ValidatedNec[String, Chain[String]]] {
-      case Validated.Invalid(e) =>
-        Applicative[F].pure(Validated.invalid(e))
-      case Validated.Valid(preps) =>
-        (compileTo match {
-          case TypescriptTarget =>
-            preps
-              .map(p => writeFile(p.target("ts"), TypescriptFile(p.model).generateTS()))
+    prepareFiles(srcPath, imports, targetPath)
+      .map(_.map(_.filter(_.hasOutput(compileTo))))
+      .flatMap[ValidatedNec[String, Chain[String]]] {
+        case Validated.Invalid(e) =>
+          Applicative[F].pure(Validated.invalid(e))
+        case Validated.Valid(preps) =>
+          (compileTo match {
+            case TypescriptTarget =>
+              preps
+                .map(p => writeFile(p.target("ts"), TypescriptFile(p.model).generateTS(bodyConfig)))
 
-          // TODO add function name to AirTarget class
-          case AirTarget =>
-            preps
-              .map(p =>
-                writeFile(
-                  p.target("air"),
-                  p.model.resolveFunctions
-                    .map(FuncAirGen)
-                    .map(g =>
-                      // add function name before body
-                      s";; function name: ${g.func.funcName}\n\n" + g.generateAir.show
+            // TODO add function name to AirTarget class
+            case AirTarget =>
+              preps
+                .flatMap(p =>
+                  p.model.resolveFunctions.map { fc =>
+                    fc.funcName -> FuncAirGen(fc).generateAir(bodyConfig).show
+                  }.map { case (n, g) =>
+                    writeFile(
+                      p.target(n + ".air"),
+                      g
                     )
-                    .toList
-                    .mkString("\n\n\n")
+                  }
                 )
-              )
 
-        }).foldLeft(
-          EitherT.rightT[F, NonEmptyChain[String]](Chain.empty[String])
-        ) { case (accET, writeET) =>
-          EitherT(for {
-            a <- accET.value
-            w <- writeET.value
-          } yield (a, w) match {
-            case (Left(errs), Left(err)) => Left(errs :+ err)
-            case (Right(res), Right(r)) => Right(res :+ r)
-            case (Left(errs), _) => Left(errs)
-            case (_, Left(err)) => Left(NonEmptyChain.of(err))
-          })
-        }.value
-          .map(Validated.fromEither)
+          }).foldLeft(
+            EitherT.rightT[F, NonEmptyChain[String]](Chain.empty[String])
+          ) { case (accET, writeET) =>
+            EitherT(for {
+              a <- accET.value
+              w <- writeET.value
+            } yield (a, w) match {
+              case (Left(errs), Left(err)) => Left(errs :+ err)
+              case (Right(res), Right(r)) => Right(res :+ r)
+              case (Left(errs), _) => Left(errs)
+              case (_, Left(err)) => Left(NonEmptyChain.of(err))
+            })
+          }.value
+            .map(Validated.fromEither)
 
-    }
+      }
 
   def writeFile[F[_]: Files: Concurrent](file: Path, content: String): EitherT[F, String, String] =
     EitherT.right[String](Files[F].deleteIfExists(file)) >>

cli/src/main/scala/aqua/Test.scala

@@ -1,5 +1,6 @@
 package aqua
 
+import aqua.model.transform.BodyConfig
 import cats.effect.{IO, IOApp}
 import cats.data.Validated
 
@@ -13,7 +14,8 @@ object Test extends IOApp.Simple {
         Paths.get("./aqua-src"),
         LazyList(Paths.get("./aqua")),
         Paths.get("./target"),
-        AquaCompiler.TypescriptTarget
+        AquaCompiler.AirTarget,
+        BodyConfig()
       )
       .map {
         case Validated.Invalid(errs) =>

model/src/main/scala/aqua/model/func/body/FuncOps.scala

@@ -3,6 +3,7 @@ package aqua.model.func.body
 import aqua.model.{LiteralModel, ValueModel}
 import aqua.model.func.Call
 import cats.data.Chain
+import cats.free.Cofree
 
 object FuncOps {
 
@@ -18,14 +19,30 @@ object FuncOps {
       )
     )
 
+  def callArrow(funcName: String, call: Call): FuncOp =
+    FuncOp.leaf(
+      CallArrowTag(
+        funcName,
+        call
+      )
+    )
+
   def onVia(on: ValueModel, via: Chain[ValueModel], wrap: FuncOp): FuncOp =
     FuncOp.wrap(
       OnTag(on, via),
       wrap
     )
 
-  def seq(op1: FuncOp, ops: FuncOp*): FuncOp =
-    FuncOp.node(SeqTag, Chain.fromSeq(op1 +: ops))
+  def seq(ops: FuncOp*): FuncOp =
+    FuncOp.node(
+      SeqTag,
+      Chain
+        .fromSeq(ops.flatMap {
+          case FuncOp(Cofree(SeqTag, subOps)) => subOps.value.toList
+          case FuncOp(cof) => cof :: Nil
+        })
+        .map(FuncOp(_))
+    )
 
   def xor(left: FuncOp, right: FuncOp): FuncOp =
     FuncOp.node(XorTag, Chain(left, right))

model/src/main/scala/aqua/model/transform/ArgsProvider.scala

@@ -0,0 +1,25 @@
+package aqua.model.transform
+
+import aqua.model.ValueModel
+import aqua.model.func.Call
+import aqua.model.func.body.{FuncOp, FuncOps}
+
+trait ArgsProvider {
+  def transform(op: FuncOp): FuncOp
+}
+
+case class ArgsFromService(dataServiceId: ValueModel, names: Seq[String]) extends ArgsProvider {
+
+  def getDataOp(name: String): FuncOp =
+    FuncOps.callService(
+      dataServiceId,
+      name,
+      Call(Nil, Some(name))
+    )
+
+  def transform(op: FuncOp): FuncOp =
+    FuncOps.seq(
+      names.map(getDataOp) :+ op: _*
+    )
+
+}

model/src/main/scala/aqua/model/transform/ErrorsCatcher.scala

@@ -0,0 +1,37 @@
+package aqua.model.transform
+
+import aqua.model.{LiteralModel, ValueModel}
+import aqua.model.func.Call
+import aqua.model.func.body.{FuncOp, FuncOps}
+import aqua.types.ScalarType.string
+
+case class ErrorsCatcher(
+  enabled: Boolean,
+  serviceId: ValueModel,
+  funcName: String,
+  callable: InitPeerCallable
+) {
+
+  def transform(op: FuncOp): FuncOp =
+    if (enabled)
+      FuncOps.xor(
+        op,
+        callable.makeCall(
+          serviceId,
+          funcName,
+          ErrorsCatcher.lastErrorCall
+        )
+      )
+    else op
+
+}
+
+object ErrorsCatcher {
+  // TODO not a string
+  val lastErrorArg: Call.Arg = Call.Arg(LiteralModel("%last_error%"), string)
+
+  val lastErrorCall: Call = Call(
+    lastErrorArg :: Nil,
+    None
+  )
+}

model/src/main/scala/aqua/model/transform/ForClient.scala

@@ -1,131 +0,0 @@
-package aqua.model.transform
-
-import aqua.model.func.body._
-import aqua.model.func.{ArgDef, ArgsCall, ArgsDef, Call, FuncCallable}
-import aqua.model.{LiteralModel, VarModel}
-import aqua.types.ScalarType.string
-import aqua.types.ArrowType
-import cats.data.Chain
-import cats.free.Cofree
-
-object ForClient {
-  // TODO not a string
-  private val lastErrorArg = Call.Arg(LiteralModel("%last_error%"), string)
-
-  // Get to init user through a relay
-  def viaRelay(op: FuncOp)(implicit conf: BodyConfig): FuncOp =
-    FuncOps.onVia(LiteralModel.initPeerId, Chain.one(VarModel(conf.relayVarName)), op)
-
-  def wrapXor(op: FuncOp)(implicit conf: BodyConfig): FuncOp =
-    if (conf.wrapWithXor)
-      FuncOp.node(
-        XorTag,
-        Chain(
-          op,
-          viaRelay(
-            FuncOps.callService(
-              conf.errorHandlingCallback,
-              conf.errorFuncName,
-              Call(
-                lastErrorArg :: Nil,
-                None
-              )
-            )
-          )
-        )
-      )
-    else op
-
-  def returnCallback(func: FuncCallable)(implicit conf: BodyConfig): Option[FuncOp] = func.ret.map {
-    retArg =>
-      viaRelay(
-        FuncOps.callService(
-          conf.callbackSrvId,
-          conf.respFuncName,
-          Call(
-            retArg :: Nil,
-            None
-          )
-        )
-      )
-  }
-
-  def initPeerCallable(name: String, arrowType: ArrowType)(implicit
-    conf: BodyConfig
-  ): FuncCallable = {
-    val (args, call, ret) = ArgsCall.arrowToArgsCallRet(arrowType)
-    FuncCallable(
-      s"init_peer_callable_$name",
-      viaRelay(
-        FuncOps.callService(
-          conf.callbackSrvId,
-          name,
-          call
-        )
-      ),
-      args,
-      ret,
-      Map.empty
-    )
-  }
-
-  // Get data with this name from a local service
-  def getDataOp(name: String)(implicit conf: BodyConfig): FuncOp =
-    FuncOps.callService(
-      conf.dataSrvId,
-      name,
-      Call(Nil, Some(name))
-    )
-
-  def resolve(func: FuncCallable, conf: BodyConfig): Cofree[Chain, OpTag] = {
-    implicit val c: BodyConfig = conf
-
-    // Like it is called from TS
-    def funcArgsCall: Call =
-      Call(
-        func.args.toCallArgs,
-        None
-      )
-
-    val funcAround: FuncCallable = FuncCallable(
-      "funcAround",
-      wrapXor(
-        viaRelay(
-          FuncOp
-            .node(
-              SeqTag,
-              (
-                func.args.dataArgNames.map(getDataOp) :+ getDataOp(conf.relayVarName)
-              )
-                .append(
-                  FuncOp.leaf(
-                    CallArrowTag(
-                      func.funcName,
-                      funcArgsCall
-                    )
-                  )
-                ) ++ Chain.fromSeq(returnCallback(func).toSeq)
-            )
-        )
-      ),
-      ArgsDef(ArgDef.Arrow(func.funcName, func.arrowType) :: Nil),
-      None,
-      func.args.arrowArgs.collect { case ArgDef.Arrow(argName, arrowType) =>
-        argName -> initPeerCallable(argName, arrowType)
-      }.toList.toMap
-    )
-
-    val body =
-      funcAround
-        .resolve(
-          Call(Call.Arg(VarModel("_func"), func.arrowType) :: Nil, None),
-          Map("_func" -> func),
-          Set.empty
-        )
-        .value
-        ._1
-        .tree
-
-    Topology.resolve(body)
-  }
-}

model/src/main/scala/aqua/model/transform/InitPeerCallable.scala

@@ -0,0 +1,23 @@
+package aqua.model.transform
+
+import aqua.model.{LiteralModel, ValueModel}
+import aqua.model.func.Call
+import aqua.model.func.body.{FuncOp, FuncOps}
+import cats.data.Chain
+
+sealed trait InitPeerCallable {
+  def transform(op: FuncOp): FuncOp
+
+  def makeCall(serviceId: ValueModel, funcName: String, call: Call): FuncOp =
+    transform(FuncOps.callService(serviceId, funcName, call))
+
+  def service(serviceId: ValueModel): (String, Call) => FuncOp = makeCall(serviceId, _, _)
+}
+
+case class InitViaRelayCallable(goThrough: Chain[ValueModel]) extends InitPeerCallable {
+
+  // Get to init user through a relay
+  override def transform(op: FuncOp): FuncOp =
+    FuncOps.onVia(LiteralModel.initPeerId, goThrough, op)
+
+}

model/src/main/scala/aqua/model/transform/ResolveFunc.scala

@@ -0,0 +1,69 @@
+package aqua.model.transform
+
+import aqua.model.VarModel
+import aqua.model.func.{ArgDef, ArgsCall, ArgsDef, Call, FuncCallable}
+import aqua.model.func.body.{FuncOp, FuncOps}
+import aqua.types.ArrowType
+import cats.Eval
+
+case class ResolveFunc(
+  transform: FuncOp => FuncOp,
+  callback: (String, Call) => FuncOp,
+  respFuncName: String,
+  wrapCallableName: String = "funcAround",
+  arrowCallbackPrefix: String = "init_peer_callable_"
+) {
+
+  def returnCallback(func: FuncCallable): Option[FuncOp] = func.ret.map { retArg =>
+    callback(
+      respFuncName,
+      Call(
+        retArg :: Nil,
+        None
+      )
+    )
+  }
+
+  def arrowToCallback(name: String, arrowType: ArrowType): FuncCallable = {
+    val (args, call, ret) = ArgsCall.arrowToArgsCallRet(arrowType)
+    FuncCallable(
+      arrowCallbackPrefix + name,
+      callback(name, call),
+      args,
+      ret,
+      Map.empty
+    )
+  }
+
+  def wrap(func: FuncCallable): FuncCallable =
+    FuncCallable(
+      wrapCallableName,
+      transform(
+        FuncOps.seq(
+          FuncOps
+            .callArrow(
+              func.funcName,
+              Call(
+                func.args.toCallArgs,
+                None
+              )
+            ) ::
+            returnCallback(func).toList: _*
+        )
+      ),
+      ArgsDef(ArgDef.Arrow(func.funcName, func.arrowType) :: Nil),
+      None,
+      func.args.arrowArgs.collect { case ArgDef.Arrow(argName, arrowType) =>
+        argName -> arrowToCallback(argName, arrowType)
+      }.toList.toMap
+    )
+
+  def resolve(func: FuncCallable, funcArgName: String = "_func"): Eval[FuncOp] =
+    wrap(func)
+      .resolve(
+        Call(Call.Arg(VarModel(funcArgName), func.arrowType) :: Nil, None),
+        Map(funcArgName -> func),
+        Set.empty
+      )
+      .map(_._1)
+}

model/src/main/scala/aqua/model/transform/Transform.scala

@@ -0,0 +1,37 @@
+package aqua.model.transform
+
+import aqua.model.func.body._
+import aqua.model.func.FuncCallable
+import aqua.model.VarModel
+import cats.data.Chain
+import cats.free.Cofree
+
+object Transform {
+
+  def forClient(func: FuncCallable, conf: BodyConfig): Cofree[Chain, OpTag] = {
+    val initCallable: InitPeerCallable = InitViaRelayCallable(
+      Chain.one(VarModel(conf.relayVarName))
+    )
+    val errorsCatcher = ErrorsCatcher(
+      enabled = conf.wrapWithXor,
+      conf.errorHandlingCallback,
+      conf.errorFuncName,
+      initCallable
+    )
+    val argsProvider: ArgsProvider =
+      ArgsFromService(conf.dataSrvId, conf.relayVarName +: func.args.dataArgNames.toList)
+
+    val transform =
+      errorsCatcher.transform _ compose initCallable.transform compose argsProvider.transform
+
+    val callback = initCallable.service(conf.callbackSrvId)
+
+    val wrapFunc = ResolveFunc(
+      transform,
+      callback,
+      conf.respFuncName
+    )
+
+    Topology.resolve(wrapFunc.resolve(func).value.tree)
+  }
+}

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

@@ -147,136 +147,4 @@ class TopologySpec extends AnyFlatSpec with Matchers {
     proc.equalsOrPrintDiff(expected) should be(true)
   }
 
-  "topology resolver" should "work well with function 1 (no calls before on)" in {
-
-    val ret = LiteralModel("\"return this\"")
-
-    val func: FuncCallable =
-      FuncCallable(
-        "ret",
-        FuncOp(on(otherPeer, Nil, call(1))),
-        ArgsDef.empty,
-        Some(Call.Arg(ret, ScalarType.string)),
-        Map.empty
-      )
-
-    val bc = BodyConfig()
-
-    val fc = ForClient.resolve(func, bc)
-
-    val procFC: Node = fc
-
-    val expectedFC =
-      xor(
-        on(
-          initPeer,
-          relayV :: Nil,
-          seq(
-            dataCall(bc, "relay", initPeer),
-            on(otherPeer, Nil, through(relayV), call(1, otherPeer)),
-            through(relayV),
-            on(initPeer, relayV :: Nil, respCall(bc, ret, initPeer))
-          )
-        ),
-        on(initPeer, relayV :: Nil, xorErrorCall(bc, initPeer))
-      )
-
-    procFC.equalsOrPrintDiff(expectedFC) should be(true)
-
-  }
-
-  "topology resolver" should "work well with function 2 (with a call before on)" in {
-
-    val ret = LiteralModel("\"return this\"")
-
-    val func: FuncCallable = FuncCallable(
-      "ret",
-      FuncOp(seq(call(0), on(otherPeer, Nil, call(1)))),
-      ArgsDef.empty,
-      Some(Call.Arg(ret, ScalarType.string)),
-      Map.empty
-    )
-
-    val bc = BodyConfig()
-
-    val fc = ForClient.resolve(func, bc)
-
-    val procFC: Node = fc
-
-    val expectedFC =
-      xor(
-        on(
-          initPeer,
-          relayV :: Nil,
-          seq(
-            dataCall(bc, "relay", initPeer),
-            seq(
-              call(0, initPeer),
-              on(otherPeer, Nil, through(relayV), call(1, otherPeer))
-            ),
-            through(relayV),
-            on(initPeer, relayV :: Nil, respCall(bc, ret, initPeer))
-          )
-        ),
-        on(initPeer, relayV :: Nil, xorErrorCall(bc, initPeer))
-      )
-
-    procFC.equalsOrPrintDiff(expectedFC) should be(true)
-
-  }
-
-  "topology resolver" should "link funcs correctly" in {
-    /*
-    func one() -> u64:
-      variable <- Demo.get42()
-      <- variable
-
-    func two() -> u64:
-      variable <- one()
-      <- variable
-     */
-
-    val f1: FuncCallable =
-      FuncCallable(
-        "f1",
-        FuncOp(Node(CallServiceTag(LiteralModel("\"srv1\""), "foo", Call(Nil, Some("v")), None))),
-        ArgsDef.empty,
-        Some(Call.Arg(VarModel("v"), ScalarType.string)),
-        Map.empty
-      )
-
-    val f2: FuncCallable =
-      FuncCallable(
-        "f2",
-        FuncOp(
-          Node(CallArrowTag("callable", Call(Nil, Some("v"))))
-        ),
-        ArgsDef.empty,
-        Some(Call.Arg(VarModel("v"), ScalarType.string)),
-        Map("callable" -> f1)
-      )
-
-    val bc = BodyConfig(wrapWithXor = false)
-
-    val res = ForClient.resolve(f2, bc): Node
-
-    res.equalsOrPrintDiff(
-      on(
-        initPeer,
-        relayV :: Nil,
-        seq(
-          dataCall(bc, "relay", initPeer),
-          Node(
-            CallServiceTag(LiteralModel("\"srv1\""), "foo", Call(Nil, Some("v")), Some(initPeer))
-          ),
-          on(
-            initPeer,
-            relayV :: Nil,
-            respCall(bc, VarModel("v"), initPeer)
-          )
-        )
-      )
-    ) should be(true)
-  }
-
 }

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

@@ -0,0 +1,144 @@
+package aqua.model.transform
+
+import aqua.model.{LiteralModel, Node, VarModel}
+import aqua.model.func.{ArgsDef, Call, FuncCallable}
+import aqua.model.func.body.{CallArrowTag, CallServiceTag, FuncOp}
+import aqua.types.ScalarType
+import org.scalatest.flatspec.AnyFlatSpec
+import org.scalatest.matchers.should.Matchers
+
+class TransformSpec extends AnyFlatSpec with Matchers {
+  import Node._
+
+  "transform.forClient" should "work well with function 1 (no calls before on)" in {
+
+    val ret = LiteralModel("\"return this\"")
+
+    val func: FuncCallable =
+      FuncCallable(
+        "ret",
+        FuncOp(on(otherPeer, Nil, call(1))),
+        ArgsDef.empty,
+        Some(Call.Arg(ret, ScalarType.string)),
+        Map.empty
+      )
+
+    val bc = BodyConfig()
+
+    val fc = Transform.forClient(func, bc)
+
+    val procFC: Node = fc
+
+    val expectedFC =
+      xor(
+        on(
+          initPeer,
+          relayV :: Nil,
+          seq(
+            dataCall(bc, "relay", initPeer),
+            on(otherPeer, Nil, through(relayV), call(1, otherPeer)),
+            through(relayV),
+            on(initPeer, relayV :: Nil, respCall(bc, ret, initPeer))
+          )
+        ),
+        on(initPeer, relayV :: Nil, xorErrorCall(bc, initPeer))
+      )
+
+    procFC.equalsOrPrintDiff(expectedFC) should be(true)
+
+  }
+
+  "transform.forClient" should "work well with function 2 (with a call before on)" in {
+
+    val ret = LiteralModel("\"return this\"")
+
+    val func: FuncCallable = FuncCallable(
+      "ret",
+      FuncOp(seq(call(0), on(otherPeer, Nil, call(1)))),
+      ArgsDef.empty,
+      Some(Call.Arg(ret, ScalarType.string)),
+      Map.empty
+    )
+
+    val bc = BodyConfig()
+
+    val fc = Transform.forClient(func, bc)
+
+    val procFC: Node = fc
+
+    val expectedFC =
+      xor(
+        on(
+          initPeer,
+          relayV :: Nil,
+          seq(
+            dataCall(bc, "relay", initPeer),
+            seq(
+              call(0, initPeer),
+              on(otherPeer, Nil, through(relayV), call(1, otherPeer))
+            ),
+            through(relayV),
+            on(initPeer, relayV :: Nil, respCall(bc, ret, initPeer))
+          )
+        ),
+        on(initPeer, relayV :: Nil, xorErrorCall(bc, initPeer))
+      )
+
+    procFC.equalsOrPrintDiff(expectedFC) should be(true)
+
+  }
+
+  "transform.forClient" should "link funcs correctly" in {
+    /*
+    func one() -> u64:
+      variable <- Demo.get42()
+      <- variable
+
+    func two() -> u64:
+      variable <- one()
+      <- variable
+     */
+
+    val f1: FuncCallable =
+      FuncCallable(
+        "f1",
+        FuncOp(Node(CallServiceTag(LiteralModel("\"srv1\""), "foo", Call(Nil, Some("v")), None))),
+        ArgsDef.empty,
+        Some(Call.Arg(VarModel("v"), ScalarType.string)),
+        Map.empty
+      )
+
+    val f2: FuncCallable =
+      FuncCallable(
+        "f2",
+        FuncOp(
+          Node(CallArrowTag("callable", Call(Nil, Some("v"))))
+        ),
+        ArgsDef.empty,
+        Some(Call.Arg(VarModel("v"), ScalarType.string)),
+        Map("callable" -> f1)
+      )
+
+    val bc = BodyConfig(wrapWithXor = false)
+
+    val res = Transform.forClient(f2, bc): Node
+
+    res.equalsOrPrintDiff(
+      on(
+        initPeer,
+        relayV :: Nil,
+        seq(
+          dataCall(bc, "relay", initPeer),
+          Node(
+            CallServiceTag(LiteralModel("\"srv1\""), "foo", Call(Nil, Some("v")), Some(initPeer))
+          ),
+          on(
+            initPeer,
+            relayV :: Nil,
+            respCall(bc, VarModel("v"), initPeer)
+          )
+        )
+      )
+    ) should be(true)
+  }
+}