From 349edecc4e468535c9b111be0839568795f0d24e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Matt=C3=A9o=20Delabre?= <spam@delab.re>
Date: Wed, 27 Jan 2021 13:48:34 +0100
Subject: [PATCH] Parallelize

---
 runall.py | 46 ++++++++++++++++++++++++++++++++--------------
 1 file changed, 32 insertions(+), 14 deletions(-)

diff --git a/runall.py b/runall.py
index 3348c54..738ffa9 100644
--- a/runall.py
+++ b/runall.py
@@ -1,4 +1,5 @@
 from itertools import product
+from multiprocessing import current_process, Pool
 import signal
 from subprocess import DEVNULL, PIPE, Popen, TimeoutExpired
 
@@ -7,6 +8,14 @@ chars = '0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ!"#$%&\'(
 # Prevent zombie processes
 signal.signal(signal.SIGCHLD, signal.SIG_IGN)
 
+# Current test case pairs used for testing (local to worker process)
+current_pairs = None
+
+
+def set_pairs(pairs):
+    global current_pairs
+    current_pairs = pairs
+
 
 def check_pair(script, instr, outstr):
     process = Popen(
@@ -29,29 +38,38 @@ def check_pair(script, instr, outstr):
         return False
 
 
-def check_pairs(script, pairs):
-    for pair in pairs:
+def check_script(script):
+    for pair in current_pairs:
         if not check_pair(script, *pair):
-            return False
-    return True
+            return script, False
+    return script, True
+
+
+def generate_scripts(max_length):
+    for length in range(max_length + 1):
+        for letters in product(chars, repeat=length):
+            yield "".join(letters)
 
 
 def find_script(pairs, max_length):
     candidates = []
+
     chars_count = len(chars)
-    total_options = int((chars_count ** (max_length + 1) - 1) / (chars_count - 1))
-    current_option = 0
+    num_tasks = int((chars_count ** (max_length + 1) - 1) / (chars_count - 1))
+    done_tasks = 0
 
-    for length in range(max_length + 1):
-        for letters in product(chars, repeat=length):
-            current_option += 1
+    with Pool(processes=8, initializer=set_pairs, initargs=(pairs,)) as pool:
+        for script, result in pool.imap_unordered(
+            check_script,
+            generate_scripts(max_length),
+            chunksize=10,
+        ):
+            done_tasks += 1
 
-            if current_option % 1000 == 0:
-                print(f"Progress: {current_option}/{total_options}")
+            if done_tasks % 10000 == 0:
+                print(f"Progress: {done_tasks}/{num_tasks}")
 
-            script = "".join(letters)
-
-            if check_pairs(script, pairs):
+            if result:
                 print("> Found candidate:", script)
                 candidates.append(script)