method pseudo-op support for ABI methods (#153)

- Add support for `method` pseudo-opcode in PyTeal.
- Add `name` field in `subroutine` to override __name__ from function implementation, for readability in generated code.

pyteal/ast/__init__.py

@@ -6,6 +6,7 @@
 from .addr import Addr
 from .bytes import Bytes
 from .int import Int, EnumInt
+from .methodsig import MethodSignature
 
 # properties
 from .arg import Arg
@@ -126,6 +127,7 @@
     "Bytes",
     "Int",
     "EnumInt",
+    "MethodSignature",
     "Arg",
     "TxnType",
     "TxnField",

pyteal/ast/methodsig.py

@@ -0,0 +1,43 @@
+from typing import TYPE_CHECKING
+from pyteal.errors import TealInputError
+
+from pyteal.types import TealType
+
+from ..types import TealType
+from ..ir import TealOp, Op, TealBlock
+from .leafexpr import LeafExpr
+
+if TYPE_CHECKING:
+    from ..compiler import CompileOptions
+
+
+class MethodSignature(LeafExpr):
+    """An expression that represents an ABI method selector"""
+
+    def __init__(self, methodName: str) -> None:
+        """Create a new method selector for ABI method call.
+
+        Args:
+            methodName: A string containing a valid ABI method signature
+        """
+        super().__init__()
+        if type(methodName) is not str:
+            raise TealInputError(
+                "invalid input type {} to Method".format(type(methodName))
+            )
+        elif len(methodName) == 0:
+            raise TealInputError("invalid input empty string to Method")
+        self.methodName = methodName
+
+    def __teal__(self, options: "CompileOptions"):
+        op = TealOp(self, Op.method_signature, '"{}"'.format(self.methodName))
+        return TealBlock.FromOp(options, op)
+
+    def __str__(self) -> str:
+        return "(method: {})".format(self.methodName)
+
+    def type_of(self) -> TealType:
+        return TealType.bytes
+
+
+MethodSignature.__module__ = "pyteal"

pyteal/ast/methodsig_test.py

@@ -0,0 +1,27 @@
+import pytest
+
+from pyteal.ast.methodsig import MethodSignature
+
+from .. import *
+
+
+def test_method():
+    expr = MethodSignature("add(uint64,uint64)uint64")
+    assert expr.type_of() == TealType.bytes
+
+    expected = TealSimpleBlock(
+        [TealOp(expr, Op.method_signature, '"add(uint64,uint64)uint64"')]
+    )
+    actual, _ = expr.__teal__(CompileOptions())
+    assert expected == actual
+
+
+def test_method_invalid():
+    with pytest.raises(TealInputError):
+        MethodSignature(114514)
+
+    with pytest.raises(TealInputError):
+        MethodSignature(['"m0()void"', '"m1()uint64"'])
+
+    with pytest.raises(TealInputError):
+        MethodSignature("")

pyteal/ast/subroutine.py

@@ -10,7 +10,6 @@
 from .scratchvar import ScratchVar
 
 if TYPE_CHECKING:
-    from ..ir import TealSimpleBlock
     from ..compiler import CompileOptions
 
 
@@ -19,7 +18,10 @@ class SubroutineDefinition:
     nextSubroutineId = 0
 
     def __init__(
-        self, implementation: Callable[..., Expr], returnType: TealType
+        self,
+        implementation: Callable[..., Expr],
+        returnType: TealType,
+        nameStr: str = None,
     ) -> None:
         super().__init__()
         self.id = SubroutineDefinition.nextSubroutineId
@@ -53,6 +55,7 @@ def __init__(
         self.returnType = returnType
 
         self.declaration: Optional["SubroutineDeclaration"] = None
+        self.__name = self.implementation.__name__ if nameStr is None else nameStr
 
     def getDeclaration(self) -> "SubroutineDeclaration":
         if self.declaration is None:
@@ -61,7 +64,7 @@ def getDeclaration(self) -> "SubroutineDeclaration":
         return self.declaration
 
     def name(self) -> str:
-        return self.implementation.__name__
+        return self.__name
 
     def argumentCount(self) -> int:
         return len(self.implementationParams)
@@ -181,17 +184,18 @@ def mySubroutine(a: Expr, b: Expr) -> Expr:
             ])
     """
 
-    def __init__(self, returnType: TealType) -> None:
+    def __init__(self, returnType: TealType, name: str = None) -> None:
         """Define a new subroutine with the given return type.
 
         Args:
             returnType: The type that the return value of this subroutine must conform to.
                 TealType.none indicates that this subroutine does not return any value.
         """
         self.returnType = returnType
+        self.name = name
 
     def __call__(self, fnImplementation: Callable[..., Expr]) -> Callable[..., Expr]:
-        subroutine = SubroutineDefinition(fnImplementation, self.returnType)
+        subroutine = SubroutineDefinition(fnImplementation, self.returnType, self.name)
 
         @wraps(fnImplementation)
         def subroutineCall(*args: Expr, **kwargs) -> Expr:

pyteal/compiler/constants.py

@@ -7,11 +7,7 @@
 from ..ir import (
     Op,
     TealOp,
-    TealLabel,
     TealComponent,
-    TealBlock,
-    TealSimpleBlock,
-    TealConditionalBlock,
 )
 from ..util import unescapeStr, correctBase32Padding
 from ..errors import TealInternalError
@@ -94,6 +90,28 @@ def extractAddrValue(op: TealOp) -> Union[str, bytes]:
     return value
 
 
+def extractMethodSigValue(op: TealOp) -> bytes:
+    """Extract the constant value being loaded by a TealOp whose op is Op.method.
+
+    Returns:
+        The bytes of method selector computed from the method signature that the op is loading.
+    """
+    if len(op.args) != 1 or type(op.args[0]) != str:
+        raise TealInternalError("Unexpected args in method opcode: {}".format(op.args))
+
+    methodSignature = cast(str, op.args[0])
+    if methodSignature[0] == methodSignature[-1] and methodSignature.startswith('"'):
+        methodSignature = methodSignature[1:-1]
+    else:
+        raise TealInternalError(
+            "Method signature opcode error: signatue {} not wrapped with double-quotes".format(
+                methodSignature
+            )
+        )
+    methodSelector = encoding.checksum(bytes(methodSignature, "utf-8"))[:4]
+    return methodSelector
+
+
 def createConstantBlocks(ops: List[TealComponent]) -> List[TealComponent]:
     """Convert TEAL code from using pseudo-ops for constants to using assembled constant blocks.
 
@@ -124,6 +142,9 @@ def createConstantBlocks(ops: List[TealComponent]) -> List[TealComponent]:
         elif basicOp == Op.addr:
             addrValue = extractAddrValue(op)
             byteFreqs[addrValue] = byteFreqs.get(addrValue, 0) + 1
+        elif basicOp == Op.method_signature:
+            methodValue = extractMethodSigValue(op)
+            byteFreqs[methodValue] = byteFreqs.get(methodValue, 0) + 1
 
     assembled: List[TealComponent] = []
 
@@ -177,12 +198,22 @@ def createConstantBlocks(ops: List[TealComponent]) -> List[TealComponent]:
                     assembled.append(TealOp(op.expr, Op.intc, index, "//", *op.args))
                 continue
 
-            if basicOp == Op.byte or basicOp == Op.addr:
-                byteValue = (
-                    extractBytesValue(op)
-                    if basicOp == Op.byte
-                    else extractAddrValue(op)
-                )
+            if (
+                basicOp == Op.byte
+                or basicOp == Op.addr
+                or basicOp == Op.method_signature
+            ):
+                if basicOp == Op.byte:
+                    byteValue = extractBytesValue(op)
+                elif basicOp == Op.addr:
+                    byteValue = extractAddrValue(op)
+                elif basicOp == Op.method_signature:
+                    byteValue = extractMethodSigValue(op)
+                else:
+                    raise TealInternalError(
+                        "Expect a byte-like constant opcode, get {}".format(op)
+                    )
+
                 if byteFreqs[byteValue] == 1:
                     encodedValue = (
                         ("0x" + byteValue.hex())

pyteal/compiler/constants_test.py

@@ -5,6 +5,7 @@
     extractBytesValue,
     extractAddrValue,
     createConstantBlocks,
+    extractMethodSigValue,
 )
 
 
@@ -63,6 +64,44 @@ def test_extractAddrValue():
         assert actual == expected
 
 
+# test case came from: https://gist.github.com/jasonpaulos/99e4f8a75f2fc2ec9b8073c064530359
+def test_extractMethodValue():
+    tests = [
+        (
+            TealOp(None, Op.method_signature, '"create(uint64)uint64"'),
+            b"\x43\x46\x41\x01",
+        ),
+        (TealOp(None, Op.method_signature, '"update()void"'), b"\xa0\xe8\x18\x72"),
+        (
+            TealOp(None, Op.method_signature, '"optIn(string)string"'),
+            b"\xcf\xa6\x8e\x36",
+        ),
+        (TealOp(None, Op.method_signature, '"closeOut()string"'), b"\xa9\xf4\x2b\x3d"),
+        (TealOp(None, Op.method_signature, '"delete()void"'), b"\x24\x37\x8d\x3c"),
+        (
+            TealOp(None, Op.method_signature, '"add(uint64,uint64)uint64"'),
+            b"\xfe\x6b\xdf\x69",
+        ),
+        (TealOp(None, Op.method_signature, '"empty()void"'), b"\xa8\x8c\x26\xa5"),
+        (
+            TealOp(None, Op.method_signature, '"payment(pay,uint64)bool"'),
+            b"\x3e\x3b\x3d\x28",
+        ),
+        (
+            TealOp(
+                None,
+                Op.method_signature,
+                '"referenceTest(account,application,account,asset,account,asset,asset,application,application)uint8[9]"',
+            ),
+            b"\x0d\xf0\x05\x0f",
+        ),
+    ]
+
+    for op, expected in tests:
+        actual = extractMethodSigValue(op)
+        assert actual == expected
+
+
 def test_createConstantBlocks_empty():
     ops = []
 
@@ -184,12 +223,14 @@ def test_createConstantBlocks_pushbytes():
     ops = [
         TealOp(None, Op.byte, "0x0102"),
         TealOp(None, Op.byte, "0x0103"),
+        TealOp(None, Op.method_signature, '"empty()void"'),
         TealOp(None, Op.concat),
     ]
 
     expected = [
         TealOp(None, Op.pushbytes, "0x0102", "//", "0x0102"),
         TealOp(None, Op.pushbytes, "0x0103", "//", "0x0103"),
+        TealOp(None, Op.pushbytes, "0xa88c26a5", "//", '"empty()void"'),
         TealOp(None, Op.concat),
     ]
 
@@ -240,6 +281,9 @@ def test_createConstantBlocks_byteblock_multiple():
             None, Op.addr, "WSJHNPJ6YCLX5K4GUMQ4ISPK3ABMS3AL3F6CSVQTCUI5F4I65PWEMCWT3M"
         ),
         TealOp(None, Op.concat),
+        TealOp(None, Op.method_signature, '"closeOut()string"'),
+        TealOp(None, Op.concat),
+        TealOp(None, Op.byte, "base64(qfQrPQ==)"),
     ]
 
     expected = [
@@ -249,6 +293,7 @@ def test_createConstantBlocks_byteblock_multiple():
             "0x0102",
             "0x74657374",
             "0xb49276bd3ec0977eab86a321c449ead802c96c0bd97c2956131511d2f11eebec",
+            "0xa9f42b3d",
         ),
         TealOp(None, Op.bytec_0, "//", "0x0102"),
         TealOp(None, Op.bytec_0, "//", "base64(AQI=)"),
@@ -273,6 +318,9 @@ def test_createConstantBlocks_byteblock_multiple():
             "WSJHNPJ6YCLX5K4GUMQ4ISPK3ABMS3AL3F6CSVQTCUI5F4I65PWEMCWT3M",
         ),
         TealOp(None, Op.concat),
+        TealOp(None, Op.bytec_3, "//", '"closeOut()string"'),
+        TealOp(None, Op.concat),
+        TealOp(None, Op.bytec_3, "//", "base64(qfQrPQ==)"),
     ]
 
     actual = createConstantBlocks(ops)

pyteal/ir/ops.py

@@ -74,6 +74,7 @@ def min_version(self) -> int:
     bytec_3             = OpType("bytec_3",             Mode.Signature | Mode.Application, 2)
     byte                = OpType("byte",                Mode.Signature | Mode.Application, 2)
     addr                = OpType("addr",                Mode.Signature | Mode.Application, 2)
+    method_signature    = OpType("method",              Mode.Signature | Mode.Application, 2)
     arg                 = OpType("arg",                 Mode.Signature,                    2)
     txn                 = OpType("txn",                 Mode.Signature | Mode.Application, 2)
     global_             = OpType("global",              Mode.Signature | Mode.Application, 2)

pyteal/ir/tealop.py

@@ -1,4 +1,4 @@
-from typing import cast, Union, List, Optional, TYPE_CHECKING
+from typing import Union, List, Optional, TYPE_CHECKING
 
 from .tealcomponent import TealComponent
 from .labelref import LabelReference