Source code for beaker.precompile

from import KeysView
from dataclasses import dataclass, field
from typing import TYPE_CHECKING

from pyteal import (
from pyteal.types import require_type

from beaker.compilation import Program
from beaker.consts import PROGRAM_DOMAIN_SEPARATOR, num_extra_program_pages
from beaker.lib.strings import EncodeUVarInt

    from algosdk.v2client.algod import AlgodClient

    from beaker.application import Application
    from beaker.logic_signature import (

__all__ = [

[docs]class PrecompiledApplication: """ AppPrecompile allows a smart contract to signal that some child Application should be fully compiled prior to constructing its own program. """
[docs] def __init__(self, app: "Application", client: "AlgodClient"): app_spec = self._global_schema = app_spec.global_state_schema self._local_schema = app_spec.local_state_schema # at this point, we should have all the dependant logic built # so we can compile the app teal self.approval_program = Program(app_spec.approval_program, client) self.clear_program = Program(app_spec.clear_program, client)
[docs] def get_create_config(self) -> dict[TxnField, Expr | list[Expr]]: """get a dictionary of the fields and values that should be set when creating this application that can be passed directly to the InnerTxnBuilder.Execute method """ result: dict[TxnField, Expr | list[Expr]] = { TxnField.type_enum: TxnType.ApplicationCall, } extra_pages = num_extra_program_pages( self.approval_program.raw_binary, self.clear_program.raw_binary ) approval_pages = self.approval_program.pages clear_pages = self.clear_program.pages if extra_pages == 0: assert len(approval_pages) == 1 assert len(clear_pages) == 1 result[TxnField.approval_program] = approval_pages[0] result[TxnField.clear_state_program] = clear_pages[0] else: result[TxnField.approval_program_pages] = approval_pages result[TxnField.clear_state_program_pages] = clear_pages result[TxnField.extra_program_pages] = Int(extra_pages) if l_nbs := self._local_schema.num_byte_slices: result[TxnField.local_num_byte_slices] = Int(l_nbs) if l_nui := self._local_schema.num_uints: result[TxnField.local_num_uints] = Int(l_nui) if g_nbs := self._global_schema.num_byte_slices: result[TxnField.global_num_byte_slices] = Int(g_nbs) if g_nui := self._global_schema.num_uints: result[TxnField.global_num_uints] = Int(g_nui) return result
[docs]class PrecompiledLogicSignature: """ LSigPrecompile allows a smart contract to signal that some child Logic Signature should be fully compiled prior to constructing its own program. """
[docs] def __init__(self, lsig: "LogicSignature", client: "AlgodClient"): self.logic_program = Program(lsig.program, client)
[docs] def address(self) -> Expr: """Get the address from this LSig program.""" return Addr(self.logic_program.binary_hash)
@dataclass class PrecompileTemplateValue: #: Whether or not this variable is bytes (if false, its uint64) is_bytes: bool = field(kw_only=True) #: The line number in the source TEAL this variable is present line: int = field(kw_only=True) #: The pc of the variable in the assembled bytecode pc: int = 0
[docs]class PrecompiledLogicSignatureTemplate: """ LSigPrecompile allows a smart contract to signal that some child Logic Signature should be fully compiled prior to constructing its own program. """
[docs] def __init__(self, lsig: "LogicSignatureTemplate", client: "AlgodClient"): self._template_values: dict[str, PrecompileTemplateValue] = {} lines = lsig.program.splitlines() # Replace the teal program TMPL_* template variables with # the 0 value for the given type and save the list of TemplateValues for rtt_var in lsig.runtime_template_variables.values(): token = rtt_var.token is_bytes = rtt_var.type_of() == TealType.bytes op = "pushbytes" if is_bytes else "pushint" statement = f"{op} {token} // {token}" idx = lines.index(statement) lines[idx] = lines[idx].replace(token, '""' if is_bytes else "0", 1) self._template_values[] = PrecompileTemplateValue( is_bytes=is_bytes, line=idx ) self.logic_program = Program("\n".join(lines), client) for tv in self._template_values.values(): # +1 to acount for the pushbytes/pushint op pcs = self.logic_program.source_map.get_pcs_for_line(tv.line) if pcs is None: tv.pc = 0 else: tv.pc = pcs[0] + 1
[docs] def address(self, **kwargs: Expr) -> Expr: """ returns an expression that will generate the expected hash given some set of values that should be included in the logic itself """ self._check_kwargs(kwargs.keys()) return Sha512_256( Concat( Bytes(PROGRAM_DOMAIN_SEPARATOR), self.populate_template_expr(**kwargs), ) )
def _check_kwargs(self, keys: KeysView[str]) -> None: if keys != self._template_values.keys(): raise ValueError( f"Expected arguments named: {', '.join(self._template_values.keys())} " f"but got: {', '.join(keys)}" )
[docs] def populate_template_expr(self, **kwargs: Expr) -> Expr: """ populate_template_expr returns the Expr that will patch a blank binary given a set of arguments. It is called by ``address`` to return an Expr that can be used to compare with a sender given some arguments. """ # To understand how this works, first look at the pure python one above # it should produce an identical output in terms of populated binary. # This function just reproduces the same effects in pyteal self._check_kwargs(kwargs.keys()) populate_program: list[Expr] = [ (last_pos := ScratchVar(TealType.uint64)).store(Int(0)), (offset := ScratchVar(TealType.uint64)).store(Int(0)), (curr_val := ScratchVar(TealType.bytes)).store(Bytes("")), (buff := ScratchVar(TealType.bytes)).store(Bytes("")), ] for name, tv in self._template_values.items(): # Add expressions to encode the values and insert # them into the working buffer arg = kwargs[name] require_type(arg, TealType.bytes if tv.is_bytes else TealType.uint64) populate_program += [, arg)) if tv.is_bytes else, Concat( buff.load(), Substring( self.logic_program.binary, last_pos.load(), Int(tv.pc), ), curr_val.load(), ) ), + Len(curr_val.load()) - Int(1)), + Int(1)), ] # append the bytes from the last template variable to the end populate_program += [ Concat(buff.load(), Suffix(self.logic_program.binary, last_pos.load())) ), buff.load(), ] return Seq(*populate_program)
[docs] def populate_template(self, **kwargs: str | bytes | int) -> bytes: """ populate_template returns the bytes resulting from patching the set of arguments passed into the blank binary The args passed should be of the same type and in the same order as the template values declared. """ self._check_kwargs(kwargs.keys()) # Get a copy of the binary so we can work on it in place populated_binary = list(self.logic_program.raw_binary) # Any time we add bytes, we need to update the offset so the rest # of the pc values can be updated to account for the difference offset = 0 for name, tv in self._template_values.items(): arg = kwargs[name] if tv.is_bytes: if type(arg) is int: raise TealTypeError(type(arg), bytes | str) if type(arg) is str: arg = arg.encode("utf-8") assert type(arg) is bytes # Bytes are encoded as uvarint(len(bytes)) + bytes curr_val = _py_encode_uvarint(len(arg)) + arg else: if type(arg) is not int: raise TealTypeError(type(arg), int) # Ints are just the uvarint encoded number curr_val = _py_encode_uvarint(arg) # update the working buffer to include the new value, # replacing the current 0 value populated_binary[tv.pc + offset : tv.pc + offset + 1] = curr_val # update the offset with the length(value) - 1 to account # for the existing 0 value and help keep track of how to shift the pc later offset += len(curr_val) - 1 return bytes(populated_binary)
class PrecompileContextError(Exception): pass def _py_encode_uvarint(integer: int) -> bytes: """Encodes an integer as an uvarint. :param integer: the integer to encode :return: bytes containing the integer encoded as an uvarint """ def to_byte(x: int) -> int: return x & 0b1111_1111 buffer: bytearray = bytearray() while integer >= 0b1000_0000: buffer.append(to_byte(integer) | 0b1000_0000) integer >>= 7 buffer.append(to_byte(integer)) return bytes(buffer)