From 40a9bef6a2029f29b4da48b94c4bef02df195770 Mon Sep 17 00:00:00 2001
From: Nikhil Parasaram <nikhilparasaram@Nikhils-MacBook-Pro.local>
Date: Tue, 26 May 2020 11:09:51 +0530
Subject: [PATCH] Update the state merge plugin

---
 mythril/analysis/symbolic.py                    |  4 ++--
 mythril/laser/plugin/builder.py                 |  1 -
 mythril/laser/plugin/plugins/__init__.py        |  2 +-
 .../plugin/plugins/state_merge/__init__.py      | 17 +++++++++++------
 4 files changed, 14 insertions(+), 10 deletions(-)

diff --git a/mythril/analysis/symbolic.py b/mythril/analysis/symbolic.py
index 07da0c79..d3c35c4a 100644
--- a/mythril/analysis/symbolic.py
+++ b/mythril/analysis/symbolic.py
@@ -23,7 +23,7 @@ from mythril.laser.plugin.plugins import (
     MutationPrunerBuilder,
     DependencyPrunerBuilder,
     CoveragePluginBuilder,
-    StateMerge,
+    StateMergePluginBuilder,
 )
 from mythril.laser.ethereum.strategy.extensions.bounded_loops import (
     BoundedLoopsStrategy,
@@ -132,7 +132,7 @@ class SymExecWrapper:
         plugin_loader.load(CoveragePluginBuilder())
         plugin_loader.load(MutationPrunerBuilder())
         if enable_state_merging:
-            plugin_loader.load(StateMerge())
+            plugin_loader.load(StateMergePluginBuilder())
 
         if not disable_dependency_pruning:
             plugin_loader.load(DependencyPrunerBuilder())
diff --git a/mythril/laser/plugin/builder.py b/mythril/laser/plugin/builder.py
index fdddb569..5d337415 100644
--- a/mythril/laser/plugin/builder.py
+++ b/mythril/laser/plugin/builder.py
@@ -1,7 +1,6 @@
 from mythril.laser.plugin.interface import LaserPlugin
 
 from abc import ABC, abstractmethod
-from typing import Optional
 
 
 class PluginBuilder(ABC):
diff --git a/mythril/laser/plugin/plugins/__init__.py b/mythril/laser/plugin/plugins/__init__.py
index 7079c489..f00f23f3 100644
--- a/mythril/laser/plugin/plugins/__init__.py
+++ b/mythril/laser/plugin/plugins/__init__.py
@@ -9,4 +9,4 @@ from mythril.laser.plugin.plugins.benchmark import BenchmarkPluginBuilder
 from mythril.laser.plugin.plugins.coverage.coverage_plugin import CoveragePluginBuilder
 from mythril.laser.plugin.plugins.dependency_pruner import DependencyPrunerBuilder
 from mythril.laser.plugin.plugins.mutation_pruner import MutationPrunerBuilder
-from mythril.laser.plugin.plugins.state_merge import StateMerge
+from mythril.laser.plugin.plugins.state_merge import StateMergePluginBuilder
diff --git a/mythril/laser/plugin/plugins/state_merge/__init__.py b/mythril/laser/plugin/plugins/state_merge/__init__.py
index 9f7ac87f..a7d9f55f 100644
--- a/mythril/laser/plugin/plugins/state_merge/__init__.py
+++ b/mythril/laser/plugin/plugins/state_merge/__init__.py
@@ -1,11 +1,10 @@
 from copy import copy
 from typing import Set, List
 from mythril.laser.ethereum.svm import LaserEVM
-from mythril.laser.ethereum.plugins.plugin import LaserPlugin
-from mythril.laser.ethereum.plugins.implementations.state_merge.merge_states import (
-    merge_states,
-)
-from mythril.laser.ethereum.plugins.implementations.state_merge.check_mergeability import (
+from mythril.laser.plugin.interface import LaserPlugin
+from mythril.laser.plugin.plugins.state_merge.merge_states import merge_states
+from mythril.laser.plugin.builder import PluginBuilder
+from mythril.laser.plugin.plugins.state_merge.check_mergeability import (
     check_ws_merge_condition,
 )
 from mythril.laser.ethereum.state.world_state import WorldState
@@ -19,6 +18,13 @@ class MergeAnnotation(StateAnnotation):
     pass
 
 
+class StateMergePluginBuilder(PluginBuilder):
+    plugin_name = "state merge"
+
+    def __call__(self, *args, **kwargs):
+        return StateMerge()
+
+
 class StateMerge(LaserPlugin):
     """
     Tries to merge states based on their similarity.
@@ -38,7 +44,6 @@ class StateMerge(LaserPlugin):
 
         @symbolic_vm.laser_hook("stop_sym_trans")
         def execute_stop_sym_trans_hook(svm: LaserEVM):
-
             open_states = svm.open_states
             if len(open_states) <= 1:
                 return