Convenience wrapper for subroutine calls

[zejoeh]

Hi there, I spoke to @tzaffi on the graviton repo who suggested I make a pull request here. I wrote a CallableSubroutine decorator to make directly testing subroutine calls a bit more convenient for me personally. It doesn't add any new functionality, it is just a way to skip along some steps in case you only care about quickly comparing some outputs.

Example use:

# can be used to wrap existing Subroutine
@CallableSubroutine(input_types=[TealType.uint64, TealType.uint64])
@Subroutine(TealType.bytes)
def bytes_subr(x: Expr, y: Expr) -> Expr:
    # (x * 1000) // y
    return BytesDiv(BytesMul(Itob(x), Bytes("base16", "0x03e8")), Itob(y))

# can also wrap a regular function and internally wrap as subroutine if output_type is specified
@CallableSubroutine(
    input_types=[TealType.uint64, TealType.uint64],
    output_type=TealType.uint64,
)
def uint64_subr(x: Expr, y: Expr) -> Expr:
    return Div(x, y)

# returns 500: reorders arguments according to order in subroutine
assert bytes_subr(y=5, x=10) == 2000 == bytes_subr(x=10, y=5)
# all equal 0
assert uint64_subr(y=10, x=5) == 0 == uint64_subr(x=5, y=10)
# x / 0: raises RuntimeError and reports error message from dryrun
bytes_subr(5, y=0)

I removed a lot of the things I had initially implemented locally, since I was unaware of how much work had been done on the feature/abi branch on PyTeal integration. I tried to strip the thing whole thing down a bit, but I use something like this in my own testing and found it super convenient, since it makes it pretty easy to write unit tests the way you would for any other Python project.

It's pretty hacky still and I'm sure it could be polished a lot by incorporating all the work happening on this branch, but it's been suitable for what I need and thought others might find it useful too. Hopefully I didn't miss something like this already existing in the repo, otherwise feel free to close!

It'd be pretty nice to infer the (Python-native) return types based on the ABI types too, but I only needed it for some numeric calculations and didn't want to spend too much time in case nobody else would want this. Otherwise happy to incorporate feedback. :)

[CLAassistant]

All committers have signed the CLA.

[zejoeh]

I changed the name to the decorator to the (IMO) more appropriate CallableSubroutine and enabled it to internally wrap a regular function as a Subroutine if the output_type is provided. For me this is nice because I avoid wrapping things as Subroutines until they really need to be so in a contract, since this allows me to control when things are inlined and lets me test components as standalone Subroutines easily without getting any dependencies on other Subroutines.

Ideas for further improvement: support writing the ABI types directly as type annotations in the original function and use the internal Subroutine wrapping to map these to the appropriate TealTypes. This would allow to skip both the output_types and input_type arguments and to infer the appropriate representation of the output (e.g. for strings or fixed-point types). To call this as a Subroutine within a contract, one would then have to do this via the subroutine attribute instead (or we do some dispatching based on argument types: e.g. if int | bytes | str, dispatch to _SubroutineRunner, else dispatch to _SubroutineRunner.subroutine. As a bonus, it would obviously make functions more readable, since the "Expr" annotations don't really tell us much, but I'm sure that now I'm getting to ideas that are overlapping heavily with the work already put into this branch. :)

[zejoeh]

6627d14

Added inlining of "subroutines" (technically not really subroutines any more) to better support the workflow I was doing manually before, where I would write functions to generate Expr and call those directly when I wanted to inline, and wrap them in the CallableSubroutine when I wanted to test it.

Now the _SubroutineRunner checks that all arguments are Expr or ScratchVar; in this case, it will either make the call as a subroutine, or it will inline and return the call as an Expr for use in another subroutine, depending on the inline flag in the decorator. If all arguments are not Expr or ScratchVar, it will try to do cast the arguments to bytes and execute a dryrun.

This allows for testing things in isolation by treating all functions as subroutines, even if they are actually just meant to be inlined as part of more complex logic. I think it is relatively safe, since (as far as I'm aware) the PyTeal compiler does not do too much optimization?

For example:

@CallableSubroutine(
    input_types=[TealType.uint64],
    output_type=TealType.uint64,
    inline=True
)
def square_subr(x: Expr) -> Expr:
    return x ** Int(2)

@CallableSubroutine(
    input_types=[TealType.uint64],
    output_type=TealType.uint64,
)
def pow4_subr(x: Expr) -> Expr:
    return square_subr(square_subr(x))

# yep, squaring works
assert square_subr(5) == 5 ** 2

# squaring twice works too, amazing!
assert pow4_subr(5) == 5 ** 4

# inner "subroutine" was inlined!
print(pow4_subr.report(5))

#        step |   PC# |   L# | Teal                   | Scratch   | Stack
# --------+-------+------+------------------------+-----------+----------------------
#       1 |     1 |    1 | #pragma version 6      |           | []
#       2 |     4 |    2 | intcblock 2            |           | []
#       3 |     7 |    3 | txna ApplicationArgs 0 |           | [0x05]
#       4 |     8 |    4 | btoi                   |           | [5]
#       5 |    20 |   12 | label1:                |           | [5]
#       6 |    22 |   13 | store 0                | 0->5      | []
#       7 |    24 |   14 | load 0                 |           | [5]
#       8 |    25 |   15 | intc_0 // 2            |           | [5, 2]
#       9 |    26 |   16 | exp                    |           | [25]
#      10 |    27 |   17 | intc_0 // 2            |           | [25, 2]
#      11 |    28 |   18 | exp                    |           | [625]
#      12 |    11 |    5 | callsub label1         |           | [625]
#      13 |    13 |    6 | store 1                | 1->625    | []
#      14 |    15 |    7 | load 1                 |           | [625]
#      15 |    16 |    8 | itob                   |           | [0x0000000000000271]
#      16 |    17 |    9 | log                    |           | []
#      17 |    19 |   10 | load 1                 |           | [625]
#      18 |    29 |   19 | retsub                 |           | [625]
#     ===============

This workflow may be too implicit for a lot of people though.