computable index test

ergoplatform · Aug 30, 2024 · 107d666 · 107d666
1 parent a07b67b
commit 107d666
Show file tree

Hide file tree

Showing 7 changed files with 50 additions and 24 deletions.
diff --git a/data/shared/src/main/scala/sigma/ast/methods.scala b/data/shared/src/main/scala/sigma/ast/methods.scala
@@ -1072,20 +1072,22 @@ case object SBoxMethods extends MonoTypeMethods {
         """.stripMargin,
         ArgInfo("regId", "zero-based identifier of the register."))
 
-  lazy val tokensMethod = SMethod(
-    this, "tokens", SFunc(SBox, ErgoBox.STokensRegType), 8, FixedCost(JitCost(15)))
-      .withIRInfo(MethodCallIrBuilder)
-      .withInfo(PropertyCall, "Secondary tokens")
-
   lazy val getRegMethodV6 = SMethod(this, "getReg",
-    SFunc(Array(SBox, SInt), SOption(tT), Array(paramT)), 7, ExtractRegisterAs.costKind, Seq(tT))
-    .withIRInfo(MethodCallIrBuilder)
+    SFunc(Array(SBox, SInt), SOption(tT), Array(paramT)), 7, FixedCost(JitCost(50)), Seq(tT))
+    .withIRInfo(MethodCallIrBuilder,
+      javaMethodOf[Box, Int, RType[_]]("getReg"),
+      { mtype => Array(mtype.tRange.asOption[SType].elemType) })
     .withInfo(MethodCall, """ Extracts register by id and type.
                             | Type param \lst{T} expected type of the register.
                             | Returns \lst{Some(value)} if the register is defined and has given type and \lst{None} otherwise
         """.stripMargin,
       ArgInfo("regId", "zero-based identifier of the register."))
 
+  lazy val tokensMethod = SMethod(
+    this, "tokens", SFunc(SBox, ErgoBox.STokensRegType), 8, FixedCost(JitCost(15)))
+      .withIRInfo(MethodCallIrBuilder)
+      .withInfo(PropertyCall, "Secondary tokens")
+
   lazy val commonBoxMethods = Array(
     ValueMethod, // see ExtractAmount
     PropositionBytesMethod, // see ExtractScriptBytes

diff --git a/sc/shared/src/main/scala/sigma/compiler/ir/GraphBuilding.scala b/sc/shared/src/main/scala/sigma/compiler/ir/GraphBuilding.scala
@@ -1020,8 +1020,8 @@ trait GraphBuilding extends Base with DefRewriting { IR: IRContext =>
               box.tokens
             case SBoxMethods.getRegMethodV6.name if VersionContext.current.isV6SoftForkActivated =>
               val c1 = asRep[Int](argsV(0))
-              val c3 = stypeToElem(typeSubst.apply(tT))
-              box.getReg(c1)(c3)
+              val c2 = stypeToElem(typeSubst.apply(tT))
+              box.getReg(c1)(c2)
             case _ => throwError
           }
           case (ctx: Ref[Context]@unchecked, SContextMethods) => method.name match {

diff --git a/sc/shared/src/main/scala/sigma/compiler/ir/IRContext.scala b/sc/shared/src/main/scala/sigma/compiler/ir/IRContext.scala
@@ -153,7 +153,7 @@ trait IRContext
   override def invokeUnlifted(e: Elem[_], mc: MethodCall, dataEnv: DataEnv): Any = e match {
     case _: CollElem[_,_] => mc match {
       case CollMethods.map(_, f) =>
-        val newMC = mc.copy(args = mc.args :+ f.elem.eRange)(mc.resultType, mc.isAdapterCall)
+        val newMC = mc.copy(args = mc.args :+ f.elem.eRange)(mc.resultType, mc.isAdapterCall, mc.typeSubst)
         super.invokeUnlifted(e, newMC, dataEnv)
       case _ =>
         super.invokeUnlifted(e, mc, dataEnv)

diff --git a/sc/shared/src/main/scala/sigma/compiler/ir/MethodCalls.scala b/sc/shared/src/main/scala/sigma/compiler/ir/MethodCalls.scala
@@ -1,6 +1,7 @@
 package sigma.compiler.ir
 
 import debox.{cfor, Buffer => DBuffer}
+import sigma.ast.{SType, STypeVar}
 import sigma.compiler.DelayInvokeException
 import sigma.reflection.RMethod
 import sigma.util.CollectionUtil.TraversableOps
@@ -26,7 +27,7 @@ trait MethodCalls extends Base { self: IRContext =>
     *                      given `method`.
     */
   case class MethodCall private[MethodCalls](receiver: Sym, method: RMethod, args: Seq[AnyRef], neverInvoke: Boolean)
-                                            (val resultType: Elem[Any], val isAdapterCall: Boolean = false) extends Def[Any] {
+                                            (val resultType: Elem[Any], val isAdapterCall: Boolean = false, val typeSubst: Map[STypeVar, SType]) extends Def[Any] {
 
     override def mirror(t: Transformer): Ref[Any] = {
       val len = args.length
@@ -100,8 +101,8 @@ trait MethodCalls extends Base { self: IRContext =>
 
   /** Creates new MethodCall node and returns its node ref. */
   def mkMethodCall(receiver: Sym, method: RMethod, args: Seq[AnyRef],
-                   neverInvoke: Boolean, isAdapterCall: Boolean, resultElem: Elem[_]): Sym = {
-    reifyObject(MethodCall(receiver, method, args, neverInvoke)(asElem[Any](resultElem), isAdapterCall))
+                   neverInvoke: Boolean, isAdapterCall: Boolean, resultElem: Elem[_], typeSubst: Map[STypeVar, SType] = Map.empty): Sym = {
+    reifyObject(MethodCall(receiver, method, args, neverInvoke)(asElem[Any](resultElem), isAdapterCall, typeSubst))
   }
 
   @tailrec

diff --git a/sc/shared/src/main/scala/sigma/compiler/ir/TreeBuilding.scala b/sc/shared/src/main/scala/sigma/compiler/ir/TreeBuilding.scala
@@ -290,13 +290,10 @@ trait TreeBuilding extends Base { IR: IRContext =>
         mkExtractAmount(box.asBox)
       case BoxM.propositionBytes(In(box)) =>
         mkExtractScriptBytes(box.asBox)
-      case BoxM.getReg(In(box), regId, _) =>
+      case BoxM.getReg(In(box), regId, _) if regId.isConst =>
         val tpe = elemToSType(s.elem).asOption
-        if (regId.isConst)
-          mkExtractRegisterAs(box.asBox, ErgoBox.allRegisters(valueFromRep(regId)), tpe)
-        else
-          error(s"Non constant expressions (${regId.node}) are not supported in getReg")
-      case BoxM.creationInfo(In(box)) =>
+        mkExtractRegisterAs(box.asBox, ErgoBox.allRegisters(valueFromRep(regId)), tpe)
+     case BoxM.creationInfo(In(box)) =>
         mkExtractCreationInfo(box.asBox)
       case BoxM.id(In(box)) =>
         mkExtractId(box.asBox)
@@ -399,13 +396,14 @@ trait TreeBuilding extends Base { IR: IRContext =>
         mkMultiplyGroup(obj.asGroupElement, arg.asGroupElement)
 
       // Fallback MethodCall rule: should be the last in this list of cases
-      case Def(MethodCall(objSym, m, argSyms, _)) =>
+      case Def(mc @ MethodCall(objSym, m, argSyms, _)) =>
         val obj = recurse[SType](objSym)
         val args = argSyms.collect { case argSym: Sym => recurse[SType](argSym) }
         MethodsContainer.getMethod(obj.tpe, m.getName) match {
           case Some(method) =>
-            val specMethod = method.specializeFor(obj.tpe, args.map(_.tpe))
-            builder.mkMethodCall(obj, specMethod, args.toIndexedSeq, Map())
+            val typeSubst = mc.typeSubst
+            val specMethod = method.specializeFor(obj.tpe, args.map(_.tpe)).withConcreteTypes(typeSubst)
+            builder.mkMethodCall(obj, specMethod, args.toIndexedSeq, typeSubst)
           case None =>
             error(s"Cannot find method ${m.getName} in object $obj")
         }

diff --git a/sc/shared/src/main/scala/sigma/compiler/ir/wrappers/sigma/impl/SigmaDslImpl.scala b/sc/shared/src/main/scala/sigma/compiler/ir/wrappers/sigma/impl/SigmaDslImpl.scala
@@ -8,6 +8,8 @@ import sigma.compiler.ir.wrappers.sigma.impl.SigmaDslDefs
 import scala.collection.compat.immutable.ArraySeq
 
 package impl {
+  import sigma.Evaluation
+  import sigma.ast.SType.tT
   import sigma.compiler.ir.meta.ModuleInfo
   import sigma.compiler.ir.wrappers.sigma.SigmaDsl
   import sigma.compiler.ir.{Base, GraphIRReflection, IRContext}
@@ -620,10 +622,11 @@ object Box extends EntityObject("Box") {
     }
 
     override def getReg[T](i: Ref[Int])(implicit cT: Elem[T]): Ref[WOption[T]] = {
+      val st = Evaluation.rtypeToSType(cT.sourceType)
       asRep[WOption[T]](mkMethodCall(self,
         BoxClass.getMethod("getReg", classOf[Sym], classOf[Elem[_]]),
         Array[AnyRef](i, cT),
-        true, false, element[WOption[T]]))
+        true, false, element[WOption[T]], Map(tT -> st) ))
     }
 
     override def tokens: Ref[Coll[(Coll[Byte], Long)]] = {
@@ -695,10 +698,11 @@ object Box extends EntityObject("Box") {
     }
 
     def getReg[T](i: Ref[Int])(implicit cT: Elem[T]): Ref[WOption[T]] = {
+      val st = Evaluation.rtypeToSType(cT.sourceType)
       asRep[WOption[T]](mkMethodCall(source,
         BoxClass.getMethod("getReg", classOf[Sym], classOf[Elem[_]]),
         Array[AnyRef](i, cT),
-        true, true, element[WOption[T]]))
+        true, true, element[WOption[T]], Map(tT -> st)))
     }
 
     def tokens: Ref[Coll[(Coll[Byte], Long)]] = {

diff --git a/sc/shared/src/test/scala/sigmastate/utxo/BasicOpsSpecification.scala b/sc/shared/src/test/scala/sigmastate/utxo/BasicOpsSpecification.scala
@@ -785,7 +785,28 @@ class BasicOpsSpecification extends CompilerTestingCommons
     } else {
       an[Exception] should be thrownBy getRegTest()
     }
+  }
 
+  property("Box.getReg - computable index") {
+    val ext: Seq[VarBinding] = Seq(
+      (intVar1, IntConstant(0))
+    )
+    def getRegTest(): Assertion = {
+      test("Box.getReg", env, ext,
+        """{
+          |   val idx = getVar[Int](1).get
+          |   val x = SELF.getReg[Long](idx).get
+          |   x == SELF.value
+          |}""".stripMargin,
+        null
+      )
+    }
+
+    if (VersionContext.current.isV6SoftForkActivated) {
+      getRegTest()
+    } else {
+      an[Exception] should be thrownBy getRegTest()
+    }
   }
 
 }