From a5321af4e41f6f6d25fceaec079d16a8786caf38 Mon Sep 17 00:00:00 2001 From: Alexander Slesarenko Date: Tue, 28 May 2024 20:44:27 +0200 Subject: [PATCH] v6.0-serialize: added CoreByteWriter.checkCostLimit --- .../sigma/serialization/CoreByteWriter.scala | 150 ++++++++++++++---- .../sigma/serialization/CoreSerializer.scala | 2 +- .../sigma/serialization/SigmaByteWriter.scala | 65 ++++++-- .../sigma/serialization/SigmaSerializer.scala | 7 +- 4 files changed, 178 insertions(+), 46 deletions(-) diff --git a/core/shared/src/main/scala/sigma/serialization/CoreByteWriter.scala b/core/shared/src/main/scala/sigma/serialization/CoreByteWriter.scala index aa4255449c..302663ba92 100644 --- a/core/shared/src/main/scala/sigma/serialization/CoreByteWriter.scala +++ b/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,49 +53,105 @@ 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 @@ -81,12 +159,18 @@ class CoreByteWriter(val w: Writer) extends 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 } @@ -94,19 +178,28 @@ class CoreByteWriter(val w: Writer) extends Writer { 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 } } @@ -114,6 +207,9 @@ class CoreByteWriter(val w: Writer) extends Writer { 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 */ diff --git a/core/shared/src/main/scala/sigma/serialization/CoreSerializer.scala b/core/shared/src/main/scala/sigma/serialization/CoreSerializer.scala index 938d3f22c1..2ec3ee5c1a 100644 --- a/core/shared/src/main/scala/sigma/serialization/CoreSerializer.scala +++ b/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 } diff --git a/data/shared/src/main/scala/sigma/serialization/SigmaByteWriter.scala b/data/shared/src/main/scala/sigma/serialization/SigmaByteWriter.scala index 35d5e0c9b9..87b115e942 100644 --- a/data/shared/src/main/scala/sigma/serialization/SigmaByteWriter.scala +++ b/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) diff --git a/data/shared/src/main/scala/sigma/serialization/SigmaSerializer.scala b/data/shared/src/main/scala/sigma/serialization/SigmaSerializer.scala index 3765adb029..e842258f80 100644 --- a/data/shared/src/main/scala/sigma/serialization/SigmaSerializer.scala +++ b/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 } }