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>

Author: StylishTriangles

pyteal/__init__.pyi

@@ -134,6 +134,7 @@ __all__ = [
     "BytesGt",
     "BytesGe",
     "BytesNot",
+    "BytesSqrt",
     "BytesZero",
     "ExtractUint16",
     "ExtractUint32",

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",

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.
 

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)

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)