Merge pull request #995 from ScorexFoundation/i994-fix-subst-constants

[6.0] Fix SubstConstants: serialize ErgoTree size

data/shared/src/main/scala/sigma/ast/SigmaPredef.scala

@@ -5,7 +5,6 @@ import org.ergoplatform.{ErgoAddressEncoder, P2PKAddress}
 import scorex.util.encode.{Base16, Base58, Base64}
 import sigma.ast.SCollection.{SByteArray, SIntArray}
 import sigma.ast.SOption.SIntOption
-import sigma.ast.SigmaPropConstant
 import sigma.ast.syntax._
 import sigma.data.Nullable
 import sigma.exceptions.InvalidArguments
@@ -544,7 +543,7 @@ object SigmaPredef {
 
     val funcs: Map[String, PredefinedFunc] = globalFuncs ++ infixFuncs ++ unaryFuncs
 
-    /** WARNING: This operations are not used in frontend, and should be be used.
+    /** WARNING: This operations are not used in frontend, and should not be used.
       * They are used in SpecGen only the source of metadata for the corresponding ErgoTree nodes.
       */
     val specialFuncs: Map[String, PredefinedFunc] = Seq(

data/shared/src/main/scala/sigma/serialization/ErgoTreeSerializer.scala

@@ -287,6 +287,9 @@ class ErgoTreeSerializer {
     * allow to use serialized scripts as pre-defined templates.
     * See [[SubstConstants]] for details.
     *
+    * Note, this operation doesn't require (de)serialization of ErgoTree expression,
+    * thus it is more efficient than serialization roundtrip.
+    *
     * @param scriptBytes serialized ErgoTree with ConstantSegregationFlag set to 1.
     * @param positions   zero based indexes in ErgoTree.constants array which
     *                    should be replaced with new values
@@ -304,39 +307,62 @@ class ErgoTreeSerializer {
       s"expected positions and newVals to have the same length, got: positions: ${positions.toSeq},\n newVals: ${newVals.toSeq}")
     val r = SigmaSerializer.startReader(scriptBytes)
     val (header, _, constants, treeBytes) = deserializeHeaderWithTreeBytes(r)
-    val w = SigmaSerializer.startWriter()
-    w.put(header)
+    val nConstants = constants.length
+
+    val resBytes = if (VersionContext.current.isJitActivated) {
+      // need to measure the serialized size of the new constants
+      // by serializing them into a separate writer
+      val constW = SigmaSerializer.startWriter()
 
-    if (VersionContext.current.isJitActivated) {
       // The following `constants.length` should not be serialized when segregation is off
       // in the `header`, because in this case there is no `constants` section in the
       // ErgoTree serialization format. Thus, applying this `substituteConstants` for
       // non-segregated trees will return non-parsable ErgoTree bytes (when
       // `constants.length` is put in `w`).
       if (ErgoTree.isConstantSegregation(header)) {
-        w.putUInt(constants.length)
+        constW.putUInt(constants.length)
       }
 
       // The following is optimized O(nConstants + position.length) implementation
-      val nConstants = constants.length
       if (nConstants > 0) {
         val backrefs = getPositionsBackref(positions, nConstants)
         cfor(0)(_ < nConstants, _ + 1) { i =>
           val c = constants(i)
           val iPos = backrefs(i) // index to `positions`
           if (iPos == -1) {
             // no position => no substitution, serialize original constant
-            constantSerializer.serialize(c, w)
+            constantSerializer.serialize(c, constW)
           } else {
-            assert(positions(iPos) == i) // INV: backrefs and positions are mutually inverse
+            require(positions(iPos) == i) // INV: backrefs and positions are mutually inverse
             val newConst = newVals(iPos)
             require(c.tpe == newConst.tpe,
               s"expected new constant to have the same ${c.tpe} tpe, got ${newConst.tpe}")
-            constantSerializer.serialize(newConst, w)
+            constantSerializer.serialize(newConst, constW)
           }
         }
       }
+
+      val constBytes = constW.toBytes // nConstants + serialized new constants
+
+      // start composing the resulting tree bytes
+      val w = SigmaSerializer.startWriter()
+      w.put(header) // header byte
+
+      if (VersionContext.current.isV6SoftForkActivated) {
+        // fix in v6.0 to save tree size to respect size bit of the original tree
+        if (ErgoTree.hasSize(header)) {
+          val size = constBytes.length + treeBytes.length
+          w.putUInt(size)     // tree size
+        }
+      }
+
+      w.putBytes(constBytes)  // constants section
+      w.putBytes(treeBytes)   // tree section
+      w.toBytes
     } else {
+      val w = SigmaSerializer.startWriter()
+      w.put(header)
+
       // for v4.x compatibility we save constants.length here (see the above comment to
       // understand the consequences)
       w.putUInt(constants.length)
@@ -357,10 +383,12 @@ class ErgoTreeSerializer {
         case (c, _) =>
           constantSerializer.serialize(c, w)
       }
+
+      w.putBytes(treeBytes)
+      w.toBytes
     }
 
-    w.putBytes(treeBytes)
-    (w.toBytes, constants.length)
+    (resBytes, nConstants)
   }
 
 }

sc/shared/src/test/scala/sigma/LanguageSpecificationBase.scala

@@ -1,10 +1,11 @@
 package sigma
 
 import org.scalatest.BeforeAndAfterAll
-import sigma.ast.JitCost
-import sigma.eval.{EvalSettings, Profiler}
+import sigma.ast.{Apply, FixedCostItem, FuncValue, GetVar, JitCost, OptionGet, ValUse}
+import sigma.eval.{CostDetails, EvalSettings, Profiler}
 import sigmastate.CompilerCrossVersionProps
 import sigmastate.interpreter.CErgoTreeEvaluator
+
 import scala.util.Success
 
 /** Base class for language test suites (one suite for each language version: 5.0, 6.0, etc.)
@@ -123,4 +124,23 @@ abstract class LanguageSpecificationBase extends SigmaDslTesting
     prepareSamples[(PreHeader, PreHeader)]
     prepareSamples[(Header, Header)]
   }
+
+  ///=====================================================
+  ///         CostDetails shared among test cases
+  ///-----------------------------------------------------
+  val traceBase = Array(
+    FixedCostItem(Apply),
+    FixedCostItem(FuncValue),
+    FixedCostItem(GetVar),
+    FixedCostItem(OptionGet),
+    FixedCostItem(FuncValue.AddToEnvironmentDesc, FuncValue.AddToEnvironmentDesc_CostKind),
+    FixedCostItem(ValUse)
+  )
+
+  /** Helper method to create the given expected results for all tree versions. */
+  def expectedSuccessForAllTreeVersions[A](value: A, cost: Int, costDetails: CostDetails) = {
+    val res = ExpectedResult(Success(value), Some(cost)) -> Some(costDetails)
+    Seq(0, 1, 2, 3).map(version => version -> res)
+  }
+
 }