Using local variable in subroutine through frame pointer

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 # Added
 * Added frame pointer support for subroutine arguments, replacing the previous usage of scratch. ([#562](https://github.com/algorand/pyteal/pull/562))
+* Added frame pointer support for local ABI variables in subroutine. ([#606](https://github.com/algorand/pyteal/pull/606))
 
 # Fixed
 * Allowing the `MethodCall` and `ExecuteMethodCall` to be passed `None` as app_id argument in the case of an app create transaction ([#592](https://github.com/algorand/pyteal/pull/592))

pyteal/ast/abi/type.py

@@ -3,6 +3,7 @@
 
 from pyteal.ast.expr import Expr
 from pyteal.ast.abstractvar import AbstractVar
+from pyteal.ast.frame import FrameVar, FRAME_LOCAL_NUM
 from pyteal.ast.scratchvar import ScratchVar
 from pyteal.ast.seq import Seq
 from pyteal.errors import TealInputError
@@ -76,9 +77,27 @@ class BaseType(ABC):
 
     def __init__(self, spec: TypeSpec) -> None:
         """Create a new BaseType."""
+        from pyteal.ast.subroutine import SubroutineEval
+
         super().__init__()
         self._type_spec: Final[TypeSpec] = spec
-        self._stored_value: AbstractVar = ScratchVar(spec.storage_type())
+        self._stored_value: AbstractVar
+
+        if SubroutineEval.Context.proto:
+            proto = SubroutineEval.Context.proto
+
+            assert proto.mem_layout
+            local_types = proto.mem_layout.local_stack_types
+
+            # NOTE: you can have at most 128 local variables.
+            # len(local_types) + 1 computes the resulting length,
+            # should be <= 128
+            if len(local_types) + 1 <= FRAME_LOCAL_NUM:
+                local_types.append(self._type_spec.storage_type())
+                self._stored_value = FrameVar(proto, len(local_types) - 1)
+                return
+
+        self._stored_value = ScratchVar(spec.storage_type())
 
     def type_spec(self) -> TypeSpec:
         """Get the TypeSpec for this ABI type instance."""
@@ -221,17 +240,23 @@ def store_into(self, output: BaseType) -> Expr:
                 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()}."
-            )
+        # HANG NOTE! This get_declaration check applies only for pre frame pointer case
+        # the post frame pointer case should not apply
+        # need to somehow expose the context of evaluation
+        try:
+            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()}."
+                )
+        except Exception:
+            pass
 
         return output._stored_value.store(self.computation)
 

pyteal/ast/frame.py

@@ -1,5 +1,5 @@
 from itertools import groupby
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING, Optional, Final
 
 from pyteal.ast.expr import Expr
 from pyteal.ast.int import Int
@@ -13,6 +13,9 @@
     from pyteal.compiler import CompileOptions
 
 
+FRAME_LOCAL_NUM: Final[int] = 128
+
+
 class LocalTypeSegment(Expr):
     """An expression that allocates stack spaces for local variable.
 
@@ -349,3 +352,31 @@ def has_return(self) -> bool:
 
 
 DupN.__module__ = "pyteal"
+
+
+class Bury(Expr):
+    def __init__(self, depth: int):
+        super(Bury, self).__init__()
+        if depth <= 0:
+            raise TealInputError("bury depth should be positive")
+        self.depth = depth
+
+    def __teal__(self, options: "CompileOptions") -> tuple[TealBlock, TealSimpleBlock]:
+        verifyProgramVersion(
+            Op.bury.min_version,
+            options.version,
+            "Program version too low to use bury",
+        )
+        return TealBlock.FromOp(options, TealOp(self, Op.bury, self.depth))
+
+    def __str__(self):
+        return f"(Bury {self.depth})"
+
+    def has_return(self) -> bool:
+        return False
+
+    def type_of(self) -> TealType:
+        return TealType.none
+
+
+Bury.__module__ = "pyteal"

pyteal/ast/subroutine.py

@@ -1,3 +1,4 @@
+from contextlib import AbstractContextManager
 from dataclasses import dataclass
 from docstring_parser import parse as parse_docstring
 from inspect import isclass, Parameter, signature, get_annotations
@@ -9,7 +10,7 @@
 from pyteal.ast.expr import Expr
 from pyteal.ast.seq import Seq
 from pyteal.ast.scratchvar import DynamicScratchVar, ScratchVar, ScratchSlot
-from pyteal.ast.frame import Proto, FrameVar, ProtoStackLayout
+from pyteal.ast.frame import Bury, Proto, FrameVar, ProtoStackLayout
 from pyteal.errors import TealInputError, TealInternalError, verifyProgramVersion
 from pyteal.ir import TealOp, Op, TealBlock
 from pyteal.types import TealType
@@ -1006,7 +1007,17 @@ def __call__(self, subroutine: SubroutineDefinition) -> SubroutineDeclaration:
             abi_output_kwargs[output_kwarg_info.name] = output_carrying_abi
 
         # Arg usage "B" supplied to build an AST from the user-defined PyTEAL function:
-        subroutine_body = subroutine.implementation(*loaded_args, **abi_output_kwargs)
+        subroutine_body: Expr
+        if not self.use_frame_pt:
+            subroutine_body = subroutine.implementation(
+                *loaded_args, **abi_output_kwargs
+            )
+        else:
+            with SubroutineEval.Context.CompileWithFrameContext(proto):
+                subroutine_body = subroutine.implementation(
+                    *loaded_args, **abi_output_kwargs
+                )
+
         if not isinstance(subroutine_body, Expr):
             raise TealInputError(
                 f"Subroutine function does not return a PyTeal expression. Got type {type(subroutine_body)}."
@@ -1025,6 +1036,14 @@ def __call__(self, subroutine: SubroutineDefinition) -> SubroutineDeclaration:
             if not self.use_frame_pt:
                 deferred_expr = output_carrying_abi._stored_value.load()
 
+        if self.use_frame_pt:
+            assert proto.mem_layout
+            depth = len(proto.mem_layout.local_stack_types)
+            # only when we have 1 return, and with other local variables
+            # we use bury to bury the result to 0 index against frame pointer
+            if not abi_output_kwargs and 0 < proto.num_returns < depth:
+                deferred_expr = Bury(depth)
+
         # Arg usage "A" to be pick up and store in scratch parameters that have been placed on the stack
         # need to reverse order of argumentVars because the last argument will be on top of the stack
 
@@ -1053,5 +1072,21 @@ def normal_evaluator(cls) -> "SubroutineEval":
     def fp_evaluator(cls) -> "SubroutineEval":
         return cls(SubroutineEval.var_n_loaded_fp, True)
 
+    class Context:
+        proto: Optional[Proto] = None
+
+        class CompileWithFrameContext(AbstractContextManager):
+            def __init__(self, _proto: Proto):
+                super().__init__()
+                self.prev_ctxt_proto = SubroutineEval.Context.proto
+                SubroutineEval.Context.proto = _proto
+
+            def __enter__(self):
+                return self
+
+            def __exit__(self, *_):
+                SubroutineEval.Context.proto = self.prev_ctxt_proto
+                return None
+
 
 SubroutineEval.__module__ = "pyteal"

tests/integration/teal/roundtrip/app_roundtrip_()_v8.teal

@@ -15,6 +15,8 @@ return
 tuplecomplement_0:
 proto 1 1
 byte ""
+int 0
+dupn 1
 byte ""
 frame_bury 0
 retsub
@@ -23,16 +25,19 @@ retsub
 roundtripper_1:
 proto 1 1
 byte ""
+dupn 2
+int 0
+dupn 1
 frame_dig -1
 callsub tuplecomplement_0
-store 2
-load 2
+frame_bury 1
+frame_dig 1
 callsub tuplecomplement_0
-store 3
+frame_bury 2
 frame_dig -1
-load 2
+frame_dig 1
 concat
-load 3
+frame_dig 2
 concat
 frame_bury 0
 retsub