Frame Ops to avm8 branch PR

pyteal/ast/frame.py

@@ -0,0 +1,132 @@
+from typing import TYPE_CHECKING
+
+from pyteal.ast.expr import Expr
+from pyteal.types import TealType, require_type
+from pyteal.errors import TealInputError, verifyProgramVersion
+from pyteal.ir import TealBlock, TealSimpleBlock, TealOp, Op
+
+if TYPE_CHECKING:
+    from pyteal.compiler import CompileOptions
+
+
+class Proto(Expr):
+    def __init__(self, num_args: int, num_returns: int):
+        super().__init__()
+        if num_args < 0:
+            raise TealInputError(
+                f"the number of arguments provided to Proto must be >= 0 but {num_args=}"
+            )
+        if num_returns < 0:
+            raise TealInputError(
+                f"the number of return values provided to Proto must be >= 0 but {num_returns=}"
+            )
+        self.num_args = num_args
+        self.num_returns = num_returns
+
+    def __teal__(self, options: "CompileOptions") -> tuple[TealBlock, TealSimpleBlock]:
+        verifyProgramVersion(
+            Op.proto.min_version,
+            options.version,
+            "Program version too low to use op proto",
+        )
+        op = TealOp(self, Op.proto, self.num_args, self.num_returns)
+        return TealBlock.FromOp(options, op)
+
+    def __str__(self) -> str:
+        return f"(proto: num_args = {self.num_args}, num_returns = {self.num_returns})"
+
+    def type_of(self) -> TealType:
+        return TealType.none
+
+    def has_return(self) -> bool:
+        return False
+
+
+Proto.__module__ = "pyteal"
+
+
+class FrameDig(Expr):
+    def __init__(self, frame_index: int):
+        super().__init__()
+        self.frame_index = frame_index
+
+    def __teal__(self, options: "CompileOptions") -> tuple[TealBlock, TealSimpleBlock]:
+        verifyProgramVersion(
+            Op.frame_dig.min_version,
+            options.version,
+            "Program version too low to use op frame_dig",
+        )
+        op = TealOp(self, Op.frame_dig, self.frame_index)
+        return TealBlock.FromOp(options, op)
+
+    def __str__(self) -> str:
+        return f"(frame_dig: dig_from = {self.frame_index})"
+
+    def type_of(self) -> TealType:
+        return TealType.anytype
+
+    def has_return(self) -> bool:
+        return False
+
+
+FrameDig.__module__ = "pyteal"
+
+
+class FrameBury(Expr):
+    def __init__(self, value: Expr, frame_index: int):
+        super().__init__()
+        require_type(value, TealType.anytype)
+        self.value = value
+        self.frame_index = frame_index
+
+    def __teal__(self, options: "CompileOptions") -> tuple[TealBlock, TealSimpleBlock]:
+        verifyProgramVersion(
+            Op.frame_bury.min_version,
+            options.version,
+            "Program version too low to use op frame_bury",
+        )
+        op = TealOp(self, Op.frame_bury, self.frame_index)
+        return TealBlock.FromOp(options, op, self.value)
+
+    def __str__(self) -> str:
+        return f"(frame_bury (bury_to = {self.frame_index}) ({self.value}))"
+
+    def type_of(self) -> TealType:
+        return TealType.none
+
+    def has_return(self) -> bool:
+        return False
+
+
+FrameBury.__module__ = "pyteal"
+
+
+class DupN(Expr):
+    def __init__(self, value: Expr, repetition: int):
+        super().__init__()
+        require_type(value, TealType.anytype)
+        if repetition < 0:
+            raise TealInputError("dupn repetition should be non negative")
+        self.value = value
+        self.repetition = repetition
+
+    def __teal__(self, options: "CompileOptions") -> tuple[TealBlock, TealSimpleBlock]:
+        verifyProgramVersion(
+            Op.dupn.min_version,
+            options.version,
+            "Program version too low to use op dupn",
+        )
+        op = TealOp(self, Op.dupn, self.repetition)
+        return TealBlock.FromOp(options, op, self.value)
+
+    def __str__(self) -> str:
+        return f"(dupn (repetition = {self.repetition}) ({self.value}))"
+
+    def type_of(self) -> TealType:
+        return self.value.type_of()
+
+    def has_return(self) -> bool:
+        return False
+
+
+DupN.__module__ = "pyteal"

pyteal/ast/frame_test.py

@@ -0,0 +1,107 @@
+import pytest
+import pyteal as pt
+from pyteal.ast.frame import FrameBury, FrameDig, Proto, DupN
+
+avm7Options = pt.CompileOptions(version=7)
+avm8Options = pt.CompileOptions(version=8)
+
+
+@pytest.mark.parametrize("input_num, output_num", [(1, 1), (1, 0), (5, 5)])
+def test_proto(input_num: int, output_num: int):
+    expr = Proto(input_num, output_num)
+    assert not expr.has_return()
+    assert expr.type_of() == pt.TealType.none
+
+    expected = pt.TealSimpleBlock([pt.TealOp(expr, pt.Op.proto, input_num, output_num)])
+    actual, _ = expr.__teal__(avm8Options)
+    actual.addIncoming()
+    actual = pt.TealBlock.NormalizeBlocks(actual)
+
+    assert actual == expected
+
+
+def test_proto_invalid():
+    with pytest.raises(pt.TealInputError):
+        Proto(-1, 1)
+
+    with pytest.raises(pt.TealInputError):
+        Proto(1, -1)
+
+    with pytest.raises(pt.TealInputError):
+        Proto(1, 1).__teal__(avm7Options)
+
+
+@pytest.mark.parametrize("depth", [-1, 0, 1, 2])
+def test_frame_dig(depth: int):
+    expr = FrameDig(depth)
+    assert not expr.has_return()
+    assert expr.type_of() == pt.TealType.anytype
+
+    expected = pt.TealSimpleBlock([pt.TealOp(expr, pt.Op.frame_dig, depth)])
+    actual, _ = expr.__teal__(avm8Options)
+    actual.addIncoming()
+    actual = pt.TealBlock.NormalizeBlocks(actual)
+
+    assert actual == expected
+
+
+def test_frame_dig_invalid():
+    with pytest.raises(pt.TealInputError):
+        FrameDig(1).__teal__(avm7Options)
+
+
+def test_frame_bury():
+    byte_expr = pt.Bytes("Astartes")
+    expr = FrameBury(byte_expr, 4)
+    assert not expr.has_return()
+    assert expr.type_of() == pt.TealType.none
+
+    expected = pt.TealSimpleBlock(
+        [
+            pt.TealOp(byte_expr, pt.Op.byte, '"Astartes"'),
+            pt.TealOp(expr, pt.Op.frame_bury, 4),
+        ]
+    )
+    actual, _ = expr.__teal__(avm8Options)
+    actual.addIncoming()
+    actual = pt.TealBlock.NormalizeBlocks(actual)
+
+    assert actual == expected
+
+
+def test_frame_bury_invalid():
+    with pytest.raises(pt.TealTypeError):
+        FrameBury(pt.Seq(), 1)
+
+    with pytest.raises(pt.TealInputError):
+        FrameBury(pt.Int(1), 1).__teal__(avm7Options)
+
+
+def test_dupn():
+    byte_expr = pt.Bytes("Astartes")
+    expr = DupN(byte_expr, 4)
+    assert not expr.has_return()
+    assert expr.type_of() == byte_expr.type_of()
+
+    expected = pt.TealSimpleBlock(
+        [
+            pt.TealOp(byte_expr, pt.Op.byte, '"Astartes"'),
+            pt.TealOp(expr, pt.Op.dupn, 4),
+        ]
+    )
+    actual, _ = expr.__teal__(avm8Options)
+    actual.addIncoming()
+    actual = pt.TealBlock.NormalizeBlocks(actual)
+
+    assert actual == expected
+
+
+def test_dupn_invalid():
+    with pytest.raises(pt.TealTypeError):
+        DupN(pt.Seq(), 1)
+
+    with pytest.raises(pt.TealInputError):
+        DupN(pt.Int(1), -1)
+
+    with pytest.raises(pt.TealInputError):
+        DupN(pt.Int(1), 1).__teal__(avm7Options)

pyteal/compiler/compiler.py

@@ -26,6 +26,7 @@
 from pyteal.compiler.constants import createConstantBlocks
 
 MAX_PROGRAM_VERSION = 8
+FRAME_POINTER_VERSION = 8
 MIN_PROGRAM_VERSION = 2
 DEFAULT_PROGRAM_VERSION = MIN_PROGRAM_VERSION
 

pyteal/ir/ops.py

@@ -33,7 +33,7 @@ def min_version(self) -> int:
     # fmt: off
     # meta
     comment             = OpType("//",                  Mode.Signature | Mode.Application, 0)
-    # avm 
+    # avm
     err                 = OpType("err",                 Mode.Signature | Mode.Application, 2)
     sha256              = OpType("sha256",              Mode.Signature | Mode.Application, 2)
     keccak256           = OpType("keccak256",           Mode.Signature | Mode.Application, 2)
@@ -197,6 +197,12 @@ def min_version(self) -> int:
     box_len             = OpType("box_len",             Mode.Application,                  8)
     box_get             = OpType("box_get",             Mode.Application,                  8)
     box_put             = OpType("box_put",             Mode.Application,                  8)
+    popn                = OpType("popn",                Mode.Signature | Mode.Application, 8)
+    dupn                = OpType("dupn",                Mode.Signature | Mode.Application, 8)
+    bury                = OpType("bury",                Mode.Signature | Mode.Application, 8)
+    frame_dig           = OpType("frame_dig",           Mode.Signature | Mode.Application, 8)
+    frame_bury          = OpType("frame_bury",          Mode.Signature | Mode.Application, 8)
+    proto               = OpType("proto",               Mode.Signature | Mode.Application, 8)
     # fmt: on