Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Changes to router clear state #386

Merged
merged 3 commits into from
Jun 7, 2022
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
56 changes: 33 additions & 23 deletions pyteal/ast/router.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@
from pyteal.ast.methodsig import MethodSignature
from pyteal.ast.naryexpr import And, Or
from pyteal.ast.txn import Txn
from pyteal.ast.return_ import Approve
from pyteal.ast.return_ import Approve, Reject


class CallConfig(IntFlag):
Expand All @@ -50,7 +50,7 @@ class CallConfig(IntFlag):
CREATE = 2
ALL = 3

def condition_under_config(self) -> Expr | int:
def approval_condition_under_config(self) -> Expr | int:
match self:
case CallConfig.NEVER:
return 0
Expand All @@ -63,6 +63,19 @@ def condition_under_config(self) -> Expr | int:
case _:
raise TealInternalError(f"unexpected CallConfig {self}")

def clear_state_condition_under_config(self) -> int:
match self:
case CallConfig.NEVER:
return 0
case CallConfig.CALL:
return 1
case CallConfig.CREATE | CallConfig.ALL:
raise TealInputError(
"Only CallConfig.CALL or CallConfig.NEVER are valid for a clear state CallConfig, since clear state can never be invoked during creation"
)
case _:
raise TealInputError(f"unexpected CallConfig {self}")


CallConfig.__module__ = "pyteal"

Expand Down Expand Up @@ -101,7 +114,7 @@ def approval_cond(self) -> Expr | int:
else:
cond_list = []
for config, oc in config_oc_pairs:
config_cond = config.condition_under_config()
config_cond = config.approval_condition_under_config()
match config_cond:
case Expr():
cond_list.append(And(Txn.on_completion() == oc, config_cond))
Expand All @@ -116,7 +129,7 @@ def approval_cond(self) -> Expr | int:
return Or(*cond_list)

def clear_state_cond(self) -> Expr | int:
return self.clear_state.condition_under_config()
return self.clear_state.clear_state_condition_under_config()


@dataclass(frozen=True)
Expand Down Expand Up @@ -214,7 +227,11 @@ def approval_construction(self) -> Optional[Expr]:
cond_body = wrapped_handler
case CallConfig.CALL | CallConfig.CREATE:
cond_body = Seq(
Assert(cast(Expr, oca.call_config.condition_under_config())),
Assert(
cast(
Expr, oca.call_config.approval_condition_under_config()
)
),
wrapped_handler,
)
case _:
Expand All @@ -233,29 +250,16 @@ def clear_state_construction(self) -> Optional[Expr]:
if self.clear_state.is_empty():
return None

wrapped_handler = ASTBuilder.wrap_handler(
# call this to make sure we error if the CallConfig is CREATE or ALL
self.clear_state.call_config.clear_state_condition_under_config()

return ASTBuilder.wrap_handler(
False,
cast(
Expr | SubroutineFnWrapper | ABIReturnSubroutine,
self.clear_state.action,
),
)
match self.clear_state.call_config:
case CallConfig.ALL:
return wrapped_handler
case CallConfig.CALL | CallConfig.CREATE:
return Seq(
Assert(
cast(
Expr, self.clear_state.call_config.condition_under_config()
)
),
wrapped_handler,
)
case _:
raise TealInternalError(
f"Unexpected CallConfig: {self.clear_state.call_config!r}"
)


BareCallActions.__module__ = "pyteal"
Expand Down Expand Up @@ -475,7 +479,7 @@ def add_method_to_ast(

def program_construction(self) -> Expr:
if not self.conditions_n_branches:
raise TealInputError("ABIRouter: Cannot build program with an empty AST")
return Reject()
return Cond(*[[n.condition, n.branch] for n in self.conditions_n_branches])


Expand Down Expand Up @@ -659,6 +663,9 @@ def build_program(self) -> tuple[Expr, Expr, sdk_abi.Contract]:
Constructs ASTs for approval and clear-state programs from the registered methods in the router,
also generates a JSON object of contract to allow client read and call the methods easily.

Note that if no methods or bare app call actions have been registered to either the approval
or clear state programs, then that program will reject all transactions.

Returns:
approval_program: AST for approval program
clear_state_program: AST for clear-state program
Expand All @@ -681,6 +688,9 @@ def compile_program(
Combining `build_program` and `compileTeal`, compiles built Approval and ClearState programs
and returns Contract JSON object for off-chain calling.

Note that if no methods or bare app call actions have been registered to either the approval
or clear state programs, then that program will reject all transactions.

Returns:
approval_program: compiled approval program
clear_state_program: compiled clear-state program
Expand Down
198 changes: 181 additions & 17 deletions pyteal/ast/router_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -266,20 +266,43 @@ def camel_to_snake(name: str) -> str:

def test_call_config():
for cc in pt.CallConfig:
cond_on_cc: pt.Expr | int = cc.condition_under_config()
match cond_on_cc:
approval_cond_on_cc: pt.Expr | int = cc.approval_condition_under_config()
match approval_cond_on_cc:
case pt.Expr():
expected_cc = (
(pt.Txn.application_id() == pt.Int(0))
if cc == pt.CallConfig.CREATE
else (pt.Txn.application_id() != pt.Int(0))
)
with pt.TealComponent.Context.ignoreExprEquality():
assert assemble_helper(cond_on_cc) == assemble_helper(expected_cc)
assert assemble_helper(approval_cond_on_cc) == assemble_helper(
expected_cc
)
case int():
assert cond_on_cc == int(cc) & 1
assert approval_cond_on_cc == int(cc) & 1
case _:
raise pt.TealInternalError(
f"unexpected approval_cond_on_cc {approval_cond_on_cc}"
)

if cc in (pt.CallConfig.CREATE, pt.CallConfig.ALL):
with pytest.raises(
pt.TealInputError,
match=r"Only CallConfig.CALL or CallConfig.NEVER are valid for a clear state CallConfig, since clear state can never be invoked during creation$",
):
cc.clear_state_condition_under_config()
continue

clear_state_cond_on_cc: int = cc.clear_state_condition_under_config()
match clear_state_cond_on_cc:
case 0:
assert cc == pt.CallConfig.NEVER
case 1:
assert cc == pt.CallConfig.CALL
case _:
raise pt.TealInternalError(f"unexpected cond_on_cc {cond_on_cc}")
raise pt.TealInternalError(
f"unexpected clear_state_cond_on_cc {clear_state_cond_on_cc}"
)


def test_method_config():
Expand All @@ -304,18 +327,14 @@ def test_method_config():
match mc.clear_state:
case pt.CallConfig.NEVER:
assert mc.clear_state_cond() == 0
case pt.CallConfig.ALL:
assert mc.clear_state_cond() == 1
case pt.CallConfig.CALL:
with pt.TealComponent.Context.ignoreExprEquality():
assert assemble_helper(
mc.clear_state_cond()
) == assemble_helper(pt.Txn.application_id() != pt.Int(0))
case pt.CallConfig.CREATE:
with pt.TealComponent.Context.ignoreExprEquality():
assert assemble_helper(
mc.clear_state_cond()
) == assemble_helper(pt.Txn.application_id() == pt.Int(0))
assert mc.clear_state_cond() == 1
case pt.CallConfig.CREATE | pt.CallConfig.ALL:
with pytest.raises(
pt.TealInputError,
match=r"Only CallConfig.CALL or CallConfig.NEVER are valid for a clear state CallConfig, since clear state can never be invoked during creation$",
):
mc.clear_state_cond()
if mc.is_never() or all(
getattr(mc, i) == pt.CallConfig.NEVER
for i, _ in approval_check_names_n_ocs
Expand All @@ -330,7 +349,9 @@ def test_method_config():
continue
list_of_cc = [
(
typing.cast(pt.CallConfig, getattr(mc, i)).condition_under_config(),
typing.cast(
pt.CallConfig, getattr(mc, i)
).approval_condition_under_config(),
oc,
)
for i, oc in approval_check_names_n_ocs
Expand Down Expand Up @@ -624,3 +645,146 @@ def test_contract_json_obj():
sdk_contract = sdk_abi.Contract(contract_name, method_list)
contract = router.contract_construct()
assert contract == sdk_contract


def test_build_program_all_empty():
router = pt.Router("test")

approval, clear_state, contract = router.build_program()

expected_empty_program = pt.TealSimpleBlock(
[
pt.TealOp(None, pt.Op.int, 0),
pt.TealOp(None, pt.Op.return_),
]
)

with pt.TealComponent.Context.ignoreExprEquality():
assert assemble_helper(approval) == expected_empty_program
assert assemble_helper(clear_state) == expected_empty_program

expected_contract = sdk_abi.Contract("test", [])
assert contract == expected_contract


def test_build_program_approval_empty():
router = pt.Router(
"test",
pt.BareCallActions(clear_state=pt.OnCompleteAction.call_only(pt.Approve())),
)

approval, clear_state, contract = router.build_program()

expected_empty_program = pt.TealSimpleBlock(
[
pt.TealOp(None, pt.Op.int, 0),
pt.TealOp(None, pt.Op.return_),
]
)

with pt.TealComponent.Context.ignoreExprEquality():
assert assemble_helper(approval) == expected_empty_program
assert assemble_helper(clear_state) != expected_empty_program

expected_contract = sdk_abi.Contract("test", [])
assert contract == expected_contract


def test_build_program_clear_state_empty():
router = pt.Router(
"test", pt.BareCallActions(no_op=pt.OnCompleteAction.always(pt.Approve()))
)

approval, clear_state, contract = router.build_program()

expected_empty_program = pt.TealSimpleBlock(
[
pt.TealOp(None, pt.Op.int, 0),
pt.TealOp(None, pt.Op.return_),
]
)

with pt.TealComponent.Context.ignoreExprEquality():
assert assemble_helper(approval) != expected_empty_program
assert assemble_helper(clear_state) == expected_empty_program

expected_contract = sdk_abi.Contract("test", [])
assert contract == expected_contract


def test_build_program_clear_state_invalid_config():
for config in (pt.CallConfig.CREATE, pt.CallConfig.ALL):
bareCalls = pt.BareCallActions(
clear_state=pt.OnCompleteAction(action=pt.Approve(), call_config=config)
)
with pytest.raises(
pt.TealInputError,
match=r"Only CallConfig.CALL or CallConfig.NEVER are valid for a clear state CallConfig, since clear state can never be invoked during creation$",
):
pt.Router("test", bareCalls)

router = pt.Router("test")

@pt.ABIReturnSubroutine
def clear_state_method():
return pt.Approve()

with pytest.raises(
pt.TealInputError,
match=r"Only CallConfig.CALL or CallConfig.NEVER are valid for a clear state CallConfig, since clear state can never be invoked during creation$",
):
router.add_method_handler(
clear_state_method,
method_config=pt.MethodConfig(clear_state=config),
)


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