def test_elseif_multiple():
args = [pt.Int(0), pt.Int(1), pt.Int(2), pt.Int(3), pt.Int(4), pt.Int(5), pt.Int(6)]
expr = (
pt.If(args[0])
.Then(args[1])
.ElseIf(args[2])
.Then(args[3])
.ElseIf(args[4])
.Then(args[5])
.Else(args[6])
)
assert expr.type_of() == pt.TealType.uint64
elseIfExpr = pt.If(args[2], args[3], pt.If(args[4], args[5], args[6]))
expected, _ = args[0].__teal__(options)
thenBlock, _ = args[1].__teal__(options)
elseStart, elseEnd = elseIfExpr.__teal__(options)
expectedBranch = pt.TealConditionalBlock([])
expectedBranch.setTrueBlock(thenBlock)
expectedBranch.setFalseBlock(elseStart)
expected.setNextBlock(expectedBranch)
end = pt.TealSimpleBlock([])
thenBlock.setNextBlock(end)
elseEnd.setNextBlock(end)
actual, _ = expr.__teal__(options)
assert actual == expected
def test_if_has_return():
exprWithReturn = pt.If(pt.Int(1), pt.Return(pt.Int(1)), pt.Return(pt.Int(0)))
assert exprWithReturn.has_return()
exprWithoutReturn = pt.If(pt.Int(1), pt.Int(1), pt.Int(0))
assert not exprWithoutReturn.has_return()
exprSemiReturn = pt.If(
pt.Int(1),
pt.Return(pt.Int(1)),
pt.App.globalPut(pt.Bytes("key"), pt.Bytes("value")),
)
assert not exprSemiReturn.has_return()
def test_if_invalid():
with pytest.raises(pt.TealTypeError):
pt.If(pt.Int(0), pt.Txn.amount(), pt.Txn.sender())
with pytest.raises(pt.TealTypeError):
pt.If(pt.Txn.sender(), pt.Int(1), pt.Int(0))
with pytest.raises(pt.TealTypeError):
pt.If(pt.Int(0), pt.Txn.sender())
with pytest.raises(pt.TealTypeError):
pt.If(pt.Int(0), pt.Int(2))
with pytest.raises(pt.TealCompileError):
expr = pt.If(pt.Int(0))
expr.__teal__(options)
def test_while_break():
i = pt.ScratchVar()
i.store(pt.Int(0))
items = [
i.load() < pt.Int(2),
i.store(i.load() + pt.Int(1)),
pt.If(i.load() == pt.Int(1), pt.Break()),
]
expr = pt.While(items[0]).Do(pt.Seq(items[1], items[2]))
assert expr.type_of() == pt.TealType.none
assert not expr.has_return()
options.enterLoop()
expected, condEnd = items[0].__teal__(options)
do, doEnd = pt.Seq([items[1], items[2]]).__teal__(options)
expectedBranch = pt.TealConditionalBlock([])
end = pt.TealSimpleBlock([])
expectedBranch.setTrueBlock(do)
expectedBranch.setFalseBlock(end)
condEnd.setNextBlock(expectedBranch)
doEnd.setNextBlock(expected)
breakBlocks, _ = options.exitLoop()
for block in breakBlocks:
block.setNextBlock(end)
actual, _ = expr.__teal__(options)
assert actual == expected
def subr_string_mult(s: pt.ScratchVar, n):
tmp = pt.ScratchVar(pt.TealType.bytes)
return (pt.If(n == pt.Int(0)).Then(s.store(pt.Bytes(""))).Else(
pt.Seq(
tmp.store(s.load()),
subr_string_mult(s, n - pt.Int(1)),
s.store(pt.Concat(s.load(), tmp.load())),
)))
def plus_one(n: pt.ScratchVar):
tmp = pt.ScratchVar(pt.TealType.uint64)
return (pt.If(n.load() == pt.Int(0)).Then(n.store(pt.Int(1))).Else(
pt.Seq(
tmp.store(n.load() - pt.Int(1)),
plus_one(tmp),
n.store(tmp.load() + pt.Int(1)),
)))
def test_continue_break():
i = pt.ScratchVar()
expr = pt.For(i.store(pt.Int(0)), pt.Int(1), i.store(i.load() + pt.Int(1))).Do(
pt.Seq([pt.If(pt.Int(1), pt.Break(), pt.Continue())])
)
assert expr.type_of() == pt.TealType.none
assert not expr.has_return()
expr.__teal__(options)
def factorial_BAD(n: pt.ScratchVar):
tmp = pt.ScratchVar(pt.TealType.uint64)
return (pt.If(n.load() <= pt.Int(1)).Then(n.store(pt.Int(1))).Else(
pt.Seq(
tmp.store(n.load() - pt.Int(1)),
factorial_BAD(tmp),
n.store(n.load() * tmp.load()),
)))
def test_optimize_subroutine_with_global_var():
global_var = pt.ScratchVar(pt.TealType.uint64)
@pt.Subroutine(pt.TealType.uint64)
def add(a1: pt.Expr) -> pt.Expr:
return pt.Seq(global_var.store(pt.Int(2)), global_var.load() + a1)
program = pt.Seq([
pt.If(pt.Txn.sender() == pt.Global.creator_address()).Then(
pt.Pop(add(pt.Int(1)))),
global_var.store(pt.Int(5)),
pt.Approve(),
])
optimize_options = OptimizeOptions()
# unoptimized
expected = """#pragma version 4
txn Sender
global CreatorAddress
==
bz main_l2
int 1
callsub add_0
pop
main_l2:
int 5
store 0
int 1
return
// add
add_0:
store 1
int 2
store 0
load 0
load 1
+
retsub
""".strip()
actual = pt.compileTeal(program,
version=4,
mode=pt.Mode.Application,
optimize=optimize_options)
assert actual == expected
# optimization should not apply to global vars
optimize_options = OptimizeOptions(scratch_slots=True)
actual = pt.compileTeal(program,
version=4,
mode=pt.Mode.Application,
optimize=optimize_options)
assert actual == expected
def tallygo():
result = pt.ScratchVar(pt.TealType.bytes)
# pt.If-Then is a hook for creating + opting in without providing any args
return (pt.If(
pt.Or(pt.App.id() == pt.Int(0),
pt.Txn.application_args.length() == pt.Int(0))).Then(
pt.Int(1)).Else(
pt.Seq(
result.store(pt.Bytes("dummy")),
tally(pt.Int(4), result),
pt.Btoi(result.load()),
)))
def test_optimize_subroutine_with_reserved_local_var():
local_var = pt.ScratchVar(pt.TealType.uint64, 0)
@pt.Subroutine(pt.TealType.uint64)
def add(a1: pt.Expr) -> pt.Expr:
return pt.Seq(local_var.store(pt.Int(2)), local_var.load() + a1)
program = pt.Seq([
pt.If(pt.Txn.sender() == pt.Global.creator_address()).Then(
pt.Pop(add(pt.Int(1)))),
pt.Approve(),
])
optimize_options = OptimizeOptions()
# unoptimized
expected = """#pragma version 4
txn Sender
global CreatorAddress
==
bz main_l2
int 1
callsub add_0
pop
main_l2:
int 1
return
// add
add_0:
store 1
int 2
store 0
load 0
load 1
+
retsub
""".strip()
actual = pt.compileTeal(program,
version=4,
mode=pt.Mode.Application,
optimize=optimize_options)
assert actual == expected
# The optimization must skip over the reserved slot id so the expected result
# hasn't changed.
optimize_options = OptimizeOptions(scratch_slots=True)
actual = pt.compileTeal(program,
version=4,
mode=pt.Mode.Application,
optimize=optimize_options)
assert actual == expected
def fac_by_ref_args():
n = pt.ScratchVar(pt.TealType.uint64)
return pt.Seq(
pt.If(
pt.Or(
pt.App.id() == pt.Int(0),
pt.Txn.application_args.length() == pt.Int(0),
)).Then(pt.Int(1)).Else(
pt.Seq(
n.store(pt.Btoi(pt.Txn.application_args[0])),
factorial(n),
n.load(),
)))
def test_break_fail():
with pytest.raises(pt.TealCompileError):
pt.Break().__teal__(options)
with pytest.raises(pt.TealCompileError):
pt.If(pt.Int(1), pt.Break()).__teal__(options)
with pytest.raises(pt.TealCompileError):
pt.Seq([pt.Break()]).__teal__(options)
with pytest.raises(TypeError):
pt.Break(pt.Int(1))
def test_elseif_multiple_with_multi():
args = [
pt.Int(0),
[pt.Pop(pt.Int(1)), pt.Int(2)],
pt.Int(3),
[pt.Pop(pt.Int(4)), pt.Int(5)],
pt.Int(6),
[pt.Pop(pt.Int(7)), pt.Int(8)],
[pt.Pop(pt.Int(9)), pt.Int(10)],
]
expr = (
pt.If(args[0])
.Then(*args[1])
.ElseIf(args[2])
.Then(*args[3])
.ElseIf(args[4])
.Then(*args[5])
.Else(*args[6])
)
elseIfExpr = pt.If(
args[2], pt.Seq(args[3]), pt.If(args[4], pt.Seq(args[5]), pt.Seq(args[6]))
)
expected, _ = args[0].__teal__(options)
thenBlock, thenBlockEnd = pt.Seq(args[1]).__teal__(options)
elseStart, elseBlockEnd = elseIfExpr.__teal__(options)
expectedBranch = pt.TealConditionalBlock([])
expectedBranch.setTrueBlock(thenBlock)
expectedBranch.setFalseBlock(elseStart)
expected.setNextBlock(expectedBranch)
end = pt.TealSimpleBlock([])
thenBlockEnd.setNextBlock(end)
elseBlockEnd.setNextBlock(end)
actual, _ = expr.__teal__(options)
assert actual == expected
def test_build_program_clear_state_valid_config():
action = pt.If(pt.Txn.fee() == pt.Int(4)).Then(pt.Approve()).Else(pt.Reject())
config = pt.CallConfig.CALL
router_with_bare_call = pt.Router(
"test",
pt.BareCallActions(
clear_state=pt.OnCompleteAction(action=action, call_config=config)
),
)
_, actual_clear_state_with_bare_call, _ = router_with_bare_call.build_program()
expected_clear_state_with_bare_call = assemble_helper(
pt.Cond([pt.Txn.application_args.length() == pt.Int(0), action])
)
with pt.TealComponent.Context.ignoreExprEquality():
assert (
assemble_helper(actual_clear_state_with_bare_call)
== expected_clear_state_with_bare_call
)
router_with_method = pt.Router("test")
@pt.ABIReturnSubroutine
def clear_state_method():
return action
router_with_method.add_method_handler(
clear_state_method, method_config=pt.MethodConfig(clear_state=config)
)
_, actual_clear_state_with_method, _ = router_with_method.build_program()
expected_clear_state_with_method = assemble_helper(
pt.Cond(
[
pt.Txn.application_args[0]
== pt.MethodSignature("clear_state_method()void"),
pt.Seq(clear_state_method(), pt.Approve()),
]
)
)
with pt.TealComponent.Context.ignoreExprEquality():
assert (
assemble_helper(actual_clear_state_with_method)
== expected_clear_state_with_method
)
def test_many_ifs():
"""
Test with many pt.If statements to trigger potential corner cases in code generation.
Previous versions of PyTeal took an exponential time to generate the TEAL code for this PyTEAL.
"""
sv = pt.ScratchVar(pt.TealType.uint64)
s = pt.Seq([
pt.If(
pt.Int(3 * i) == pt.Int(3 * i),
sv.store(pt.Int(3 * i + 1)),
sv.store(pt.Int(3 * i + 2)),
) for i in range(30)
] + [pt.Return(sv.load())])
pt.compileTeal(s, mode=pt.Mode.Signature, version=2)
def test_if_single():
args = [pt.Int(1), pt.Pop(pt.Int(1))]
expr = pt.If(args[0], args[1])
assert expr.type_of() == pt.TealType.none
expected, _ = args[0].__teal__(options)
thenBlockStart, thenBlockEnd = args[1].__teal__(options)
end = pt.TealSimpleBlock([])
expectedBranch = pt.TealConditionalBlock([])
expectedBranch.setTrueBlock(thenBlockStart)
expectedBranch.setFalseBlock(end)
expected.setNextBlock(expectedBranch)
thenBlockEnd.setNextBlock(end)
actual, _ = expr.__teal__(options)
assert actual == expected
def test_else_alt_multi():
args = [pt.Int(0), pt.Int(1), [pt.Pop(pt.Int(2)), pt.Int(3)]]
expr = pt.If(args[0]).Then(args[1]).Else(*args[2])
expected, _ = args[0].__teal__(options)
thenBlockStart, thenBlockEnd = args[1].__teal__(options)
elseBlockStart, elseBlockEnd = pt.Seq(*args[2]).__teal__(options)
expectedBranch = pt.TealConditionalBlock([])
expectedBranch.setTrueBlock(thenBlockStart)
expectedBranch.setFalseBlock(elseBlockStart)
expected.setNextBlock(expectedBranch)
end = pt.TealSimpleBlock([])
thenBlockEnd.setNextBlock(end)
elseBlockEnd.setNextBlock(end)
actual, _ = expr.__teal__(options)
assert actual == expected
def test_scratch_store():
for value in (
pt.Int(1),
pt.Bytes("test"),
pt.App.globalGet(pt.Bytes("key")),
pt.If(pt.Int(1), pt.Int(2), pt.Int(3)),
):
slot = pt.ScratchSlot()
expr = pt.ScratchStore(slot, value)
assert expr.type_of() == pt.TealType.none
expected, valueEnd = value.__teal__(options)
storeBlock = pt.TealSimpleBlock([pt.TealOp(expr, pt.Op.store, slot)])
valueEnd.setNextBlock(storeBlock)
actual, _ = expr.__teal__(options)
assert actual == expected
def test_if_bytes():
args = [pt.Int(1), pt.Txn.sender(), pt.Txn.receiver()]
expr = pt.If(args[0], args[1], args[2])
assert expr.type_of() == pt.TealType.bytes
expected, _ = args[0].__teal__(options)
thenBlock, _ = args[1].__teal__(options)
elseBlock, _ = args[2].__teal__(options)
expectedBranch = pt.TealConditionalBlock([])
expectedBranch.setTrueBlock(thenBlock)
expectedBranch.setFalseBlock(elseBlock)
expected.setNextBlock(expectedBranch)
end = pt.TealSimpleBlock([])
thenBlock.setNextBlock(end)
elseBlock.setNextBlock(end)
actual, _ = expr.__teal__(options)
assert actual == expected
def test_if_alt_int():
args = [pt.Int(0), pt.Int(1), pt.Int(2)]
expr = pt.If(args[0]).Then(args[1]).Else(args[2])
assert expr.type_of() == pt.TealType.uint64
expected, _ = args[0].__teal__(options)
thenBlock, _ = args[1].__teal__(options)
elseBlock, _ = args[2].__teal__(options)
expectedBranch = pt.TealConditionalBlock([])
expectedBranch.setTrueBlock(thenBlock)
expectedBranch.setFalseBlock(elseBlock)
expected.setNextBlock(expectedBranch)
end = pt.TealSimpleBlock([])
thenBlock.setNextBlock(end)
elseBlock.setNextBlock(end)
actual, _ = expr.__teal__(options)
assert actual == expected
def test_if_alt_none():
args = [pt.Int(0), pt.Pop(pt.Txn.sender()), pt.Pop(pt.Txn.receiver())]
expr = pt.If(args[0]).Then(args[1]).Else(args[2])
assert expr.type_of() == pt.TealType.none
expected, _ = args[0].__teal__(options)
thenBlockStart, thenBlockEnd = args[1].__teal__(options)
elseBlockStart, elseBlockEnd = args[2].__teal__(options)
expectedBranch = pt.TealConditionalBlock([])
expectedBranch.setTrueBlock(thenBlockStart)
expectedBranch.setFalseBlock(elseBlockStart)
expected.setNextBlock(expectedBranch)
end = pt.TealSimpleBlock([])
thenBlockEnd.setNextBlock(end)
elseBlockEnd.setNextBlock(end)
actual, _ = expr.__teal__(options)
assert actual == expected
def test_scratch_store_index_expression():
for value in (
pt.Int(1),
pt.Bytes("test"),
pt.App.globalGet(pt.Bytes("key")),
pt.If(pt.Int(1), pt.Int(2), pt.Int(3)),
):
expr = pt.ScratchStore(slot=None, value=value, index_expression=pt.Int(1337))
assert expr.type_of() == pt.TealType.none
expected = pt.TealSimpleBlock([pt.TealOp(None, pt.Op.int, 1337)])
valueStart, valueEnd = value.__teal__(options)
expected.setNextBlock(valueStart)
storeBlock = pt.TealSimpleBlock([pt.TealOp(expr, pt.Op.stores)])
valueEnd.setNextBlock(storeBlock)
actual, _ = expr.__teal__(options)
with pt.TealComponent.Context.ignoreExprEquality():
assert actual == expected
def test_for_continue():
i = pt.ScratchVar()
items = [
(i.store(pt.Int(0))),
i.load() < pt.Int(10),
i.store(i.load() + pt.Int(1)),
pt.If(i.load() < pt.Int(4), pt.Continue()),
pt.App.globalPut(pt.Itob(i.load()), i.load() * pt.Int(2)),
]
expr = pt.For(items[0], items[1], items[2]).Do(pt.Seq([items[3], items[4]]))
assert expr.type_of() == pt.TealType.none
assert not expr.has_return()
options.enterLoop()
expected, varEnd = items[0].__teal__(options)
condStart, condEnd = items[1].__teal__(options)
stepStart, stepEnd = items[2].__teal__(options)
do, doEnd = pt.Seq([items[3], items[4]]).__teal__(options)
expectedBranch = pt.TealConditionalBlock([])
end = pt.TealSimpleBlock([])
doEnd.setNextBlock(stepStart)
stepEnd.setNextBlock(condStart)
expectedBranch.setTrueBlock(do)
expectedBranch.setFalseBlock(end)
condEnd.setNextBlock(expectedBranch)
varEnd.setNextBlock(condStart)
_, continueBlocks = options.exitLoop()
for block in continueBlocks:
block.setNextBlock(stepStart)
actual, _ = expr.__teal__(options)
assert actual == expected
def __test_single_conditional(expr: pt.MultiValue, op, args: List[pt.Expr],
iargs, reducer):
__test_single(expr)
expected_call = pt.TealSimpleBlock([
pt.TealOp(expr, op, *iargs),
pt.TealOp(expr.output_slots[1].store(), pt.Op.store,
expr.output_slots[1]),
pt.TealOp(expr.output_slots[0].store(), pt.Op.store,
expr.output_slots[0]),
])
ifExpr = (pt.If(expr.output_slots[1].load()).Then(
expr.output_slots[0].load()).Else(pt.App.globalGet(pt.Bytes("None"))))
ifBlockStart, _ = ifExpr.__teal__(options)
expected_call.setNextBlock(ifBlockStart)
if len(args) == 0:
expected: pt.TealBlock = expected_call
elif len(args) == 1:
expected, after_arg = args[0].__teal__(options)
after_arg.setNextBlock(expected_call)
elif len(args) == 2:
expected, after_arg_1 = args[0].__teal__(options)
arg_2, after_arg_2 = args[1].__teal__(options)
after_arg_1.setNextBlock(arg_2)
after_arg_2.setNextBlock(expected_call)
expected.addIncoming()
expected = pt.TealBlock.NormalizeBlocks(expected)
actual, _ = expr.outputReducer(reducer).__teal__(options)
actual.addIncoming()
actual = pt.TealBlock.NormalizeBlocks(actual)
with pt.TealComponent.Context.ignoreExprEquality():
assert actual == expected
def test_multi_value(op, type, iargs, args):
reducer = (lambda value, hasValue: pt.If(hasValue).Then(value).Else(
pt.App.globalGet(pt.Bytes("None"))))
expr = pt.MultiValue(op, [type, pt.TealType.uint64],
immediate_args=iargs,
args=args)
__test_single_conditional(expr, op, args, iargs, reducer)
reducer = lambda value, hasValue: pt.Seq(pt.Assert(hasValue), value
) # noqa: E731
expr = pt.MultiValue(op, [type, pt.TealType.uint64],
immediate_args=iargs,
args=args)
__test_single_assert(expr, op, args, iargs, reducer)
hasValueVar = pt.ScratchVar(pt.TealType.uint64)
valueVar = pt.ScratchVar(type)
reducer = lambda value, hasValue: pt.Seq( # noqa: E731
hasValueVar.store(hasValue), valueVar.store(value))
expr = pt.MultiValue(op, [type, pt.TealType.uint64],
immediate_args=iargs,
args=args)
__test_single_with_vars(expr, op, args, iargs, hasValueVar, valueVar,
reducer)
def recursiveIsEven(i):
return (pt.If(i == pt.Int(0)).Then(pt.Int(1)).ElseIf(i == pt.Int(1)).Then(
pt.Int(0)).Else(recursiveIsEven(i - pt.Int(2))))
def slow_fibonacci(n):
return (pt.If(n <= pt.Int(1)).Then(n).Else(
slow_fibonacci(n - pt.Int(2)) + slow_fibonacci(n - pt.Int(1))))
def test_continue_break():
expr = pt.While(pt.Int(0)).Do(
pt.Seq([pt.If(pt.Int(1), pt.Break(), pt.Continue())]))
assert expr.type_of() == pt.TealType.none
assert not expr.has_return()
expr.__teal__(options)
def safe_clear_state_delete():
return (
pt.If(pt.Txn.sender() == pt.Global.creator_address())
.Then(pt.Approve())
.Else(pt.Reject())
)
