From ab719e325837015474ded910e821776aecbfde3c Mon Sep 17 00:00:00 2001
From: InversionSpaces <InversionSpaces@vivaldi.net>
Date: Tue, 20 Jun 2023 11:42:24 +0000
Subject: [PATCH] Fix wrapping

---
 .../model/inline/CollectionRawInlinerSpec.scala    | 10 +++++-----
 model/raw/src/main/scala/aqua/raw/ops/RawTag.scala |  2 +-
 model/src/main/scala/aqua/model/OpModel.scala      |  6 +++---
 model/tree/src/main/scala/aqua/tree/TreeNode.scala | 14 +++++++++-----
 4 files changed, 18 insertions(+), 14 deletions(-)

diff --git a/model/inline/src/test/scala/aqua/model/inline/CollectionRawInlinerSpec.scala b/model/inline/src/test/scala/aqua/model/inline/CollectionRawInlinerSpec.scala
index 9d9e13e62..07a23e025 100644
--- a/model/inline/src/test/scala/aqua/model/inline/CollectionRawInlinerSpec.scala
+++ b/model/inline/src/test/scala/aqua/model/inline/CollectionRawInlinerSpec.scala
@@ -7,6 +7,7 @@ import aqua.raw.ops.*
 import aqua.raw.value.{CollectionRaw, LiteralRaw, MakeStructRaw, VarRaw}
 import aqua.types.{CanonStreamType, OptionType, ScalarType, StreamType, StructType}
 import cats.data.{NonEmptyList, NonEmptyMap}
+import cats.syntax.show.*
 import org.scalatest.flatspec.AnyFlatSpec
 import org.scalatest.matchers.should.Matchers
 
@@ -25,15 +26,14 @@ class CollectionRawInlinerSpec extends AnyFlatSpec with Matchers {
     val raw = CollectionRaw(NonEmptyList.of(makeStruct), OptionType(nestedType))
 
     val (v, tree) =
-      RawValueInliner.valueToModel[InliningState](raw, false).run(InliningState()).value._2
+      RawValueInliner.valueToModel[InliningState](raw, false).runA(InliningState()).value
 
     val resultValue = VarModel("option-inline-0", CanonStreamType(nestedType))
 
     v shouldBe resultValue
 
-    tree.get.equalsOrShowDiff(
-      // create a stream
-      RestrictionModel("option-inline", StreamType(nestedType)).wrap(
+    val expected =
+      RestrictionModel("option-inline", StreamType(nestedType)).wrap( // create a stream
         SeqModel.wrap(
           // create an object
           CallServiceModel(
@@ -61,8 +61,8 @@ class CollectionRawInlinerSpec extends AnyFlatSpec with Matchers {
           ).leaf
         )
       )
-    ) shouldBe true
 
+    tree.get.equalsOrShowDiff(expected) shouldBe true
   }
 
 }
diff --git a/model/raw/src/main/scala/aqua/raw/ops/RawTag.scala b/model/raw/src/main/scala/aqua/raw/ops/RawTag.scala
index 8089bd253..9a8c84ce3 100644
--- a/model/raw/src/main/scala/aqua/raw/ops/RawTag.scala
+++ b/model/raw/src/main/scala/aqua/raw/ops/RawTag.scala
@@ -61,7 +61,7 @@ sealed trait ParGroupTag extends GroupTag
 case object SeqTag extends SeqGroupTag {
 
   override def wrap(children: Tree*): Tree =
-    super.wrapNonEmpty(children.filterNot(_.head == EmptyTag).toList, RawTag.empty)
+    super.wrapNonEmpty(Chain.fromSeq(children).filterNot(_.head == EmptyTag), RawTag.empty)
 }
 
 case object ParTag extends ParGroupTag {
diff --git a/model/src/main/scala/aqua/model/OpModel.scala b/model/src/main/scala/aqua/model/OpModel.scala
index 838bbabe1..bc62a42b4 100644
--- a/model/src/main/scala/aqua/model/OpModel.scala
+++ b/model/src/main/scala/aqua/model/OpModel.scala
@@ -71,12 +71,12 @@ sealed trait ParGroupModel extends GroupOpModel
 
 case object SeqModel extends SeqGroupModel {
 
-  override def wrap(children: Tree*): Tree =
-    super.wrapNonEmpty(children.filterNot(_.head == EmptyModel).toList, EmptyModel.leaf)
+  override def wrap(children: Chain[Tree]): Tree =
+    super.wrapNonEmpty(children.filterNot(_.head == EmptyModel), EmptyModel.leaf)
 
   // EmptyModel allowed – useful for tests
   def wrapWithEmpty(children: Tree*): Tree =
-    super.wrapNonEmpty(children.toList, EmptyModel.leaf)
+    super.wrapNonEmpty(Chain.fromSeq(children), EmptyModel.leaf)
 
 }
 
diff --git a/model/tree/src/main/scala/aqua/tree/TreeNode.scala b/model/tree/src/main/scala/aqua/tree/TreeNode.scala
index aff67a341..dbab2f386 100644
--- a/model/tree/src/main/scala/aqua/tree/TreeNode.scala
+++ b/model/tree/src/main/scala/aqua/tree/TreeNode.scala
@@ -1,6 +1,7 @@
 package aqua.tree
 
 import cats.data.Chain
+import cats.data.Chain.==:
 import cats.free.Cofree
 import cats.Eval
 
@@ -16,10 +17,13 @@ trait TreeNode[T <: TreeNode[T]] {
 
   def wrap(children: Chain[Tree]): Tree = Cofree(self, Eval.now(children))
 
-  protected def wrapNonEmpty(children: List[Tree], empty: Tree): Tree = children match {
-    case Nil => empty
-    case x :: Nil => x
-    case _ => Cofree(self, Eval.now(Chain.fromSeq(children)))
-  }
+  protected def wrapNonEmpty(children: Chain[Tree], empty: Tree): Tree =
+    children match {
+      case Chain.nil => empty
+      case x ==: Chain.nil => x
+      // Do not use `wrap` here as children
+      // could redefine `wrap` through this method
+      case _ => Cofree(self, Eval.now(children))
+    }
 
 }