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()) )