def fac_by_ref_BAD():
n = pt.ScratchVar(pt.TealType.uint64)
return pt.Seq(
n.store(pt.Int(10)),
factorial_BAD(n),
n.load(),
)
def test_for():
i = pt.ScratchVar()
items = [
(i.store(pt.Int(0))),
i.load() < pt.Int(10),
i.store(i.load() + pt.Int(1)),
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]]))
assert expr.type_of() == pt.TealType.none
assert not expr.has_return()
expected, varEnd = items[0].__teal__(options)
condStart, condEnd = items[1].__teal__(options)
stepStart, stepEnd = items[2].__teal__(options)
do, doEnd = pt.Seq([items[3]]).__teal__(options)
expectedBranch = pt.TealConditionalBlock([])
end = pt.TealSimpleBlock([])
varEnd.setNextBlock(condStart)
doEnd.setNextBlock(stepStart)
expectedBranch.setTrueBlock(do)
expectedBranch.setFalseBlock(end)
condEnd.setNextBlock(expectedBranch)
stepEnd.setNextBlock(condStart)
actual, _ = expr.__teal__(options)
assert actual == expected
def concat_strings(
b: pt.abi.DynamicArray[pt.abi.String], *, output: pt.abi.String
) -> pt.Expr:
idx = pt.ScratchVar()
buff = pt.ScratchVar()
init = idx.store(pt.Int(0))
cond = idx.load() < b.length()
_iter = idx.store(idx.load() + pt.Int(1))
return pt.Seq(
buff.store(pt.Bytes("")),
pt.For(init, cond, _iter).Do(
b[idx.load()].use(lambda s: buff.store(pt.Concat(buff.load(), s.get())))
),
output.set(buff.load()),
)
def test_while_compiles():
i = pt.ScratchVar()
expr = pt.While(pt.Int(2)).Do(pt.Seq([i.store(pt.Int(0))]))
assert expr.type_of() == pt.TealType.none
assert not expr.has_return()
expr.__teal__(options)
def swap(x: pt.ScratchVar, y: pt.ScratchVar):
z = pt.ScratchVar(pt.TealType.anytype)
return pt.Seq(
z.store(x.load()),
x.store(y.load()),
y.store(z.load()),
)
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 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 sub_logcat_dynamic():
first = pt.ScratchVar(pt.TealType.bytes)
return pt.Seq(
first.store(pt.Bytes("hello")),
logcat_dynamic(first, pt.Int(42)),
pt.Assert(pt.Bytes("hello42") == first.load()),
pt.Int(1),
)
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 string_mult():
s = pt.ScratchVar(pt.TealType.bytes)
return pt.Seq(
s.store(pt.Txn.application_args[0]),
subr_string_mult(s, pt.Btoi(pt.Txn.application_args[1])),
pt.Log(s.load()),
pt.Int(100),
)
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 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_for_compiles():
i = pt.ScratchVar()
expr = pt.For(i.store(pt.Int(0)), pt.Int(1), i.store(i.load() + pt.Int(1))).Do(
pt.App.globalPut(pt.Itob(pt.Int(0)), pt.Itob(pt.Int(2)))
)
assert expr.type_of() == pt.TealType.none
assert not expr.has_return()
expr.__teal__(options)
def test_invalid_for():
with pytest.raises(TypeError):
expr = pt.For()
with pytest.raises(TypeError):
expr = pt.For(pt.Int(2))
with pytest.raises(TypeError):
expr = pt.For(pt.Int(1), pt.Int(2))
with pytest.raises(pt.TealTypeError):
i = pt.ScratchVar()
expr = pt.For(i.store(pt.Int(0)), pt.Int(1), pt.Int(2))
expr.__teal__(options)
with pytest.raises(pt.TealCompileError):
i = pt.ScratchVar()
expr = pt.For(i.store(pt.Int(0)), pt.Int(1), i.store(i.load() + pt.Int(1)))
expr.type_of()
with pytest.raises(pt.TealCompileError):
i = pt.ScratchVar()
expr = pt.For(i.store(pt.Int(0)), pt.Int(1), i.store(i.load() + pt.Int(1)))
expr.__str__()
with pytest.raises(pt.TealTypeError):
i = pt.ScratchVar()
expr = pt.For(i.store(pt.Int(0)), pt.Int(1), i.store(i.load() + pt.Int(1))).Do(
pt.Int(0)
)
with pytest.raises(pt.TealTypeError):
i = pt.ScratchVar()
expr = pt.For(i.store(pt.Int(0)), pt.Int(1), i.store(i.load() + pt.Int(1))).Do(
pt.Pop(pt.Int(1)), pt.Int(2)
)
with pytest.raises(pt.TealCompileError):
expr = (
pt.For(i.store(pt.Int(0)), pt.Int(1), i.store(i.load() + pt.Int(1)))
.Do(pt.Continue())
.Do(pt.Continue())
)
expr.__str__()
def test_dynamic_scratchvar_cannot_set_index_to_another_dynamic():
myvar = pt.DynamicScratchVar()
myvar.load()
regvar = pt.ScratchVar()
myvar.set_index(regvar)
dynvar = pt.DynamicScratchVar()
with pytest.raises(pt.TealInputError):
myvar.set_index(dynvar)
def test_scratchvar_index():
myvar = pt.ScratchVar()
expr = myvar.index()
expected = pt.TealSimpleBlock([pt.TealOp(expr, pt.Op.int, myvar.slot)])
actual, _ = expr.__teal__(options)
actual.addIncoming()
actual = pt.TealBlock.NormalizeBlocks(actual)
assert actual == expected
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 wilt_the_stilt():
player_score = pt.DynamicScratchVar(pt.TealType.uint64)
wilt = pt.ScratchVar(pt.TealType.uint64, 129)
kobe = pt.ScratchVar(pt.TealType.uint64)
dt = pt.ScratchVar(pt.TealType.uint64, 131)
return pt.Seq(
player_score.set_index(wilt),
player_score.store(pt.Int(100)),
player_score.set_index(kobe),
player_score.store(pt.Int(81)),
player_score.set_index(dt),
player_score.store(pt.Int(73)),
pt.Assert(player_score.load() == pt.Int(73)),
pt.Assert(player_score.index() == pt.Int(131)),
player_score.set_index(wilt),
pt.Assert(player_score.load() == pt.Int(100)),
pt.Assert(player_score.index() == pt.Int(129)),
pt.Int(100),
)
def test_scratchvar_assign_load():
slotId = 5
myvar = pt.ScratchVar(slotId=slotId)
expr = myvar.load()
expected = pt.TealSimpleBlock([pt.TealOp(expr, pt.Op.load, myvar.slot)])
actual, _ = expr.__teal__(options)
actual.addIncoming()
actual = pt.TealBlock.NormalizeBlocks(actual)
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 test_nested_for_compiles():
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.For(i.store(pt.Int(0)), pt.Int(1), i.store(i.load() + pt.Int(1))).Do(
pt.Seq([i.store(pt.Int(0))])
)
]
)
)
assert expr.type_of() == pt.TealType.none
assert not expr.has_return()
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_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 should_it_work() -> pt.Expr:
xs = [
pt.ScratchVar(pt.TealType.uint64),
pt.ScratchVar(pt.TealType.uint64),
]
def store_initial_values():
return [s.store(pt.Int(i + 1)) for i, s in enumerate(xs)]
d = pt.DynamicScratchVar(pt.TealType.uint64)
@pt.Subroutine(pt.TealType.none)
def retrieve_and_increment(s: pt.ScratchVar):
return pt.Seq(d.set_index(s), d.store(d.load() + pt.Int(1)))
def asserts():
return [pt.Assert(x.load() == pt.Int(i + 2)) for i, x in enumerate(xs)]
return pt.Seq(
pt.Seq(store_initial_values()),
pt.Seq([retrieve_and_increment(x) for x in xs]),
pt.Seq(asserts()),
pt.Int(1),
)
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_scratchvar_store():
myvar = pt.ScratchVar(pt.TealType.bytes)
arg = pt.Bytes("value")
expr = myvar.store(arg)
expected = pt.TealSimpleBlock(
[
pt.TealOp(arg, pt.Op.byte, '"value"'),
pt.TealOp(expr, pt.Op.store, myvar.slot),
]
)
actual, _ = expr.__teal__(options)
actual.addIncoming()
actual = pt.TealBlock.NormalizeBlocks(actual)
assert actual == expected
def test_scratchvar_assign_store():
slotId = 2
myvar = pt.ScratchVar(pt.TealType.uint64, slotId)
arg = pt.Int(10)
expr = myvar.store(arg)
expected = pt.TealSimpleBlock(
[
pt.TealOp(arg, pt.Op.int, 10),
pt.TealOp(expr, pt.Op.store, myvar.slot),
]
)
actual, _ = expr.__teal__(options)
actual.addIncoming()
actual = pt.TealBlock.NormalizeBlocks(actual)
assert actual == expected
def test_while():
i = pt.ScratchVar()
i.store(pt.Int(0))
items = [i.load() < pt.Int(2), [i.store(i.load() + pt.Int(1))]]
expr = pt.While(items[0]).Do(pt.Seq(items[1]))
assert expr.type_of() == pt.TealType.none
assert not expr.has_return()
expected, condEnd = items[0].__teal__(options)
do, doEnd = pt.Seq(items[1]).__teal__(options)
expectedBranch = pt.TealConditionalBlock([])
end = pt.TealSimpleBlock([])
expectedBranch.setTrueBlock(do)
expectedBranch.setFalseBlock(end)
condEnd.setNextBlock(expectedBranch)
doEnd.setNextBlock(expected)
actual, _ = expr.__teal__(options)
assert actual == expected
def test_optimize_dynamic_var():
myvar = pt.DynamicScratchVar()
regvar = pt.ScratchVar()
program = pt.Seq(
regvar.store(pt.Int(1)),
myvar.set_index(regvar),
regvar.store(pt.Int(2)),
pt.Pop(regvar.load()),
pt.Approve(),
)
optimize_options = OptimizeOptions()
# unoptimized
expected = """#pragma version 4
int 1
store 1
int 1
store 0
int 2
store 1
load 1
pop
int 1
return""".strip()
actual = pt.compileTeal(program,
version=4,
mode=pt.Mode.Application,
optimize=optimize_options)
assert actual == expected
# optimization should not change the code because the candidate slot
# is used by the dynamic slot variable.
optimize_options = OptimizeOptions(scratch_slots=True)
actual = pt.compileTeal(program,
version=4,
mode=pt.Mode.Application,
optimize=optimize_options)
assert actual == expected