From e2586894b0bced4910c836af298eef9f4b014a73 Mon Sep 17 00:00:00 2001
From: Nick Drozd <nicholasdrozd@gmail.com>
Date: Tue, 16 Jan 2024 23:11:35 -0600
Subject: [PATCH] Centralize macro cache

---
 tm/macro.py | 28 ++++++++++++++++++----------
 1 file changed, 18 insertions(+), 10 deletions(-)

diff --git a/tm/macro.py b/tm/macro.py
index f548a6c4..dff6f420 100644
--- a/tm/macro.py
+++ b/tm/macro.py
@@ -2,6 +2,7 @@
 
 from abc import abstractmethod
 from typing import TYPE_CHECKING
+from collections import defaultdict
 
 from tm.parse import tcompile
 
@@ -13,6 +14,17 @@
 
 ########################################
 
+COLOR_TO_TAPE_CACHE: dict[
+    int,
+    dict[Color, tuple[Color, ...]],
+] = defaultdict(dict)
+
+TAPE_TO_COLOR_CACHE: dict[
+    int,
+    dict[tuple[Color, ...], Color],
+] = defaultdict(dict)
+
+
 class MacroProg:
     program: str | GetInstr
     comp: GetInstr
@@ -29,9 +41,6 @@ class MacroProg:
 
     cells: int
 
-    color_to_tape_cache: dict[Color, tuple[Color, ...]]
-    tape_to_color_cache: dict[tuple[Color, ...], Color]
-
     def __init__(self, program: str | GetInstr):
         self.program = program
 
@@ -56,9 +65,6 @@ def __init__(self, program: str | GetInstr):
 
         self.instrs = {}
 
-        self.color_to_tape_cache = {}
-        self.tape_to_color_cache = {}
-
     def __getitem__(self, slot: Slot) -> Instr | None:
         try:
             instr = self.instrs[slot]
@@ -139,7 +145,7 @@ def run_simulator(self, config: Config) -> Config | None:
         return state, (cells <= pos, tape)
 
     def tape_to_color(self, tape: Tape) -> Color:
-        if (cached := self.tape_to_color_cache.get(
+        if (cached := TAPE_TO_COLOR_CACHE[cells := self.cells].get(
                 tuple_tape := tuple(tape))) is not None:
             return cached
 
@@ -148,8 +154,8 @@ def tape_to_color(self, tape: Tape) -> Color:
             for place, value in enumerate(reversed(tape))
         )
 
-        self.tape_to_color_cache[tuple_tape] = color
-        self.color_to_tape_cache[color] = tuple_tape
+        TAPE_TO_COLOR_CACHE[cells][tuple_tape] = color
+        COLOR_TO_TAPE_CACHE[cells][color] = tuple_tape
 
         return color
 
@@ -157,7 +163,9 @@ def color_to_tape(self, color: Color) -> Tape:
         if color == 0:
             return [0] * self.cells
 
-        assert (prev := self.color_to_tape_cache.get(color)) is not None
+        assert (
+            prev := COLOR_TO_TAPE_CACHE[self.cells].get(color)
+        ) is not None
 
         return list(prev)