Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
57 changes: 54 additions & 3 deletions pyteal/ast/router_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -418,7 +418,7 @@ def test_wrap_handler_method_call():
)
for abi_subroutine in ONLY_ABI_SUBROUTINE_CASES:
wrapped: pt.Expr = ASTBuilder.wrap_handler(True, abi_subroutine)
assembled_wrapped: pt.TealBlock = assemble_helper(wrapped)
actual: pt.TealBlock = assemble_helper(wrapped)

args: list[pt.abi.BaseType] = [
spec.new_instance()
Expand Down Expand Up @@ -466,9 +466,60 @@ def test_wrap_handler_method_call():
else:
evaluate = abi_subroutine(*args)

actual = assemble_helper(pt.Seq(*loading, evaluate, pt.Approve()))
expected = assemble_helper(pt.Seq(*loading, evaluate, pt.Approve()))
with pt.TealComponent.Context.ignoreScratchSlotEquality(), pt.TealComponent.Context.ignoreExprEquality():
assert actual == assembled_wrapped
assert actual == expected

assert pt.TealBlock.MatchScratchSlotReferences(
pt.TealBlock.GetReferencedScratchSlots(actual),
pt.TealBlock.GetReferencedScratchSlots(expected),
)


def test_wrap_handler_method_call_many_args():
wrapped: pt.Expr = ASTBuilder.wrap_handler(True, many_args)
actual: pt.TealBlock = assemble_helper(wrapped)

args = [pt.abi.Uint64() for _ in range(20)]
last_arg = pt.abi.TupleTypeSpec(
*[pt.abi.Uint64TypeSpec() for _ in range(6)]
).new_instance()

output_temp = pt.abi.Uint64()
expected_ast = pt.Seq(
args[0].decode(pt.Txn.application_args[1]),
args[1].decode(pt.Txn.application_args[2]),
args[2].decode(pt.Txn.application_args[3]),
args[3].decode(pt.Txn.application_args[4]),
args[4].decode(pt.Txn.application_args[5]),
args[5].decode(pt.Txn.application_args[6]),
args[6].decode(pt.Txn.application_args[7]),
args[7].decode(pt.Txn.application_args[8]),
args[8].decode(pt.Txn.application_args[9]),
args[9].decode(pt.Txn.application_args[10]),
args[10].decode(pt.Txn.application_args[11]),
args[11].decode(pt.Txn.application_args[12]),
args[12].decode(pt.Txn.application_args[13]),
args[13].decode(pt.Txn.application_args[14]),
last_arg.decode(pt.Txn.application_args[15]),
last_arg[0].store_into(args[14]),
last_arg[1].store_into(args[15]),
last_arg[2].store_into(args[16]),
last_arg[3].store_into(args[17]),
last_arg[4].store_into(args[18]),
last_arg[5].store_into(args[19]),
many_args(*args).store_into(output_temp),
pt.abi.MethodReturn(output_temp),
pt.Approve(),
)
expected = assemble_helper(expected_ast)
with pt.TealComponent.Context.ignoreScratchSlotEquality(), pt.TealComponent.Context.ignoreExprEquality():
assert actual == expected

assert pt.TealBlock.MatchScratchSlotReferences(
pt.TealBlock.GetReferencedScratchSlots(actual),
pt.TealBlock.GetReferencedScratchSlots(expected),
)


def test_contract_json_obj():
Expand Down
15 changes: 15 additions & 0 deletions pyteal/ir/tealcomponent.py
Original file line number Diff line number Diff line change
Expand Up @@ -67,6 +67,21 @@ def __exit__(self, *args):

@classmethod
def ignoreScratchSlotEquality(cls):
"""When comparing TealOps, do not verify the equality of any ScratchSlot arguments.

This is commonly used in testing to verify the that two control flow graphs contains the
same operations, but may use different ScratchSlots in them. In this case, you will most
likely want to also use use the following code after comparing with this option enabled:

.. code-block:: python
TealBlock.MatchScratchSlotReferences(
TealBlock.GetReferencedScratchSlots(actual),
TealBlock.GetReferencedScratchSlots(expected),
)

This ensures that the ScratchSlot usages between the two control flow graphs is
equivalent. See :any:`TealBlock.MatchScratchSlotReferences` for more info.
"""
return cls.ScratchSlotEqualityContext()


Expand Down