fix(compiler): Generate stream restriction for scoped exprs [fixes LNG-222] (#841)

* Add show for AST

* Update ForSem

* Fix if and try

* Fix else, otherwise, catch, add tests

* Add integration tests

build.sbt

@@ -182,7 +182,7 @@ lazy val parser = crossProject(JVMPlatform, JSPlatform)
       "org.typelevel" %%% "cats-free"  % catsV
     )
   )
-  .dependsOn(types)
+  .dependsOn(types, helpers)
 
 lazy val linker = crossProject(JVMPlatform, JSPlatform)
   .withoutSuffixFor(JVMPlatform)
@@ -200,6 +200,7 @@ lazy val tree = crossProject(JVMPlatform, JSPlatform)
       "org.typelevel" %%% "cats-free" % catsV
     )
   )
+  .dependsOn(helpers)
 
 lazy val raw = crossProject(JVMPlatform, JSPlatform)
   .withoutSuffixFor(JVMPlatform)
@@ -212,7 +213,7 @@ lazy val model = crossProject(JVMPlatform, JSPlatform)
   .withoutSuffixFor(JVMPlatform)
   .crossType(CrossType.Pure)
   .settings(commons)
-  .dependsOn(types, tree, raw)
+  .dependsOn(types, tree, raw, helpers)
 
 lazy val res = crossProject(JVMPlatform, JSPlatform)
   .withoutSuffixFor(JVMPlatform)
@@ -228,7 +229,6 @@ lazy val inline = crossProject(JVMPlatform, JSPlatform)
   .settings(commons)
   .dependsOn(raw, model)
 
-
 lazy val transform = crossProject(JVMPlatform, JSPlatform)
   .withoutSuffixFor(JVMPlatform)
   .crossType(CrossType.Pure)
@@ -304,6 +304,18 @@ lazy val constants = crossProject(JVMPlatform, JSPlatform)
   )
   .dependsOn(parser, raw)
 
+lazy val helpers = crossProject(JVMPlatform, JSPlatform)
+  .withoutSuffixFor(JVMPlatform)
+  .crossType(CrossType.Pure)
+  .in(file("utils/helpers"))
+  .settings(commons)
+  .settings(
+    libraryDependencies ++= Seq(
+      "org.typelevel" %%% "cats-core" % catsV,
+      "org.typelevel" %%% "cats-free" % catsV
+    )
+  )
+
 lazy val `backend-air` = crossProject(JVMPlatform, JSPlatform)
   .withoutSuffixFor(JVMPlatform)
   .crossType(CrossType.Pure)

integration-tests/aqua/examples/streamScopes.aqua

@@ -0,0 +1,83 @@
+aqua StreamExports
+
+export FailureSrv, streamIf, streamTry, streamFor, streamComplex
+
+service FailureSrv("failure"):
+    fail(msg: string)
+
+func streamIf() -> i8:
+    on HOST_PEER_ID:
+        if true:
+            stream: *i8
+            stream <<- 1
+        else:
+            stream: *i8
+            stream <<- 2
+        
+        if false:
+            stream: *i8
+            stream <<- 3
+        else:
+            stream: *i8
+            stream <<- 4
+
+        stream: *i8
+        stream <<- 5
+
+    <- stream!
+
+func streamTry() -> i8:
+    on HOST_PEER_ID:
+        try:
+            stream: *i8
+            stream <<- 1
+            FailureSrv.fail("try")
+        catch e:
+            stream: *i8
+            stream <<- 2
+            FailureSrv.fail("catch")
+        otherwise:
+            stream: *i8
+            stream <<- 3
+
+        stream: *i8
+        stream <<- 4
+
+    <- stream!
+
+func streamFor() -> i8:
+    on HOST_PEER_ID:
+        for i <- [1, 2, 3]:
+            stream: *i8
+            stream <<- i
+        
+        stream: *i8
+        stream <<- 4
+
+    <- stream!
+
+func streamComplex() -> i8:
+    on HOST_PEER_ID:
+        for i <- [1, 2, 3]:
+            try:
+                if i == 2:
+                    stream: *i8
+                    stream <<- i
+                    FailureSrv.fail("if")
+                else:
+                    stream: *i8
+                    stream <<- i
+                
+                stream: *i8
+                stream <<- i + 3
+            catch e:
+                stream: *i8
+                stream <<- i + 6
+            
+            stream: *i8
+            stream <<- i + 9
+
+        stream: *i8
+        stream <<- 13
+
+    <- stream!

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

@@ -39,6 +39,7 @@ import { coCall } from '../examples/coCall.js';
 import { bugLNG60Call, passArgsCall } from '../examples/passArgsCall.js';
 import { streamArgsCall } from '../examples/streamArgsCall.js';
 import { streamResultsCall } from '../examples/streamResultsCall.js';
+import { streamIfCall, streamForCall, streamTryCall, streamComplexCall } from '../examples/streamScopes.js';
 import { pushToStreamCall } from '../examples/pushToStreamCall.js';
 import { literalCall } from '../examples/returnLiteralCall.js';
 import { multiReturnCall } from '../examples/multiReturnCall.js';
@@ -158,6 +159,30 @@ describe('Testing examples', () => {
         expect(streamResResult).toEqual([[], ['a', 'b', 'c']]);
     });
 
+    it('streamScopes.aqua streamIf', async () => {
+        let streamIfResult = await streamIfCall();
+
+        expect(streamIfResult).toEqual(5);
+    });
+
+    it('streamScopes.aqua streamTry', async () => {
+        let streamTryResult = await streamTryCall();
+
+        expect(streamTryResult).toEqual(4);
+    });
+
+    it('streamScopes.aqua streamFor', async () => {
+        let streamTryResult = await streamForCall();
+
+        expect(streamTryResult).toEqual(4);
+    });
+
+    it('streamScopes.aqua streamComplex', async () => {
+        let streamTryResult = await streamComplexCall();
+
+        expect(streamTryResult).toEqual(13);
+    });
+
     it('if.aqua', async () => {
         await ifCall();
     });

integration-tests/src/examples/streamScopes.ts

@@ -0,0 +1,35 @@
+import {
+    streamIf,
+    streamTry,
+    streamFor,
+    streamComplex,
+    registerFailureSrv,
+} from '../compiled/examples/streamScopes.js';
+
+export async function streamIfCall() {
+    return await streamIf();
+}
+
+export async function streamTryCall() {
+    registerFailureSrv({
+        fail: (msg) => {
+            return Promise.reject(msg);
+        },
+    });
+
+    return await streamTry();
+}
+
+export async function streamForCall() {
+    return await streamFor();
+}
+
+export async function streamComplexCall() {
+    registerFailureSrv({
+        fail: (msg) => {
+            return Promise.reject(msg);
+        },
+    });
+
+    return await streamComplex();
+}

model/tree/src/main/scala/aqua/tree/TreeNodeCompanion.scala

@@ -9,12 +9,15 @@ import cats.syntax.apply.*
 
 import scala.annotation.tailrec
 
+import aqua.helpers.Tree
+
 trait TreeNodeCompanion[T <: TreeNode[T]] {
 
   given showTreeLabel: Show[T]
 
   type Tree = Cofree[Chain, T]
 
+  // TODO: Use helpers.Tree istead of this function
   private def showOffset(what: Tree, offset: Int): String = {
     val spaces = "| " * offset
     spaces + what.head.show + what.tail.map {
@@ -98,8 +101,7 @@ trait TreeNodeCompanion[T <: TreeNode[T]] {
 
   given Show[Tree] with
 
-    override def show(t: Tree): String =
+    override def show(t: Tree): String = Tree.show(t)
-      showOffset(t, 0)
 
   given Show[(Tree, Tree)] with
 

parser/src/main/scala/aqua/parser/Ast.scala

@@ -4,10 +4,13 @@ import aqua.parser.expr.*
 import aqua.parser.head.{HeadExpr, HeaderExpr}
 import aqua.parser.lift.{LiftParser, Span}
 import aqua.parser.lift.LiftParser.*
+import aqua.helpers.Tree
+
 import cats.data.{Chain, Validated, ValidatedNec}
 import cats.free.Cofree
 import cats.{Comonad, Eval}
 import cats.~>
+import cats.Show
 
 case class Ast[S[_]](head: Ast.Head[S], tree: Ast.Tree[S]) {
 
@@ -19,6 +22,16 @@ case class Ast[S[_]](head: Ast.Head[S], tree: Ast.Tree[S]) {
 }
 
 object Ast {
-  type Tree[S[_]] = Cofree[Chain, Expr[S]]
   type Head[S[_]] = Cofree[Chain, HeaderExpr[S]]
+  type Tree[S[_]] = Cofree[Chain, Expr[S]]
+
+  given [S[_]]: Show[Ast[S]] with {
+
+    def show(ast: Ast[S]): String = {
+      val head = Tree.show(ast.head)
+      val body = Tree.show(ast.tree)
+
+      s"$head\n$body"
+    }
+  }
 }

parser/src/main/scala/aqua/parser/Expr.scala

@@ -7,16 +7,18 @@ import aqua.parser.expr.func.ReturnExpr
 import aqua.parser.lift.LiftParser.*
 import aqua.parser.lift.Span.{P0ToSpan, PToSpan}
 import aqua.parser.lift.{LiftParser, Span}
+import aqua.parser.Ast.Tree
+import aqua.parser.ListToTreeConverter
+
 import cats.data.Chain.:==
 import cats.data.{Chain, NonEmptyChain, Validated, ValidatedNec}
 import cats.free.Cofree
+import cats.Show
 import cats.data.Validated.{invalid, invalidNec, invalidNel, valid, validNec, validNel}
 import cats.parse.{Parser as P, Parser0 as P0}
 import cats.syntax.comonad.*
 import cats.{~>, Comonad, Eval}
 import scribe.Logging
-import aqua.parser.Ast.Tree
-import aqua.parser.ListToTreeConverter
 
 abstract class Expr[F[_]](val companion: Expr.Companion, val token: Token[F]) {
 
@@ -109,4 +111,9 @@ object Expr {
           .result
       }
   }
+
+  given [S[_]]: Show[Expr[S]] with {
+    // TODO: Make it better
+    def show(e: Expr[S]): String = e.toString
+  }
 }

parser/src/main/scala/aqua/parser/head/HeaderExpr.scala

@@ -3,17 +3,19 @@ package aqua.parser.head
 import aqua.parser.Ast
 import aqua.parser.lexer.Token
 import aqua.parser.lift.LiftParser
-import cats.{Comonad, Eval}
-import cats.data.Chain
-import cats.free.Cofree
-import cats.parse.Parser as P
-import cats.~>
 import aqua.parser.lift.Span
 import aqua.parser.lift.Span.{P0ToSpan, PToSpan}
 
+import cats.{Comonad, Eval}
+import cats.data.Chain
+import cats.free.Cofree
+import cats.Show
+import cats.parse.Parser as P
+import cats.~>
+
 trait HeaderExpr[S[_]] {
   def token: Token[S]
-  
+
   def mapK[K[_]: Comonad](fk: S ~> K): HeaderExpr[K]
 }
 
@@ -30,4 +32,9 @@ object HeaderExpr {
     override def ast: P[Ast.Head[Span.S]] =
       p.map(Cofree[Chain, HeaderExpr[Span.S]](_, Eval.now(Chain.empty)))
   }
+
+  given [S[_]]: Show[HeaderExpr[S]] with {
+    // TODO: Make it better
+    def show(e: HeaderExpr[S]): String = e.toString
+  }
 }

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

@@ -22,26 +22,26 @@ class CatchSem[S[_]](val expr: CatchExpr[S]) extends AnyVal {
   ): Prog[Alg, Raw] =
     Prog
       .around(
-        N.beginScope(expr.name) >>
+        N.define(expr.name, ValueRaw.lastError.baseType),
-          L.beginScope() >>
-          N.define(expr.name, ValueRaw.lastError.baseType),
         (_, g: Raw) =>
-          N.endScope() >> L.endScope() as (
+          g match {
-            g match {
+            case FuncOp(op) =>
-              case FuncOp(op) =>
+              for {
-                TryTag.Catch
+                restricted <- FuncOpSem.restrictStreamsInScope(op)
+                tag = TryTag.Catch
                   .wrap(
                     SeqTag.wrap(
                       AssignmentTag(ValueRaw.lastError, expr.name.value).leaf,
-                      op
+                      restricted
                     )
                   )
-                  .toFuncOp
+              } yield tag.toFuncOp
-              case _ =>
+            case _ =>
-                Raw.error("Wrong body of the `catch` expression")
+              Raw.error("Wrong body of the `catch` expression").pure
-            }
+          }
-          )
       )
       .abilitiesScope[S](expr.token)
+      .namesScope(expr.token)
+      .locationsScope()
 
 }

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

@@ -9,6 +9,7 @@ import aqua.semantics.rules.locations.LocationsAlgebra
 import aqua.semantics.rules.names.NamesAlgebra
 
 import cats.syntax.applicative.*
+import cats.syntax.functor.*
 import cats.Monad
 
 class ElseOtherwiseSem[S[_]](val expr: ElseOtherwiseExpr[S]) extends AnyVal {
@@ -19,20 +20,23 @@ class ElseOtherwiseSem[S[_]](val expr: ElseOtherwiseExpr[S]) extends AnyVal {
     L: LocationsAlgebra[S, Alg]
   ): Prog[Alg, Raw] =
     Prog
-      .after[Alg, Raw] {
+      .after((ops: Raw) =>
-        case FuncOp(op) =>
+        ops match {
-          expr.kind
+          case FuncOp(op) =>
-            .fold(
+            for {
-              ifElse = IfTag.Else,
+              restricted <- FuncOpSem.restrictStreamsInScope(op)
-              ifOtherwise = TryTag.Otherwise
+              tag = expr.kind
-            )
+                .fold(
-            .wrap(op)
+                  ifElse = IfTag.Else,
-            .toFuncOp
+                  ifOtherwise = TryTag.Otherwise
-            .pure
+                )
-        case _ =>
+                .wrap(restricted)
-          val name = expr.kind.fold("`else`", "`otherwise`")
+            } yield tag.toFuncOp
-          Raw.error(s"Wrong body of the $name expression").pure
+          case _ =>
-      }
+            val name = expr.kind.fold("`else`", "`otherwise`")
+            Raw.error(s"Wrong body of the $name expression").pure
+        }
+      )
       .abilitiesScope(expr.token)
       .namesScope(expr.token)
       .locationsScope()

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

@@ -11,6 +11,7 @@ import aqua.semantics.rules.abilities.AbilitiesAlgebra
 import aqua.semantics.rules.names.NamesAlgebra
 import aqua.semantics.rules.types.TypesAlgebra
 import aqua.types.{ArrayType, BoxType, StreamType}
+import aqua.semantics.expr.func.FuncOpSem
 
 import cats.Monad
 import cats.data.Chain
@@ -18,6 +19,7 @@ import cats.syntax.applicative.*
 import cats.syntax.apply.*
 import cats.syntax.flatMap.*
 import cats.syntax.functor.*
+import cats.syntax.option.*
 
 class ForSem[S[_]](val expr: ForExpr[S]) extends AnyVal {
 
@@ -29,49 +31,44 @@ class ForSem[S[_]](val expr: ForExpr[S]) extends AnyVal {
   ): Prog[F, Raw] =
     Prog
       .around(
-        V.valueToRaw(expr.iterable).flatMap[Option[ValueRaw]] {
+        V.valueToRaw(expr.iterable).flatMap {
           case Some(vm) =>
             vm.`type` match {
               case t: BoxType =>
-                N.define(expr.item, t.element).as(Option(vm))
+                N.define(expr.item, t.element).as(vm.some)
               case dt =>
-                T.ensureTypeMatches(expr.iterable, ArrayType(dt), dt).as(Option.empty[ValueRaw])
+                T.ensureTypeMatches(expr.iterable, ArrayType(dt), dt).as(none)
             }
 
-          case _ => None.pure[F]
+          case _ => none.pure
         },
-        (stOpt: Option[ValueRaw], ops: Raw) =>
+        // Without type of ops specified
-          N.streamsDefinedWithinScope()
+        // scala compiler fails to compile this
-            .map(streams =>
+        (iterable, ops: Raw) =>
-              (stOpt, ops) match {
+          (iterable, ops) match {
-                case (Some(vm), FuncOp(op)) =>
+            case (Some(vm), FuncOp(op)) =>
-                  val innerTag = expr.mode.fold(SeqTag) {
+              FuncOpSem.restrictStreamsInScope(op).map { restricted =>
-                    case ForExpr.Mode.ParMode => ParTag
+                val innerTag = expr.mode.fold(SeqTag) {
-                    case ForExpr.Mode.TryMode => TryTag
+                  case ForExpr.Mode.ParMode => ParTag
-                  }
+                  case ForExpr.Mode.TryMode => TryTag
+                }
 
-                  val mode = expr.mode.collect { case ForExpr.Mode.ParMode => WaitMode }
+                val mode = expr.mode.collect { case ForExpr.Mode.ParMode => WaitMode }
 
-                  val forTag =
+                val forTag = ForTag(expr.item.value, vm, mode).wrap(
-                    ForTag(expr.item.value, vm, mode).wrap(
+                  innerTag.wrap(
-                      innerTag
+                    restricted,
-                        .wrap(
+                    NextTag(expr.item.value).leaf
-                          // Restrict the streams created within this scope
+                  )
-                          streams.toList.foldLeft(op) { case (tree, (streamName, streamType)) =>
+                )
-                            RestrictionTag(streamName, streamType).wrap(tree)
-                          },
-                          NextTag(expr.item.value).leaf
-                        )
-                    )
 
-                  // Fix: continue execution after fold par immediately, without finding a path out from par branches
+                // Fix: continue execution after fold par immediately, without finding a path out from par branches
-                  if (innerTag == ParTag) ParTag.Detach.wrap(forTag).toFuncOp
+                if (innerTag == ParTag) ParTag.Detach.wrap(forTag).toFuncOp
-                  else forTag.toFuncOp
+                else forTag.toFuncOp
-                case _ =>
-                  Raw.error("Wrong body of the `for` expression")
               }
-            )
+            case _ => Raw.error("Wrong body of the `for` expression").pure[F]
+          }
       )
-      .namesScope[S](expr.token)
+      .namesScope(expr.token)
-      .abilitiesScope[S](expr.token)
+      .abilitiesScope(expr.token)
 }

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

@@ -0,0 +1,22 @@
+package aqua.semantics.expr.func
+
+import cats.Monad
+import cats.syntax.functor.*
+
+import aqua.semantics.rules.names.NamesAlgebra
+import aqua.raw.Raw
+import aqua.raw.ops.{RawTag, RestrictionTag}
+
+object FuncOpSem {
+
+  def restrictStreamsInScope[S[_], Alg[_]: Monad](tree: RawTag.Tree)(using
+    N: NamesAlgebra[S, Alg]
+  ): Alg[RawTag.Tree] = N
+    .streamsDefinedWithinScope()
+    .map(streams =>
+      streams.toList
+        .foldLeft(tree) { case (tree, (streamName, streamType)) =>
+          RestrictionTag(streamName, streamType).wrap(tree)
+        }
+    )
+}

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

@@ -41,17 +41,18 @@ class IfSem[S[_]](val expr: IfExpr[S]) extends AnyVal {
               ).map(Option.when(_)(raw))
             )
           ),
-        (value: Option[ValueRaw], ops: Raw) =>
+        // Without type of ops specified
-          value
+        // scala compiler fails to compile this
-            .fold(
+        (value, ops: Raw) =>
-              Raw.error("`if` expression errored in matching types")
+          (value, ops) match {
-            )(raw =>
+            case (Some(vr), FuncOp(op)) =>
-              ops match {
+              for {
-                case FuncOp(op) => IfTag(raw).wrap(op).toFuncOp
+                restricted <- FuncOpSem.restrictStreamsInScope(op)
-                case _ => Raw.error("Wrong body of the `if` expression")
+                tag = IfTag(vr).wrap(restricted)
-              }
+              } yield tag.toFuncOp
-            )
+            case (None, _) => Raw.error("`if` expression errored in matching types").pure
-            .pure
+            case _ => Raw.error("Wrong body of the `if` expression").pure
+          }
       )
       .abilitiesScope[S](expr.token)
       .namesScope[S](expr.token)

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

@@ -14,6 +14,7 @@ import cats.data.Chain
 import cats.syntax.applicative.*
 import cats.syntax.apply.*
 import cats.syntax.flatMap.*
+import cats.syntax.traverse.*
 import cats.syntax.functor.*
 import cats.{Monad, Traverse}
 
@@ -27,8 +28,8 @@ class OnSem[S[_]](val expr: OnExpr[S]) extends AnyVal {
     Prog.around(
       (
         V.ensureIsString(expr.peerId),
-        Traverse[List]
+        expr.via
-          .traverse(expr.via)(v =>
+          .traverse(v =>
             V.valueToRaw(v).flatTap {
               case Some(vm) =>
                 vm.`type` match {

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

@@ -9,7 +9,9 @@ import aqua.semantics.rules.abilities.AbilitiesAlgebra
 import aqua.semantics.rules.locations.LocationsAlgebra
 import aqua.semantics.rules.names.NamesAlgebra
 import aqua.semantics.rules.types.TypesAlgebra
+
 import cats.syntax.applicative.*
+import cats.syntax.functor.*
 import cats.Monad
 
 class TrySem[S[_]](val expr: TryExpr[S]) extends AnyVal {
@@ -20,12 +22,19 @@ class TrySem[S[_]](val expr: TryExpr[S]) extends AnyVal {
     L: LocationsAlgebra[S, Alg]
   ): Prog[Alg, Raw] =
     Prog
-      .after[Alg, Raw] {
+      // Without type of ops specified
-        case FuncOp(o) =>
+      // scala compiler fails to compile this
-          TryTag.wrap(o).toFuncOp.pure[Alg]
+      .after((ops: Raw) =>
-        case _ =>
+        ops match {
-          Raw.error("Wrong body of the `try` expression").pure[Alg]
+          case FuncOp(op) =>
-      }
+            for {
+              restricted <- FuncOpSem.restrictStreamsInScope(op)
+              tag = TryTag.wrap(restricted)
+            } yield tag.toFuncOp
+          case _ =>
+            Raw.error("Wrong body of the `try` expression").pure[Alg]
+        }
+      )
       .abilitiesScope(expr.token)
       .namesScope(expr.token)
       .locationsScope()

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

@@ -5,7 +5,7 @@ import aqua.parser.Ast
 import aqua.raw.ops.{Call, CallArrowRawTag, FuncOp, OnTag, ParTag, RawTag, SeqGroupTag, SeqTag}
 import aqua.parser.Parser
 import aqua.parser.lift.{LiftParser, Span}
-import aqua.raw.value.{ApplyBinaryOpRaw, LiteralRaw, ValueRaw}
+import aqua.raw.value.{ApplyBinaryOpRaw, LiteralRaw, ValueRaw, VarRaw}
 import aqua.types.*
 import aqua.raw.ops.*
 
@@ -72,6 +72,24 @@ class SemanticsSpec extends AnyFlatSpec with Matchers with Inside {
   def neq(left: ValueRaw, right: ValueRaw): ApplyBinaryOpRaw =
     ApplyBinaryOpRaw(ApplyBinaryOpRaw.Op.Neq, left, right)
 
+  def declareStreamPush(
+    name: String,
+    value: String
+  ): RawTag.Tree = {
+    val streamType = StreamType(ScalarType.string)
+    val stream = VarRaw(name, streamType)
+
+    RestrictionTag(stream.name, streamType).wrap(
+      SeqTag.wrap(
+        DeclareStreamTag(stream).leaf,
+        PushToStreamTag(
+          LiteralRaw.quote(value),
+          Call.Export(name, streamType)
+        ).leaf
+      )
+    )
+  }
+
   // use it to fix https://github.com/fluencelabs/aqua/issues/90
   "semantics" should "create right model" in {
     val script =
@@ -459,4 +477,53 @@ class SemanticsSpec extends AnyFlatSpec with Matchers with Inside {
       body.equalsOrShowDiff(expected) should be(true)
     }
   }
+
+  it should "restrict streams inside `if`" in {
+    val script = """
+                   |func test():           
+                   |   if "a" != "b":
+                   |      stream: *string
+                   |      stream <<- "a"
+                   |   else:
+                   |      stream: *string
+                   |      stream <<- "b"
+                   |""".stripMargin
+
+    insideBody(script) { body =>
+      val expected = IfTag(neq(LiteralRaw.quote("a"), LiteralRaw.quote("b"))).wrap(
+        declareStreamPush("stream", "a"),
+        declareStreamPush("stream", "b")
+      )
+
+      body.equalsOrShowDiff(expected) should be(true)
+    }
+  }
+
+  it should "restrict streams inside `try`" in {
+    val script = """
+                   |func test():           
+                   |   try:
+                   |      stream: *string
+                   |      stream <<- "a"
+                   |   catch e:
+                   |      stream: *string
+                   |      stream <<- "b"
+                   |   otherwise:
+                   |      stream: *string
+                   |      stream <<- "c"
+                   |""".stripMargin
+
+    insideBody(script) { body =>
+      val expected = TryTag.wrap(
+        declareStreamPush("stream", "a"),
+        SeqTag.wrap(
+          AssignmentTag(ValueRaw.lastError, "e").leaf,
+          declareStreamPush("stream", "b")
+        ),
+        declareStreamPush("stream", "c")
+      )
+
+      body.equalsOrShowDiff(expected) should be(true)
+    }
+  }
 }

utils/helpers/src/main/scala/aqua/tree/Tree.scala

@@ -0,0 +1,30 @@
+package aqua.helpers
+
+import cats.data.Chain
+import cats.free.Cofree
+import cats.Traverse
+import cats.Show
+import cats.Eval
+import cats.syntax.show.*
+import cats.syntax.traverse.*
+import cats.syntax.foldable.*
+
+object Tree {
+
+  def show[F[_]: Traverse, A: Show](
+    what: Cofree[F, A]
+  ): String =
+    Cofree
+      .cata[F, A, List[String]](what) { case (head, tail) =>
+        Eval.later {
+          val children = tail.combineAll.map("| " + _)
+          val parent = head.show
+
+          if (children.isEmpty) List(parent)
+          else (parent + ":") +: children
+        }
+      }
+      .value
+      .mkString("\n")
+
+}