Prepare for Testing the ABI-Router in PyTeal

CHANGELOG.md

@@ -1,6 +1,15 @@
 <!-- markdownlint-disable MD024 -->
 # Changelog
 
+## `v0.9.0` (_aka_ 🐐)
+
+### Breaking changes
+
+* `class ABIContractExecutor` renamed to `ABICallStrategy` and moved from `graviton/blackbox.py` to `graviton/abi_strategy.py`. Some of the methods have been renamed as well.
+
+### Added
+* `class Simulation` in `graviton/sim.py` unifies the ability to run an argument strategy and check that invariants hold using its `run_and_assert()` method
+* `class DryRunTransactionParameters` has a new method `update_fields()`
 
 ## `v0.8.0` (_aka_ 🦛)
 

graviton/abi_strategy.py

@@ -4,16 +4,28 @@
 TODO: Leverage Hypothesis!
 """
 from abc import ABC, abstractmethod
+from enum import Enum, auto
 import random
 import string
-from typing import List, Optional, Sequence, cast
+from typing import List, Optional, Sequence, Type, cast
 
 from algosdk import abi, encoding
 
+
 from graviton.models import PyTypes
 
 
 class ABIStrategy(ABC):
+    """
+    TODO: when incorporating hypothesis strategies, we'll need a more holistic
+    approach that looks at relationships amongst various args.
+    Current approach only looks at each argument as a completely independent entity.
+    """
+
+    @abstractmethod
+    def __init__(self, abi_instance: abi.ABIType, dynamic_length: Optional[int] = None):
+        pass
+
     @abstractmethod
     def get(self) -> PyTypes:
         pass
@@ -167,6 +179,15 @@ def address_logic(x):
 
 
 class RandomABIStrategyHalfSized(RandomABIStrategy):
+    """
+    This strategy only generates data that is half the size that _ought_ to be possible.
+    This is useful in the case that operations involving the generated arguments
+    could overflow due to multiplication.
+
+    Since this only makes sense for `abi.UintType`, it degenerates to the standard
+    `RandomABIStrategy` for other types.
+    """
+
     def __init__(
         self,
         abi_instance: abi.ABIType,
@@ -183,3 +204,162 @@ def get(self) -> PyTypes:
         return cast(int, full_random) % (
             1 << (cast(abi.UintType, self.abi_type).bit_size // 2)
         )
+
+
+class ABIArgsMod(Enum):
+    selector_byte_insert = auto()
+    selector_byte_delete = auto()
+    selector_byte_replace = auto()
+    parameter_delete = auto()
+    parameter_append = auto()
+
+
+class ABICallStrategy:
+    """
+    TODO: refactor to comport with ABIStrategy + Hypothesis
+    TODO: make this generic on the strategy type
+    """
+
+    append_args_type: abi.ABIType = abi.ByteType()
+
+    def __init__(
+        self,
+        contract: str,
+        argument_strategy: Type[ABIStrategy] = RandomABIStrategy,
+        *,
+        num_dryruns: int = 1,
+        handle_selector: bool = True,
+        abi_args_mod: Optional[ABIArgsMod] = None,
+    ):
+        """
+        contract - ABI Contract JSON
+
+        argument_strategy (default=RandomABIStrategy) - ABI strategy for generating arguments
+
+        num_dry_runs (default=1) - the number of dry runs to run
+            (generates different inputs each time)
+
+        handle_selector (default=True) - usually we'll want to let
+            `ABIContractExecutor.run_sequence()`
+            handle adding the method selector so this param.
+            But if set False: when providing `inputs`
+            ensure that the 0'th argument for method calls is the selector.
+            And when set True: when NOT providing `inputs`, the selector arg
+            at index 0 will be added automatically.
+
+        abi_args_mod (optional) - when desiring to mutate the args, provide an ABIArgsMod value
+        """
+        self.contract: abi.Contract = abi.Contract.from_json(contract)
+        self.argument_strategy: Type[ABIStrategy] = argument_strategy
+        self.num_dryruns = num_dryruns
+        self.handle_selector = handle_selector
+        self.abi_args_mod = abi_args_mod
+
+    def abi_method(self, method: Optional[str]) -> abi.Method:
+        assert method, "cannot get abi.Method for bare app call"
+
+        return self.contract.get_method_by_name(method)
+
+    def method_signature(self, method: Optional[str]) -> Optional[str]:
+        """Returns None, for a bare app call (method=None signals this)"""
+        if method is None:
+            return None
+
+        return self.abi_method(method).get_signature()
+
+    def method_selector(self, method: Optional[str]) -> bytes:
+        assert method, "cannot get method_selector for bare app call"
+
+        return self.abi_method(method).get_selector()
+
+    def argument_types(self, method: Optional[str]) -> List[abi.ABIType]:
+        """
+        Argument types (excluding selector)
+        """
+        if method is None:
+            return []
+
+        return [cast(abi.ABIType, arg.type) for arg in self.abi_method(method).args]
+
+    def num_args(self, method: Optional[str]) -> int:
+        return len(self.argument_types(method))
+
+    def generate_inputs(self, method: Optional[str]) -> List[Sequence[PyTypes]]:
+        """
+        Generates inputs appropriate for bare app calls and method calls
+        according to available argument_strategy.
+        """
+        assert (
+            self.argument_strategy
+        ), "cannot generate inputs without an argument_strategy"
+
+        mutating = self.abi_args_mod is not None
+
+        if not (method or mutating):
+            # bare calls receive no arguments (unless mutating)
+            return [tuple() for _ in range(self.num_dryruns)]
+
+        arg_types = self.argument_types(method)
+
+        prefix: List[bytes] = []
+        if self.handle_selector and method:
+            prefix = [self.method_selector(method)]
+
+        modify_selector = False
+        if (action := self.abi_args_mod) in (
+            ABIArgsMod.selector_byte_delete,
+            ABIArgsMod.selector_byte_insert,
+            ABIArgsMod.selector_byte_replace,
+        ):
+            assert (
+                prefix
+            ), f"{self.abi_args_mod=} which means we need to modify the selector, but we don't have one available to modify"
+            modify_selector = True
+
+        def selector_mod(prefix):
+            assert isinstance(prefix, list) and len(prefix) <= 1
+            if not (prefix and modify_selector):
+                return prefix
+
+            selector = prefix[0]
+            idx = random.randint(0, 4)
+            x, y = selector[:idx], selector[idx:]
+            if action == ABIArgsMod.selector_byte_insert:
+                selector = x + random.randbytes(1) + y
+            elif action == ABIArgsMod.selector_byte_delete:
+                selector = (x[:-1] + y) if x else y[:-1]
+            else:
+                assert (
+                    action == ABIArgsMod.selector_byte_replace
+                ), f"expected action={ABIArgsMod.selector_byte_replace} but got [{action}]"
+                idx = random.randint(0, 3)
+                selector = (
+                    selector[:idx]
+                    + bytes([(selector[idx] + 1) % 256])
+                    + selector[idx + 1 :]
+                )
+            return [selector]
+
+        def args_mod(args):
+            if action not in (ABIArgsMod.parameter_append, ABIArgsMod.parameter_delete):
+                return args
+
+            if action == ABIArgsMod.parameter_delete:
+                return args if not args else tuple(args[:-1])
+
+            assert action == ABIArgsMod.parameter_append
+            return args + (self.generate_value(self.append_args_type),)
+
+        def gen_args():
+            # TODO: when incorporating hypothesis strategies, we'll need a more holistic
+            # approach that looks at relationships amongst various args
+            args = tuple(
+                selector_mod(prefix)
+                + [self.generate_value(atype) for atype in arg_types]
+            )
+            return args_mod(args)
+
+        return [gen_args() for _ in range(self.num_dryruns)]
+
+    def generate_value(self, gen_type: abi.ABIType) -> PyTypes:
+        return cast(Type[ABIStrategy], self.argument_strategy)(gen_type).get()