From 6212a7a513cfe9f36a1ef3302eab503ee4884e52 Mon Sep 17 00:00:00 2001
From: Boyan MILANOV <boyan.milanov@trailofbits.com>
Date: Thu, 17 Feb 2022 19:18:57 +0100
Subject: [PATCH] Optimise forking when there is only 1 solution (#2527)

* Optimise forking when there is only 1 solution

* Update etherum test

* Add test for forking on a unique solution

* Replace strlen test with dedicated fork test

* Fix typo

* Fix output filename in test
---
 manticore/core/manticore.py    | 54 ++++++++++++++++++++++------------
 tests/ethereum/test_general.py |  2 +-
 tests/native/test_models.py    |  1 +
 tests/other/test_fork.py       | 35 ++++++++++++++++++++++
 4 files changed, 73 insertions(+), 19 deletions(-)
 create mode 100644 tests/other/test_fork.py

diff --git a/manticore/core/manticore.py b/manticore/core/manticore.py
index 579068221..3d5b8cf0e 100644
--- a/manticore/core/manticore.py
+++ b/manticore/core/manticore.py
@@ -523,30 +523,48 @@ def setstate(x, y):
 
         self._publish("will_fork_state", state, expression, solutions, policy)
 
-        # Build and enqueue a state for each solution
+        # Create new states
         children = []
-        for new_value in solutions:
-            with state as new_state:
-                new_state.constrain(expression == new_value)
+        if len(solutions) == 1:
+            # If only one solution don't copy the state but update the current
+            # state instead
+            # Add state constraint
+            new_value = solutions[0]
+            state.constrain(expression == new_value)
+            setstate(state, new_value)
+
+            # Put the state back in ready list
+            self._put_state(state)  # Doesn't change state id
+            self._publish("did_fork_state", state, expression, solutions, policy, [])
+
+            # Remove the state from busy list
+            with self._lock:
+                self._busy_states.remove(state.id)
+                self._lock.notify_all()
+        else:
+            # Build and enqueue a state for each solution
+            for new_value in solutions:
+                with state as new_state:
+                    new_state.constrain(expression == new_value)
 
-                # and set the PC of the new state to the concrete pc-dest
-                # (or other register or memory address to concrete)
-                setstate(new_state, new_value)
+                    # and set the PC of the new state to the concrete pc-dest
+                    # (or other register or memory address to concrete)
+                    setstate(new_state, new_value)
 
-                # enqueue new_state, assign new state id
-                new_state_id = self._put_state(new_state)
+                    # enqueue new_state, assign new state id
+                    new_state_id = self._put_state(new_state)
 
-                # maintain a list of children for logging purpose
-                children.append(new_state_id)
+                    # maintain a list of children for logging purpose
+                    children.append(new_state_id)
 
-        self._publish("did_fork_state", state, expression, solutions, policy, children)
-        logger.debug("Forking current state %r into states %r", state.id, children)
+            self._publish("did_fork_state", state, expression, solutions, policy, children)
+            logger.debug("Forking current state %r into states %r", state.id, children)
 
-        with self._lock:
-            self._busy_states.remove(state.id)
-            self._remove(state.id)
-            state._id = None
-            self._lock.notify_all()
+            with self._lock:
+                self._busy_states.remove(state.id)
+                self._remove(state.id)
+                state._id = None
+                self._lock.notify_all()
 
     @staticmethod
     @deprecated("Use utils.log.set_verbosity instead.")
diff --git a/tests/ethereum/test_general.py b/tests/ethereum/test_general.py
index 00308a73f..661b8cde2 100644
--- a/tests/ethereum/test_general.py
+++ b/tests/ethereum/test_general.py
@@ -779,7 +779,7 @@ def test_gen_testcase_only_if(self):
             for ext in ("summary", "constraints", "pkl", "tx.json", "tx", "trace", "logs")
         }
 
-        expected_files.add("state_00000001.pkl")
+        expected_files.add("state_00000000.pkl")
 
         actual_files = set((fn for fn in os.listdir(self.mevm.workspace) if not fn.startswith(".")))
         self.assertEqual(actual_files, expected_files)
diff --git a/tests/native/test_models.py b/tests/native/test_models.py
index 2b41cafbf..964b805b7 100644
--- a/tests/native/test_models.py
+++ b/tests/native/test_models.py
@@ -12,6 +12,7 @@
     issymbolic,
     ArraySelect,
     BitVecITE,
+    ArrayProxy,
 )
 from manticore.native.state import State
 from manticore.platforms import linux
diff --git a/tests/other/test_fork.py b/tests/other/test_fork.py
new file mode 100644
index 000000000..2189a150d
--- /dev/null
+++ b/tests/other/test_fork.py
@@ -0,0 +1,35 @@
+import unittest
+import tempfile
+from manticore.native import Manticore
+from manticore.core.state import Concretize
+from pathlib import Path
+from glob import glob
+
+
+class TestFork(unittest.TestCase):
+    def test_fork_unique_solution(self):
+        binary = str(
+            Path(__file__).parent.parent.parent.joinpath(
+                "tests", "native", "binaries", "hello_world"
+            )
+        )
+        tmp_dir = tempfile.TemporaryDirectory(prefix="mcore_test_fork_")
+        m = Manticore(binary, stdin_size=10, workspace_url=str(tmp_dir.name))
+
+        @m.hook(0x3E50)  # Entrypoint
+        def concretize_var(state):
+            # Concretize symbolic var that has only 1 solution
+            var = BitVecVariable(size=32, name="bar")
+            state.constrain(var == 5)
+            raise Concretize(var)
+
+        m.run()
+        m.finalize()
+
+        # Check that no additional state was created when forking
+        states = f"{str(m.workspace)}/test_*.pkl"
+        self.assertEqual(len(glob(states)), 1)
+
+
+if __name__ == "__main__":
+    unittest.main()