From 97b70e2ef09b2e47506b97f2166b059c20bc89f8 Mon Sep 17 00:00:00 2001
From: Forest Belton <65484+forestbelton@users.noreply.github.com>
Date: Tue, 31 Aug 2021 21:35:44 -0400
Subject: [PATCH] Minor improvements

---
 gbso/cpu/state.py         | 13 ++++++++++++-
 gbso/optimize.py          | 33 +++++++++++++++++++--------------
 gbso/program/test_case.py |  3 +++
 3 files changed, 34 insertions(+), 15 deletions(-)

diff --git a/gbso/cpu/state.py b/gbso/cpu/state.py
index 8a928bc..be7e261 100644
--- a/gbso/cpu/state.py
+++ b/gbso/cpu/state.py
@@ -1,9 +1,11 @@
 from collections import defaultdict
 from dataclasses import dataclass, field
-from typing import Dict
+from typing import Dict, List, Tuple, Union
 
 from gbso.cpu.regs import R8
 
+Loc = Union[R8, int]
+
 
 @dataclass
 class CPUState:
@@ -19,3 +21,12 @@ class CPUState:
             reg8=self.reg8.copy(),
             memory=self.memory.copy(),
         )
+
+    def with_vals(self, *init_vals: List[Tuple[Loc, int]]) -> "CPUState":
+        cpu = self.copy()
+        for loc, val in init_vals:
+            if type(loc) == R8:
+                cpu.reg8[loc] = val
+            else:
+                cpu.memory[loc] = val
+        return cpu
diff --git a/gbso/optimize.py b/gbso/optimize.py
index 309137b..1d8fe22 100644
--- a/gbso/optimize.py
+++ b/gbso/optimize.py
@@ -53,6 +53,7 @@ def cost_noperf(
 class OptimizationParameters:
     max_size: int
     beta: float = DEFAULT_ANNEALING_CONSTANT
+    synthesize: bool = True
     synthesis_iters: int = DEFAULT_SYNTHESIS_ITERS
     optimize_iters: int = DEFAULT_OPTIMIZE_ITERS
     num_candidates: int = DEFAULT_NUM_CANDIDATES
@@ -118,21 +119,25 @@ def optimize(
     outputs: List[Output],
     params: OptimizationParameters,
 ) -> Program:
-    print("Synthesizing candidates...")
-    candidates = [
-        _optimize(
-            target_prgm,
-            test_cases,
-            outputs,
-            replace(params, cost_fn=cost_noperf),
-            num_iters=params.synthesis_iters,
-            init_prgm=create_random_program(params.max_size),
+    best_candidate = target_prgm
+
+    if params.synthesize:
+        print("Synthesizing candidates...")
+        candidates = [
+            _optimize(
+                target_prgm,
+                test_cases,
+                outputs,
+                replace(params, cost_fn=cost_noperf),
+                num_iters=params.synthesis_iters,
+                init_prgm=create_random_program(params.max_size),
+            )
+            for _ in range(params.num_candidates)
+        ]
+        best_candidate = min(
+            candidates, key=lambda p: cost(target_prgm, test_cases, outputs, p)[0]
         )
-        for _ in range(params.num_candidates)
-    ]
-    best_candidate = min(
-        candidates, key=lambda p: cost(target_prgm, test_cases, outputs, p)[0]
-    )
+
     print("Optimizing...")
     return _optimize(
         target_prgm,
diff --git a/gbso/program/test_case.py b/gbso/program/test_case.py
index 88b35e3..d5cc2a9 100644
--- a/gbso/program/test_case.py
+++ b/gbso/program/test_case.py
@@ -39,6 +39,9 @@ def eq_on_output(o: Output, p_cpu: CPU, q_cpu: CPU) -> int:
 
 # Counts differing bits between two 8-bit values
 def eq_8bit(x: int, y: int) -> int:
+    if x == y:
+        return 0
+
     delta = 0
     for i in range(8):
         mask = 1 << i