Add generic class for dataflow analysis and class to find constraints on group size fields

.github/workflows/pytest.yml

@@ -54,3 +54,8 @@ jobs:
     - name: Run detectors tests
       run: |
         pytest tests/test_detectors.py
+    
+    - name: Run dataflow analysis tests
+      run: |
+        pytest tests/transaction_context/test_group_sizes.py
+        pytest tests/transaction_context/test_group_indices.py

tealer/analyses/dataflow/all_constraints.py

@@ -0,0 +1,2 @@
+# pylint: disable=unused-import
+from tealer.analyses.dataflow.int_fields import GroupIndices

tealer/analyses/dataflow/int_fields.py

@@ -0,0 +1,195 @@
+from typing import TYPE_CHECKING, List, Set, Tuple, Dict
+
+from tealer.analyses.dataflow.generic import DataflowTransactionContext
+from tealer.teal.instructions.instructions import (
+    Global,
+    Eq,
+    Neq,
+    Greater,
+    GreaterE,
+    Less,
+    LessE,
+    Txn,
+)
+from tealer.teal.global_field import GroupSize
+from tealer.teal.instructions.transaction_field import GroupIndex
+from tealer.utils.analyses import is_int_push_ins
+from tealer.utils.algorand_constants import MAX_GROUP_SIZE
+
+if TYPE_CHECKING:
+    from tealer.teal.instructions.instructions import Instruction
+
+group_size_key = "GroupSize"
+group_index_key = "GroupIndex"
+analysis_keys = [group_size_key, group_index_key]
+universal_sets = {}
+universal_sets[group_size_key] = list(range(1, MAX_GROUP_SIZE + 1))
+universal_sets[group_index_key] = list(range(0, MAX_GROUP_SIZE))
+
+
+class GroupIndices(DataflowTransactionContext):  # pylint: disable=too-few-public-methods
+
+    GROUP_SIZE_KEY = group_size_key
+    GROUP_INDEX_KEY = group_index_key
+    BASE_KEYS: List[str] = analysis_keys
+    KEYS_WITH_GTXN: List[str] = []  # gtxn information is not collected for any of the keys
+    UNIVERSAL_SETS: Dict[str, List] = universal_sets
+
+    def _universal_set(self, key: str) -> Set:
+        return set(self.UNIVERSAL_SETS[key])
+
+    def _null_set(self, key: str) -> Set:
+        return set()
+
+    def _union(self, key: str, a: Set, b: Set) -> Set:
+        return a | b
+
+    def _intersection(self, key: str, a: Set, b: Set) -> Set:
+        return a & b
+
+    @staticmethod
+    def _get_asserted_int_values(
+        comparison_ins: "Instruction", compared_int: int, universal_set: List[int]
+    ) -> List[int]:
+        """return list of ints from universal set(U) that will satisfy the comparison.
+
+        if the given condition uses ==, return compared int list.
+        if the condition uses != then return {U - compared_int}
+        if the given condition uses <, return U[ : U.index(compared_int)]
+        if the given condition uses <=, return U[ : U.index(compared_int) + 1]
+        if the given condition uses >, return U[U.index(compared_int) + 1:]
+        if the given condition uses >=, return U[U.index(compared_int): ]
+
+        Args:
+            comparison_ins: comparison instruction used. can be [==, !=, <, <=, >, >=]
+            compared_int: integer value compared.
+            universal_set: list of all possible integer values for the field.
+
+        Returns:
+            list of ints that will satisfy the comparison
+        """
+        U = list(universal_set)
+
+        if isinstance(comparison_ins, Eq):  # pylint: disable=no-else-return
+            return [compared_int]
+        elif isinstance(comparison_ins, Neq):
+            if compared_int in U:
+                U.remove(compared_int)
+            return U
+        elif isinstance(comparison_ins, Less):
+            return [i for i in U if i < compared_int]
+        elif isinstance(comparison_ins, LessE):
+            return [i for i in U if i <= compared_int]
+        elif isinstance(comparison_ins, Greater):
+            return [i for i in U if i > compared_int]
+        elif isinstance(comparison_ins, GreaterE):
+            return [i for i in U if i >= compared_int]
+        else:
+            return U
+
+    def _get_asserted_groupsizes(self, ins_stack: List["Instruction"]) -> Tuple[Set[int], Set[int]]:
+        """return set of values for groupsize that will make the comparison true and false
+
+        checks for instruction sequence and returns group size values that will make the comparison true.
+
+        [ Global GroupSize | (int | pushint)]
+        [ (int | pushint) | Global GroupSize]
+        [ == | != | < | <= | > | >=]
+
+        Args:
+            ins_stack: list of instructions that are executed up until the comparison instruction (including the comparison instruction).
+
+        Returns:
+            set of groupsize values that will make the comparison true, set of groupsize values that will make the comparison false.
+        """
+        U = list(self.UNIVERSAL_SETS[self.GROUP_SIZE_KEY])
+        if len(ins_stack) < 3:
+            return set(U), set(U)
+
+        if isinstance(ins_stack[-1], (Eq, Neq, Less, LessE, Greater, GreaterE)):
+            ins1 = ins_stack[-2]
+            ins2 = ins_stack[-3]
+            compared_value = None
+
+            if isinstance(ins1, Global) and isinstance(ins1.field, GroupSize):
+                is_int, value = is_int_push_ins(ins2)
+                if is_int:
+                    compared_value = value
+            elif isinstance(ins2, Global) and isinstance(ins2.field, GroupSize):
+                is_int, value = is_int_push_ins(ins1)
+                if is_int:
+                    compared_value = value
+
+            if compared_value is None or not isinstance(compared_value, int):
+                # if the comparison does not check groupsize, return U as values that make the comparison false
+                return set(U), set(U)
+
+            ins = ins_stack[-1]
+            asserted_values = self._get_asserted_int_values(ins, compared_value, U)
+            return set(asserted_values), set(U) - set(asserted_values)
+        return set(U), set(U)
+
+    def _get_asserted_groupindices(
+        self, ins_stack: List["Instruction"]
+    ) -> Tuple[Set[int], Set[int]]:
+        """return list of values for group index that will make the comparison true and false
+
+        checks for instruction sequence and returns group index values that will make the comparison true.
+
+        [ txn GroupIndex | (int | pushint)]
+        [ (int | pushint) | txn GroupIndex]
+        [ == | != | < | <= | > | >=]
+
+        Args:
+            ins_stack: list of instructions that are executed up until the comparison instruction (including the comparison instruction).
+
+        Returns:
+            List of groupindex values that will make the comparison true.
+        """
+        U = list(self.UNIVERSAL_SETS[self.GROUP_INDEX_KEY])
+        if len(ins_stack) < 3:
+            return set(U), set(U)
+
+        if isinstance(ins_stack[-1], (Eq, Neq, Less, LessE, Greater, GreaterE)):
+            ins1 = ins_stack[-2]
+            ins2 = ins_stack[-3]
+            compared_value = None
+
+            if isinstance(ins1, Txn) and isinstance(ins1.field, GroupIndex):
+                is_int, value = is_int_push_ins(ins2)
+                if is_int:
+                    compared_value = value
+            elif isinstance(ins2, Txn) and isinstance(ins2.field, GroupIndex):
+                is_int, value = is_int_push_ins(ins1)
+                if is_int:
+                    compared_value = value
+
+            if compared_value is None or not isinstance(compared_value, int):
+                return set(U), set(U)
+
+            ins = ins_stack[-1]
+            asserted_values = self._get_asserted_int_values(ins, compared_value, U)
+            return set(asserted_values), set(U) - set(asserted_values)
+        return set(U), set(U)
+
+    def _get_asserted(self, key: str, ins_stack: List["Instruction"]) -> Tuple[Set, Set]:
+        if key == self.GROUP_SIZE_KEY:
+            return self._get_asserted_groupsizes(ins_stack)
+        return self._get_asserted_groupindices(ins_stack)
+
+    def _store_results(self) -> None:
+        # use group_sizes to update group_indices
+        group_sizes_context = self._block_contexts[self.GROUP_SIZE_KEY]
+        group_indices_context = self._block_contexts[self.GROUP_INDEX_KEY]
+        for bi in self._teal.bbs:
+            group_indices_context[bi] = group_indices_context[bi] & set(
+                range(0, max(group_sizes_context[bi], default=0))
+            )
+
+        group_size_block_context = self._block_contexts[self.GROUP_SIZE_KEY]
+        for block in self._teal.bbs:
+            block.transaction_context.group_sizes = list(group_size_block_context[block])
+
+        group_index_block_context = self._block_contexts[self.GROUP_INDEX_KEY]
+        for block in self._teal.bbs:
+            block.transaction_context.group_indices = list(group_index_block_context[block])

tealer/teal/basic_blocks.py

@@ -16,12 +16,14 @@
 from typing import List, Optional, TYPE_CHECKING
 
 from tealer.teal.instructions.instructions import Instruction
+from tealer.teal.context.block_transaction_context import BlockTransactionContext
+
 
 if TYPE_CHECKING:
     from tealer.teal.teal import Teal
 
 
-class BasicBlock:
+class BasicBlock:  # pylint: disable=too-many-instance-attributes
     """Class to represent basic blocks of the teal contract.
 
     A basic block is a sequence of instructions with a single entry
@@ -37,6 +39,9 @@ def __init__(self) -> None:
         self._next: List[BasicBlock] = []
         self._idx: int = 0
         self._teal: Optional["Teal"] = None
+        self._transaction_context = BlockTransactionContext()
+        self._callsub_block: Optional[BasicBlock] = None
+        self._sub_return_point: Optional[BasicBlock] = None
 
     def add_instruction(self, instruction: Instruction) -> None:
         """Append instruction to this basic block.
@@ -113,6 +118,28 @@ def idx(self) -> int:
     def idx(self, i: int) -> None:
         self._idx = i
 
+    @property
+    def callsub_block(self) -> Optional["BasicBlock"]:
+        """If this block is the return point of a subroutine, `callsub_block` is the block
+        that called the subroutine.
+        """
+        return self._callsub_block
+
+    @callsub_block.setter
+    def callsub_block(self, b: "BasicBlock") -> None:
+        self._callsub_block = b
+
+    @property
+    def sub_return_point(self) -> Optional["BasicBlock"]:
+        """If a subroutine is executed after this block i.e exit instruction is callsub.
+        then, sub_return_point will be basic block that will be executed after the subroutine.
+        """
+        return self._sub_return_point
+
+    @sub_return_point.setter
+    def sub_return_point(self, b: "BasicBlock") -> None:
+        self._sub_return_point = b
+
     @property
     def cost(self) -> int:
         """cost of executing all instructions in this basic block"""
@@ -127,6 +154,10 @@ def teal(self) -> Optional["Teal"]:
     def teal(self, teal_instance: "Teal") -> None:
         self._teal = teal_instance
 
+    @property
+    def transaction_context(self) -> "BlockTransactionContext":
+        return self._transaction_context
+
     def __str__(self) -> str:
         ret = ""
         for ins in self._instructions:

tealer/teal/context/block_transaction_context.py

@@ -0,0 +1,30 @@
+from typing import List, Optional
+
+from tealer.exceptions import TealerException
+from tealer.utils.algorand_constants import MAX_GROUP_SIZE
+
+
+class BlockTransactionContext:  # pylint: disable=too-few-public-methods
+
+    _group_transactions_context: Optional[List["BlockTransactionContext"]] = None
+
+    def __init__(self, tail: bool = False) -> None:
+        if not tail:
+            self._group_transactions_context = [BlockTransactionContext(True) for _ in range(16)]
+
+        # set default values
+        if tail:
+            # information from gtxn {i} instructions.
+            self.group_indices = []
+            self.group_sizes = []
+        else:
+            self.group_sizes = list(range(1, MAX_GROUP_SIZE + 1))
+            self.group_indices = list(range(0, MAX_GROUP_SIZE))
+
+    def gtxn_context(self, txn_index: int) -> "BlockTransactionContext":
+        """context information collected from gtxn {txn_index} field instructions"""
+        if self._group_transactions_context is None:
+            raise TealerException()
+        if txn_index >= MAX_GROUP_SIZE:
+            raise TealerException()
+        return self._group_transactions_context[txn_index]

tealer/teal/instructions/instructions.py

@@ -49,7 +49,7 @@ class ContractType(ComparableEnum):
 }
 
 
-class Instruction:
+class Instruction:  # pylint: disable=too-many-instance-attributes
     """Base class for Teal instructions.
 
     Any class that represents a teal instruction must inherit from
@@ -66,6 +66,7 @@ def __init__(self) -> None:
         self._bb: Optional["BasicBlock"] = None
         self._version: int = 1
         self._mode: ContractType = ContractType.ANY
+        self._callsub_ins: Optional["Instruction"] = None
 
     def add_prev(self, prev_ins: "Instruction") -> None:
         """Add instruction that may execute just before this instruction.
@@ -137,6 +138,25 @@ def bb(self) -> Optional["BasicBlock"]:
     def bb(self, b: "BasicBlock") -> None:
         self._bb = b
 
+    @property
+    def callsub_ins(self) -> Optional["Instruction"]:
+        """if this instruction is a return point to a callsub instruction i.e callsub instruction is
+        present right before this instruction, then callsub_ins returns a reference to that callsub
+        instruction object.
+
+        e.g
+        callsub main
+        int 1
+        return
+
+        callsub_ins of `int 1` will be instruction obj of `callsub main`.
+        """
+        return self._callsub_ins
+
+    @callsub_ins.setter
+    def callsub_ins(self, ins: "Instruction") -> None:
+        self._callsub_ins = ins
+
     @property
     def version(self) -> int:
         """Teal version this instruction is introduced in and supported from."""