Move iprof to plugin (#1406)

* Add instruction level plugin support

* Add iprof register hooks

* Finish moving iprof to plugins

* Add comments and refactor

* Revert changes to pruner

* Handle changes to tests

Co-authored-by: Nikhil Parasaram <nikhilparasaram@Nikhils-MacBook-Pro.local>

mythril/analysis/symbolic.py

@@ -3,7 +3,6 @@ purposes."""
 
 from mythril.analysis.module import EntryPoint, ModuleLoader, get_detection_module_hooks
 from mythril.laser.ethereum import svm
-from mythril.laser.ethereum.iprof import InstructionProfiler
 from mythril.laser.ethereum.state.account import Account
 from mythril.laser.ethereum.state.world_state import WorldState
 from mythril.laser.ethereum.strategy.basic import (
@@ -24,6 +23,7 @@ from mythril.laser.plugin.plugins import (
     DependencyPrunerBuilder,
     CoveragePluginBuilder,
     CallDepthLimitBuilder,
+    InstructionProfilerBuilder,
 )
 from mythril.laser.ethereum.strategy.extensions.bounded_loops import (
     BoundedLoopsStrategy,
@@ -55,7 +55,6 @@ class SymExecWrapper:
         transaction_count: int = 2,
         modules: Optional[List[str]] = None,
         compulsory_statespace: bool = True,
-        iprof: Optional[InstructionProfiler] = None,
         disable_dependency_pruning: bool = False,
         run_analysis_modules: bool = True,
         custom_modules_directory: str = "",
@@ -121,7 +120,6 @@ class SymExecWrapper:
             create_timeout=create_timeout,
             transaction_count=transaction_count,
             requires_statespace=requires_statespace,
-            iprof=iprof,
         )
 
         if loop_bound is not None:
@@ -131,9 +129,11 @@ class SymExecWrapper:
         plugin_loader.load(CoveragePluginBuilder())
         plugin_loader.load(MutationPrunerBuilder())
         plugin_loader.load(CallDepthLimitBuilder())
+        plugin_loader.load(InstructionProfilerBuilder())
         plugin_loader.add_args(
             "call-depth-limit", call_depth_limit=args.call_depth_limit
         )
+
         if not disable_dependency_pruning:
             plugin_loader.load(DependencyPrunerBuilder())
 

mythril/laser/ethereum/instructions.py

@@ -4,7 +4,7 @@ import binascii
 import logging
 
 from copy import copy, deepcopy
-from typing import cast, Callable, List, Union
+from typing import cast, Callable, List, Union, Optional
 from datetime import datetime
 
 from mythril.laser.smt import (
@@ -202,7 +202,13 @@ class Instruction:
     """Instruction class is used to mutate a state according to the current
     instruction."""
 
-    def __init__(self, op_code: str, dynamic_loader: DynLoader, iprof=None) -> None:
+    def __init__(
+        self,
+        op_code: str,
+        dynamic_loader: DynLoader,
+        pre_hooks: List[Callable] = None,
+        post_hooks: List[Callable] = None,
+    ) -> None:
         """
 
         :param op_code:
@@ -211,7 +217,16 @@ class Instruction:
         """
         self.dynamic_loader = dynamic_loader
         self.op_code = op_code.upper()
-        self.iprof = iprof
+        self.pre_hook = pre_hooks if pre_hooks else []
+        self.post_hook = post_hooks if post_hooks else []
+
+    def _execute_pre_hooks(self, global_state: GlobalState):
+        for hook in self.pre_hook:
+            hook(global_state)
+
+    def _execute_post_hooks(self, global_state: GlobalState):
+        for hook in self.post_hook:
+            hook(global_state)
 
     def evaluate(self, global_state: GlobalState, post=False) -> List[GlobalState]:
         """Performs the mutation for this instruction.
@@ -241,13 +256,10 @@ class Instruction:
         if instruction_mutator is None:
             raise NotImplementedError
 
-        if self.iprof is None:
-            result = instruction_mutator(global_state)
-        else:
-            start_time = datetime.now()
+        self._execute_pre_hooks(global_state)
+
         result = instruction_mutator(global_state)
-            end_time = datetime.now()
-            self.iprof.record(op, start_time, end_time)
+        self._execute_post_hooks(global_state)
 
         return result
 

mythril/laser/ethereum/svm.py

@@ -5,6 +5,7 @@ from copy import copy
 from datetime import datetime, timedelta
 from typing import Callable, Dict, DefaultDict, List, Tuple, Optional
 
+from mythril.support.opcodes import opcodes as OPCODES
 from mythril.analysis.potential_issues import check_potential_issues
 from mythril.laser.ethereum.cfg import NodeFlags, Node, Edge, JumpType
 from mythril.laser.ethereum.evm_exceptions import StackUnderflowException
@@ -105,7 +106,11 @@ class LaserEVM:
         self._stop_sym_exec_hooks = []  # type: List[Callable]
 
         self.iprof = iprof
-
+        self.instr_pre_hook = {}  # type: Dict[str, List[Callable]]
+        self.instr_post_hook = {}  # type: Dict[str, List[Callable]]
+        for op, _, _, _ in OPCODES.values():
+            self.instr_pre_hook[op] = []
+            self.instr_post_hook[op] = []
         log.info("LASER EVM initialized with dynamic loader: " + str(dynamic_loader))
 
     def extend_strategy(self, extension: ABCMeta, *args) -> None:
@@ -176,9 +181,6 @@ class LaserEVM:
                 self.total_states,
             )
 
-        if self.iprof is not None:
-            log.info("Instruction Statistics:\n{}".format(self.iprof))
-
         for hook in self._stop_sym_exec_hooks:
             hook()
 
@@ -331,7 +333,10 @@ class LaserEVM:
 
         try:
             new_global_states = Instruction(
-                op_code, self.dynamic_loader, self.iprof
+                op_code,
+                self.dynamic_loader,
+                pre_hooks=self.instr_pre_hook[op_code],
+                post_hooks=self.instr_post_hook[op_code],
             ).evaluate(global_state)
 
         except VmException as e:
@@ -443,7 +448,10 @@ class LaserEVM:
 
         # Execute the post instruction handler
         new_global_states = Instruction(
-            op_code, self.dynamic_loader, self.iprof
+            op_code,
+            self.dynamic_loader,
+            pre_hooks=self.instr_pre_hook[op_code],
+            post_hooks=self.instr_post_hook[op_code],
         ).evaluate(return_global_state, True)
 
         # In order to get a nice call graph we need to set the nodes here
@@ -579,6 +587,37 @@ class LaserEVM:
                 "Invalid hook type %s. Must be one of {add_world_state}", hook_type
             )
 
+    def register_instr_hooks(self, hook_type: str, opcode: str, hook: Callable):
+        """Registers instructions hooks from plugins"""
+        if hook_type == "pre":
+            if opcode is None:
+                for op, _, _, _ in OPCODES.values():
+                    self.instr_pre_hook[op].append(hook(op))
+            else:
+                self.instr_pre_hook[opcode].append(hook)
+        else:
+            if opcode is None:
+                for op, _, _, _ in OPCODES.values():
+                    self.instr_post_hook[op].append(hook(op))
+            else:
+                self.instr_post_hook[opcode].append(hook)
+
+    def instr_hook(self, hook_type, opcode) -> Callable:
+        """Registers the annoted function with register_instr_hooks
+
+        :param hook_type: Type of hook pre/post
+        :param opcode: The opcode related to the function
+        """
+
+        def hook_decorator(func: Callable):
+            """ Hook decorator generated by laser_hook
+
+            :param func: Decorated function
+            """
+            self.register_instr_hooks(hook_type, opcode, func)
+
+        return hook_decorator
+
     def laser_hook(self, hook_type: str) -> Callable:
         """Registers the annotated function with register_laser_hooks
 

mythril/laser/plugin/plugins/__init__.py

@@ -10,3 +10,4 @@ from mythril.laser.plugin.plugins.coverage.coverage_plugin import CoveragePlugin
 from mythril.laser.plugin.plugins.dependency_pruner import DependencyPrunerBuilder
 from mythril.laser.plugin.plugins.mutation_pruner import MutationPrunerBuilder
 from mythril.laser.plugin.plugins.call_depth_limiter import CallDepthLimitBuilder
+from mythril.laser.plugin.plugins.instruction_profiler import InstructionProfilerBuilder

mythril/laser/plugin/plugins/instruction_profiler.py

@@ -0,0 +1,116 @@
+from collections import namedtuple
+from datetime import datetime
+from typing import Dict, List, Tuple
+from mythril.laser.plugin.builder import PluginBuilder
+from mythril.laser.plugin.interface import LaserPlugin
+from mythril.laser.ethereum.svm import LaserEVM
+from mythril.laser.ethereum.state.global_state import GlobalState
+from datetime import datetime
+import logging
+
+# Type annotations:
+#   start_time: datetime
+#   end_time: datetime
+_InstrExecRecord = namedtuple("_InstrExecRecord", ["start_time", "end_time"])
+
+# Type annotations:
+#   total_time: float
+#   total_nr: float
+#   min_time: float
+#   max_time: float
+_InstrExecStatistic = namedtuple(
+    "_InstrExecStatistic", ["total_time", "total_nr", "min_time", "max_time"]
+)
+
+# Map the instruction opcode to its records if all execution times
+_InstrExecRecords = Dict[str, List[_InstrExecRecord]]
+
+# Map the instruction opcode to the statistic of its execution times
+_InstrExecStatistics = Dict[str, _InstrExecStatistic]
+
+log = logging.getLogger(__name__)
+
+
+class InstructionProfilerBuilder(PluginBuilder):
+    plugin_name = "dependency-pruner"
+
+    def __call__(self, *args, **kwargs):
+        return InstructionProfiler()
+
+
+class InstructionProfiler(LaserPlugin):
+    """Performance profile for the execution of each instruction.
+    """
+
+    def __init__(self):
+        self._reset()
+
+    def _reset(self):
+        self.records = dict()
+        self.start_time = None
+
+    def initialize(self, symbolic_vm: LaserEVM) -> None:
+        @symbolic_vm.instr_hook("pre", None)
+        def get_start_time(op_code: str):
+            def start_time_wrapper(global_state: GlobalState):
+                self.start_time = datetime.now()
+
+            return start_time_wrapper
+
+        @symbolic_vm.instr_hook("post", None)
+        def record(op_code: str):
+            def record_opcode(global_state: GlobalState):
+                end_time = datetime.now()
+                try:
+                    self.records[op_code].append(
+                        _InstrExecRecord(self.start_time, end_time)
+                    )
+                except KeyError:
+                    self.records[op_code] = [
+                        _InstrExecRecord(self.start_time, end_time)
+                    ]
+
+            return record_opcode
+
+        @symbolic_vm.laser_hook("stop_sym_exec")
+        def print_stats():
+            total, stats = self._make_stats()
+
+            s = "Total: {} s\n".format(total)
+
+            for op in sorted(stats):
+                stat = stats[op]
+                s += "[{:12s}] {:>8.4f} %,  nr {:>6},  total {:>8.4f} s,  avg {:>8.4f} s,  min {:>8.4f} s,  max {:>8.4f} s\n".format(
+                    op,
+                    stat.total_time * 100 / total,
+                    stat.total_nr,
+                    stat.total_time,
+                    stat.total_time / stat.total_nr,
+                    stat.min_time,
+                    stat.max_time,
+                )
+
+            log.info(s)
+
+    def _make_stats(self) -> Tuple[float, _InstrExecStatistics]:
+        periods = {
+            op: list(
+                map(lambda r: r.end_time.timestamp() - r.start_time.timestamp(), rs)
+            )
+            for op, rs in self.records.items()
+        }
+
+        stats = dict()
+        total_time = 0
+
+        for _, (op, times) in enumerate(periods.items()):
+            stat = _InstrExecStatistic(
+                total_time=sum(times),
+                total_nr=len(times),
+                min_time=min(times),
+                max_time=max(times),
+            )
+            total_time += stat.total_time
+            stats[op] = stat
+
+        return total_time, stats

mythril/mythril/mythril_analyzer.py

@@ -65,7 +65,6 @@ class MythrilAnalyzer:
         self.execution_timeout = execution_timeout
         self.loop_bound = loop_bound
         self.create_timeout = create_timeout
-        self.iprof = InstructionProfiler() if enable_iprof else None
         self.disable_dependency_pruning = disable_dependency_pruning
         self.custom_modules_directory = custom_modules_directory
         args.sparse_pruning = sparse_pruning
@@ -73,6 +72,7 @@ class MythrilAnalyzer:
         args.parallel_solving = parallel_solving
         args.unconstrained_storage = unconstrained_storage
         args.call_depth_limit = call_depth_limit
+        args.iprof = enable_iprof
 
     def dump_statespace(self, contract: EVMContract = None) -> str:
         """
@@ -88,7 +88,6 @@ class MythrilAnalyzer:
             max_depth=self.max_depth,
             execution_timeout=self.execution_timeout,
             create_timeout=self.create_timeout,
-            iprof=self.iprof,
             disable_dependency_pruning=self.disable_dependency_pruning,
             run_analysis_modules=False,
             custom_modules_directory=self.custom_modules_directory,
@@ -121,7 +120,6 @@ class MythrilAnalyzer:
             execution_timeout=self.execution_timeout,
             transaction_count=transaction_count,
             create_timeout=self.create_timeout,
-            iprof=self.iprof,
             disable_dependency_pruning=self.disable_dependency_pruning,
             run_analysis_modules=False,
             custom_modules_directory=self.custom_modules_directory,
@@ -156,7 +154,6 @@ class MythrilAnalyzer:
                     transaction_count=transaction_count,
                     modules=modules,
                     compulsory_statespace=False,
-                    iprof=self.iprof,
                     disable_dependency_pruning=self.disable_dependency_pruning,
                     custom_modules_directory=self.custom_modules_directory,
                 )
@@ -166,8 +163,6 @@ class MythrilAnalyzer:
                 raise e
             except KeyboardInterrupt:
                 log.critical("Keyboard Interrupt")
-                if self.iprof is not None:
-                    log.info("Instruction Statistics:\n{}".format(self.iprof))
                 issues = retrieve_callback_issues(modules)
             except Exception:
                 log.critical(
@@ -176,8 +171,6 @@ class MythrilAnalyzer:
                 )
                 issues = retrieve_callback_issues(modules)
                 exceptions.append(traceback.format_exc())
-                if self.iprof is not None:
-                    log.info("Instruction Statistics:\n{}".format(self.iprof))
             for issue in issues:
                 issue.add_code_info(contract)
 

mythril/support/support_args.py

@@ -10,6 +10,7 @@ class Args:
         self.unconstrained_storage = False
         self.parallel_solving = False
         self.call_depth_limit = 3
+        self.iprof = True
 
 
 args = Args()