feat(compiler): Restrict exporting functions that return arrow types or ability types [fixes LNG-209] (#815)

aqua-src/antithesis.aqua

@@ -1,50 +1,11 @@
 aqua Main
 
-use DECLARE_CONST, decl_bar from "declare.aqua" as Declare
+export a
 
-export handleAb
+alias CL: string -> ()
 
-service SomeService("wed"):
+func a(cl: string -> ()) -> CL:
-  getStr(s: string) -> string
+  <- cl
 
-ability SomeAb:
+func b(c: u32, d: u32) -> u32:
-  someArrow(s: string) -> string, string
+  <- c + d
-  str: string
-
-ability SecondAb:
-  arrow(s: string) -> string
-  num: u32
-
-func funcStr(s: string) -> string, string:
-  strInFunc <- SomeService.getStr(Declare.DECLARE_CONST)
-  -- SomeService.getStr(s)
-  <- strInFunc, s
-
---
--- func diffFunc(s: string) -> string:
---   differentStr <- SomeService.different(s)
---   <- differentStr
---
--- func unit():
---   funcStr("")
-
--- func bbbbbbb()
---
--- func aaaaaa():
---   closure = (a: string) -> string:
---     <- SomeService.str()
-
-func handleSecAb {SomeAb, SecondAb}() -> string, string:
-  SomeAb.someArrow("eferfrfrf")
-  b, c <- SomeAb.someArrow("efre")
-  <- b, c
-
-func returnAb(s: string) -> SomeAb:
-  SomeAb = SomeAb(someArrow = funcStr, str = s)
-  <- SomeAb
-
-func handleAb(fff: string) -> string, string:
-  SomeAb = returnAb(fff)
-  SecondAb = SecondAb(arrow = funcStr, num = 12)
-  d, g <- handleSecAb{SomeAb, SecondAb}()
-  <- d, g

compiler/src/main/scala/aqua/compiler/CompilerAPI.scala

@@ -8,19 +8,19 @@ import aqua.parser.{Ast, ParserError}
 import aqua.raw.RawPart.Parts
 import aqua.raw.{RawContext, RawPart}
 import aqua.res.AquaRes
-import aqua.semantics.{CompilerState, RawSemantics, Semantics}
 import aqua.semantics.header.{HeaderHandler, HeaderSem}
+import aqua.semantics.{CompilerState, RawSemantics, Semantics}
 import cats.data.*
-import cats.data.Validated.{invalid, validNec, Invalid, Valid}
+import cats.data.Validated.{Invalid, Valid, invalid, validNec}
 import cats.parse.Parser0
 import cats.syntax.applicative.*
 import cats.syntax.flatMap.*
+import cats.syntax.foldable.*
 import cats.syntax.functor.*
 import cats.syntax.monoid.*
-import cats.syntax.traverse.*
 import cats.syntax.semigroup.*
-import cats.syntax.foldable.*
+import cats.syntax.traverse.*
-import cats.{~>, Comonad, Monad, Monoid, Order}
+import cats.{Comonad, Monad, Monoid, Order, ~>}
 import scribe.Logging
 
 import scala.collection.MapView
@@ -63,20 +63,6 @@ object CompilerAPI extends Logging {
       )
       .rawContextMonoid
 
-    implicit val headerSemMonoid: Monoid[HeaderSem[S, RawContext]] =
-      new Monoid[HeaderSem[S, RawContext]] {
-        override def empty: HeaderSem[S, RawContext] = HeaderSem(rc.empty, (c, _) => validNec(c))
-
-        override def combine(
-          a: HeaderSem[S, RawContext],
-          b: HeaderSem[S, RawContext]
-        ): HeaderSem[S, RawContext] =
-          HeaderSem(
-            a.initCtx |+| b.initCtx,
-            (c, i) => a.finInitCtx(c, i).andThen(b.finInitCtx(_, i))
-          )
-      }
-
     val semantics = new RawSemantics[S]()
 
     new AquaCompiler[F, E, I, S, RawContext](new HeaderHandler[S, RawContext](), semantics)

language-server/language-server-api/src/main/scala/aqua/lsp/LspContext.scala

@@ -56,6 +56,9 @@ object LspContext {
 
     override def blank: LspContext[S] = LspContext[S](Picker[RawContext].blank, Map.empty)
     override def exports(ctx: LspContext[S]): Option[Map[String, Option[String]]] = ops(ctx).exports
+
+    override def funcReturnAbilityOrArrow(ctx: LspContext[S], name: String): Boolean =
+      ops(ctx).funcReturnAbilityOrArrow(name)
     override def funcNames(ctx: LspContext[S]): List[String] = ops(ctx).funcNames
 
     override def addPart(ctx: LspContext[S], part: (LspContext[S], RawPart)): LspContext[S] =
@@ -104,7 +107,6 @@ object LspContext {
         }
       }.getOrElse(ctx.tokens)
 
-
       ops(ctx)
         .pick(name, rename, declared)
         .map(rc =>
@@ -124,6 +126,7 @@ object LspContext {
 
     override def pickDeclared(
       ctx: LspContext[S]
-    )(implicit semi: Semigroup[LspContext[S]]): LspContext[S] = ctx.copy(raw = ops(ctx).pickDeclared)
+    )(using Semigroup[LspContext[S]]): LspContext[S] =
+      ctx.copy(raw = ops(ctx).pickDeclared)
   }
 }

semantics/src/main/scala/aqua/semantics/header/HeaderHandler.scala

@@ -3,18 +3,18 @@ package aqua.semantics.header
 import aqua.parser.Ast
 import aqua.parser.head.*
 import aqua.parser.lexer.{Ability, Token}
-import aqua.raw.RawContext
 import aqua.semantics.header.Picker.*
 import aqua.semantics.{HeaderError, SemanticError}
-import cats.data.Validated.{invalidNec, validNec, Invalid, Valid}
 import cats.data.*
+import cats.data.Validated.{invalidNec, validNec}
 import cats.free.Cofree
 import cats.instances.list.*
 import cats.instances.option.*
 import cats.kernel.Semigroup
 import cats.syntax.foldable.*
-import cats.syntax.monoid
+import cats.syntax.functor.*
 import cats.syntax.semigroup.*
+import cats.syntax.validated.*
 import cats.{Comonad, Eval, Monoid}
 
 class HeaderHandler[S[_]: Comonad, C](implicit
@@ -115,15 +115,15 @@ class HeaderHandler[S[_]: Comonad, C](implicit
                   ctx
                     .pick(n, None, ctx.module.nonEmpty)
                     // We just validate, nothing more
-                    .map(_ => validNec(1))
+                    .as(validNec(1))
                     .getOrElse(
                       error(
                         t,
-                        s"`${n}` is expected to be declared, but declaration is not found in the file"
+                        s"`$n` is expected to be declared, but declaration is not found in the file"
                       )
                     )
                 }.combineAll
-                  .map(_ =>
+                  .as(
                     // TODO: why module name and declares is lost? where is it lost?
                     ctx.setModule(name.value, declares = shouldDeclare)
                   )
@@ -173,17 +173,24 @@ class HeaderHandler[S[_]: Comonad, C](implicit
                   )
                 )
                 .map { case (token, name, rename) =>
-                  (initCtx |+| ctx)
+                  val sumCtx = initCtx |+| ctx
-                    .pick(name, rename, declared = false)
+
-                    .map(_ => Map(name -> rename))
+                  if (sumCtx.funcReturnAbilityOrArrow(name))
-                    .map(validNec)
+                    error(
-                    .getOrElse(
+                      token,
-                      error(
+                      s"The function '$name' cannot be exported, because it returns arrow type or ability type"
-                        token,
-                        s"File has no $name declaration or import, cannot export, available funcs: ${(initCtx |+| ctx).funcNames
-                          .mkString(", ")}"
-                      )
                     )
+                  else
+                    sumCtx
+                      .pick(name, rename, declared = false)
+                      .as(Map(name -> rename).validNec)
+                      .getOrElse(
+                        error(
+                          token,
+                          s"File has no $name declaration or import, cannot export, available funcs: ${sumCtx.funcNames
+                            .mkString(", ")}"
+                        )
+                      )
                 }
                 .foldLeft[ResT[S, Map[String, Option[String]]]](
                   validNec(ctx.exports.getOrElse(Map.empty))

semantics/src/main/scala/aqua/semantics/header/HeaderSem.scala

@@ -1,7 +1,11 @@
 package aqua.semantics.header
 
+import aqua.raw.RawContext
 import aqua.semantics.SemanticError
+import cats.{Comonad, Monoid}
 import cats.data.*
+import cats.syntax.monoid.*
+import cats.data.Validated.validNec
 
 case class HeaderSem[S[_], C](
   initCtx: C,
@@ -11,3 +15,21 @@ case class HeaderSem[S[_], C](
   def finCtx: C => ValidatedNec[SemanticError[S], C] =
     finInitCtx(_, initCtx)
 }
+
+object HeaderSem {
+
+  given [S[_]: Comonad](using
+    rc: Monoid[RawContext]
+  ): Monoid[HeaderSem[S, RawContext]] with {
+    override def empty: HeaderSem[S, RawContext] = HeaderSem(rc.empty, (c, _) => validNec(c))
+
+    override def combine(
+      a: HeaderSem[S, RawContext],
+      b: HeaderSem[S, RawContext]
+    ): HeaderSem[S, RawContext] =
+      HeaderSem(
+        a.initCtx |+| b.initCtx,
+        (c, i) => a.finInitCtx(c, i).andThen(b.finInitCtx(_, i))
+      )
+  }
+}

semantics/src/main/scala/aqua/semantics/header/Picker.scala

@@ -1,11 +1,8 @@
 package aqua.semantics.header
 
 import aqua.raw.{RawContext, RawPart}
-import aqua.semantics.CompilerState
+import aqua.types.{AbilityType, ArrowType}
-import aqua.semantics.rules.abilities.AbilitiesState
+import cats.Semigroup
-import aqua.semantics.rules.names.NamesState
-import aqua.semantics.rules.types.TypesState
-import cats.{Comonad, Semigroup}
 import cats.syntax.semigroup.*
 
 // Able to pick info from different contexts
@@ -19,6 +16,7 @@ trait Picker[A] {
   def pickHeader(ctx: A): A
   def module(ctx: A): Option[String]
   def exports(ctx: A): Option[Map[String, Option[String]]]
+  def funcReturnAbilityOrArrow(ctx: A, name: String): Boolean
   def declares(ctx: A): Set[String]
   def setAbility(ctx: A, name: String, ctxAb: A): A
   def setModule(ctx: A, name: Option[String], declares: Set[String]): A
@@ -39,6 +37,7 @@ final class PickerOps[A: Picker](p: A) {
   def pickHeader: A = Picker[A].pickHeader(p)
   def module: Option[String] = Picker[A].module(p)
   def exports: Option[Map[String, Option[String]]] = Picker[A].exports(p)
+  def funcReturnAbilityOrArrow(name: String): Boolean = Picker[A].funcReturnAbilityOrArrow(p, name)
   def declares: Set[String] = Picker[A].declares(p)
   def setAbility(name: String, ctx: A): A = Picker[A].setAbility(p, name, ctx)
   def setInit(ctx: Option[A]): A = Picker[A].setInit(p, ctx)
@@ -56,6 +55,14 @@ final class PickerOps[A: Picker](p: A) {
 
 object Picker {
 
+  def returnsAbilityOrArrow(arrowType: ArrowType): Boolean = {
+    arrowType.codomain.toList.exists {
+      case _: AbilityType => true
+      case _: ArrowType => true
+      case _ => false
+    }
+  }
+
   implicit final def apply[A](implicit ev: Picker[A]): Picker[A] = ev
 
   implicit final def syntaxPicker[A: Picker](a: A): PickerOps[A] =
@@ -65,6 +72,8 @@ object Picker {
 
     override def blank: RawContext = RawContext.blank
     override def exports(ctx: RawContext): Option[Map[String, Option[String]]] = ctx.exports
+    override def funcReturnAbilityOrArrow(ctx: RawContext, name: String): Boolean =
+      ctx.funcs.get(name).map(_.arrow.`type`).exists(returnsAbilityOrArrow)
     override def funcNames(ctx: RawContext): List[String] = ctx.funcs.keys.toList
 
     override def addPart(ctx: RawContext, part: (RawContext, RawPart)): RawContext =

semantics/src/main/scala/aqua/semantics/rules/names/NamesInterpreter.scala

@@ -2,14 +2,13 @@ package aqua.semantics.rules.names
 
 import aqua.parser.lexer.{Name, Token}
 import aqua.semantics.Levenshtein
-import aqua.semantics.rules.locations.LocationsAlgebra
 import aqua.semantics.rules.StackInterpreter
 import aqua.semantics.rules.errors.ReportErrors
-import aqua.types.{ArrowType, StreamType, Type}
+import aqua.semantics.rules.locations.LocationsAlgebra
+import aqua.types.{AbilityType, ArrowType, StreamType, Type}
 import cats.data.{OptionT, State}
 import cats.syntax.flatMap.*
 import cats.syntax.functor.*
-import cats.~>
 import monocle.Lens
 import monocle.macros.GenLens
 
@@ -129,7 +128,7 @@ class NamesInterpreter[S[_], X](implicit
         ).as(true)
     }.flatTap(_ => locations.addToken(name.value, name))
 
-  override def defineArrow(name: Name[S], gen: ArrowType, isRoot: Boolean): SX[Boolean] =
+  override def defineArrow(name: Name[S], arrowType: ArrowType, isRoot: Boolean): SX[Boolean] =
     readName(name.value).flatMap {
       case Some(_) =>
         getState.map(_.definitions.get(name.value).exists(_ == name)).flatMap {
@@ -142,7 +141,7 @@ class NamesInterpreter[S[_], X](implicit
           if (isRoot)
             modify(st =>
               st.copy(
-                rootArrows = st.rootArrows.updated(name.value, gen),
+                rootArrows = st.rootArrows.updated(name.value, arrowType),
                 definitions = st.definitions.updated(name.value, name)
               )
             )
@@ -150,7 +149,7 @@ class NamesInterpreter[S[_], X](implicit
           else
             report(name, "Cannot define a variable in the root scope")
               .as(false)
-        )(fr => fr.addArrow(name, gen) -> true)
+        )(fr => fr.addArrow(name, arrowType) -> true)
     }.flatTap(_ => locations.addToken(name.value, name))
 
   override def streamsDefinedWithinScope(): SX[Map[String, StreamType]] =

semantics/src/test/scala/aqua/semantics/HeaderSpec.scala

@@ -0,0 +1,50 @@
+package aqua.semantics
+
+import aqua.parser.head.{ExportExpr, FromExpr, HeaderExpr}
+import aqua.parser.lexer.Name
+import aqua.raw.RawContext
+import aqua.raw.arrow.{ArrowRaw, FuncRaw}
+import aqua.raw.ops.RawTag
+import aqua.raw.value.VarRaw
+import aqua.semantics.header.{HeaderHandler, HeaderSem}
+import aqua.types.{ArrowType, NilType, ProductType}
+import cats.data.{Chain, NonEmptyList, Validated}
+import cats.free.Cofree
+import cats.{Eval, Id, Monoid}
+import org.scalatest.Inside
+import org.scalatest.flatspec.AnyFlatSpec
+import org.scalatest.matchers.should.Matchers
+
+class HeaderSpec extends AnyFlatSpec with Matchers with Inside {
+
+  "header handler" should "generate an error on exported function that returns arrow or ability" in {
+    implicit val rc: Monoid[RawContext] = RawContext.implicits(RawContext.blank).rawContextMonoid
+
+    val handler = new HeaderHandler[Id, RawContext]()
+
+    val funcName = "funcName"
+
+    val exp: FromExpr.NameOrAbAs[Id] = Left((Name[Id](funcName), None))
+    val ast =
+      Cofree[Chain, HeaderExpr[Id]](ExportExpr[Id](NonEmptyList.of(exp)), Eval.now(Chain.empty))
+
+    val retArrowType = ArrowType(NilType, NilType)
+    val arrowType = ArrowType(NilType, ProductType.apply(retArrowType :: Nil))
+
+    val initCtx = RawContext(parts =
+      Chain.one(
+        (
+          RawContext.blank,
+          FuncRaw(funcName, ArrowRaw(arrowType, VarRaw("", retArrowType) :: Nil, RawTag.empty))
+        )
+      )
+    )
+
+    val result = handler.sem(Map.empty, ast).andThen(_.finCtx(initCtx))
+    
+    inside(result) {
+      case Validated.Invalid(errors) =>
+        errors.head shouldBe a [HeaderError[Id]]   
+    }
+  }
+}

utils/logging/src/main/scala/aqua/logging/LogLevels.scala

@@ -1,9 +1,8 @@
 package aqua.logging
 
+import cats.data.Validated.{invalidNel, validNel}
+import cats.data.{NonEmptyList, Validated, ValidatedNel}
 import scribe.Level
-import cats.data.Validated.{invalid, invalidNec, invalidNel, valid, validNec, validNel}
-import cats.data.{Validated, ValidatedNel}
-import cats.data.NonEmptyList
 
 case class LogLevels(
   compiler: Level = Level.Error,