Add BytesSqrt (#163)

* Add BytesSqrt * Update pyteal/ast/unaryexpr_test.py Co-authored-by: Jason Paulos <jasonpaulos@users.noreply.github.com> Co-authored-by: Jason Paulos <jasonpaulos@users.noreply.github.com>
algorand · Jan 19, 2022 · 6bebe92 · 6bebe92
1 parent 7b73148
commit 6bebe92
Show file tree

Hide file tree

Showing 5 changed files with 36 additions and 0 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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)