ABI Type subroutine return (#256)

* wtf * update * update to f-str * define void type * update instantiated computed type returnedType for ABI return in subroutine * minor * update stuffs to help infer type annotation of return ABI * minor * minor * minor * minor * minor * try my best to save stuffs * simplify decorator to single function * tear it down * minor * sheeeesh emacs * updates * minor, renaming something * new design, start over * updates * abi fn wrapper for now * minor * minor update on subroutine def * minor fixes * minor fixes * changes * more constraint on void ret * update comment examples * import from abi * update some error msg and comments * testcases partial * upgrade testscripts * Bundle optional refactorings to subroutine.py (#308) * Bundle optional refactorings to subroutine.py * Refactor to remove branching * storing local changes * pr review partly * pr review partly * update test script * Abi subroutine feature merge (#315) * resolve conflicts * lint * minor reconstruct * missing imports * missing requirements from bad merge * cooperate with typespec n var def change * update comments * some comments resolving? * trim * update some comments * bring testcase back * restriction on output kwarg name * stop ci! i am reformatting * simplify name constrain * resolving comments * Use deferred subroutine expression for ABI returns (#328) * Add subroutine deferred expr * Allow multiple deferred blocks * Remove error * flake8 * Make test more realistic * Add second test * clean up code and comments * remove deferred blocks function * Add coverage for multiple ops in block error * hack return from abi var * a hack * better returned type * simplify * add some compiler test * flake break no way * PR comment resolve partial * polishing Co-authored-by: Michael Diamant <michaeldiamant@users.noreply.github.com> Co-authored-by: Zeph Grunschlag <tzaffi@users.noreply.github.com> Co-authored-by: Zeph Grunschlag <zeph@algorand.com> Co-authored-by: Jason Paulos <jasonpaulos@users.noreply.github.com>
algorand · May 11, 2022 · c72a7ef · c72a7ef
1 parent 0bd9c01
commit c72a7ef
Show file tree

Hide file tree

Showing 10 changed files with 1,244 additions and 177 deletions.
diff --git a/pyteal/__init__.pyi b/pyteal/__init__.pyi
@@ -23,6 +23,7 @@ from pyteal.errors import (
 from pyteal.config import MAX_GROUP_SIZE, NUM_SLOTS
 
 __all__ = [
+    "ABIReturnSubroutine",
     "AccountParam",
     "Add",
     "Addr",

diff --git a/pyteal/ast/__init__.py b/pyteal/ast/__init__.py
@@ -102,7 +102,7 @@
 from pyteal.ast.substring import Substring, Extract, Suffix
 
 # more ops
-from pyteal.ast.naryexpr import NaryExpr, And, Add, Mul, Or, Concat
+from pyteal.ast.naryexpr import NaryExpr, Add, And, Mul, Or, Concat
 from pyteal.ast.widemath import WideRatio
 
 # control flow
@@ -118,6 +118,7 @@
     SubroutineDeclaration,
     SubroutineCall,
     SubroutineFnWrapper,
+    ABIReturnSubroutine,
 )
 from pyteal.ast.while_ import While
 from pyteal.ast.for_ import For
@@ -242,6 +243,7 @@
     "SubroutineDeclaration",
     "SubroutineCall",
     "SubroutineFnWrapper",
+    "ABIReturnSubroutine",
     "ScratchIndex",
     "ScratchLoad",
     "ScratchSlot",

diff --git a/pyteal/ast/abi/__init__.py b/pyteal/ast/abi/__init__.py
@@ -4,7 +4,7 @@
     Address,
     AddressLength,
 )
-from pyteal.ast.abi.type import TypeSpec, BaseType, ComputedValue
+from pyteal.ast.abi.type import TypeSpec, BaseType, ComputedValue, ReturnedValue
 from pyteal.ast.abi.bool import BoolTypeSpec, Bool
 from pyteal.ast.abi.uint import (
     UintTypeSpec,
@@ -47,6 +47,7 @@
     "TypeSpec",
     "BaseType",
     "ComputedValue",
+    "ReturnedValue",
     "BoolTypeSpec",
     "Bool",
     "UintTypeSpec",

diff --git a/pyteal/ast/abi/type.py b/pyteal/ast/abi/type.py
@@ -147,7 +147,7 @@ class ComputedValue(ABC, Generic[T]):
     """Represents an ABI Type whose value must be computed by an expression."""
 
     @abstractmethod
-    def produced_type_spec(cls) -> TypeSpec:
+    def produced_type_spec(self) -> TypeSpec:
         """Get the ABI TypeSpec that this object produces."""
         pass
 
@@ -182,3 +182,40 @@ def use(self, action: Callable[[T], Expr]) -> Expr:
 
 
 ComputedValue.__module__ = "pyteal"
+
+
+class ReturnedValue(ComputedValue):
+    def __init__(self, type_spec: TypeSpec, computation_expr: Expr):
+        from pyteal.ast.subroutine import SubroutineCall
+
+        self.type_spec = type_spec
+        if not isinstance(computation_expr, SubroutineCall):
+            raise TealInputError(
+                f"Expecting computation_expr to be SubroutineCall but get {type(computation_expr)}"
+            )
+        self.computation = computation_expr
+
+    def produced_type_spec(self) -> TypeSpec:
+        return self.type_spec
+
+    def store_into(self, output: BaseType) -> Expr:
+        if output.type_spec() != self.produced_type_spec():
+            raise TealInputError(
+                f"expected type_spec {self.produced_type_spec()} but get {output.type_spec()}"
+            )
+
+        declaration = self.computation.subroutine.get_declaration()
+
+        if declaration.deferred_expr is None:
+            raise TealInputError(
+                "ABI return subroutine must have deferred_expr to be not-None."
+            )
+        if declaration.deferred_expr.type_of() != output.type_spec().storage_type():
+            raise TealInputError(
+                f"ABI return subroutine deferred_expr is expected to be typed {output.type_spec().storage_type()}, "
+                f"but has type {declaration.deferred_expr.type_of()}."
+            )
+        return output.stored_value.slot.store(self.computation)
+
+
+ReturnedValue.__module__ = "pyteal"
diff --git a/pyteal/ast/abi/util.py b/pyteal/ast/abi/util.py
@@ -12,7 +12,7 @@ def substringForDecoding(
     *,
     startIndex: Expr = None,
     endIndex: Expr = None,
-    length: Expr = None
+    length: Expr = None,
 ) -> Expr:
     """A helper function for getting the substring to decode according to the rules of BaseType.decode."""
     if length is not None and endIndex is not None: