Skip to content

Commit

Permalink
Add a handful of type-hint annotations.
Browse files Browse the repository at this point in the history
This is not exhaustive, just a handful of annotations that
improved my quality-of-life while developing. Note that
I occasionally import type dependencies using the
`if typing.TYPE_CHECKING` construct .. this avoids runtime
circular imports.
  • Loading branch information
robmcl4 committed Dec 18, 2023
1 parent 4475f51 commit 71c695f
Show file tree
Hide file tree
Showing 4 changed files with 48 additions and 14 deletions.
8 changes: 8 additions & 0 deletions greed/factory.py
Original file line number Diff line number Diff line change
@@ -1,14 +1,22 @@
import typing
import logging

from greed.block import Block
from greed.sim_manager import SimulationManager
from greed.state import SymbolicEVMState

if typing.TYPE_CHECKING:
from greed.project import Project
from greed.TAC import TAC_Statement
from greed.function import TAC_Function


log = logging.getLogger(__name__)


class Factory:
project: 'Project'

def __init__(self, project: "Project"):
self.project = project

Expand Down
21 changes: 12 additions & 9 deletions greed/sim_manager.py
Original file line number Diff line number Diff line change
@@ -1,11 +1,14 @@
import logging
import os
import sys
from typing import Callable
from typing import Callable, List, Optional, TYPE_CHECKING

from greed import options
from greed.state import SymbolicEVMState

if TYPE_CHECKING:
from greed.exploration_techniques import ExplorationTechnique

log = logging.getLogger(__name__)


Expand Down Expand Up @@ -35,35 +38,35 @@ def set_error(self, s: str):
self.error += [s]

@property
def states(self):
def states(self) -> List[SymbolicEVMState]:
"""
:return: All the states
"""
return sum(self.stashes.values(), [])

@property
def active(self):
def active(self) -> List[SymbolicEVMState]:
"""
:return: Active stash
"""
return self.stashes['active']

@property
def deadended(self):
def deadended(self) -> List[SymbolicEVMState]:
"""
:return: Deadended stash
"""
return self.stashes['deadended']

@property
def found(self):
def found(self) -> List[SymbolicEVMState]:
"""
:return: Found stash
"""
return self.stashes['found']

@property
def one_active(self):
def one_active(self) -> Optional[SymbolicEVMState]:
"""
:return: First element of the active stash, or None if the stash is empty
"""
Expand All @@ -73,7 +76,7 @@ def one_active(self):
return None

@property
def one_deadended(self):
def one_deadended(self) -> Optional[SymbolicEVMState]:
"""
:return: First element of the deadended stash, or None if the stash is empty
"""
Expand All @@ -83,7 +86,7 @@ def one_deadended(self):
return None

@property
def one_found(self):
def one_found(self) -> Optional[SymbolicEVMState]:
"""
:return: First element of the found stash, or None if the stash is empty
"""
Expand Down Expand Up @@ -146,7 +149,7 @@ def step(self, find: Callable[[SymbolicEVMState], bool] = lambda s: False,
for s in self.stashes['pruned'] + self.stashes['unsat'] + self.stashes['errored']:
s.solver.dispose_context()

def single_step_state(self, state: SymbolicEVMState):
def single_step_state(self, state: SymbolicEVMState) -> List[SymbolicEVMState]:
log.debug(f"Stepping {state}")
log.debug(state.curr_stmt)

Expand Down
24 changes: 22 additions & 2 deletions greed/state.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import datetime
import logging
import typing

from greed import options as opt
from greed.memory import LambdaMemory, PartialConcreteStorage
Expand All @@ -11,9 +12,28 @@

log = logging.getLogger(__name__)

if typing.TYPE_CHECKING:
from greed.project import Project


class SymbolicEVMState:
uuid_generator = UUIDGenerator()
xid: int
project: "Project"
code: bytes
uuid: int
active_plugins: typing.Dict[str, SimStatePlugin]
_pc: str
trace: typing.List[str]
memory: LambdaMemory
options: typing.Dict[str, typing.Any]
registers: typing.Dict[str, typing.Any]

# default plugins
solver: SimStateSolver
globals: SimStateGlobals
inspect: SimStateInspect


def __init__(self, xid, project, partial_init=False, init_ctx=None, options=None, max_calldatasize=None, partial_concrete_storage=False):
self.xid = xid
Expand Down Expand Up @@ -158,7 +178,7 @@ def set_init_ctx(self, init_ctx=None):
self.ctx["CALLVALUE"] = BVV(init_ctx["CALLVALUE"], 256)

@property
def pc(self):
def pc(self) -> str:
return self._pc

@property
Expand All @@ -170,7 +190,7 @@ def constraints(self):
return self.solver.constraints

@pc.setter
def pc(self, value):
def pc(self, value: str):
self._pc = value

def set_next_pc(self):
Expand Down
9 changes: 6 additions & 3 deletions greed/state_plugins/solver.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,11 +3,14 @@
from greed import options
from greed.solver.shortcuts import *
from greed.state_plugins import SimStatePlugin
from greed.solver.solver import Solver

log = logging.getLogger(__name__)


class SimStateSolver(SimStatePlugin):
_solver: Solver

def __init__(self, partial_init=False):
super(SimStateSolver, self).__init__()

Expand Down Expand Up @@ -41,14 +44,14 @@ def _add_assertions(self, assertions):
# Adding the constraints to the backend
self._solver.add_assertions(assertions)

def push(self):
def push(self) -> int:
self._curr_frame_level += 1
self._path_constraints[self._curr_frame_level] = set()
self._memory_constraints[self._curr_frame_level] = set()
self._solver.push()
return self._curr_frame_level

def pop(self):
def pop(self) -> int:
if self._curr_frame_level == 0:
log.fatal("Cannot pop() level 0 frame")
return self._curr_frame_level
Expand Down Expand Up @@ -184,7 +187,7 @@ def eval_memory_at(self, array, offset, length, raw=False):
else:
return f"{int_value:0{bv_unsigned_value(length)*2}x}"

def copy(self):
def copy(self) -> 'SimStateSolver':
new_solver = SimStateSolver(partial_init=True)
new_solver._solver = self._solver.copy()

Expand Down

0 comments on commit 71c695f

Please sign in to comment.