v6.0-serialize: added CoreByteWriter.checkCostLimit

core/shared/src/main/scala/sigma/serialization/CoreByteWriter.scala

@@ -3,25 +3,47 @@ package sigma.serialization
 import scorex.util.serialization.Writer.Aux
 import scorex.util.serialization.{VLQByteBufferWriter, Writer}
 import sigma.ast.SType
-import sigma.serialization.CoreByteWriter.{Bits, DataInfo, U, Vlq, ZigZag}
+import sigma.serialization.CoreByteWriter.{Bits, CostLimitChecker, DataInfo, U, Vlq, ZigZag}
 
 /** Implementation of [[Writer]] provided by `sigma-core` module.
-  * @param w destination [[Writer]] to which all the call got delegated.
+  *
+  * @param w              destination [[Writer]] to which all the call got delegated.
+  * @param checkCostLimitOpt callback to check if the cost limit at current writing position
+  *                       is reached. The callback will throw an exception if the limit is reached.
+  *                       Note, the cost of serialization is approximated using formula
+  *                       `this.length * CostPerByte + InitialCost`
   */
-class CoreByteWriter(val w: Writer) extends Writer {
+class CoreByteWriter(val w: Writer, val checkCostLimitOpt: Option[CostLimitChecker]) extends Writer {
   type CH = w.CH
 
+  /** Check the current writer length against the cost limit. */
+  @inline protected def checkCostLimit(): Unit = {
+    if (checkCostLimitOpt.isDefined)
+      checkCostLimitOpt.get(this.length())
+  }
+
   @inline override def length(): Int = w.length()
 
   @inline override def newWriter(): Aux[CH] = w.newWriter()
 
-  @inline override def putChunk(chunk: CH): this.type = { w.putChunk(chunk); this }
+  @inline override def putChunk(chunk: CH): this.type = {
+    w.putChunk(chunk);
+    checkCostLimit()
+    this
+  }
 
   @inline override def result(): CH = w.result()
 
-  @inline def put(x: Byte): this.type = { w.put(x); this }
+  @inline def put(x: Byte): this.type = {
+    w.put(x);
+    checkCostLimit()
+    this
+  }
+
   @inline def put(x: Byte, info: DataInfo[Byte]): this.type = {
-    w.put(x); this
+    w.put(x);
+    checkCostLimit()
+    this
   }
 
   override def putUByte(x: Int): this.type = {
@@ -31,89 +53,163 @@ class CoreByteWriter(val w: Writer) extends Writer {
     super.putUByte(x)
   }
 
-  @inline def putBoolean(x: Boolean): this.type = { w.putBoolean(x); this }
+  @inline def putBoolean(x: Boolean): this.type = {
+    w.putBoolean(x);
+    checkCostLimit()
+    this
+  }
+
   @inline def putBoolean(x: Boolean, info: DataInfo[Boolean]): this.type = {
-    w.putBoolean(x); this
+    w.putBoolean(x);
+    checkCostLimit()
+    this
+  }
+
+  @inline def putShort(x: Short): this.type = {
+    w.putShort(x);
+    checkCostLimit()
+    this
   }
 
-  @inline def putShort(x: Short): this.type = { w.putShort(x); this }
   @inline def putShort(x: Short, info: DataInfo[Short]): this.type = {
-    w.putShort(x); this
+    w.putShort(x);
+    checkCostLimit()
+    this
+  }
+
+  @inline def putUShort(x: Int): this.type = {
+    w.putUShort(x);
+    checkCostLimit()
+    this
   }
 
-  @inline def putUShort(x: Int): this.type = { w.putUShort(x); this }
   @inline def putUShort(x: Int, info: DataInfo[Vlq[U[Short]]]): this.type = {
-    w.putUShort(x); this
+    w.putUShort(x);
+    checkCostLimit()
+    this
+  }
+
+  @inline def putInt(x: Int): this.type = {
+    w.putInt(x);
+    checkCostLimit()
+    this
   }
 
-  @inline def putInt(x: Int): this.type = { w.putInt(x); this }
   @inline def putInt(x: Int, info: DataInfo[Int]): this.type = {
-    w.putInt(x); this
+    w.putInt(x);
+    checkCostLimit()
+    this
   }
 
-  @inline def putUInt(x: Long): this.type = { w.putUInt(x); this }
+  @inline def putUInt(x: Long): this.type = {
+    w.putUInt(x);
+    checkCostLimit()
+    this
+  }
   @inline def putUInt(x: Long, info: DataInfo[Vlq[U[Int]]]): this.type = {
-    w.putUInt(x); this
+    w.putUInt(x);
+    checkCostLimit()
+    this
   }
 
-  @inline def putLong(x: Long): this.type = { w.putLong(x); this }
+  @inline def putLong(x: Long): this.type = {
+    w.putLong(x);
+    checkCostLimit()
+    this
+  }
   @inline def putLong(x: Long, info: DataInfo[Vlq[ZigZag[Long]]]): this.type = {
-    w.putLong(x); this
+    w.putLong(x);
+    checkCostLimit()
+    this
+  }
+
+  @inline def putULong(x: Long): this.type = {
+    w.putULong(x);
+    checkCostLimit()
+    this
   }
 
-  @inline def putULong(x: Long): this.type = { w.putULong(x); this }
   @inline def putULong(x: Long, info: DataInfo[Vlq[U[Long]]]): this.type = {
-    w.putULong(x); this
+    w.putULong(x);
+    checkCostLimit()
+    this
   }
 
   override def putBytes(xs: Array[Byte],
                         offset: Int,
                         length: Int): this.type = {
-    w.putBytes(xs, offset, length); this
+    w.putBytes(xs, offset, length);
+    checkCostLimit()
+    this
+  }
+
+  @inline def putBytes(xs: Array[Byte]): this.type = {
+    w.putBytes(xs);
+    checkCostLimit()
+    this
   }
-  @inline def putBytes(xs: Array[Byte]): this.type = { w.putBytes(xs); this }
   @inline def putBytes(xs: Array[Byte], info: DataInfo[Array[Byte]]): this.type = {
-    w.putBytes(xs); this
+    w.putBytes(xs);
+    checkCostLimit()
+    this
   }
 
   /** Put the two bytes of the big-endian representation of the Short value into the
    * writer. */
   @inline def putShortBytes(value: Short): this.type = {
     w.put((value >> 8).toByte)
     w.put(value.toByte)
+    checkCostLimit()
     this
   }
 
-  @inline def putBits(xs: Array[Boolean]): this.type = { w.putBits(xs); this }
+  @inline def putBits(xs: Array[Boolean]): this.type = {
+    w.putBits(xs);
+    checkCostLimit()
+    this
+  }
   @inline def putBits(xs: Array[Boolean], info: DataInfo[Bits]): this.type = {
     w.putBits(xs);
+    checkCostLimit()
     this
   }
 
   @inline def putOption[T](x: Option[T])(putValueC: (this.type, T) => Unit): this.type = {
     w.putOption(x) { (_, v) =>
       putValueC(this, v)
     }
+    checkCostLimit()
     this
   }
 
-  @inline def putShortString(s: String): this.type = { w.putShortString(s); this }
+  @inline def putShortString(s: String): this.type = {
+    w.putShortString(s);
+    checkCostLimit()
+    this
+  }
 
   // TODO refactor: move to Writer
   @inline def toBytes: Array[Byte] = w match {
     case wr: VLQByteBufferWriter => wr.toBytes
   }
 
-  @inline def putType[T <: SType](x: T): this.type = { TypeSerializer.serialize(x, this); this }
+  @inline def putType[T <: SType](x: T): this.type = {
+    TypeSerializer.serialize(x, this); // the cost is checked in TypeSerializer
+    this
+  }
   @inline def putType[T <: SType](x: T, info: DataInfo[SType]): this.type = {
-    TypeSerializer.serialize(x, this); this
+    TypeSerializer.serialize(x, this); // the cost is checked in TypeSerializer
+    this
   }
 
 }
 
 object CoreByteWriter {
   import scala.language.implicitConversions
 
+  /** Callback type of cost limit checker. */
+  type CostLimitChecker = Int => Unit
+
   /** Format descriptor type family. */
   trait FormatDescriptor[T] {
     /** Size formula associated with this format */

core/shared/src/main/scala/sigma/serialization/CoreSerializer.scala

@@ -66,7 +66,7 @@ object CoreSerializer {
   def startWriter(): CoreByteWriter = {
     val b = new ByteArrayBuilder()
     val wi = new VLQByteBufferWriter(b)
-    val w = new CoreByteWriter(wi)
+    val w = new CoreByteWriter(wi, None)
     w
   }
 

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

@@ -4,15 +4,28 @@ import scorex.util.serialization.Writer
 import sigma.ast.syntax._
 import sigma.ast._
 import sigma.serialization.CoreByteWriter.{ArgInfo, DataInfo, FormatDescriptor, SeqFmt}
-
-class SigmaByteWriter(override val w: Writer,
-                      val constantExtractionStore: Option[ConstantStore])
-    extends CoreByteWriter(w) {
+import CoreByteWriter.CostLimitChecker
+
+/** Implementation of [[Writer]] provided by `sigma-data` module.
+  *
+  * @param w                       destination [[Writer]] to which all the call got delegated.
+  * @param constantExtractionStore optional store to segregate constants to while
+  *                                replacing them with placeholders.
+  * @param checkCostLimitOpt          callback to check if the cost limit at current writing position
+  *                                is reached. The callback will throw an exception if the limit is reached.
+  */
+class SigmaByteWriter(
+  override val w: Writer,
+  val constantExtractionStore: Option[ConstantStore],
+  override val checkCostLimitOpt: Option[CostLimitChecker]
+)
+    extends CoreByteWriter(w, checkCostLimitOpt) {
   import CoreByteWriter._
   import ValueSerializer._
 
   override def put(x: Byte, info: DataInfo[Byte]): this.type = {
     ValueSerializer.addArgInfo(info)
+    checkCostLimit()
     w.put(x); this
   }
 
@@ -23,59 +36,81 @@ class SigmaByteWriter(override val w: Writer,
 
   @inline override def putBoolean(x: Boolean, info: DataInfo[Boolean]): this.type = {
     ValueSerializer.addArgInfo(info)
-    w.putBoolean(x); this
+    w.putBoolean(x);
+    checkCostLimit()
+    this
   }
 
   @inline override def putShort(x: Short, info: DataInfo[Short]): this.type = {
     ValueSerializer.addArgInfo(info)
-    w.putShort(x); this
+    w.putShort(x);
+    checkCostLimit()
+    this
   }
 
   @inline override def putUShort(x: Int, info: DataInfo[Vlq[U[Short]]]): this.type = {
     ValueSerializer.addArgInfo(info)
-    w.putUShort(x); this
+    w.putUShort(x);
+    checkCostLimit()
+    this
   }
 
   @inline override def putInt(x: Int, info: DataInfo[Int]): this.type = {
     ValueSerializer.addArgInfo(info)
-    w.putInt(x); this
+    w.putInt(x);
+    checkCostLimit()
+    this
   }
 
   @inline override def putUInt(x: Long, info: DataInfo[Vlq[U[Int]]]): this.type = {
     ValueSerializer.addArgInfo(info)
-    w.putUInt(x); this
+    w.putUInt(x);
+    checkCostLimit()
+    this
   }
 
   @inline override def putLong(x: Long, info: DataInfo[Vlq[ZigZag[Long]]]): this.type = {
     ValueSerializer.addArgInfo(info)
-    w.putLong(x); this
+    w.putLong(x);
+    checkCostLimit()
+    this
   }
 
   @inline override def putULong(x: Long, info: DataInfo[Vlq[U[Long]]]): this.type = {
     ValueSerializer.addArgInfo(info)
-    w.putULong(x); this
+    w.putULong(x);
+    checkCostLimit()
+    this
   }
 
   @inline override def putBytes(xs: Array[Byte], info: DataInfo[Array[Byte]]): this.type = {
     ValueSerializer.addArgInfo(info)
-    w.putBytes(xs); this
+    w.putBytes(xs);
+    checkCostLimit()
+    this
   }
 
   @inline override def putBits(xs: Array[Boolean], info: DataInfo[Bits]): this.type = {
     ValueSerializer.addArgInfo(info)
     w.putBits(xs);
+    checkCostLimit()
     this
   }
 
   @inline override def putType[T <: SType](x: T, info: DataInfo[SType]): this.type = {
     ValueSerializer.addArgInfo(info)
-    TypeSerializer.serialize(x, this); this
+    TypeSerializer.serialize(x, this); // the cost is checked in TypeSerializer
+    this
   }
 
-  @inline def putValue[T <: SType](x: Value[T]): this.type = { ValueSerializer.serialize(x, this); this }
+  @inline def putValue[T <: SType](x: Value[T]): this.type = {
+    ValueSerializer.serialize(x, this) // the cost is checked in ValueSerializer
+    this
+  }
   @inline def putValue[T <: SType](x: Value[T], info: DataInfo[SValue]): this.type = {
     ValueSerializer.addArgInfo(info)
-    ValueSerializer.serialize(x, this); this
+    ValueSerializer.serialize(x, this); // the cost is checked in ValueSerializer
+    this
   }
   @inline def putValues[T <: SType](xs: Seq[Value[T]]): this.type = {
     putUInt(xs.length)

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

@@ -4,6 +4,7 @@ import java.nio.ByteBuffer
 import scorex.util.ByteArrayBuilder
 import scorex.util.serialization._
 import sigma.data.SigmaConstants
+import sigma.serialization.CoreByteWriter.CostLimitChecker
 import sigma.serialization.ValueCodes.OpCode
 
 object SigmaSerializer {
@@ -51,14 +52,14 @@ object SigmaSerializer {
   def startWriter(): SigmaByteWriter = {
     val b = new ByteArrayBuilder()
     val wi = new VLQByteBufferWriter(b)
-    val w = new SigmaByteWriter(wi, constantExtractionStore = None)
+    val w = new SigmaByteWriter(wi, constantExtractionStore = None, checkCostLimitOpt = None)
     w
   }
 
-  def startWriter(constantExtractionStore: ConstantStore): SigmaByteWriter = {
+  def startWriter(constantExtractionStore: ConstantStore, checkCostLimit: Option[CostLimitChecker] = None): SigmaByteWriter = {
     val b = new ByteArrayBuilder()
     val wi = new VLQByteBufferWriter(b)
-    val w = new SigmaByteWriter(wi, constantExtractionStore = Some(constantExtractionStore))
+    val w = new SigmaByteWriter(wi, constantExtractionStore = Some(constantExtractionStore), checkCostLimit)
     w
   }
 }