def test_assignScratchSlotsToSubroutines_slot_used_before_assignment():
def sub1Impl():
return None
def sub2Impl(a1):
return None
def sub3Impl(a1, a2, a3):
return None
subroutine1 = pt.SubroutineDefinition(sub1Impl, pt.TealType.uint64)
subroutine2 = pt.SubroutineDefinition(sub2Impl, pt.TealType.bytes)
subroutine3 = pt.SubroutineDefinition(sub3Impl, pt.TealType.none)
globalSlot1 = pt.ScratchSlot()
subroutine1Slot1 = pt.ScratchSlot()
subroutine1Slot2 = pt.ScratchSlot()
subroutine1Ops = [
pt.TealOp(None, pt.Op.int, 1),
pt.TealOp(None, pt.Op.store, subroutine1Slot1),
pt.TealOp(None, pt.Op.int, 3),
pt.TealOp(None, pt.Op.store, subroutine1Slot2),
pt.TealOp(None, pt.Op.load, globalSlot1),
pt.TealOp(None, pt.Op.retsub),
]
subroutine2Slot1 = pt.ScratchSlot()
subroutine2Ops = [
pt.TealOp(None, pt.Op.byte, '"value"'),
pt.TealOp(None, pt.Op.store, subroutine2Slot1),
pt.TealOp(None, pt.Op.load, subroutine2Slot1),
pt.TealOp(None, pt.Op.retsub),
]
subroutine3Ops = [
pt.TealOp(None, pt.Op.retsub),
]
mainSlot1 = pt.ScratchSlot()
mainSlot2 = pt.ScratchSlot()
mainOps = [
pt.TealOp(None, pt.Op.int, 7),
pt.TealOp(None, pt.Op.store, globalSlot1),
pt.TealOp(None, pt.Op.int, 2),
pt.TealOp(None, pt.Op.store, mainSlot2),
pt.TealOp(None, pt.Op.load, mainSlot1),
pt.TealOp(None, pt.Op.return_),
]
subroutineBlocks = {
None: pt.TealSimpleBlock(mainOps),
subroutine1: pt.TealSimpleBlock(subroutine1Ops),
subroutine2: pt.TealSimpleBlock(subroutine2Ops),
subroutine3: pt.TealSimpleBlock(subroutine3Ops),
}
with pytest.raises(pt.TealInternalError):
assignScratchSlotsToSubroutines(subroutineBlocks)
def test_evaluate_subroutine_no_args():
cases = (
(pt.TealType.none, pt.Return()),
(pt.TealType.uint64, pt.Int(1) + pt.Int(2)),
(pt.TealType.uint64, pt.Return(pt.Int(1) + pt.Int(2))),
(pt.TealType.bytes, pt.Bytes("value")),
(pt.TealType.bytes, pt.Return(pt.Bytes("value"))),
)
for (returnType, returnValue) in cases:
def mySubroutine():
return returnValue
definition = pt.SubroutineDefinition(mySubroutine, returnType)
declaration = evaluate_subroutine(definition)
assert isinstance(declaration, pt.SubroutineDeclaration)
assert declaration.subroutine is definition
assert declaration.type_of() == returnValue.type_of()
assert declaration.has_return() == returnValue.has_return()
options.setSubroutine(definition)
expected, _ = pt.Seq([returnValue]).__teal__(options)
actual, _ = declaration.__teal__(options)
options.setSubroutine(None)
assert actual == expected
def test_subroutine_return_value():
cases = (
(pt.TealType.uint64, pt.Int(1), pt.Op.int, 1),
(pt.TealType.bytes, pt.Bytes("value"), pt.Op.byte, '"value"'),
(pt.TealType.anytype, pt.Int(1), pt.Op.int, 1),
(pt.TealType.anytype, pt.Bytes("value"), pt.Op.byte, '"value"'),
)
for (tealType, value, op, opValue) in cases:
expr = pt.Return(value)
def mySubroutine():
return expr
subroutine = pt.SubroutineDefinition(mySubroutine, tealType)
assert expr.type_of() == pt.TealType.none
assert expr.has_return()
expected = pt.TealSimpleBlock(
[pt.TealOp(value, op, opValue), pt.TealOp(expr, pt.Op.retsub)]
)
options.setSubroutine(subroutine)
actual, _ = expr.__teal__(options)
options.setSubroutine(None)
actual.addIncoming()
actual = pt.TealBlock.NormalizeBlocks(actual)
assert actual == expected
def mock_subroutine_definition(implementation, has_abi_output=False):
mock = pt.SubroutineDefinition(lambda: pt.Return(pt.Int(1)),
pt.TealType.uint64)
mock._validate() # haven't failed with dummy implementation
mock.implementation = implementation
mock.has_abi_output = has_abi_output
return mock
def test_subroutine_call():
def mySubroutine():
return pt.Return()
returnTypes = (
pt.TealType.uint64,
pt.TealType.bytes,
pt.TealType.anytype,
pt.TealType.none,
)
argCases = (
[],
[pt.Int(1)],
[pt.Int(1), pt.Bytes("value")],
)
for returnType in returnTypes:
definition = pt.SubroutineDefinition(mySubroutine, returnType)
for args in argCases:
expr = pt.SubroutineCall(definition, args)
assert expr.type_of() == returnType
assert not expr.has_return()
expected, _ = pt.TealBlock.FromOp(
options, pt.TealOp(expr, pt.Op.callsub, definition), *args)
actual, _ = expr.__teal__(options)
assert actual == expected
def test_flattenSubroutines_1_subroutine():
subroutine = pt.SubroutineDefinition(
lambda: pt.Int(1) + pt.Int(2) + pt.Int(3), pt.TealType.uint64)
subroutineToLabel = OrderedDict()
subroutineToLabel[subroutine] = "sub0"
subroutineLabel = pt.LabelReference(subroutineToLabel[subroutine])
subroutineOps = [
pt.TealOp(None, pt.Op.int, 1),
pt.TealOp(None, pt.Op.int, 2),
pt.TealOp(None, pt.Op.int, 3),
pt.TealOp(None, pt.Op.add),
pt.TealOp(None, pt.Op.add),
pt.TealOp(None, pt.Op.retsub),
]
l1Label = pt.LabelReference("l1")
mainOps = [
pt.TealOp(None, pt.Op.txn, "Fee"),
pt.TealOp(None, pt.Op.int, 0),
pt.TealOp(None, pt.Op.eq),
pt.TealOp(None, pt.Op.bz, l1Label),
pt.TealOp(None, pt.Op.callsub, subroutineLabel),
pt.TealOp(None, pt.Op.return_),
pt.TealLabel(None, l1Label),
pt.TealOp(None, pt.Op.int, 0),
pt.TealOp(None, pt.Op.return_),
]
subroutineMapping = {None: mainOps, subroutine: subroutineOps}
expectedL1Label = pt.LabelReference("main_l1")
expectedSubroutineLabel = pt.LabelReference("sub0")
expected = [
pt.TealOp(None, pt.Op.txn, "Fee"),
pt.TealOp(None, pt.Op.int, 0),
pt.TealOp(None, pt.Op.eq),
pt.TealOp(None, pt.Op.bz, expectedL1Label),
pt.TealOp(None, pt.Op.callsub, expectedSubroutineLabel),
pt.TealOp(None, pt.Op.return_),
pt.TealLabel(None, expectedL1Label),
pt.TealOp(None, pt.Op.int, 0),
pt.TealOp(None, pt.Op.return_),
pt.TealLabel(None, expectedSubroutineLabel, "<lambda>"),
pt.TealOp(None, pt.Op.int, 1),
pt.TealOp(None, pt.Op.int, 2),
pt.TealOp(None, pt.Op.int, 3),
pt.TealOp(None, pt.Op.add),
pt.TealOp(None, pt.Op.add),
pt.TealOp(None, pt.Op.retsub),
]
actual = flattenSubroutines(subroutineMapping, subroutineToLabel)
assert actual == expected
def test_subroutine_return_none_invalid():
for value in (pt.Int(1), pt.Bytes("value")):
expr = pt.Return(value)
def mySubroutine():
return expr
subroutine = pt.SubroutineDefinition(mySubroutine, pt.TealType.none)
options.setSubroutine(subroutine)
with pytest.raises(pt.TealCompileError):
expr.__teal__(options)
options.setSubroutine(None)
def test_evaluate_subroutine_2_args():
cases = (
(pt.TealType.none, pt.Return()),
(pt.TealType.uint64, pt.Int(1) + pt.Int(2)),
(pt.TealType.uint64, pt.Return(pt.Int(1) + pt.Int(2))),
(pt.TealType.bytes, pt.Bytes("value")),
(pt.TealType.bytes, pt.Return(pt.Bytes("value"))),
)
for (returnType, returnValue) in cases:
argSlots: List[pt.ScratchSlot] = []
def mySubroutine(a1, a2):
assert isinstance(a1, pt.ScratchLoad)
argSlots.append(a1.slot)
assert isinstance(a2, pt.ScratchLoad)
argSlots.append(a2.slot)
return returnValue
definition = pt.SubroutineDefinition(mySubroutine, returnType)
declaration = evaluate_subroutine(definition)
assert isinstance(declaration, pt.SubroutineDeclaration)
assert declaration.subroutine is definition
assert declaration.type_of() == returnValue.type_of()
assert declaration.has_return() == returnValue.has_return()
assert isinstance(declaration.body, pt.Seq)
assert len(declaration.body.args) == 3
assert isinstance(declaration.body.args[0], pt.ScratchStackStore)
assert isinstance(declaration.body.args[1], pt.ScratchStackStore)
assert declaration.body.args[0].slot is argSlots[-1]
assert declaration.body.args[1].slot is argSlots[-2]
options.setSubroutine(definition)
expected, _ = pt.Seq(
[declaration.body.args[0], declaration.body.args[1],
returnValue]).__teal__(options)
actual, _ = declaration.__teal__(options)
options.setSubroutine(None)
assert actual == expected
def test_subroutine_return_value_invalid():
cases = (
(pt.TealType.bytes, pt.Int(1)),
(pt.TealType.uint64, pt.Bytes("value")),
)
for (tealType, value) in cases:
expr = pt.Return(value)
def mySubroutine():
return expr
subroutine = pt.SubroutineDefinition(mySubroutine, tealType)
options.setSubroutine(subroutine)
with pytest.raises(pt.TealCompileError):
expr.__teal__(options)
options.setSubroutine(None)
def test_evaluate_subroutine_10_args():
cases = (
(pt.TealType.none, pt.Return()),
(pt.TealType.uint64, pt.Int(1) + pt.Int(2)),
(pt.TealType.uint64, pt.Return(pt.Int(1) + pt.Int(2))),
(pt.TealType.bytes, pt.Bytes("value")),
(pt.TealType.bytes, pt.Return(pt.Bytes("value"))),
)
for (returnType, returnValue) in cases:
argSlots: List[pt.ScratchSlot] = []
def mySubroutine(a1, a2, a3, a4, a5, a6, a7, a8, a9, a10):
for a in (a1, a2, a3, a4, a5, a6, a7, a8, a9, a10):
assert isinstance(a, pt.ScratchLoad)
argSlots.append(a.slot)
return returnValue
definition = pt.SubroutineDefinition(mySubroutine, returnType)
declaration = evaluate_subroutine(definition)
assert isinstance(declaration, pt.SubroutineDeclaration)
assert declaration.subroutine is definition
assert declaration.type_of() == returnValue.type_of()
assert declaration.has_return() == returnValue.has_return()
assert isinstance(declaration.body, pt.Seq)
assert len(declaration.body.args) == 11
for i in range(10):
assert isinstance(declaration.body.args[i], pt.ScratchStackStore)
for i in range(10):
assert declaration.body.args[i].slot is argSlots[-i - 1]
options.setSubroutine(definition)
expected, _ = pt.Seq(declaration.body.args[:10] +
[returnValue]).__teal__(options)
actual, _ = declaration.__teal__(options)
options.setSubroutine(None)
assert actual == expected
def test_subroutine_return_none():
expr = pt.Return()
def mySubroutine():
return expr
subroutine = pt.SubroutineDefinition(mySubroutine, pt.TealType.none)
assert expr.type_of() == pt.TealType.none
assert expr.has_return()
expected = pt.TealSimpleBlock([pt.TealOp(expr, pt.Op.retsub)])
options.setSubroutine(subroutine)
actual, _ = expr.__teal__(options)
options.setSubroutine(None)
actual.addIncoming()
actual = pt.TealBlock.NormalizeBlocks(actual)
assert actual == expected
def test_subroutine_declaration():
cases = (
(pt.TealType.none, pt.Return()),
(pt.TealType.uint64, pt.Return(pt.Int(1))),
(pt.TealType.uint64, pt.Int(1)),
(pt.TealType.bytes, pt.Bytes("value")),
(pt.TealType.anytype, pt.App.globalGet(pt.Bytes("key"))),
)
for (returnType, value) in cases:
def mySubroutine():
return value
definition = pt.SubroutineDefinition(mySubroutine, returnType)
declaration = pt.SubroutineDeclaration(definition, value)
assert declaration.type_of() == value.type_of()
assert declaration.has_return() == value.has_return()
options.currentSubroutine = definition
assert declaration.__teal__(options) == value.__teal__(options)
options.setSubroutine(None)
def test_subroutine_definition_invalid():
def fnWithDefaults(a, b=None):
return pt.Return()
def fnWithKeywordArgs(a, *, output):
return pt.Return()
def fnWithKeywordArgsWrongKWName(a, *, b: pt.abi.Uint64):
return pt.Return()
def fnWithMultipleABIKeywordArgs(a, *, b: pt.abi.Byte, c: pt.abi.Bool):
return pt.Return()
def fnWithVariableArgs(a, *b):
return pt.Return()
def fnWithNonExprReturnAnnotation(a, b) -> pt.TealType.uint64:
return pt.Return()
def fnWithNonExprParamAnnotation(a, b: pt.TealType.uint64):
return pt.Return()
def fnWithScratchVarSubclass(a, b: pt.DynamicScratchVar):
return pt.Return()
def fnReturningExprSubclass(a: pt.ScratchVar, b: pt.Expr) -> pt.Return:
return pt.Return()
def fnWithMixedAnns4AndBytesReturn(a: pt.Expr,
b: pt.ScratchVar) -> pt.Bytes:
return pt.Bytes("hello uwu")
def fnWithMixedAnnsABIRet1(
a: pt.Expr, b: pt.ScratchVar, c: pt.abi.Uint16
) -> pt.abi.StaticArray[pt.abi.Uint32, Literal[10]]:
return pt.abi.StaticArray(
pt.abi.StaticArrayTypeSpec(pt.abi.Uint32TypeSpec(), 10))
def fnWithMixedAnnsABIRet2(a: pt.Expr, b: pt.abi.Byte,
c: pt.ScratchVar) -> pt.abi.Uint64:
return pt.abi.Uint64()
cases = (
(
1,
"TealInputError('Input to SubroutineDefinition is not callable'",
"TealInputError('Input to ABIReturnSubroutine is not callable'",
),
(
None,
"TealInputError('Input to SubroutineDefinition is not callable'",
"TealInputError('Input to ABIReturnSubroutine is not callable'",
),
(
fnWithDefaults,
"TealInputError('Function has a parameter with a default value, which is not allowed in a subroutine: b'",
"TealInputError('Function has a parameter with a default value, which is not allowed in a subroutine: b'",
),
(
fnWithKeywordArgs,
"TealInputError('Function has a parameter type that is not allowed in a subroutine: parameter output with type",
"TealInputError('ABI return subroutine output-kwarg output must specify ABI type')",
),
(
fnWithKeywordArgsWrongKWName,
"TealInputError('Function has a parameter type that is not allowed in a subroutine: parameter b with type",
"TealInputError('ABI return subroutine output-kwarg name must be `output` at this moment",
),
(
fnWithMultipleABIKeywordArgs,
"TealInputError('Function has a parameter type that is not allowed in a subroutine: parameter b with type",
"multiple output arguments (2) with type annotations",
),
(
fnWithVariableArgs,
"TealInputError('Function has a parameter type that is not allowed in a subroutine: parameter b with type",
"Function has a parameter type that is not allowed in a subroutine: parameter b with type VAR_POSITIONAL",
),
(
fnWithNonExprReturnAnnotation,
"Function has return of disallowed type TealType.uint64. Only Expr is allowed",
"Function has return of disallowed type TealType.uint64. Only Expr is allowed",
),
(
fnWithNonExprParamAnnotation,
"Function has parameter b of declared type TealType.uint64 which is not a class",
"Function has parameter b of declared type TealType.uint64 which is not a class",
),
(
fnWithScratchVarSubclass,
"Function has parameter b of disallowed type <class 'pyteal.DynamicScratchVar'>",
"Function has parameter b of disallowed type <class 'pyteal.DynamicScratchVar'>",
),
(
fnReturningExprSubclass,
"Function has return of disallowed type <class 'pyteal.Return'>",
"Function has return of disallowed type <class 'pyteal.Return'>. Only Expr is allowed",
),
(
fnWithMixedAnns4AndBytesReturn,
"Function has return of disallowed type <class 'pyteal.Bytes'>",
"Function has return of disallowed type <class 'pyteal.Bytes'>. Only Expr is allowed",
),
(
fnWithMixedAnnsABIRet1,
"Function has return of disallowed type pyteal.abi.StaticArray[pyteal.abi.Uint32, typing.Literal[10]]. "
"Only Expr is allowed",
"Function has return of disallowed type pyteal.abi.StaticArray[pyteal.abi.Uint32, typing.Literal[10]]. "
"Only Expr is allowed",
),
(
fnWithMixedAnnsABIRet2,
"Function has return of disallowed type <class 'pyteal.abi.Uint64'>. Only Expr is allowed",
"Function has return of disallowed type <class 'pyteal.abi.Uint64'>. Only Expr is allowed",
),
)
for fn, sub_def_msg, abi_sub_def_msg in cases:
with pytest.raises(pt.TealInputError) as e:
print(f"case=[{sub_def_msg}]")
pt.SubroutineDefinition(fn, pt.TealType.none)
assert sub_def_msg in str(e), f"failed for case [{fn.__name__}]"
with pytest.raises(pt.TealInputError) as e:
print(f"case=[{abi_sub_def_msg}]")
pt.ABIReturnSubroutine(fn)
assert abi_sub_def_msg in str(e), f"failed for case[{fn.__name__}]"
def test_subroutine_invocation_param_types():
def fnWithNoAnnotations(a, b):
return pt.Return()
def fnWithExprAnnotations(a: pt.Expr, b: pt.Expr) -> pt.Expr:
return pt.Return()
def fnWithSVAnnotations(a: pt.ScratchVar, b: pt.ScratchVar):
return pt.Return()
def fnWithABIAnnotations(
a: pt.abi.Byte,
b: pt.abi.StaticArray[pt.abi.Uint32, Literal[10]],
c: pt.abi.DynamicArray[pt.abi.Bool],
):
return pt.Return()
def fnWithMixedAnns1(a: pt.ScratchVar, b: pt.Expr) -> pt.Expr:
return pt.Return()
def fnWithMixedAnns2(a: pt.ScratchVar, b) -> pt.Expr:
return pt.Return()
def fnWithMixedAnns3(a: pt.Expr, b: pt.ScratchVar):
return pt.Return()
def fnWithMixedAnns4(a: pt.ScratchVar, b, c: pt.abi.Uint16) -> pt.Expr:
return pt.Return()
sv = pt.ScratchVar()
x = pt.Int(42)
s = pt.Bytes("hello")
av_u16 = pt.abi.Uint16()
av_bool_dym_arr = pt.abi.DynamicArray(
pt.abi.DynamicArrayTypeSpec(pt.abi.BoolTypeSpec()))
av_u32_static_arr = pt.abi.StaticArray(
pt.abi.StaticArrayTypeSpec(pt.abi.Uint32TypeSpec(), 10))
av_bool = pt.abi.Bool()
av_byte = pt.abi.Byte()
cases = [
("vanilla 1", fnWithNoAnnotations, [x, s], None),
("vanilla 2", fnWithNoAnnotations, [x, x], None),
("vanilla no sv's allowed 1", fnWithNoAnnotations, [x, sv],
pt.TealInputError),
("exprs 1", fnWithExprAnnotations, [x, s], None),
("exprs 2", fnWithExprAnnotations, [x, x], None),
("exprs no sv's allowed 1", fnWithExprAnnotations, [x, sv],
pt.TealInputError),
("all sv's 1", fnWithSVAnnotations, [sv, sv], None),
("all sv's but strings", fnWithSVAnnotations, [s,
s], pt.TealInputError),
("all sv's but ints", fnWithSVAnnotations, [x, x], pt.TealInputError),
(
"all abi's 1",
fnWithABIAnnotations,
[av_byte, av_u32_static_arr, av_bool_dym_arr],
None,
),
(
"all abi's but ints 1",
fnWithABIAnnotations,
[x, av_u32_static_arr, av_bool_dym_arr],
pt.TealInputError,
),
(
"all abi's but ints 2",
fnWithABIAnnotations,
[x, av_u32_static_arr, x],
pt.TealInputError,
),
("all abi's but ints 3", fnWithABIAnnotations, [x, x,
x], pt.TealInputError),
(
"all abi's but sv's 1",
fnWithABIAnnotations,
[sv, av_u32_static_arr, av_bool_dym_arr],
pt.TealInputError,
),
(
"all abi's but sv's 2",
fnWithABIAnnotations,
[av_byte, av_u32_static_arr, sv],
pt.TealInputError,
),
(
"all abi's but sv's 3",
fnWithABIAnnotations,
[av_byte, sv, av_u32_static_arr],
pt.TealInputError,
),
(
"all abi's but wrong typed 1",
fnWithABIAnnotations,
[av_u32_static_arr, av_u32_static_arr, av_bool_dym_arr],
pt.TealInputError,
),
(
"all abi's but wrong typed 2",
fnWithABIAnnotations,
[av_bool, av_bool_dym_arr, av_u16],
pt.TealInputError,
),
(
"all abi's but wrong typed 3",
fnWithABIAnnotations,
[av_u16, av_bool, av_byte],
pt.TealInputError,
),
("mixed1 copacetic", fnWithMixedAnns1, [sv, x], None),
("mixed1 flipped", fnWithMixedAnns1, [x, sv], pt.TealInputError),
("mixed1 missing the sv", fnWithMixedAnns1, [x, s], pt.TealInputError),
("mixed1 missing the non-sv", fnWithMixedAnns1, [sv, sv],
pt.TealInputError),
("mixed2 copacetic", fnWithMixedAnns2, [sv, x], None),
("mixed2 flipped", fnWithMixedAnns2, [x, sv], pt.TealInputError),
("mixed2 missing the sv", fnWithMixedAnns2, [x, s], pt.TealInputError),
("mixed2 missing the non-sv", fnWithMixedAnns2, [sv, sv],
pt.TealInputError),
("mixed3 copacetic", fnWithMixedAnns3, [s, sv], None),
("mixed3 flipped", fnWithMixedAnns3, [sv, x], pt.TealInputError),
("mixed3 missing the sv", fnWithMixedAnns3, [x, s], pt.TealInputError),
("mixed anno", fnWithMixedAnns4, [sv, x, av_u16], None),
(
"mixed anno but wrong typed 1",
fnWithMixedAnns4,
[av_byte, x, av_u16],
pt.TealInputError,
),
(
"mixed anno but wrong typed 2",
fnWithMixedAnns4,
[sv, av_byte, sv],
pt.TealInputError,
),
(
"mixed anno but wrong typed 3",
fnWithMixedAnns4,
[sv, x, av_byte],
pt.TealInputError,
),
]
for case_name, fn, args, err in cases:
definition = pt.SubroutineDefinition(fn, pt.TealType.none)
assert definition.argument_count() == len(args), case_name
assert definition.name() == fn.__name__, case_name
if err is None:
assert len(definition.by_ref_args) == len(
[x for x in args if isinstance(x, pt.ScratchVar)]), case_name
invocation = definition.invoke(args)
assert isinstance(invocation, pt.SubroutineCall), case_name
assert invocation.subroutine is definition, case_name
assert invocation.args == args, case_name
assert invocation.has_return() is False, case_name
else:
try:
with pytest.raises(err):
definition.invoke(args)
except Exception as e:
assert (
not e
), f"EXPECTED ERROR of type {err}. encountered unexpected error during invocation case <{case_name}>: {e}"
def test_assignScratchSlotsToSubroutines_invalid_requested_id():
def sub1Impl():
return None
def sub2Impl(a1):
return None
def sub3Impl(a1, a2, a3):
return None
subroutine1 = pt.SubroutineDefinition(sub1Impl, pt.TealType.uint64)
subroutine2 = pt.SubroutineDefinition(sub2Impl, pt.TealType.bytes)
subroutine3 = pt.SubroutineDefinition(sub3Impl, pt.TealType.none)
globalSlot1 = pt.ScratchSlot(requestedSlotId=8)
subroutine1Slot1 = pt.ScratchSlot()
subroutine1Slot2 = pt.ScratchSlot(requestedSlotId=5)
subroutine1Ops = [
pt.TealOp(None, pt.Op.int, 1),
pt.TealOp(None, pt.Op.store, subroutine1Slot1),
pt.TealOp(None, pt.Op.int, 3),
pt.TealOp(None, pt.Op.store, subroutine1Slot2),
pt.TealOp(None, pt.Op.load, globalSlot1),
pt.TealOp(None, pt.Op.retsub),
]
subroutine2Slot1 = pt.ScratchSlot(requestedSlotId=100)
subroutine2Ops = [
pt.TealOp(None, pt.Op.byte, '"value"'),
pt.TealOp(None, pt.Op.store, subroutine2Slot1),
pt.TealOp(None, pt.Op.load, subroutine2Slot1),
pt.TealOp(None, pt.Op.retsub),
]
subroutine3Ops = [
pt.TealOp(None, pt.Op.retsub),
]
mainSlot1 = pt.ScratchSlot()
mainSlot2 = pt.ScratchSlot(requestedSlotId=100)
mainOps = [
pt.TealOp(None, pt.Op.int, 7),
pt.TealOp(None, pt.Op.store, globalSlot1),
pt.TealOp(None, pt.Op.int, 1),
pt.TealOp(None, pt.Op.store, mainSlot1),
pt.TealOp(None, pt.Op.int, 2),
pt.TealOp(None, pt.Op.store, mainSlot2),
pt.TealOp(None, pt.Op.load, mainSlot1),
pt.TealOp(None, pt.Op.return_),
]
subroutineBlocks = {
None: pt.TealSimpleBlock(mainOps),
subroutine1: pt.TealSimpleBlock(subroutine1Ops),
subroutine2: pt.TealSimpleBlock(subroutine2Ops),
subroutine3: pt.TealSimpleBlock(subroutine3Ops),
}
# mainSlot2 and subroutine2Slot1 request the same ID, 100
with pytest.raises(pt.TealInternalError):
assignScratchSlotsToSubroutines(subroutineBlocks)
def test_assignScratchSlotsToSubroutines_no_requested_ids():
def sub1Impl():
return None
def sub2Impl(a1):
return None
def sub3Impl(a1, a2, a3):
return None
subroutine1 = pt.SubroutineDefinition(sub1Impl, pt.TealType.uint64)
subroutine2 = pt.SubroutineDefinition(sub2Impl, pt.TealType.bytes)
subroutine3 = pt.SubroutineDefinition(sub3Impl, pt.TealType.none)
globalSlot1 = pt.ScratchSlot()
subroutine1Slot1 = pt.ScratchSlot()
subroutine1Slot2 = pt.ScratchSlot()
subroutine1Ops = [
pt.TealOp(None, pt.Op.int, 1),
pt.TealOp(None, pt.Op.store, subroutine1Slot1),
pt.TealOp(None, pt.Op.int, 3),
pt.TealOp(None, pt.Op.store, subroutine1Slot2),
pt.TealOp(None, pt.Op.load, globalSlot1),
pt.TealOp(None, pt.Op.retsub),
]
subroutine2Slot1 = pt.ScratchSlot()
subroutine2Ops = [
pt.TealOp(None, pt.Op.byte, '"value"'),
pt.TealOp(None, pt.Op.store, subroutine2Slot1),
pt.TealOp(None, pt.Op.load, subroutine2Slot1),
pt.TealOp(None, pt.Op.retsub),
]
subroutine3Ops = [
pt.TealOp(None, pt.Op.retsub),
]
mainSlot1 = pt.ScratchSlot()
mainSlot2 = pt.ScratchSlot()
mainOps = [
pt.TealOp(None, pt.Op.int, 7),
pt.TealOp(None, pt.Op.store, globalSlot1),
pt.TealOp(None, pt.Op.int, 1),
pt.TealOp(None, pt.Op.store, mainSlot1),
pt.TealOp(None, pt.Op.int, 2),
pt.TealOp(None, pt.Op.store, mainSlot2),
pt.TealOp(None, pt.Op.load, mainSlot1),
pt.TealOp(None, pt.Op.return_),
]
subroutineBlocks = {
None: pt.TealSimpleBlock(mainOps),
subroutine1: pt.TealSimpleBlock(subroutine1Ops),
subroutine2: pt.TealSimpleBlock(subroutine2Ops),
subroutine3: pt.TealSimpleBlock(subroutine3Ops),
}
expectedAssignments = {
globalSlot1: 0,
subroutine1Slot1: 1,
subroutine1Slot2: 2,
subroutine2Slot1: 3,
mainSlot1: 4,
mainSlot2: 5,
}
expected = {
None: {expectedAssignments[mainSlot1], expectedAssignments[mainSlot2]},
subroutine1: {
expectedAssignments[subroutine1Slot1],
expectedAssignments[subroutine1Slot2],
},
subroutine2: {expectedAssignments[subroutine2Slot1]},
subroutine3: set(),
}
actual = assignScratchSlotsToSubroutines(subroutineBlocks)
assert actual == expected
assert subroutine1Ops == [
pt.TealOp(None, pt.Op.int, 1),
pt.TealOp(None, pt.Op.store, expectedAssignments[subroutine1Slot1]),
pt.TealOp(None, pt.Op.int, 3),
pt.TealOp(None, pt.Op.store, expectedAssignments[subroutine1Slot2]),
pt.TealOp(None, pt.Op.load, expectedAssignments[globalSlot1]),
pt.TealOp(None, pt.Op.retsub),
]
assert subroutine2Ops == [
pt.TealOp(None, pt.Op.byte, '"value"'),
pt.TealOp(None, pt.Op.store, expectedAssignments[subroutine2Slot1]),
pt.TealOp(None, pt.Op.load, expectedAssignments[subroutine2Slot1]),
pt.TealOp(None, pt.Op.retsub),
]
assert subroutine3Ops == [
pt.TealOp(None, pt.Op.retsub),
]
assert mainOps == [
pt.TealOp(None, pt.Op.int, 7),
pt.TealOp(None, pt.Op.store, expectedAssignments[globalSlot1]),
pt.TealOp(None, pt.Op.int, 1),
pt.TealOp(None, pt.Op.store, expectedAssignments[mainSlot1]),
pt.TealOp(None, pt.Op.int, 2),
pt.TealOp(None, pt.Op.store, expectedAssignments[mainSlot2]),
pt.TealOp(None, pt.Op.load, expectedAssignments[mainSlot1]),
pt.TealOp(None, pt.Op.return_),
]
def test_subroutine_definition():
def fn0Args():
return pt.Return()
def fn1Args(a1):
return pt.Return()
def fn2Args(a1, a2):
return pt.Return()
def fn10Args(a1, a2, a3, a4, a5, a6, a7, a8, a9, a10):
return pt.Return()
lam0Args = lambda: pt.Return() # noqa: E731
lam1Args = lambda a1: pt.Return() # noqa: E731
lam2Args = lambda a1, a2: pt.Return() # noqa: E731
lam10Args = (lambda a1, a2, a3, a4, a5, a6, a7, a8, a9, a10: pt.Return()
) # noqa: E731
def fnWithExprAnnotations(a: pt.Expr, b: pt.Expr) -> pt.Expr:
return pt.Return()
def fnWithOnlyReturnExprAnnotations(a, b) -> pt.Expr:
return pt.Return()
def fnWithOnlyArgExprAnnotations(a: pt.Expr, b: pt.Expr):
return pt.Return()
def fnWithPartialExprAnnotations(a, b: pt.Expr) -> pt.Expr:
return pt.Return()
cases = (
(fn0Args, 0, "fn0Args"),
(fn1Args, 1, "fn1Args"),
(fn2Args, 2, "fn2Args"),
(fn10Args, 10, "fn10Args"),
(lam0Args, 0, "<lambda>"),
(lam1Args, 1, "<lambda>"),
(lam2Args, 2, "<lambda>"),
(lam10Args, 10, "<lambda>"),
(fnWithExprAnnotations, 2, "fnWithExprAnnotations"),
(fnWithOnlyReturnExprAnnotations, 2,
"fnWithOnlyReturnExprAnnotations"),
(fnWithOnlyArgExprAnnotations, 2, "fnWithOnlyArgExprAnnotations"),
(fnWithPartialExprAnnotations, 2, "fnWithPartialExprAnnotations"),
)
for (fn, numArgs, name) in cases:
definition = pt.SubroutineDefinition(fn, pt.TealType.none)
assert definition.argument_count() == numArgs
assert definition.name() == name
if numArgs > 0:
with pytest.raises(pt.TealInputError):
definition.invoke([pt.Int(1)] * (numArgs - 1))
with pytest.raises(pt.TealInputError):
definition.invoke([pt.Int(1)] * (numArgs + 1))
if numArgs > 0:
with pytest.raises(pt.TealInputError):
definition.invoke([1] * numArgs)
args = [pt.Int(1)] * numArgs
invocation = definition.invoke(args)
assert isinstance(invocation, pt.SubroutineCall)
assert invocation.subroutine is definition
assert invocation.args == args
def test_flattenSubroutines_multiple_subroutines():
def sub1Impl():
return None
def sub2Impl(a1):
return None
def sub3Impl(a1, a2, a3):
return None
subroutine1 = pt.SubroutineDefinition(sub1Impl, pt.TealType.uint64)
subroutine2 = pt.SubroutineDefinition(sub2Impl, pt.TealType.bytes)
subroutine3 = pt.SubroutineDefinition(sub3Impl, pt.TealType.none)
subroutineToLabel = OrderedDict()
subroutineToLabel[subroutine1] = "sub0"
subroutineToLabel[subroutine2] = "sub1"
subroutineToLabel[subroutine3] = "sub2"
subroutine1Label = pt.LabelReference(subroutineToLabel[subroutine1])
subroutine1Ops = [
pt.TealOp(None, pt.Op.int, 1),
pt.TealOp(None, pt.Op.int, 2),
pt.TealOp(None, pt.Op.int, 3),
pt.TealOp(None, pt.Op.add),
pt.TealOp(None, pt.Op.add),
pt.TealOp(None, pt.Op.retsub),
]
subroutine2Label = pt.LabelReference(subroutineToLabel[subroutine2])
subroutine2L1Label = pt.LabelReference("l1")
subroutine2L2Label = pt.LabelReference("l2")
subroutine2L3Label = pt.LabelReference("l3")
subroutine2L4Label = pt.LabelReference("l4")
subroutine2Ops = [
pt.TealOp(None, pt.Op.dup),
pt.TealOp(None, pt.Op.int, 1),
pt.TealOp(None, pt.Op.eq),
pt.TealOp(None, pt.Op.bnz, subroutine2L1Label),
pt.TealOp(None, pt.Op.dup),
pt.TealOp(None, pt.Op.int, 2),
pt.TealOp(None, pt.Op.eq),
pt.TealOp(None, pt.Op.bnz, subroutine2L2Label),
pt.TealOp(None, pt.Op.dup),
pt.TealOp(None, pt.Op.int, 3),
pt.TealOp(None, pt.Op.eq),
pt.TealOp(None, pt.Op.bnz, subroutine2L3Label),
pt.TealOp(None, pt.Op.dup),
pt.TealOp(None, pt.Op.int, 4),
pt.TealOp(None, pt.Op.eq),
pt.TealOp(None, pt.Op.bnz, subroutine2L4Label),
pt.TealOp(None, pt.Op.err),
pt.TealLabel(None, subroutine2L1Label),
pt.TealOp(None, pt.Op.pop),
pt.TealOp(None, pt.Op.byte, '"1"'),
pt.TealOp(None, pt.Op.retsub),
pt.TealLabel(None, subroutine2L2Label),
pt.TealOp(None, pt.Op.pop),
pt.TealOp(None, pt.Op.byte, '"2"'),
pt.TealOp(None, pt.Op.retsub),
pt.TealLabel(None, subroutine2L3Label),
pt.TealOp(None, pt.Op.pop),
pt.TealOp(None, pt.Op.byte, '"3"'),
pt.TealOp(None, pt.Op.retsub),
pt.TealLabel(None, subroutine2L4Label),
pt.TealOp(None, pt.Op.pop),
pt.TealOp(None, pt.Op.byte, '"4"'),
pt.TealOp(None, pt.Op.retsub),
]
subroutine3Label = pt.LabelReference(subroutineToLabel[subroutine3])
subroutine3L1Label = pt.LabelReference("l1")
subroutine3Ops = [
pt.TealLabel(None, subroutine3L1Label),
pt.TealOp(None, pt.Op.app_local_put),
pt.TealOp(None, pt.Op.retsub),
pt.TealOp(None, pt.Op.b, subroutine3L1Label),
]
l1Label = pt.LabelReference("l1")
mainOps = [
pt.TealOp(None, pt.Op.byte, '"account"'),
pt.TealOp(None, pt.Op.byte, '"key"'),
pt.TealOp(None, pt.Op.byte, '"value"'),
pt.TealOp(None, pt.Op.callsub, subroutine3Label),
pt.TealOp(None, pt.Op.txn, "Fee"),
pt.TealOp(None, pt.Op.int, 0),
pt.TealOp(None, pt.Op.eq),
pt.TealOp(None, pt.Op.bz, l1Label),
pt.TealOp(None, pt.Op.int, 3),
pt.TealOp(None, pt.Op.callsub, subroutine2Label),
pt.TealOp(None, pt.Op.pop),
pt.TealOp(None, pt.Op.callsub, subroutine1Label),
pt.TealOp(None, pt.Op.return_),
pt.TealLabel(None, l1Label),
pt.TealOp(None, pt.Op.int, 0),
pt.TealOp(None, pt.Op.return_),
]
subroutineMapping = {
None: mainOps,
subroutine1: subroutine1Ops,
subroutine2: subroutine2Ops,
subroutine3: subroutine3Ops,
}
expectedL1Label = pt.LabelReference("main_l1")
expectedSubroutine1Label = pt.LabelReference("sub0")
expectedSubroutine2Label = pt.LabelReference("sub1")
expectedSubroutine2L1Label = pt.LabelReference("sub1_l1")
expectedSubroutine2L2Label = pt.LabelReference("sub1_l2")
expectedSubroutine2L3Label = pt.LabelReference("sub1_l3")
expectedSubroutine2L4Label = pt.LabelReference("sub1_l4")
expectedSubroutine3Label = pt.LabelReference("sub2")
expectedSubroutine3L1Label = pt.LabelReference("sub2_l1")
expected = [
pt.TealOp(None, pt.Op.byte, '"account"'),
pt.TealOp(None, pt.Op.byte, '"key"'),
pt.TealOp(None, pt.Op.byte, '"value"'),
pt.TealOp(None, pt.Op.callsub, subroutine3Label),
pt.TealOp(None, pt.Op.txn, "Fee"),
pt.TealOp(None, pt.Op.int, 0),
pt.TealOp(None, pt.Op.eq),
pt.TealOp(None, pt.Op.bz, expectedL1Label),
pt.TealOp(None, pt.Op.int, 3),
pt.TealOp(None, pt.Op.callsub, subroutine2Label),
pt.TealOp(None, pt.Op.pop),
pt.TealOp(None, pt.Op.callsub, subroutine1Label),
pt.TealOp(None, pt.Op.return_),
pt.TealLabel(None, expectedL1Label),
pt.TealOp(None, pt.Op.int, 0),
pt.TealOp(None, pt.Op.return_),
pt.TealLabel(None, expectedSubroutine1Label, "sub1Impl"),
pt.TealOp(None, pt.Op.int, 1),
pt.TealOp(None, pt.Op.int, 2),
pt.TealOp(None, pt.Op.int, 3),
pt.TealOp(None, pt.Op.add),
pt.TealOp(None, pt.Op.add),
pt.TealOp(None, pt.Op.retsub),
pt.TealLabel(None, expectedSubroutine2Label, "sub2Impl"),
pt.TealOp(None, pt.Op.dup),
pt.TealOp(None, pt.Op.int, 1),
pt.TealOp(None, pt.Op.eq),
pt.TealOp(None, pt.Op.bnz, expectedSubroutine2L1Label),
pt.TealOp(None, pt.Op.dup),
pt.TealOp(None, pt.Op.int, 2),
pt.TealOp(None, pt.Op.eq),
pt.TealOp(None, pt.Op.bnz, expectedSubroutine2L2Label),
pt.TealOp(None, pt.Op.dup),
pt.TealOp(None, pt.Op.int, 3),
pt.TealOp(None, pt.Op.eq),
pt.TealOp(None, pt.Op.bnz, expectedSubroutine2L3Label),
pt.TealOp(None, pt.Op.dup),
pt.TealOp(None, pt.Op.int, 4),
pt.TealOp(None, pt.Op.eq),
pt.TealOp(None, pt.Op.bnz, expectedSubroutine2L4Label),
pt.TealOp(None, pt.Op.err),
pt.TealLabel(None, expectedSubroutine2L1Label),
pt.TealOp(None, pt.Op.pop),
pt.TealOp(None, pt.Op.byte, '"1"'),
pt.TealOp(None, pt.Op.retsub),
pt.TealLabel(None, expectedSubroutine2L2Label),
pt.TealOp(None, pt.Op.pop),
pt.TealOp(None, pt.Op.byte, '"2"'),
pt.TealOp(None, pt.Op.retsub),
pt.TealLabel(None, expectedSubroutine2L3Label),
pt.TealOp(None, pt.Op.pop),
pt.TealOp(None, pt.Op.byte, '"3"'),
pt.TealOp(None, pt.Op.retsub),
pt.TealLabel(None, expectedSubroutine2L4Label),
pt.TealOp(None, pt.Op.pop),
pt.TealOp(None, pt.Op.byte, '"4"'),
pt.TealOp(None, pt.Op.retsub),
pt.TealLabel(None, expectedSubroutine3Label, "sub3Impl"),
pt.TealLabel(None, expectedSubroutine3L1Label),
pt.TealOp(None, pt.Op.app_local_put),
pt.TealOp(None, pt.Op.retsub),
pt.TealOp(None, pt.Op.b, expectedSubroutine3L1Label),
]
actual = flattenSubroutines(subroutineMapping, subroutineToLabel)
assert actual == expected
def test_collectScratchSlots():
def sub1Impl():
return None
def sub2Impl(a1):
return None
def sub3Impl(a1, a2, a3):
return None
subroutine1 = pt.SubroutineDefinition(sub1Impl, pt.TealType.uint64)
subroutine2 = pt.SubroutineDefinition(sub2Impl, pt.TealType.bytes)
subroutine3 = pt.SubroutineDefinition(sub3Impl, pt.TealType.none)
globalSlot1 = pt.ScratchSlot()
subroutine1Slot1 = pt.ScratchSlot()
subroutine1Slot2 = pt.ScratchSlot()
subroutine1Ops = [
pt.TealOp(None, pt.Op.int, 1),
pt.TealOp(None, pt.Op.store, subroutine1Slot1),
pt.TealOp(None, pt.Op.int, 3),
pt.TealOp(None, pt.Op.store, subroutine1Slot2),
pt.TealOp(None, pt.Op.load, globalSlot1),
pt.TealOp(None, pt.Op.retsub),
]
subroutine2Slot1 = pt.ScratchSlot()
subroutine2Ops = [
pt.TealOp(None, pt.Op.byte, '"value"'),
pt.TealOp(None, pt.Op.store, subroutine2Slot1),
pt.TealOp(None, pt.Op.load, subroutine2Slot1),
pt.TealOp(None, pt.Op.retsub),
]
subroutine3Ops = [
pt.TealOp(None, pt.Op.retsub),
]
mainSlot1 = pt.ScratchSlot()
mainSlot2 = pt.ScratchSlot()
mainOps = [
pt.TealOp(None, pt.Op.int, 7),
pt.TealOp(None, pt.Op.store, globalSlot1),
pt.TealOp(None, pt.Op.int, 1),
pt.TealOp(None, pt.Op.store, mainSlot1),
pt.TealOp(None, pt.Op.int, 2),
pt.TealOp(None, pt.Op.store, mainSlot2),
pt.TealOp(None, pt.Op.load, mainSlot1),
pt.TealOp(None, pt.Op.return_),
]
subroutineBlocks = {
None: pt.TealSimpleBlock(mainOps),
subroutine1: pt.TealSimpleBlock(subroutine1Ops),
subroutine2: pt.TealSimpleBlock(subroutine2Ops),
subroutine3: pt.TealSimpleBlock(subroutine3Ops),
}
expected_global = {globalSlot1}
expected_local = {
None: {mainSlot1, mainSlot2},
subroutine1: {subroutine1Slot1, subroutine1Slot2},
subroutine2: {subroutine2Slot1},
subroutine3: set(),
}
actual_global, actual_local = collectScratchSlots(subroutineBlocks)
assert actual_global == expected_global
assert actual_local == expected_local