Add generic dataflow framework + example + stack value analysis

tealer/analyses/dataflow/abstract.py

@@ -0,0 +1,60 @@
+from abc import ABC, abstractmethod
+from typing import List, Any, TypeVar, Generic, TYPE_CHECKING
+
+from tealer.teal.basic_blocks import BasicBlock
+
+AbstractValues = TypeVar("AbstractValues")
+
+if TYPE_CHECKING:
+ from tealer.teal.teal import Teal
+
+
+class AbstractDataflow(ABC, Generic[AbstractValues]):
+ def __init__(self, teal: "Teal"):
+ self.teal = teal
+
+ @abstractmethod
+ def _merge_predecessor(self, bb: BasicBlock) -> AbstractValues:
+ pass
+
+ @abstractmethod
+ def _is_fix_point(self, bb: BasicBlock, values: AbstractValues) -> bool:
+ pass
+
+ @abstractmethod
+ def _transfer_function(self, bb: BasicBlock) -> AbstractValues:
+ pass
+
+ @abstractmethod
+ def _store_values_in(self, bb: BasicBlock, values: AbstractValues) -> None:
+ pass
+
+ @abstractmethod
+ def _store_values_out(self, bb: BasicBlock, values: AbstractValues) -> None:
+ pass
+
+ @abstractmethod
+ def _filter_successors(self, bb: BasicBlock) -> List[BasicBlock]:
+ pass
+
+ @abstractmethod
+ def result(self) -> Any:
+ pass
+
+ def explore(self, bb: BasicBlock, is_entry_node: bool = False) -> None:
+
+ values = self._merge_predecessor(bb)
+
+ if not is_entry_node and self._is_fix_point(bb, values):
+ return
+
+ self._store_values_in(bb, values)
+ values = self._transfer_function(bb)
+ self._store_values_out(bb, values)
+
+ successors = self._filter_successors(bb)
+ for successor in successors:
+ self.explore(successor)
+
+ def run_analysis(self) -> None:
+ self.explore(self.teal.bbs[0], is_entry_node=True)

tealer/analyses/dataflow/collect_instructions.py

@@ -0,0 +1,102 @@
+from collections import defaultdict
+from typing import Union, Set, Any, List, Dict, TYPE_CHECKING
+
+from tealer.analyses.dataflow.abstract import AbstractDataflow
+from tealer.teal.basic_blocks import BasicBlock
+from tealer.teal.instructions.instructions import Instruction
+
+if TYPE_CHECKING:
+ from tealer.teal.teal import Teal
+
+MAX_ELEMS = 35
+
+
+class InstructionSet:
+ def __init__(self, values: Union[str, Instruction, Set[Instruction]]) -> None:
+
+ if isinstance(values, str):
+ assert values in ["TOP", "BOTTOM"]
+
+ if isinstance(values, set) and len(values) > MAX_ELEMS:
+ values = "TOP"
+ if isinstance(values, Instruction):
+ values = {values}
+
+ self.values: Union[str, Set[Instruction]] = values
+
+ @property
+ def is_top(self) -> bool:
+ return isinstance(self.values, str) and self.values == "TOP"
+
+ @property
+ def is_bottom(self) -> bool:
+ return isinstance(self.values, str) and self.values == "BOTTOM"
+
+ def union(self, instructions: "InstructionSet") -> "InstructionSet":
+ v0 = self.values
+ v1 = instructions.values
+ if v0 == "TOP" or v1 == "TOP":
+ return InstructionSet("TOP")
+
+ if v0 == "BOTTOM":
+ return InstructionSet(v1)
+
+ if v1 == "BOTTOM":
+ return InstructionSet(v0)
+
+ assert isinstance(v0, set)
+ assert isinstance(v1, set)
+
+ return InstructionSet(v0.union(v1))
+
+ def __eq__(self, other: Any) -> bool:
+ if isinstance(other, InstructionSet):
+ if isinstance(self.values, str) and isinstance(other.values, str):
+ return self.values == other.values
+ if isinstance(self.values, set) and isinstance(other.values, set):
+ return self.values == other.values
+ return False
+
+ def __str__(self) -> str:
+ if isinstance(self.values, str):
+ return self.values
+ return "[" + ",".join([str(x) for x in self.values]) + "]"
+
+
+class CollectInstructions(AbstractDataflow[InstructionSet]):
+ def __init__(self, teal: "Teal") -> None:
+ super().__init__(teal)
+ self.bb_in: Dict[BasicBlock, InstructionSet] = defaultdict(lambda: InstructionSet("BOTTOM"))
+ self.bb_out: Dict[BasicBlock, InstructionSet] = defaultdict(
+ lambda: InstructionSet("BOTTOM")
+ )
+
+ def _merge_predecessor(self, bb: BasicBlock) -> InstructionSet:
+ s = InstructionSet("BOTTOM")
+ for bb_prev in bb.prev:
+ s = s.union(self.bb_out[bb_prev])
+
+ return s
+
+ def _is_fix_point(self, bb: BasicBlock, values: InstructionSet) -> bool:
+ return self.bb_in[bb] == values
+
+ def _transfer_function(self, bb: BasicBlock) -> InstructionSet:
+ bb_out = self.bb_in[bb]
+
+ for ins in bb.instructions:
+ bb_out = bb_out.union(InstructionSet(ins))
+
+ return bb_out
+
+ def _store_values_in(self, bb: BasicBlock, values: InstructionSet) -> None:
+ self.bb_in[bb] = values
+
+ def _store_values_out(self, bb: BasicBlock, values: InstructionSet) -> None:
+ self.bb_out[bb] = values
+
+ def _filter_successors(self, bb: BasicBlock) -> List[BasicBlock]:
+ return bb.next
+
+ def result(self) -> Dict[BasicBlock, InstructionSet]:
+ return self.bb_out

tealer/analyses/dataflow/stack_value.py

@@ -0,0 +1,229 @@
+from collections import defaultdict
+from typing import Union, Set, Any, List, Dict, TYPE_CHECKING, Type, Callable, Tuple
+
+from tealer.analyses.dataflow.abstract import AbstractDataflow
+from tealer.teal.basic_blocks import BasicBlock
+from tealer.teal.global_field import GlobalField
+from tealer.teal.instructions import instructions
+from tealer.teal.instructions.instructions import Instruction
+from tealer.teal.instructions.transaction_field import TransactionField
+
+if TYPE_CHECKING:
+ from tealer.teal.teal import Teal
+
+MAX_ELEMS_PER_STACK_VALUE = 35
+MAX_STACK_DEPTH = 100
+
+
+# pylint: disable=too-few-public-methods
+class TOP:
+ def __eq__(self, other: Any) -> bool:
+ return isinstance(other, TOP)
+
+ def __str__(self) -> str:
+ return "TOP"
+
+
+# pylint: disable=too-few-public-methods
+class BOTTOM:
+ def __eq__(self, other: Any) -> bool:
+ return isinstance(other, BOTTOM)
+
+ def __str__(self) -> str:
+ return "BOTTOM"
+
+
+VALUES_TRACKED = Union[Set[Union[GlobalField, TransactionField, int, str]], TOP, BOTTOM]
+
+
+class StackValue:
+ """
+ StackValue represent an abstract value on the stack
+ It can be either a set of int/str/fields, or TOP/BOTTOM
+ The set's size is limited to MAX_ELEMS_PER_STACK_VALUE, if above, it becomes TOP
+
+ """
+
+ def __init__(self, values: VALUES_TRACKED):
+ if isinstance(values, set) and len(values) > MAX_ELEMS_PER_STACK_VALUE:
+ values = TOP()
+ self.values = values
+
+ def union(self, other_stack_value: "StackValue") -> "StackValue":
+ self_values = self.values
+ other_values = other_stack_value.values
+ if isinstance(self_values, TOP) or isinstance(other_values, TOP):
+ return StackValue(TOP())
+ if isinstance(self_values, BOTTOM):
+ return StackValue(other_values)
+ if isinstance(other_values, BOTTOM):
+ return StackValue(self_values)
+ assert isinstance(self_values, set)
+ assert isinstance(other_values, set)
+ return StackValue(self_values | other_values)
+
+ def __eq__(self, other: Any) -> bool:
+ if isinstance(other, StackValue):
+ return self.values == other.values
+ return False
+
+ def __str__(self) -> str:
+ values = self.values
+ if isinstance(values, (TOP, BOTTOM)):
+ return str(values)
+ assert isinstance(values, set)
+ return str({str(x) for x in values})
+
+
+class Stack:
+ """
+ Represent an abstract stack
+ The length is limited by MAX_STACK_DEPTH
+ self.values contains the abstract element
+
+ If there is two paths merged, where one path push 1 (concrete stack [1])
+ and the other push 3; (concrete stack [3])
+ then the abstract stack is
+ [ [1;3] ]x
+ Ie: the top can either be 1 or 3
+
+ If we pop beyond the known values, we return TOP().
+ As a result, if the most left elements are TOP in the stack, we can stop tracking them
+
+ If there is two paths merged, and the stack size has a different size, then
+ we use TOP for the elements in the difference. Ex:
+ - one path push 1; push 2; (concrete stack [2;1])
+ - and the other push 3; (concrete stack [3])
+ - then the abstract stack is
+ - [ [TOP] ; [1;3] ]
+ Ie: the top can either be 1 or 3, and the second one is TOP
+ Because the most left element of the stack can be removed, this can be simplied as
+ [ [1;3] ]
+
+
+
+ """
+
+ def __init__(self, values: List[StackValue]) -> None:
+ if len(values) > MAX_STACK_DEPTH:
+ values = values[:-MAX_STACK_DEPTH]
+
+ while values and values[0] == StackValue(TOP()):
+ values.pop()
+
+ self.values = values
+
+ def union(self, stack: "Stack") -> "Stack":
+
+ v1 = self.values
+ v2 = stack.values
+
+ min_length = min(len(v1), len(v2))
+ v1 = v1[:-min_length]
+ v2 = v2[:-min_length]
+
+ v3 = []
+ for i in range(min_length):
+ v3.append(v1[i].union(v2[i]))
+
+ return Stack(v3)
+
+ def pop(self, number: int) -> Tuple["Stack", List[StackValue]]:
+ if number == 0:
+ return Stack(self.values), []
+ if len(self.values) < number:
+ diff = number - len(self.values)
+ poped_values = list(self.values)
+ return Stack([]), [StackValue(TOP()) for _ in range(diff)] + poped_values
+
+ poped_values = self.values[:-number]
+ return Stack(self.values[: len(self.values) - number]), poped_values
+
+ def push(self, pushed_values: List[StackValue]) -> "Stack":
+ return Stack(self.values + pushed_values)
+
+ def __eq__(self, other: Any) -> bool:
+ if isinstance(other, Stack):
+ return self.values == other.values
+ return False
+
+ def __str__(self) -> str:
+ return str([str(x) for x in self.values])
+
+
+# pylint: disable=unused-argument
+def handle_int(ins: Instruction, stack: Stack) -> List[StackValue]:
+ assert isinstance(ins, instructions.Int)
+ return [StackValue({ins.value})]
+
+
+# pylint: disable=unused-argument
+def handle_pushint(ins: Instruction, stack: Stack) -> List[StackValue]:
+ assert isinstance(ins, instructions.PushInt)
+ return [StackValue({ins.value})]
+
+
+# pylint: disable=unused-argument
+def handle_txn(ins: Instruction, stack: Stack) -> List[StackValue]:
+ assert isinstance(ins, instructions.Txn)
+ return [StackValue({ins.field})]
+
+
+# pylint: disable=unused-argument
+def handle_global(ins: Instruction, stack: Stack) -> List[StackValue]:
+ assert isinstance(ins, instructions.Global)
+ return [StackValue({ins.field})]
+
+
+special_handling: Dict[Type[Instruction], Callable[[Instruction, Stack], List[StackValue]]] = {
+ instructions.Int: handle_int,
+ instructions.PushInt: handle_pushint,
+ instructions.Txn: handle_txn,
+ instructions.Global: handle_global,
+}
+
+
+class StackValueAnalysis(AbstractDataflow[Stack]):
+ def __init__(self, teal: "Teal") -> None:
+ super().__init__(teal)
+ self.bb_in: Dict[BasicBlock, Stack] = defaultdict(lambda: Stack([StackValue(BOTTOM())]))
+ self.bb_out: Dict[BasicBlock, Stack] = defaultdict(lambda: Stack([]))
+
+ self.ins_in: Dict[Instruction, Stack] = defaultdict(lambda: Stack([]))
+ self.ins_out: Dict[Instruction, Stack] = defaultdict(lambda: Stack([]))
+
+ def _merge_predecessor(self, bb: BasicBlock) -> Stack:
+ s = Stack([])
+ for bb_prev in bb.prev:
+ s = s.union(self.bb_out[bb_prev])
+ return s
+
+ def _is_fix_point(self, bb: BasicBlock, values: Stack) -> bool:
+ return self.bb_in[bb] == values
+
+ def _transfer_function(self, bb: BasicBlock) -> Stack:
+ bb_out = self.bb_in[bb]
+
+ for ins in bb.instructions:
+ self.ins_in[ins] = Stack(bb_out.values)
+ if type(ins) in special_handling:
+ pushed_elems = special_handling[type(ins)](ins, bb_out)
+ else:
+ pushed_elems = [StackValue(TOP()) for _ in range(ins.stack_push_size)]
+ bb_out, _ = bb_out.pop(ins.stack_pop_size)
+ bb_out = bb_out.push(pushed_elems)
+ self.ins_out[ins] = Stack(bb_out.values)
+
+ return bb_out
+
+ def _store_values_in(self, bb: BasicBlock, values: Stack) -> None:
+ self.bb_in[bb] = values
+
+ def _store_values_out(self, bb: BasicBlock, values: Stack) -> None:
+ self.bb_out[bb] = values
+
+ def _filter_successors(self, bb: BasicBlock) -> List[BasicBlock]:
+ return bb.next
+
+ def result(self) -> Dict[BasicBlock, Stack]:
+ return self.bb_out