Refactoring Router taking CSP expression (#636)

Author: ahangsu

CHANGELOG.md

@@ -3,6 +3,8 @@
 ## Added
 
 ## Changed
+* Improved Router clear state program generation. ([#636](https://github.com/algorand/pyteal/pull/636))
+  * NOTE: a backwards incompatable change was imposed in this PR: previous Clear State Program (CSP) can be constructed in router by registering ABI methods or bare app calls, now one has to use `clear_state` argument in `Router.__init__` to construct the CSP.
 
 ## Fixed
 

docs/abi.rst

@@ -732,13 +732,16 @@ The AVM supports 6 types of OnCompletion options that may be specified on an app
 #. **Update application**, which updates an app, represented by :any:`OnComplete.UpdateApplication`
 #. **Delete application**, which deletes an app, represented by :any:`OnComplete.DeleteApplication`
 
-In PyTeal, you have the ability to register a bare app call handler for each of these actions. Additionally, a bare app call handler must also specify whether the handler can be invoking during an **app creation transaction** (:any:`CallConfig.CREATE`), during a **non-creation app call** (:any:`CallConfig.CALL`), or during **either** (:any:`CallConfig.ALL`).
+.. note::
+    While **clear state** is a valid OnCompletion action, its behavior differs significantly from the others. For this reason, the :any:`Router` does not support bare app calls or methods to be called during clear state. Instead, you may use the :code:`clear_state` argument in :any:`Router.__init__` to do work during clear state.
+
+In PyTeal, you have the ability to register a bare app call handler for each of these actions, except for clear state. Additionally, a bare app call handler must also specify whether the handler can be invoking during an **app creation transaction** (:any:`CallConfig.CREATE`), during a **non-creation app call** (:any:`CallConfig.CALL`), or during **either** (:any:`CallConfig.ALL`).
 
 The :any:`BareCallActions` class is used to define a bare app call handler for on completion actions. Each bare app call handler must be an instance of the :any:`OnCompleteAction` class.
 
 The :any:`OnCompleteAction` class is responsible for holding the actual code for the bare app call handler (an instance of either :code:`Expr` or a subroutine that takes no args and returns nothing) as well as a :any:`CallConfig` option that indicates whether the action is able to be called during a creation app call, a non-creation app call, or either.
 
-All the bare app calls that an application wishes to support must be provided to the :any:`Router.__init__` method.
+All the bare app calls that an application wishes to support must be provided to the :any:`Router.__init__` method. Additionally, if you wish to perform actions during clear state, you can specify the :code:`clear_state` argument.
 
 A brief example is below:
 
@@ -785,6 +788,7 @@ A brief example is below:
                 action=assert_sender_is_creator, call_config=CallConfig.CALL
             ),
         ),
+        clear_state=Approve(),
     )
 
 .. note::
@@ -820,6 +824,7 @@ The first way to register a method is with the :any:`Router.add_method_handler`
             opt_in=OnCompleteAction(action=Approve(), call_config=CallConfig.CALL),
             close_out=OnCompleteAction(action=Approve(), call_config=CallConfig.CALL),
         ),
+        clear_state=Approve(),
     )
 
     @ABIReturnSubroutine

examples/application/abi/algobank.py

@@ -25,9 +25,6 @@ def assert_sender_is_creator() -> Expr:
         close_out=OnCompleteAction(
             action=transfer_balance_to_lost, call_config=CallConfig.CALL
         ),
-        clear_state=OnCompleteAction(
-            action=transfer_balance_to_lost, call_config=CallConfig.CALL
-        ),
         # only the creator can update or delete the app
         update_application=OnCompleteAction(
             action=assert_sender_is_creator, call_config=CallConfig.CALL
@@ -36,6 +33,7 @@ def assert_sender_is_creator() -> Expr:
             action=assert_sender_is_creator, call_config=CallConfig.CALL
         ),
     ),
+    clear_state=transfer_balance_to_lost,
 )
 
 

examples/application/abi/algobank_clear_state.teal

@@ -1,10 +1,4 @@
 #pragma version 6
-txn NumAppArgs
-int 0
-==
-bnz main_l2
-err
-main_l2:
 byte "lost"
 byte "lost"
 app_global_get

pyteal/ast/router.py

@@ -63,19 +63,6 @@ def approval_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"
 
@@ -96,6 +83,15 @@ class MethodConfig:
     update_application: CallConfig = field(kw_only=True, default=CallConfig.NEVER)
     delete_application: CallConfig = field(kw_only=True, default=CallConfig.NEVER)
 
+    def __post_init__(self):
+        if self.clear_state != CallConfig.NEVER:
+            raise TealInputError(
+                "Attempt to construct clear state program from MethodConfig: "
+                "Use Router top level argument `clear_state` instead. "
+                "For more details please refer to "
+                "https://pyteal.readthedocs.io/en/latest/abi.html#registering-bare-app-calls"
+            )
+
     def is_never(self) -> bool:
         return all(map(lambda cc: cc == CallConfig.NEVER, astuple(self)))
 
@@ -128,8 +124,11 @@ def approval_cond(self) -> Expr | int:
                         )
             return Or(*cond_list)
 
-    def clear_state_cond(self) -> Expr | int:
-        return self.clear_state.clear_state_condition_under_config()
+
+MethodConfig.__module__ = "pyteal"
+
+
+ActionType = Expr | SubroutineFnWrapper | ABIReturnSubroutine
 
 
 @dataclass(frozen=True)
@@ -138,9 +137,7 @@ class OnCompleteAction:
     OnComplete Action, registers bare calls to one single OnCompletion case.
     """
 
-    action: Optional[Expr | SubroutineFnWrapper | ABIReturnSubroutine] = field(
-        kw_only=True, default=None
-    )
+    action: Optional[ActionType] = field(kw_only=True, default=None)
     call_config: CallConfig = field(kw_only=True, default=CallConfig.NEVER)
 
     def __post_init__(self):
@@ -154,21 +151,15 @@ def never() -> "OnCompleteAction":
         return OnCompleteAction()
 
     @staticmethod
-    def create_only(
-        f: Expr | SubroutineFnWrapper | ABIReturnSubroutine,
-    ) -> "OnCompleteAction":
+    def create_only(f: ActionType) -> "OnCompleteAction":
         return OnCompleteAction(action=f, call_config=CallConfig.CREATE)
 
     @staticmethod
-    def call_only(
-        f: Expr | SubroutineFnWrapper | ABIReturnSubroutine,
-    ) -> "OnCompleteAction":
+    def call_only(f: ActionType) -> "OnCompleteAction":
         return OnCompleteAction(action=f, call_config=CallConfig.CALL)
 
     @staticmethod
-    def always(
-        f: Expr | SubroutineFnWrapper | ABIReturnSubroutine,
-    ) -> "OnCompleteAction":
+    def always(f: ActionType) -> "OnCompleteAction":
         return OnCompleteAction(action=f, call_config=CallConfig.ALL)
 
     def is_empty(self) -> bool:
@@ -197,6 +188,15 @@ class BareCallActions:
         kw_only=True, default=OnCompleteAction.never()
     )
 
+    def __post_init__(self):
+        if not self.clear_state.is_empty():
+            raise TealInputError(
+                "Attempt to construct clear state program from bare app call: "
+                "Use Router top level argument `clear_state` instead. "
+                "For more details please refer to "
+                "https://pyteal.readthedocs.io/en/latest/abi.html#registering-bare-app-calls"
+            )
+
     def is_empty(self) -> bool:
         for action_field in fields(self):
             action: OnCompleteAction = getattr(self, action_field.name)
@@ -220,7 +220,7 @@ def approval_construction(self) -> Optional[Expr]:
                 continue
             wrapped_handler = ASTBuilder.wrap_handler(
                 False,
-                cast(Expr | SubroutineFnWrapper | ABIReturnSubroutine, oca.action),
+                cast(ActionType, oca.action),
             )
             match oca.call_config:
                 case CallConfig.ALL:
@@ -246,21 +246,6 @@ def approval_construction(self) -> Optional[Expr]:
             )
         return Cond(*[[n.condition, n.branch] for n in conditions_n_branches])
 
-    def clear_state_construction(self) -> Optional[Expr]:
-        if self.clear_state.is_empty():
-            return None
-
-        # 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,
-            ),
-        )
-
 
 BareCallActions.__module__ = "pyteal"
 
@@ -280,7 +265,7 @@ class ASTBuilder:
 
     @staticmethod
     def wrap_handler(
-        is_method_call: bool, handler: ABIReturnSubroutine | SubroutineFnWrapper | Expr
+        is_method_call: bool, handler: ActionType, wrap_to_name: str | None = None
     ) -> Expr:
         """This is a helper function that handles transaction arguments passing in bare-app-call/abi-method handlers.
 
@@ -302,11 +287,13 @@ def wrap_handler(
                   passed in ABIReturnSubroutine and logged, then approve.
         """
         if not is_method_call:
+            wrap_to_name = "bare appcall" if wrap_to_name is None else wrap_to_name
+
             match handler:
                 case Expr():
                     if handler.type_of() != TealType.none:
                         raise TealInputError(
-                            f"bare appcall handler should be TealType.none not {handler.type_of()}."
+                            f"{wrap_to_name} handler should be TealType.none not {handler.type_of()}."
                         )
                     return handler if handler.has_return() else Seq(handler, Approve())
                 case SubroutineFnWrapper():
@@ -316,7 +303,7 @@ def wrap_handler(
                         )
                     if handler.subroutine.argument_count() != 0:
                         raise TealInputError(
-                            f"subroutine call should take 0 arg for bare-app call. "
+                            f"subroutine call should take 0 arg for {wrap_to_name}. "
                             f"this subroutine takes {handler.subroutine.argument_count()}."
                         )
                     return Seq(handler(), Approve())
@@ -327,22 +314,23 @@ def wrap_handler(
                         )
                     if handler.subroutine.argument_count() != 0:
                         raise TealInputError(
-                            f"abi-returning subroutine call should take 0 arg for bare-app call. "
+                            f"abi-returning subroutine call should take 0 arg for {wrap_to_name}. "
                             f"this abi-returning subroutine takes {handler.subroutine.argument_count()}."
                         )
                     return Seq(cast(Expr, handler()), Approve())
                 case _:
                     raise TealInputError(
-                        "bare appcall can only accept: none type Expr, or Subroutine/ABIReturnSubroutine with none return and no arg"
+                        f"{wrap_to_name} can only accept: none type Expr, or Subroutine/ABIReturnSubroutine with none return and no arg"
                     )
         else:
+            wrap_to_name = "method call" if wrap_to_name is None else wrap_to_name
             if not isinstance(handler, ABIReturnSubroutine):
                 raise TealInputError(
-                    f"method call should be only registering ABIReturnSubroutine, got {type(handler)}."
+                    f"{wrap_to_name} should be only registering ABIReturnSubroutine, got {type(handler)}."
                 )
             if not handler.is_abi_routable():
                 raise TealInputError(
-                    f"method call ABIReturnSubroutine is not routable "
+                    f"{wrap_to_name} ABIReturnSubroutine is not routable "
                     f"got {handler.subroutine.argument_count()} args with {len(handler.subroutine.abi_args)} ABI args."
                 )
 
@@ -508,19 +496,28 @@ def __init__(
         name: str,
         bare_calls: BareCallActions | None = None,
         descr: str | None = None,
+        *,
+        clear_state: Optional[ActionType] = None,
     ) -> None:
         """
         Args:
             name: the name of the smart contract, used in the JSON object.
             bare_calls: the bare app call registered for each on_completion.
             descr: a description of the smart contract, used in the JSON object.
+            clear_state: an expression describing the behavior of clear state program. This
+                expression will be the entirety of the clear state program; no additional code is
+                inserted by the Router. If not provided, the clear state program will always reject.
         """
 
         self.name: str = name
         self.descr = descr
 
         self.approval_ast = ASTBuilder()
-        self.clear_state_ast = ASTBuilder()
+        self.clear_state: Expr = (
+            Reject()
+            if clear_state is None
+            else ASTBuilder.wrap_handler(False, clear_state, "clear state call")
+        )
 
         self.methods: list[sdk_abi.Method] = []
         self.method_sig_to_selector: dict[str, bytes] = dict()
@@ -535,14 +532,6 @@ def __init__(
                         cast(Expr, bare_call_approval),
                     )
                 )
-            bare_call_clear = bare_calls.clear_state_construction()
-            if bare_call_clear:
-                self.clear_state_ast.conditions_n_branches.append(
-                    CondNode(
-                        Txn.application_args.length() == Int(0),
-                        cast(Expr, bare_call_clear),
-                    )
-                )
 
     def add_method_handler(
         self,
@@ -594,13 +583,9 @@ def add_method_handler(
         self.method_selector_to_sig[method_selector] = method_signature
 
         method_approval_cond = method_config.approval_cond()
-        method_clear_state_cond = method_config.clear_state_cond()
         self.approval_ast.add_method_to_ast(
             method_signature, method_approval_cond, method_call
         )
-        self.clear_state_ast.add_method_to_ast(
-            method_signature, method_clear_state_cond, method_call
-        )
         return method_call
 
     def method(
@@ -637,21 +622,31 @@ def method(
             opt_in (optional): The allowed calls during :code:`OnComplete.OptIn`.
             close_out (optional): The allowed calls during :code:`OnComplete.CloseOut`.
             clear_state (optional): The allowed calls during :code:`OnComplete.ClearState`.
+                This argument has been deprecated, and will error on compile time if one wants to access it.
+                Use Router top level argument `clear_state` instead.
             update_application (optional): The allowed calls during :code:`OnComplete.UpdateApplication`.
             delete_application (optional): The allowed calls during :code:`OnComplete.DeleteApplication`.
         """
         # we use `is None` extensively for CallConfig to distinguish 2 following cases
         # - None
         # - CallConfig.Never
         # both cases evaluate to False in if statement.
+
+        if clear_state is not None:
+            raise TealInputError(
+                "Attempt to register ABI method for clear state program: "
+                "Use Router top level argument `clear_state` instead. "
+                "For more details please refer to "
+                "https://pyteal.readthedocs.io/en/latest/abi.html#registering-bare-app-calls"
+            )
+
         def wrap(_func) -> ABIReturnSubroutine:
             wrapped_subroutine = ABIReturnSubroutine(_func, overriding_name=name)
             call_configs: MethodConfig
             if (
                 no_op is None
                 and opt_in is None
                 and close_out is None
-                and clear_state is None
                 and update_application is None
                 and delete_application is None
             ):
@@ -664,15 +659,13 @@ def none_to_never(x: None | CallConfig):
                 _no_op = none_to_never(no_op)
                 _opt_in = none_to_never(opt_in)
                 _close_out = none_to_never(close_out)
-                _clear_state = none_to_never(clear_state)
                 _update_app = none_to_never(update_application)
                 _delete_app = none_to_never(delete_application)
 
                 call_configs = MethodConfig(
                     no_op=_no_op,
                     opt_in=_opt_in,
                     close_out=_close_out,
-                    clear_state=_clear_state,
                     update_application=_update_app,
                     delete_application=_delete_app,
                 )
@@ -714,7 +707,7 @@ def build_program(self) -> tuple[Expr, Expr, sdk_abi.Contract]:
         """
         return (
             self.approval_ast.program_construction(),
-            self.clear_state_ast.program_construction(),
+            self.clear_state,
             self.contract_construct(),
         )