insertOrUpdate

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

@@ -1643,23 +1643,63 @@ case object SAvlTreeMethods extends MonoTypeMethods {
     OperationCostInfo(m.costKind.asInstanceOf[FixedCost], m.opDesc)
   }
 
-  protected override def getMethods(): Seq[SMethod] = super.getMethods() ++ Seq(
-    digestMethod,
-    enabledOperationsMethod,
-    keyLengthMethod,
-    valueLengthOptMethod,
-    isInsertAllowedMethod,
-    isUpdateAllowedMethod,
-    isRemoveAllowedMethod,
-    updateOperationsMethod,
-    containsMethod,
-    getMethod,
-    getManyMethod,
-    insertMethod,
-    updateMethod,
-    removeMethod,
-    updateDigestMethod
-  )
+  // 6.0 methods below
+  lazy val insertOrUpdateMethod = SMethod(this, "insertOrUpdate",
+    SFunc(Array(SAvlTree, CollKeyValue, SByteArray), SAvlTreeOption), 16, DynamicCost)
+    .withIRInfo(MethodCallIrBuilder)
+    .withInfo(MethodCall,
+      """
+        |  /** Perform insertions of key-value entries into this tree using proof `proof`.
+        |    * Throws exception if proof is incorrect
+        |    *
+        |    * @note CAUTION! Pairs must be ordered the same way they were in insert ops before proof was generated.
+        |    * Return Some(newTree) if successful
+        |    * Return None if operations were not performed.
+        |    * @param operations   collection of key-value pairs to insert in this authenticated dictionary.
+        |    * @param proof
+        |    */
+        |
+        """.stripMargin)
+
+  /** Implements evaluation of AvlTree.insert method call ErgoTree node.
+    * Called via reflection based on naming convention.
+    * @see SMethod.evalMethod
+    */
+  def insertOrUpdate_eval(mc: MethodCall, tree: AvlTree, entries: KeyValueColl, proof: Coll[Byte])
+                         (implicit E: ErgoTreeEvaluator): Option[AvlTree] = {
+    E.insertOrUpdate_eval(mc, tree, entries, proof)
+  }
+
+  lazy val v5Methods = {
+    super.getMethods() ++ Seq(
+      digestMethod,
+      enabledOperationsMethod,
+      keyLengthMethod,
+      valueLengthOptMethod,
+      isInsertAllowedMethod,
+      isUpdateAllowedMethod,
+      isRemoveAllowedMethod,
+      updateOperationsMethod,
+      containsMethod,
+      getMethod,
+      getManyMethod,
+      insertMethod,
+      updateMethod,
+      removeMethod,
+      updateDigestMethod
+    )
+  }
+
+  lazy val v6Methods = v5Methods ++ Seq(insertOrUpdateMethod)
+
+  protected override def getMethods(): Seq[SMethod] = {
+    if (VersionContext.current.isV6SoftForkActivated) {
+      v6Methods
+    } else {
+      v5Methods
+    }
+  }
+
 }
 
 /** Type descriptor of `Context` type of ErgoTree. */

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

@@ -8,7 +8,7 @@ import sigma.ast.TypeCodes.ConstantCode
 import sigma.ast.syntax._
 import sigma.crypto.{CryptoConstants, EcPointType}
 import sigma.data.OverloadHack.Overloaded1
-import sigma.data.{CSigmaDslBuilder, CSigmaProp, Nullable, RType, SigmaBoolean}
+import sigma.data.{AvlTreeData, CAvlTree, CSigmaDslBuilder, CSigmaProp, Nullable, RType, SigmaBoolean}
 import sigma.eval.ErgoTreeEvaluator.DataEnv
 import sigma.eval.{ErgoTreeEvaluator, SigmaDsl}
 import sigma.exceptions.InterpreterException
@@ -536,6 +536,7 @@ object SigmaPropConstant {
 
 object AvlTreeConstant {
   def apply(value: AvlTree): Constant[SAvlTree.type] = Constant[SAvlTree.type](value, SAvlTree)
+  def apply(value: AvlTreeData): Constant[SAvlTree.type] = Constant[SAvlTree.type](CAvlTree(value), SAvlTree)
 }
 
 object PreHeaderConstant {

data/shared/src/main/scala/sigma/eval/AvlTreeVerifier.scala

@@ -50,6 +50,18 @@ trait AvlTreeVerifier {
     */
   def performUpdate(key: Array[Byte], value: Array[Byte]): Try[Option[Array[Byte]]]
 
+  /**
+    * Returns Failure if the proof does not verify.
+    * Otherwise, successfully modifies tree and so returns Success.
+    * After one failure, all subsequent operations with this verifier will fail and digest
+    * is None.
+    *
+    * @param key   key to look up
+    * @param value value to check it was updated
+    * @return Success(Some(value)), Success(None), or Failure
+    */
+  def performInsertOrUpdate(key: Array[Byte], value: Array[Byte]): Try[Option[Array[Byte]]]
+
   /** Check the key has been removed in the tree.
     * If `key` exists in the tree and the operation succeeds,
     * returns `Success(Some(v))`, where v is old value associated with `key`.

data/shared/src/main/scala/sigma/eval/ErgoTreeEvaluator.scala

@@ -134,6 +134,13 @@ abstract class ErgoTreeEvaluator {
       mc: MethodCall, tree: AvlTree,
       operations: KeyValueColl, proof: Coll[Byte]): Option[AvlTree]
 
+  /** Implements evaluation of AvlTree.insert method call ErgoTree node. */
+  def insertOrUpdate_eval(
+                   mc: MethodCall,
+                   tree: AvlTree,
+                   entries: KeyValueColl,
+                   proof: Coll[Byte]): Option[AvlTree]
+
   /** Implements evaluation of AvlTree.remove method call ErgoTree node. */
   def remove_eval(
       mc: MethodCall, tree: AvlTree,

interpreter/shared/src/main/scala/sigmastate/eval/CAvlTreeVerifier.scala

@@ -1,6 +1,6 @@
 package sigmastate.eval
 
-import scorex.crypto.authds.avltree.batch.{BatchAVLVerifier, Insert, Lookup, Remove, Update}
+import scorex.crypto.authds.avltree.batch.{BatchAVLVerifier, Insert, InsertOrUpdate, Lookup, Remove, Update}
 import scorex.crypto.authds.{ADDigest, ADKey, ADValue, SerializedAdProof}
 import scorex.crypto.hash.{Blake2b256, Digest32}
 import sigma.data.CAvlTree
@@ -32,6 +32,9 @@ class CAvlTreeVerifier private(
   override def performUpdate(key: Array[Byte], value: Array[Byte]): Try[Option[Array[Byte]]] =
     performOneOperation(Update(ADKey @@ key, ADValue @@ value))
 
+  override def performInsertOrUpdate(key: Array[Byte], value: Array[Byte]): Try[Option[Array[Byte]]] =
+    performOneOperation(InsertOrUpdate(ADKey @@ key, ADValue @@ value))
+
   override def performRemove(key: Array[Byte]): Try[Option[Array[Byte]]] =
     performOneOperation(Remove(ADKey @@ key))
 

interpreter/shared/src/main/scala/sigmastate/interpreter/CErgoTreeEvaluator.scala

@@ -173,7 +173,7 @@ class CErgoTreeEvaluator(
       // when the tree is empty we still need to add the insert cost
       val nItems = Math.max(bv.treeHeight, 1)
 
-      // here we use forall as looping with fast break on first failed tree oparation
+      // here we use forall as looping with fast break on first failed tree operation
       operations.forall { case (key, value) =>
         var res = true
         // the cost of tree update is O(bv.treeHeight)
@@ -192,6 +192,37 @@ class CErgoTreeEvaluator(
     }
   }
 
+  override def insertOrUpdate_eval(
+                            mc: MethodCall, tree: AvlTree,
+                            operations: KeyValueColl, proof: Coll[Byte]): Option[AvlTree] = {
+    addCost(isUpdateAllowed_Info)
+    addCost(isInsertAllowed_Info)
+    if (!(tree.isUpdateAllowed && tree.isInsertAllowed)) {
+      None
+    } else {
+      val bv     = createVerifier(tree, proof)
+      // when the tree is empty we still need to add the insert cost
+      val nItems = Math.max(bv.treeHeight, 1)
+
+      // here we use forall as looping with fast break on first failed tree operation
+      operations.forall { case (key, value) =>
+        var res = true
+        // the cost of tree update is O(bv.treeHeight)
+        addSeqCost(UpdateAvlTree_Info, nItems) { () =>
+          val updateRes = bv.performInsertOrUpdate(key.toArray, value.toArray)
+          res = updateRes.isSuccess
+        }
+        res
+      }
+      bv.digest match {
+        case Some(d) =>
+          addCost(updateDigest_Info)
+          Some(tree.updateDigest(Colls.fromArray(d)))
+        case _ => None
+      }
+    }
+  }
+
   override def remove_eval(
       mc: MethodCall, tree: AvlTree,
       operations: Coll[Coll[Byte]], proof: Coll[Byte]): Option[AvlTree] = {

sc/shared/src/main/scala/sigma/compiler/ir/GraphBuilding.scala

@@ -1119,6 +1119,10 @@ trait GraphBuilding extends Base with DefRewriting { IR: IRContext =>
               val operations = asRep[Coll[(Coll[Byte], Coll[Byte])]](argsV(0))
               val proof = asRep[Coll[Byte]](argsV(1))
               tree.update(operations, proof)
+            case SAvlTreeMethods.insertOrUpdateMethod.name =>
+              val operations = asRep[Coll[(Coll[Byte], Coll[Byte])]](argsV(0))
+              val proof = asRep[Coll[Byte]](argsV(1))
+              tree.insertOrUpdate(operations, proof)
             case _ => throwError()
           }
           case (ph: Ref[PreHeader]@unchecked, SPreHeaderMethods) => method.name match {

sc/shared/src/main/scala/sigma/compiler/ir/wrappers/sigma/SigmaDslUnit.scala

@@ -51,6 +51,7 @@ import scalan._
       def getMany(keys: Ref[Coll[Coll[Byte]]], proof: Ref[Coll[Byte]]): Ref[Coll[WOption[Coll[Byte]]]];
       def insert(operations: Ref[Coll[scala.Tuple2[Coll[Byte], Coll[Byte]]]], proof: Ref[Coll[Byte]]): Ref[WOption[AvlTree]];
       def update(operations: Ref[Coll[scala.Tuple2[Coll[Byte], Coll[Byte]]]], proof: Ref[Coll[Byte]]): Ref[WOption[AvlTree]];
+      def insertOrUpdate(operations: Ref[Coll[scala.Tuple2[Coll[Byte], Coll[Byte]]]], proof: Ref[Coll[Byte]]): Ref[WOption[AvlTree]];
       def remove(operations: Ref[Coll[Coll[Byte]]], proof: Ref[Coll[Byte]]): Ref[WOption[AvlTree]]
     };
     trait PreHeader extends Def[PreHeader] {

sc/shared/src/main/scala/sigma/compiler/ir/wrappers/sigma/impl/SigmaDslImpl.scala

@@ -932,6 +932,13 @@ object AvlTree extends EntityObject("AvlTree") {
         true, false, element[WOption[AvlTree]]))
     }
 
+    override def insertOrUpdate(operations: Ref[Coll[(Coll[Byte], Coll[Byte])]], proof: Ref[Coll[Byte]]): Ref[WOption[AvlTree]] = {
+      asRep[WOption[AvlTree]](mkMethodCall(self,
+        AvlTreeClass.getMethod("insertOrUpdate", classOf[Sym], classOf[Sym]),
+        Array[AnyRef](operations, proof),
+        true, false, element[WOption[AvlTree]]))
+    }
+
     override def remove(operations: Ref[Coll[Coll[Byte]]], proof: Ref[Coll[Byte]]): Ref[WOption[AvlTree]] = {
       asRep[WOption[AvlTree]](mkMethodCall(self,
         AvlTreeClass.getMethod("remove", classOf[Sym], classOf[Sym]),
@@ -1056,6 +1063,13 @@ object AvlTree extends EntityObject("AvlTree") {
         true, true, element[WOption[AvlTree]]))
     }
 
+    def insertOrUpdate(operations: Ref[Coll[(Coll[Byte], Coll[Byte])]], proof: Ref[Coll[Byte]]): Ref[WOption[AvlTree]] = {
+      asRep[WOption[AvlTree]](mkMethodCall(source,
+        AvlTreeClass.getMethod("insertOrUpdate", classOf[Sym], classOf[Sym]),
+        Array[AnyRef](operations, proof),
+        true, true, element[WOption[AvlTree]]))
+    }
+
     def remove(operations: Ref[Coll[Coll[Byte]]], proof: Ref[Coll[Byte]]): Ref[WOption[AvlTree]] = {
       asRep[WOption[AvlTree]](mkMethodCall(source,
         AvlTreeClass.getMethod("remove", classOf[Sym], classOf[Sym]),

sc/shared/src/test/scala/sigmastate/utxo/BasicOpsSpecification.scala

@@ -4,7 +4,7 @@ import org.ergoplatform.ErgoBox.{AdditionalRegisters, R6, R8}
 import org.ergoplatform._
 import org.scalatest.Assertion
 import scorex.crypto.authds.{ADKey, ADValue}
-import scorex.crypto.authds.avltree.batch.{BatchAVLProver, Insert}
+import scorex.crypto.authds.avltree.batch.{BatchAVLProver, Insert, InsertOrUpdate}
 import scorex.crypto.hash.{Blake2b256, Digest32}
 import scorex.util.ByteArrayBuilder
 import scorex.util.encode.Base16
@@ -13,14 +13,12 @@ import scorex.util.encode.Base16
 import scorex.utils.Ints
 import scorex.util.serialization.VLQByteBufferWriter
 import scorex.utils.Longs
-import sigma.{Colls, SigmaTestingData}
+import sigma.{Coll, Colls, SigmaTestingData, VersionContext}
 import sigma.Extensions.ArrayOps
-import sigma.{SigmaTestingData, VersionContext}
 import sigma.VersionContext.{V6SoftForkVersion, withVersions}
 import sigma.ast.SCollection.SByteArray
-import sigma.ast.SType.AnyOps
-import sigma.data.{AvlTreeData, CAnyValue, CHeader, CSigmaDslBuilder}
-import sigma.data.{AvlTreeData, AvlTreeFlags, CAND, CAnyValue, CHeader, CSigmaDslBuilder, CSigmaProp}
+import sigma.ast.SType.{AnyOps, tD}
+import sigma.data.{AvlTreeData, AvlTreeFlags, CAND, CAnyValue, CAvlTree, CHeader, CSigmaDslBuilder, CSigmaProp}
 import sigma.util.StringUtil._
 import sigma.ast._
 import sigma.ast.syntax._
@@ -34,10 +32,9 @@ import sigmastate.interpreter.Interpreter._
 import sigma.ast.Apply
 import sigma.eval.EvalSettings
 import sigma.exceptions.InvalidType
-import sigma.serialization.ErgoTreeSerializer
+import sigma.serialization.{DataSerializer, ErgoTreeSerializer, SigmaByteWriter, SigmaSerializer, ValueSerializer}
 import sigma.interpreter.{ContextExtension, ProverResult}
 import sigma.util.NBitsUtils
-import sigma.serialization.{DataSerializer, ErgoTreeSerializer, SigmaByteWriter}
 import sigma.util.Extensions
 import sigma.validation.ValidationException
 import sigmastate.utils.Helpers
@@ -2364,4 +2361,51 @@ class BasicOpsSpecification extends CompilerTestingCommons
     }
   }
 
+  property("avltree.insertOrUpdate") {
+    val avlProver = new BatchAVLProver[Digest32, Blake2b256.type](keyLength = 32, None)
+
+    val elements = Seq(123, 22)
+    val treeElements = elements.map(i => Longs.toByteArray(i)).map(s => (ADKey @@@ Blake2b256(s), ADValue @@ s))
+    treeElements.foreach(s => avlProver.performOneOperation(Insert(s._1, s._2)))
+    avlProver.generateProof()
+    val treeData = new AvlTreeData(avlProver.digest.toColl, AvlTreeFlags.AllOperationsAllowed, 32, None)
+
+    val elements2 = Seq(1, 22)
+    val treeElements2 = elements2.map(i => Longs.toByteArray(i)).map(s => (ADKey @@@ Blake2b256(s), ADValue @@ s))
+    treeElements2.foreach(s => avlProver.performOneOperation(InsertOrUpdate(s._1, s._2)))
+    val updateProof = avlProver.generateProof()
+    val treeData2 = new AvlTreeData(avlProver.digest.toColl, AvlTreeFlags.AllOperationsAllowed, 32, None)
+
+    val v: Coll[(Coll[Byte], Coll[Byte])] = treeElements2.map(t => t._1.toColl -> t._2.toColl).toArray.toColl
+    val ops = IR.builder.mkConstant[SType](v.asWrappedType, SCollection(STuple(SByteArray, SByteArray)))
+
+    val customExt = Seq(
+      21.toByte -> AvlTreeConstant(treeData),
+      22.toByte -> AvlTreeConstant(treeData2),
+      23.toByte -> ops,
+      24.toByte -> ByteArrayConstant(updateProof)
+    )
+
+    def deserTest() = test("deserializeTo", env, customExt,
+      s"""{
+            val tree1 = getVar[AvlTree](21).get
+            val tree2 = getVar[AvlTree](22).get
+
+            val toInsert = getVar[Coll[(Coll[Byte], Coll[Byte])]](23).get
+            val proof = getVar[Coll[Byte]](24).get
+
+            val tree1Updated = tree1.insertOrUpdate(toInsert, proof).get
+            tree2.digest == tree1Updated.digest
+          }""",
+      null,
+      true
+    )
+
+    if (VersionContext.current.isV6SoftForkActivated) {
+      deserTest()
+    } else {
+      an[ValidationException] should be thrownBy deserTest()
+    }
+  }
+
 }