diff --git a/pyteal/__init__.pyi b/pyteal/__init__.pyi index 5d69ac4b2..1c9d96dac 100644 --- a/pyteal/__init__.pyi +++ b/pyteal/__init__.pyi @@ -134,6 +134,7 @@ __all__ = [ "BytesGt", "BytesGe", "BytesNot", + "BytesSqrt", "BytesZero", "ExtractUint16", "ExtractUint32", diff --git a/pyteal/ast/__init__.py b/pyteal/ast/__init__.py index 624d5fdb6..375671dd1 100644 --- a/pyteal/ast/__init__.py +++ b/pyteal/ast/__init__.py @@ -44,6 +44,7 @@ Balance, MinBalance, BytesNot, + BytesSqrt, BytesZero, Log, ) @@ -239,6 +240,7 @@ "BytesGt", "BytesGe", "BytesNot", + "BytesSqrt", "BytesZero", "ExtractUint16", "ExtractUint32", diff --git a/pyteal/ast/unaryexpr.py b/pyteal/ast/unaryexpr.py index 5a3c79086..055339d1f 100644 --- a/pyteal/ast/unaryexpr.py +++ b/pyteal/ast/unaryexpr.py @@ -153,6 +153,16 @@ def BytesNot(arg: Expr) -> UnaryExpr: return UnaryExpr(Op.b_not, TealType.bytes, TealType.bytes, arg) +def BytesSqrt(arg: Expr) -> UnaryExpr: + """Get the bytes square root of bytes. + + This will return the largest integer X such that X^2 <= arg. + + Requires TEAL version 6 or higher. + """ + return UnaryExpr(Op.bsqrt, TealType.bytes, TealType.bytes, arg) + + def BytesZero(arg: Expr) -> UnaryExpr: """Get a byte-array of a specified length, containing all zero bytes. diff --git a/pyteal/ast/unaryexpr_test.py b/pyteal/ast/unaryexpr_test.py index 7da54a3a4..b79f3e0b6 100644 --- a/pyteal/ast/unaryexpr_test.py +++ b/pyteal/ast/unaryexpr_test.py @@ -9,6 +9,7 @@ teal3Options = CompileOptions(version=3) teal4Options = CompileOptions(version=4) teal5Options = CompileOptions(version=5) +teal6Options = CompileOptions(version=6) def test_btoi(): @@ -357,6 +358,27 @@ def test_b_not_invalid(): BytesNot(Int(2)) +def test_bsqrt(): + arg = Bytes("base16", "0xFEDCBA9876543210") + expr = BytesSqrt(arg) + assert expr.type_of() == TealType.bytes + + expected = TealSimpleBlock( + [TealOp(arg, Op.byte, "0xFEDCBA9876543210"), TealOp(expr, Op.bsqrt)] + ) + + actual, _ = expr.__teal__(teal6Options) + actual.addIncoming() + actual = TealBlock.NormalizeBlocks(actual) + + assert actual == expected + + +def test_bsqrt_invalid(): + with pytest.raises(TealTypeError): + BytesSqrt(Int(2 ** 64 - 1)) + + def test_b_zero(): arg = Int(8) expr = BytesZero(arg) diff --git a/pyteal/ir/ops.py b/pyteal/ir/ops.py index 738e5a7eb..1c78b166c 100644 --- a/pyteal/ir/ops.py +++ b/pyteal/ir/ops.py @@ -170,6 +170,7 @@ def min_version(self) -> int: gtxnas = OpType("gtxnas", Mode.Signature | Mode.Application, 5) gtxnsas = OpType("gtxnsas", Mode.Signature | Mode.Application, 5) args = OpType("args", Mode.Signature, 5) + bsqrt = OpType("bsqrt", Mode.Signature | Mode.Application, 6) itxn_next = OpType("itxn_next", Mode.Application, 6) gitxn = OpType("gitxn", Mode.Application, 6) gitxna = OpType("gitxna", Mode.Application, 6)