Fix gate inlining

2024-12-04 22:50:18 +00:00 · 2023-10-06 10:00:05 +00:00 · 2023-10-06 10:00:05 +00:00 · 405e4d330f
commit 405e4d330f
parent feb7a167a2
4 changed files with 85 additions and 51 deletions
--- a/compiler/src/test/scala/aqua/compiler/AquaCompilerSpec.scala
+++ b/compiler/src/test/scala/aqua/compiler/AquaCompilerSpec.scala
@ -1,14 +1,6 @@
 package aqua.compiler

-import aqua.model.{
-  CallModel,
-  ForModel,
-  FunctorModel,
-  IntoIndexModel,
-  LiteralModel,
-  ValueModel,
-  VarModel
-}
+import aqua.model.{CallModel, ForModel, FunctorModel, LiteralModel, ValueModel, VarModel}
 import aqua.model.transform.ModelBuilder
 import aqua.model.transform.TransformConfig
 import aqua.model.transform.Transform
--- a/model/inline/src/main/scala/aqua/model/inline/raw/ApplyGateRawInliner.scala
+++ b/model/inline/src/main/scala/aqua/model/inline/raw/ApplyGateRawInliner.scala
@ -1,5 +1,6 @@
 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}
@ -11,6 +12,7 @@ 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 {
@ -108,6 +110,11 @@ object ApplyGateRawInliner extends RawInliner[ApplyGateRaw] with Logging {
      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,
@ -121,15 +128,23 @@ object ApplyGateRawInliner extends RawInliner[ApplyGateRaw] with Logging {
        resultName = uniqueResultName
      )

-      val tree = SeqModel.wrap(idxInline.predo.toList :+ gate)
-
-      val treeInline = Inline(predo = Chain.one(tree))
-
-      (
-        VarModel(uniqueResultName, ArrayType(afr.streamType.element)),
-        treeInline
+      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) =
--- a/model/inline/src/main/scala/aqua/model/inline/raw/ApplyPropertiesRawInliner.scala
+++ b/model/inline/src/main/scala/aqua/model/inline/raw/ApplyPropertiesRawInliner.scala
@ -18,6 +18,7 @@ import cats.syntax.bifunctor.*
 import cats.syntax.foldable.*
 import cats.syntax.monoid.*
 import cats.syntax.traverse.*
+import cats.syntax.option.*
 import scribe.Logging

 object ApplyPropertiesRawInliner extends RawInliner[ApplyPropertyRaw] with Logging {
@ -48,17 +49,6 @@ object ApplyPropertiesRawInliner extends RawInliner[ApplyPropertyRaw] with Loggi
    }
  }

-  private def removeProperties[S: Mangler](
-    varModel: VarModel
-  ): State[S, (VarModel, Inline)] = {
-    for {
-      nn <- Mangler[S].findAndForbidName(varModel.name + "_flat")
-    } yield {
-      val flatten = VarModel(nn, varModel.`type`)
-      flatten -> Inline.tree(FlattenModel(varModel, flatten.name).leaf)
-    }
-  }
-
  private def unfoldAbilityProperty[S: Mangler: Exports: Arrows](
    varModel: VarModel,
    abilityType: NamedType,
@ -143,9 +133,7 @@ object ApplyPropertiesRawInliner extends RawInliner[ApplyPropertyRaw] with Loggi

      case f @ FunctorRaw(_, _) =>
        for {
-          flattenVI <-
-            if (varModel.properties.nonEmpty) removeProperties(varModel)
-            else State.pure(varModel, Inline.empty)
+          flattenVI <- removeProperties(varModel)
          (flatten, inline) = flattenVI
          newVI <- ApplyFunctorRawInliner(flatten, f)
        } yield {
@ -157,9 +145,7 @@ object ApplyPropertiesRawInliner extends RawInliner[ApplyPropertyRaw] with Loggi

      case ic @ IntoCopyRaw(_, _) =>
        for {
-          flattenVI <-
-            if (varModel.properties.nonEmpty) removeProperties(varModel)
-            else State.pure(varModel, Inline.empty)
+          flattenVI <- removeProperties(varModel)
          (flatten, inline) = flattenVI
          newVI <- ApplyIntoCopyRawInliner(flatten, ic)
        } yield {
@ -182,15 +168,31 @@ object ApplyPropertiesRawInliner extends RawInliner[ApplyPropertyRaw] with Loggi
        unfold(vr, propertiesAllowed = false).flatMap {
          case (vm @ VarModel(_, _, _), inline) if vm.properties.nonEmpty =>
            removeProperties(vm).map { case (vf, inlf) =>
-              PropertyRawWithModel(iir, Option(IntoIndexModel(vf.name, t))) -> Inline(
+              PropertyRawWithModel(
+                iir,
+                IntoIndexModel
+                  .fromValueModel(vf, t)
+                  .getOrElse(
+                    internalError(s"Unexpected: could not convert ($vf) to IntoIndexModel")
+                  )
+                  .some
+              ) -> Inline(
                inline.predo ++ inlf.predo,
                mergeMode = SeqMode
              )
            }
-          case (VarModel(name, _, _), inline) =>
-            State.pure(PropertyRawWithModel(iir, Option(IntoIndexModel(name, t))) -> inline)
-          case (LiteralModel(literal, _), inline) =>
-            State.pure(PropertyRawWithModel(iir, Option(IntoIndexModel(literal, t))) -> inline)
+          case (vm, inline) =>
+            (
+              PropertyRawWithModel(
+                iir,
+                IntoIndexModel
+                  .fromValueModel(vm, t)
+                  .getOrElse(
+                    internalError(s"Unexpected: could not convert ($vm) to IntoIndexModel")
+                  )
+                  .some
+              ) -> inline
+            ).pure
        }

      case p => State.pure(PropertyRawWithModel(p, None) -> Inline.empty)
@ -251,15 +253,19 @@ object ApplyPropertiesRawInliner extends RawInliner[ApplyPropertyRaw] with Loggi
    properties: Chain[PropertyRaw],
    propertiesAllowed: Boolean
  ): State[S, (ValueModel, Inline)] = {
-    ((raw, properties.headOption) match {
-      case (vr @ VarRaw(_, st @ StreamType(_)), Some(IntoIndexRaw(idx, _))) =>
+    ((raw, properties.uncons) match {
+      case (vr @ VarRaw(_, st @ StreamType(_)), Some(IntoIndexRaw(idx, _), otherProperties)) =>
        unfold(vr).flatMap {
          case (VarModel(nameVM, _, _), inl) =>
            val gateRaw = ApplyGateRaw(nameVM, st, idx)
            unfold(gateRaw).flatMap {
              case (gateResVal: VarModel, gateResInline) =>
-                unfoldProperties(gateResInline, gateResVal, properties, propertiesAllowed).map {
-                  case (v, i) =>
+                unfoldProperties(
+                  gateResInline,
+                  gateResVal,
+                  otherProperties,
+                  propertiesAllowed
+                ).map { case (v, i) =>
                  v -> Inline(
                    inl.predo ++ i.predo,
                    mergeMode = SeqMode
@ -299,6 +305,19 @@ object ApplyPropertiesRawInliner extends RawInliner[ApplyPropertyRaw] with Loggi

  }

+  /**
+   * Remove properties from the var and return a new var without them
+   */
+  def removeProperties[S: Mangler](
+    varModel: VarModel
+  ): State[S, (VarModel, Inline)] =
+    if (varModel.properties.isEmpty) (varModel, Inline.empty).pure
+    else
+      for {
+        nn <- Mangler[S].findAndForbidName(varModel.name + "_flat")
+        flatten = VarModel(nn, varModel.`type`)
+      } yield flatten -> Inline.tree(FlattenModel(varModel, flatten.name).leaf)
+
  override def apply[S: Mangler: Exports: Arrows](
    apr: ApplyPropertyRaw,
    propertiesAllowed: Boolean
--- a/model/src/main/scala/aqua/model/ValueModel.scala
+++ b/model/src/main/scala/aqua/model/ValueModel.scala
@ -157,6 +157,15 @@ case class IntoIndexModel(idx: String, `type`: Type) extends PropertyModel {
  )
 }

+object IntoIndexModel {
+
+  def fromValueModel(vm: ValueModel, `type`: Type): Option[IntoIndexModel] = vm match {
+    case VarModel(name, _, Chain.nil) => IntoIndexModel(name, `type`).some
+    case LiteralModel(value, _) => IntoIndexModel(value, `type`).some
+    case _ => none
+  }
+}
+
 case class VarModel(name: String, baseType: Type, properties: Chain[PropertyModel] = Chain.empty)
    extends ValueModel with Logging {

@ -173,14 +182,13 @@ case class VarModel(name: String, baseType: Type, properties: Chain[PropertyMode
  private def deriveFrom(vm: VarModel): VarModel =
    vm.copy(properties = vm.properties ++ properties)

+  def withProperty(p: PropertyModel): VarModel =
+    copy(properties = properties :+ p)
+
  def intoField(field: String): Option[VarModel] = `type` match {
    case StructType(_, fields) =>
      fields(field)
-        .map(fieldType =>
-          copy(
-            properties = properties :+ IntoFieldModel(field, fieldType)
-          )
-        )
+        .map(fieldType => withProperty(IntoFieldModel(field, fieldType)))
    case _ => none
  }