ABI Router implementation and test

pyteal/__init__.py

@@ -17,7 +17,12 @@
     TealInputError,
     TealCompileError,
 )
-from pyteal.config import MAX_GROUP_SIZE, NUM_SLOTS
+from pyteal.config import (
+    MAX_GROUP_SIZE,
+    NUM_SLOTS,
+    RETURN_HASH_PREFIX,
+    METHOD_ARG_NUM_CUTOFF,
+)
 
 # begin __all__
 __all__ = (
@@ -37,6 +42,8 @@
         "TealCompileError",
         "MAX_GROUP_SIZE",
         "NUM_SLOTS",
+        "RETURN_HASH_PREFIX",
+        "METHOD_ARG_NUM_CUTOFF",
     ]
 )
 # end __all__

pyteal/__init__.pyi

@@ -20,7 +20,12 @@ from pyteal.errors import (
     TealInputError,
     TealCompileError,
 )
-from pyteal.config import MAX_GROUP_SIZE, NUM_SLOTS
+from pyteal.config import (
+    MAX_GROUP_SIZE,
+    NUM_SLOTS,
+    RETURN_HASH_PREFIX,
+    METHOD_ARG_NUM_CUTOFF,
+)
 
 __all__ = [
     "ABIReturnSubroutine",
@@ -38,6 +43,7 @@ __all__ = [
     "AssetHolding",
     "AssetParam",
     "Balance",
+    "BareCallActions",
     "BinaryExpr",
     "BitLen",
     "BitwiseAnd",
@@ -64,6 +70,7 @@ __all__ = [
     "BytesSqrt",
     "BytesXor",
     "BytesZero",
+    "CallConfig",
     "CompileOptions",
     "Concat",
     "Cond",
@@ -117,8 +124,10 @@ __all__ = [
     "Lt",
     "MAX_GROUP_SIZE",
     "MAX_TEAL_VERSION",
+    "METHOD_ARG_NUM_CUTOFF",
     "MIN_TEAL_VERSION",
     "MaybeValue",
+    "MethodConfig",
     "MethodSignature",
     "MinBalance",
     "Minus",
@@ -132,14 +141,17 @@ __all__ = [
     "Nonce",
     "Not",
     "OnComplete",
+    "OnCompleteAction",
     "Op",
     "OpUp",
     "OpUpMode",
     "OptimizeOptions",
     "Or",
     "Pop",
+    "RETURN_HASH_PREFIX",
     "Reject",
     "Return",
+    "Router",
     "ScratchIndex",
     "ScratchLoad",
     "ScratchSlot",

pyteal/ast/__init__.py

@@ -125,7 +125,6 @@
 from pyteal.ast.break_ import Break
 from pyteal.ast.continue_ import Continue
 
-
 # misc
 from pyteal.ast.scratch import (
     ScratchIndex,
@@ -139,6 +138,13 @@
 from pyteal.ast.multi import MultiValue
 from pyteal.ast.opup import OpUp, OpUpMode
 from pyteal.ast.ecdsa import EcdsaCurve, EcdsaVerify, EcdsaDecompress, EcdsaRecover
+from pyteal.ast.router import (
+    Router,
+    CallConfig,
+    MethodConfig,
+    OnCompleteAction,
+    BareCallActions,
+)
 
 # abi
 import pyteal.ast.abi as abi  # noqa: I250
@@ -280,6 +286,11 @@
     "For",
     "Break",
     "Continue",
+    "Router",
+    "CallConfig",
+    "MethodConfig",
+    "OnCompleteAction",
+    "BareCallActions",
     "abi",
     "EcdsaCurve",
     "EcdsaVerify",

pyteal/ast/abi/__init__.py

@@ -49,6 +49,7 @@
     make,
     size_of,
     type_spec_from_annotation,
+    contains_type_spec,
 )
 
 __all__ = [
@@ -103,4 +104,5 @@
     "size_of",
     "algosdk_from_annotation",
     "algosdk_from_type_spec",
+    "contains_type_spec",
 ]

pyteal/ast/abi/method_return.py

@@ -4,7 +4,7 @@
 from pyteal.errors import TealInputError
 from pyteal.ast import Expr, Log, Concat, Bytes
 from pyteal.ir import TealBlock, TealSimpleBlock, Op
-from pyteal.config import RETURN_METHOD_SELECTOR
+from pyteal.config import RETURN_HASH_PREFIX
 
 if TYPE_CHECKING:
     from pyteal.compiler import CompileOptions
@@ -22,9 +22,9 @@ def __teal__(self, options: "CompileOptions") -> Tuple[TealBlock, TealSimpleBloc
             raise TealInputError(
                 f"current version {options.version} is lower than log's min version {Op.log.min_version}"
             )
-        return Log(
-            Concat(Bytes("base16", RETURN_METHOD_SELECTOR), self.arg.encode())
-        ).__teal__(options)
+        return Log(Concat(Bytes(RETURN_HASH_PREFIX), self.arg.encode())).__teal__(
+            options
+        )
 
     def __str__(self) -> str:
         return f"(MethodReturn {self.arg.type_spec()})"

pyteal/ast/abi/string.py

@@ -1,4 +1,4 @@
-from typing import Union, TypeVar, Sequence, cast
+from typing import Union, Sequence, cast
 from collections.abc import Sequence as CollectionSequence
 
 from pyteal.ast.abi.uint import Byte
@@ -20,9 +20,6 @@ def encoded_string(s: Expr):
     return Concat(Suffix(Itob(Len(s)), Int(6)), s)
 
 
-T = TypeVar("T", bound=BaseType)
-
-
 class StringTypeSpec(DynamicArrayTypeSpec):
     def __init__(self) -> None:
         super().__init__(ByteTypeSpec())
@@ -70,12 +67,7 @@ def set(
 
         match value:
             case ComputedValue():
-                if value.produced_type_spec() == StringTypeSpec():
-                    return value.store_into(self)
-
-                raise TealInputError(
-                    f"Got ComputedValue with type spec {value.produced_type_spec()}, expected StringTypeSpec"
-                )
+                return self._set_with_computed_type(value)
             case BaseType():
                 if value.type_spec() == StringTypeSpec() or (
                     value.type_spec() == DynamicArrayTypeSpec(ByteTypeSpec())

pyteal/ast/abi/util.py

@@ -1,4 +1,4 @@
-from typing import TypeVar, Any, Literal, get_origin, get_args, cast
+from typing import Sequence, TypeVar, Any, Literal, get_origin, get_args, cast
 
 import algosdk.abi
 
@@ -235,6 +235,29 @@ def type_spec_from_annotation(annotation: Any) -> TypeSpec:
 T = TypeVar("T", bound=BaseType)
 
 
+def contains_type_spec(ts: TypeSpec, targets: Sequence[TypeSpec]) -> bool:
+    from pyteal.ast.abi.array_dynamic import DynamicArrayTypeSpec
+    from pyteal.ast.abi.array_static import StaticArrayTypeSpec
+    from pyteal.ast.abi.tuple import TupleTypeSpec
+
+    stack: list[TypeSpec] = [ts]
+
+    while stack:
+        current = stack.pop()
+        if current in targets:
+            return True
+
+        match current:
+            case TupleTypeSpec():
+                stack.extend(current.value_type_specs())
+            case DynamicArrayTypeSpec():
+                stack.append(current.value_type_spec())
+            case StaticArrayTypeSpec():
+                stack.append(current.value_type_spec())
+
+    return False
+
+
 def size_of(t: type[T]) -> int:
     """Get the size in bytes of an ABI type. Must be a static type"""
 

pyteal/ast/return_.py

@@ -73,10 +73,7 @@ def __teal__(self, options: "CompileOptions"):
                 )
             op = Op.return_
 
-        args = []
-        if self.value is not None:
-            args.append(self.value)
-
+        args = [] if self.value is None else [self.value]
         return TealBlock.FromOp(options, TealOp(self, op), *args)
 
     def __str__(self):