NODE-2549 Fixes for global variable in estimator

lang/jvm/src/main/scala/com/wavesplatform/lang/Lang.scala

@@ -26,7 +26,7 @@ object Lang {
 
   def compileDApp(input: String): DAppWithMeta =
     API
-      .compile(input, ScriptEstimatorV3(fixOverflow = true, overhead = false))
+      .compile(input, ScriptEstimatorV3.latest)
       .flatMap {
         case r: CompileResult.DApp =>
           val javaMeta = Meta(

lang/shared/src/main/scala/com/wavesplatform/lang/FileCompiler.scala

@@ -6,7 +6,7 @@ import java.nio.file.{Files, Paths}
 import scala.jdk.CollectionConverters.*
 
 object FileCompiler extends App {
-  private val estimator = ScriptEstimatorV3(true, false)
+  private val estimator = ScriptEstimatorV3.latest
   args
     .foreach { path =>
       val script = Files.readAllLines(Paths.get(path)).asScala.reduce(_ + "\n" + _)

lang/shared/src/main/scala/com/wavesplatform/lang/v1/estimator/ScriptEstimator.scala

@@ -18,5 +18,5 @@ trait ScriptEstimator {
 
 object ScriptEstimator {
   def all(fixOverflow: Boolean): List[ScriptEstimator] =
-    List(ScriptEstimatorV1, ScriptEstimatorV2, ScriptEstimatorV3(fixOverflow, overhead = false))
+    List(ScriptEstimatorV1, ScriptEstimatorV2, ScriptEstimatorV3(fixOverflow, overhead = false, letFixes = true))
 }

lang/shared/src/main/scala/com/wavesplatform/lang/v1/estimator/v3/EstimatorContext.scala

@@ -9,7 +9,8 @@ import shapeless.{Lens, lens}
 
 private[v3] case class EstimatorContext(
     funcs: Map[FunctionHeader, (Coeval[Long], Set[String])],
-    usedRefs: Set[String] = Set.empty,
+    usedRefs: Set[String] = Set(),
+    refsCosts: Map[String, Long] = Map(),
     globalFunctionsCosts: Map[String, Long] = Map(), //
     globalLetsCosts: Map[String, Long] = Map(),      // only for globalDeclarationsMode
     globalLetEvals: Map[String, EvalM[Long]] = Map() //

lang/shared/src/main/scala/com/wavesplatform/lang/v1/estimator/v3/ScriptEstimatorV3.scala

@@ -1,6 +1,6 @@
 package com.wavesplatform.lang.v1.estimator.v3
 
-import cats.implicits.*
+import cats.implicits.{toBifunctorOps, toFoldableOps, toTraverseOps}
 import cats.{Id, Monad}
 import com.wavesplatform.lang.v1.FunctionHeader
 import com.wavesplatform.lang.v1.FunctionHeader.User
@@ -13,7 +13,7 @@ import monix.eval.Coeval
 
 import scala.util.Try
 
-case class ScriptEstimatorV3(fixOverflow: Boolean, overhead: Boolean) extends ScriptEstimator {
+case class ScriptEstimatorV3(fixOverflow: Boolean, overhead: Boolean, letFixes: Boolean) extends ScriptEstimator {
   private val overheadCost: Long = if (overhead) 1 else 0
 
   override val version: Int = 3
@@ -39,55 +39,45 @@ case class ScriptEstimatorV3(fixOverflow: Boolean, overhead: Boolean) extends Sc
       globalDeclarationsMode: Boolean
   ): (EstimatorContext, Either[EstimationError, Long]) = {
     val ctxFuncs = funcs.view.mapValues((_, Set[String]())).toMap
-    evalExpr(expr, globalDeclarationsMode).run(EstimatorContext(ctxFuncs)).value
+    evalExpr(expr, Set(), globalDeclarationsMode).run(EstimatorContext(ctxFuncs)).value
   }
 
-  private def evalExpr(t: EXPR, globalDeclarationsMode: Boolean = false): EvalM[Long] =
+  private def evalExpr(t: EXPR, activeFuncArgs: Set[String], globalDeclarationsMode: Boolean = false): EvalM[Long] =
     if (Thread.currentThread().isInterrupted)
       raiseError("Script estimation was interrupted")
     else
       t match {
-        case LET_BLOCK(let, inner)       => evalLetBlock(let, inner, globalDeclarationsMode)
-        case BLOCK(let: LET, inner)      => evalLetBlock(let, inner, globalDeclarationsMode)
-        case BLOCK(f: FUNC, inner)       => evalFuncBlock(f, inner, globalDeclarationsMode)
+        case LET_BLOCK(let, inner)       => evalLetBlock(let, inner, activeFuncArgs, globalDeclarationsMode)
+        case BLOCK(let: LET, inner)      => evalLetBlock(let, inner, activeFuncArgs, globalDeclarationsMode)
+        case BLOCK(f: FUNC, inner)       => evalFuncBlock(f, inner, activeFuncArgs, globalDeclarationsMode)
         case BLOCK(_: FAILED_DEC, _)     => zero
-        case REF(str)                    => markRef(str)
+        case REF(str)                    => evalRef(str, activeFuncArgs)
         case _: EVALUATED                => const(overheadCost)
-        case IF(cond, t1, t2)            => evalIF(cond, t1, t2)
-        case GETTER(expr, _)             => evalGetter(expr)
-        case FUNCTION_CALL(header, args) => evalFuncCall(header, args)
+        case IF(cond, t1, t2)            => evalIF(cond, t1, t2, activeFuncArgs)
+        case GETTER(expr, _)             => evalGetter(expr, activeFuncArgs)
+        case FUNCTION_CALL(header, args) => evalFuncCall(header, args, activeFuncArgs)
         case _: FAILED_EXPR              => zero
       }
 
-  private def evalHoldingFuncs(expr: EXPR): EvalM[Long] =
+  private def evalLetBlock(let: LET, nextExpr: EXPR, activeFuncArgs: Set[String], globalDeclarationsMode: Boolean): EvalM[Long] =
     for {
+      _        <- if (globalDeclarationsMode) saveGlobalLetCost(let, activeFuncArgs) else doNothing
       startCtx <- get[Id, EstimatorContext, EstimationError]
-      cost     <- evalExpr(expr)
-      _        <- update(funcs.set(_)(startCtx.funcs))
-    } yield cost
-
-  private def evalLetBlock(let: LET, inner: EXPR, globalDeclarationsMode: Boolean): EvalM[Long] =
-    for {
-      startCtx <- get[Id, EstimatorContext, EstimationError]
-      overlap = startCtx.usedRefs.contains(let.name)
-      _ <- update(usedRefs.modify(_)(_ - let.name))
-      letEval = evalHoldingFuncs(let.value)
-      _        <- if (globalDeclarationsMode) saveGlobalLetCost(let) else doNothing
-      nextCost <- evalExpr(inner, globalDeclarationsMode)
-      ctx      <- get[Id, EstimatorContext, EstimationError]
-      letCost  <- if (ctx.usedRefs.contains(let.name)) letEval else zero
-      _        <- update(usedRefs.modify(_)(r => if (overlap) r + let.name else r - let.name))
-      result   <- sum(nextCost, letCost)
+      letEval = evalHoldingFuncs(let.value, activeFuncArgs)
+      _            <- beforeNextExprEval(let, letEval)
+      nextExprCost <- evalExpr(nextExpr, activeFuncArgs, globalDeclarationsMode)
+      nextExprCtx  <- get[Id, EstimatorContext, EstimationError]
+      _            <- afterNextExprEval(let, startCtx)
+      letCost      <- if (nextExprCtx.usedRefs.contains(let.name)) letEval else const(0L)
+      result       <- sum(nextExprCost, letCost)
     } yield result
 
-  private def saveGlobalLetCost(let: LET): EvalM[Unit] = {
+  private def saveGlobalLetCost(let: LET, activeFuncArgs: Set[String]): EvalM[Unit] = {
     val costEvaluation =
       for {
-        startCtx    <- get[Id, EstimatorContext, EstimationError]
-        bodyCost    <- evalExpr(let.value)
-        bodyEvalCtx <- get[Id, EstimatorContext, EstimationError]
-        usedRefs = bodyEvalCtx.usedRefs diff startCtx.usedRefs
-        letCosts <- usedRefs.toSeq.traverse(bodyEvalCtx.globalLetEvals.getOrElse(_, zero))
+        (bodyCost, usedRefs) <- withUsedRefs(evalExpr(let.value, activeFuncArgs))
+        ctx                  <- get[Id, EstimatorContext, EstimationError]
+        letCosts             <- usedRefs.toSeq.traverse(ctx.globalLetEvals.getOrElse(_, zero))
       } yield bodyCost + letCosts.sum
     for {
       cost <- local(costEvaluation)
@@ -100,26 +90,47 @@ case class ScriptEstimatorV3(fixOverflow: Boolean, overhead: Boolean) extends Sc
     } yield ()
   }
 
-  private def evalFuncBlock(func: FUNC, inner: EXPR, globalDeclarationsMode: Boolean): EvalM[Long] =
+  private def beforeNextExprEval(let: LET, eval: EvalM[Long]): EvalM[Unit] =
+    for {
+      cost <- local(eval)
+      _ <- update(ctx =>
+        usedRefs
+          .modify(ctx)(_ - let.name)
+          .copy(refsCosts = ctx.refsCosts + (let.name -> cost))
+      )
+    } yield ()
+
+  private def afterNextExprEval(let: LET, startCtx: EstimatorContext): EvalM[Unit] =
+    update(ctx =>
+      usedRefs
+        .modify(ctx)(r => if (startCtx.usedRefs.contains(let.name)) r + let.name else r - let.name)
+        .copy(refsCosts =
+          if (startCtx.refsCosts.contains(let.name))
+            ctx.refsCosts + (let.name -> startCtx.refsCosts(let.name))
+          else
+            ctx.refsCosts - let.name
+        )
+    )
+
+  private def evalFuncBlock(func: FUNC, nextExpr: EXPR, activeFuncArgs: Set[String], globalDeclarationsMode: Boolean): EvalM[Long] =
     for {
-      startCtx    <- get[Id, EstimatorContext, EstimationError]
-      _           <- checkShadowing(func, startCtx)
-      funcCost    <- evalHoldingFuncs(func.body)
-      bodyEvalCtx <- get[Id, EstimatorContext, EstimationError]
-      refsUsedInBody = bodyEvalCtx.usedRefs diff startCtx.usedRefs
-      _        <- if (globalDeclarationsMode) saveGlobalFuncCost(func.name, funcCost, bodyEvalCtx, refsUsedInBody) else doNothing
-      _        <- handleUsedRefs(func.name, funcCost, startCtx, refsUsedInBody)
-      nextCost <- evalExpr(inner, globalDeclarationsMode)
-    } yield nextCost
+      startCtx                   <- get[Id, EstimatorContext, EstimationError]
+      _                          <- checkShadowing(func, startCtx)
+      (funcCost, refsUsedInBody) <- withUsedRefs(evalHoldingFuncs(func.body, activeFuncArgs ++ func.args))
+      _                          <- if (globalDeclarationsMode) saveGlobalFuncCost(func.name, funcCost, refsUsedInBody) else doNothing
+      _                          <- handleUsedRefs(func.name, funcCost, startCtx, refsUsedInBody)
+      nextExprCost               <- evalExpr(nextExpr, activeFuncArgs, globalDeclarationsMode)
+    } yield nextExprCost
 
   private def checkShadowing(func: FUNC, startCtx: EstimatorContext): EvalM[Any] =
     if (fixOverflow && startCtx.funcs.contains(FunctionHeader.User(func.name)))
       raiseError(s"Function '${func.name}${func.args.mkString("(", ", ", ")")}' shadows preceding declaration")
     else
       doNothing
 
-  private def saveGlobalFuncCost(name: String, funcCost: Long, ctx: EstimatorContext, refsUsedInBody: Set[String]): EvalM[Unit] =
+  private def saveGlobalFuncCost(name: String, funcCost: Long, refsUsedInBody: Set[String]): EvalM[Unit] =
     for {
+      ctx      <- get[Id, EstimatorContext, EstimationError]
       letCosts <- local(refsUsedInBody.toSeq.traverse(ctx.globalLetEvals.getOrElse(_, zero)))
       totalCost = math.max(1, funcCost + letCosts.sum)
       _ <- set[Id, EstimatorContext, EstimationError](ctx.copy(globalFunctionsCosts = ctx.globalFunctionsCosts + (name -> totalCost)))
@@ -135,46 +146,75 @@ case class ScriptEstimatorV3(fixOverflow: Boolean, overhead: Boolean) extends Sc
       }
     )
 
-  private def evalIF(cond: EXPR, ifTrue: EXPR, ifFalse: EXPR): EvalM[Long] =
+  private def evalIF(cond: EXPR, ifTrue: EXPR, ifFalse: EXPR, activeFuncArgs: Set[String]): EvalM[Long] =
     for {
-      cond  <- evalHoldingFuncs(cond)
-      right <- evalHoldingFuncs(ifTrue)
-      left  <- evalHoldingFuncs(ifFalse)
+      cond  <- evalHoldingFuncs(cond, activeFuncArgs)
+      right <- evalHoldingFuncs(ifTrue, activeFuncArgs)
+      left  <- evalHoldingFuncs(ifFalse, activeFuncArgs)
       r1    <- sum(cond, Math.max(right, left))
       r2    <- sum(r1, overheadCost)
     } yield r2
 
-  private def markRef(key: String): EvalM[Long] =
-    update(usedRefs.modify(_)(_ + key)).map(_ => overheadCost)
+  private def evalRef(key: String, activeFuncArgs: Set[String]): EvalM[Long] =
+    if (activeFuncArgs.contains(key) && letFixes)
+      const(overheadCost)
+    else
+      update(usedRefs.modify(_)(_ + key)).map(_ => overheadCost)
 
-  private def evalGetter(expr: EXPR): EvalM[Long] =
-    evalExpr(expr).flatMap(sum(_, overheadCost))
+  private def evalGetter(expr: EXPR, activeFuncArgs: Set[String]): EvalM[Long] =
+    evalExpr(expr, activeFuncArgs).flatMap(sum(_, overheadCost))
 
-  private def evalFuncCall(header: FunctionHeader, args: List[EXPR]): EvalM[Long] =
+  private def evalFuncCall(header: FunctionHeader, args: List[EXPR], activeFuncArgs: Set[String]): EvalM[Long] =
     for {
-      ctx <- get[Id, EstimatorContext, EstimationError]
-      (bodyCost, bodyUsedRefs) <- funcs
-        .get(ctx)
-        .get(header)
-        .map(const)
-        .getOrElse(
-          raiseError[Id, EstimatorContext, EstimationError, (Coeval[Long], Set[String])](s"function '$header' not found")
-        )
-      _ <- update(
-        (funcs ~ usedRefs).modify(_) { case (funcs, usedRefs) =>
-          (
-            funcs + ((header, (bodyCost, Set[String]()))),
-            usedRefs ++ bodyUsedRefs
-          )
-        }
-      )
-      argsCosts    <- args.traverse(evalHoldingFuncs)
-      argsCostsSum <- argsCosts.foldM(0L)(sum)
-      bodyCostV         = bodyCost.value()
-      correctedBodyCost = if (!overhead && bodyCostV == 0) 1 else bodyCostV
+      ctx                       <- get[Id, EstimatorContext, EstimationError]
+      (bodyCost, bodyUsedRefs)  <- getFuncCost(header, ctx)
+      _                         <- setFuncToCtx(header, bodyCost, bodyUsedRefs)
+      (argsCosts, argsUsedRefs) <- withUsedRefs(args.traverse(evalHoldingFuncs(_, activeFuncArgs)))
+      argsCostsSum              <- argsCosts.foldM(0L)(sum)
+      bodyCostV = bodyCost.value()
+      correctedBodyCost =
+        if (!overhead && !letFixes && bodyCostV == 0) 1
+        else if (letFixes && bodyCostV == 0 && isBlankFunc(bodyUsedRefs ++ argsUsedRefs, ctx.refsCosts)) 1
+        else bodyCostV
       result <- sum(argsCostsSum, correctedBodyCost)
     } yield result
 
+  private def setFuncToCtx(header: FunctionHeader, bodyCost: Coeval[Long], bodyUsedRefs: Set[EstimationError]): EvalM[Unit] =
+    update(
+      (funcs ~ usedRefs).modify(_) { case (funcs, usedRefs) =>
+        (
+          funcs + (header -> (bodyCost, Set())),
+          usedRefs ++ bodyUsedRefs
+        )
+      }
+    )
+
+  private def getFuncCost(header: FunctionHeader, ctx: EstimatorContext): EvalM[(Coeval[Long], Set[EstimationError])] =
+    funcs
+      .get(ctx)
+      .get(header)
+      .map(const)
+      .getOrElse(
+        raiseError[Id, EstimatorContext, EstimationError, (Coeval[Long], Set[EstimationError])](s"function '$header' not found")
+      )
+
+  private def isBlankFunc(usedRefs: Set[String], refsCosts: Map[String, Long]): Boolean =
+    !usedRefs.exists(refsCosts.get(_).exists(_ > 0))
+
+  private def evalHoldingFuncs(expr: EXPR, activeFuncArgs: Set[String]): EvalM[Long] =
+    for {
+      startCtx <- get[Id, EstimatorContext, EstimationError]
+      cost     <- evalExpr(expr, activeFuncArgs)
+      _        <- update(funcs.set(_)(startCtx.funcs))
+    } yield cost
+
+  private def withUsedRefs[A](eval: EvalM[A]): EvalM[(A, Set[String])] =
+    for {
+      ctxBefore <- get[Id, EstimatorContext, EstimationError]
+      result    <- eval
+      ctxAfter  <- get[Id, EstimatorContext, EstimationError]
+    } yield (result, ctxAfter.usedRefs diff ctxBefore.usedRefs)
+
   private def update(f: EstimatorContext => EstimatorContext): EvalM[Unit] =
     modify[Id, EstimatorContext, EstimationError](f)
 
@@ -192,3 +232,7 @@ case class ScriptEstimatorV3(fixOverflow: Boolean, overhead: Boolean) extends Sc
     liftEither(Try(r).toEither.leftMap(_ => "Illegal script"))
   }
 }
+
+object ScriptEstimatorV3 {
+  val latest = ScriptEstimatorV3(fixOverflow = true, overhead = false, letFixes = true)
+}

lang/tests/src/test/scala/com/wavesplatform/lang/compiler/ContractCompilerTest.scala

@@ -883,7 +883,7 @@ class ContractCompilerTest extends PropSpec {
       NoLibraries,
       dAppV4Ctx,
       V4,
-      ScriptEstimatorV3(fixOverflow = true, overhead = true),
+      ScriptEstimatorV3(fixOverflow = true, overhead = true, letFixes = true),
       needCompaction = false,
       removeUnusedCode = false
     ) should produce("Script is too large: 165187 bytes > 163840 bytes")
@@ -1042,7 +1042,7 @@ class ContractCompilerTest extends PropSpec {
       NoLibraries,
       getTestContext(V4).compilerContext,
       V4,
-      ScriptEstimatorV3(fixOverflow = true, overhead = true),
+      ScriptEstimatorV3(fixOverflow = true, overhead = true, letFixes = true),
       needCompaction = false,
       removeUnusedCode = false
     ) shouldBe Symbol("right")

lang/tests/src/test/scala/com/wavesplatform/lang/compiler/ExpressionCompilerV1Test.scala

@@ -323,7 +323,7 @@ class ExpressionCompilerV1Test extends PropSpec {
       )
       .compilerContext
 
-    val e = ScriptEstimatorV3(fixOverflow = true, overhead = true)
+    val e = ScriptEstimatorV3(fixOverflow = true, overhead = true, letFixes = true)
     Global.compileExpression(expr, NoLibraries, ctx, V4, Account, e) should produce("Script is too large: 8756 bytes > 8192 bytes")
   }
 

lang/tests/src/test/scala/com/wavesplatform/lang/doc/FunctionComplexityDocTest.scala

@@ -46,7 +46,7 @@ class FunctionComplexityDocTest extends PropSpec {
           val arg  = CONST_STRING("throw").explicitGet()
           val expr = FUNCTION_CALL(function.header, List.fill(function.args.size)(arg))
           val estimatedCost =
-            ScriptEstimatorV3(fixOverflow = true, overhead = false)(
+            ScriptEstimatorV3(fixOverflow = true, overhead = false, letFixes = true)(
               varNames(ds.stdLibVersion, ds.contentType),
               functionCosts(ds.stdLibVersion, ds.contentType),
               expr

lang/tests/src/test/scala/com/wavesplatform/lang/estimator/CommonScriptEstimatorTest.scala

@@ -11,10 +11,11 @@ class CommonScriptEstimatorTest
     extends ScriptEstimatorTestBase(
       ScriptEstimatorV1,
       ScriptEstimatorV2,
-      ScriptEstimatorV3(fixOverflow = true, overhead = true),
-      ScriptEstimatorV3(fixOverflow = true, overhead = false),
-      ScriptEstimatorV3(fixOverflow = false, overhead = true),
-      ScriptEstimatorV3(fixOverflow = false, overhead = false),
+      ScriptEstimatorV3(fixOverflow = true, overhead = true, letFixes = false),
+      ScriptEstimatorV3(fixOverflow = true, overhead = false, letFixes = false),
+      ScriptEstimatorV3(fixOverflow = false, overhead = true, letFixes = false),
+      ScriptEstimatorV3(fixOverflow = false, overhead = false, letFixes = false),
+      ScriptEstimatorV3(fixOverflow = true, overhead = false, letFixes = true),   // <- all fixes
       evaluatorV2AsEstimator(overhead = true),
       evaluatorV2AsEstimator(overhead = false)
     ) {