Make the pyc.visit() cache local to a visitor.

A global cache creates a memory leak in a persistent worker; in its current form it's also incorrect, since object ids can be recycled. Visitors are a more appropriate place to store the cache, as they're not reused across calls to pyc.visit().

Note: a WeakKeyDictionary would have been more ergonomic than a dict keyed by object id, but it would require the keys to be hashable, which they currently aren't. It would also create more work for the GC, which isn't necessarily free.
PiperOrigin-RevId: 652460329

pytype/blocks/process_blocks.py

@@ -25,10 +25,11 @@ def _is_function_def(fn_code):
   return True
 
 
-class CollectAnnotationTargetsVisitor:
+class CollectAnnotationTargetsVisitor(pyc.CodeVisitor):
   """Collect opcodes that might have annotations attached."""
 
   def __init__(self):
+    super().__init__()
     # A mutable map of line: opcode for STORE_* opcodes. This is modified as the
     # visitor runs, and contains the last opcode for each line.
     self.store_ops = {}
@@ -68,10 +69,11 @@ def visit_code(self, code):
     return code
 
 
-class FunctionDefVisitor:
+class FunctionDefVisitor(pyc.CodeVisitor):
   """Add metadata to function definition opcodes."""
 
   def __init__(self, param_annotations):
+    super().__init__()
     self.annots = param_annotations
 
   def visit_code(self, code):

pytype/constant_folding.py

@@ -265,7 +265,7 @@ def resolve(self, op):
     return f
 
 
-class _FoldConstants:
+class _FoldConstants(pyc.CodeVisitor):
   """Fold constant literals in pyc code."""
 
   def visit_code(self, code):

pytype/pyc/pyc.py

@@ -1,5 +1,6 @@
 """Functions for generating, reading and parsing pyc."""
 
+import abc
 import copy
 
 from pycnite import pyc
@@ -11,6 +12,20 @@
 CompileError = compiler.CompileError
 
 
+# The abstract base class for a code visitor passed to pyc.visit.
+class CodeVisitor(abc.ABC):
+
+  def __init__(self):
+    # This cache, used by pyc.visit below, is needed to avoid visiting the same
+    # code object twice, since some visitors mutate the input object.
+    # It maps CodeType object id to the result of visiting that object.
+    self.visitor_cache = {}
+
+  @abc.abstractmethod
+  def visit_code(self, code):
+    ...
+
+
 def parse_pyc_string(data):
   """Parse pyc data from a string.
 
@@ -23,10 +38,11 @@ def parse_pyc_string(data):
   return pyc.loads(data)
 
 
-class AdjustFilename:
+class AdjustFilename(CodeVisitor):
   """Visitor for changing co_filename in a code object."""
 
   def __init__(self, filename):
+    super().__init__()
     self.filename = filename
 
   def visit_code(self, code):
@@ -66,18 +82,12 @@ def compile_src(src, filename, python_version, python_exe, mode="exec"):
   return code
 
 
-# This cache is needed to avoid visiting the same code object twice, since some
-# visitors mutate the input object.
-_VISIT_CACHE = {}
-
-
-def visit(c, visitor):
+def visit(c, visitor: CodeVisitor):
   """Recursively process constants in a pyc using a visitor."""
   if hasattr(c, "co_consts"):
     # This is a CodeType object (because it has co_consts). Visit co_consts,
     # and then the CodeType object itself.
-    k = (id(c), visitor)
-    if k not in _VISIT_CACHE:
+    if id(c) not in visitor.visitor_cache:
       new_consts = []
       changed = False
       for const in c.co_consts:
@@ -87,7 +97,7 @@ def visit(c, visitor):
       if changed:
         c = copy.copy(c)
         c.co_consts = new_consts
-      _VISIT_CACHE[k] = visitor.visit_code(c)
-    return _VISIT_CACHE[k]
+      visitor.visitor_cache[id(c)] = visitor.visit_code(c)
+    return visitor.visitor_cache[id(c)]
   else:
     return c

pytype/vm_utils.py

@@ -21,6 +21,7 @@
 from pytype.errors import error_types
 from pytype.overlays import metaclass
 from pytype.pyc import opcodes
+from pytype.pyc import pyc
 from pytype.pytd import mro
 from pytype.pytd import pytd
 from pytype.pytd import slots
@@ -72,17 +73,17 @@ def __repr__(self):
     return f"Block({self.type}: {self.op_index} level={self.level})"
 
 
-class FindIgnoredTypeComments:
+class FindIgnoredTypeComments(pyc.CodeVisitor):
   """A visitor that finds type comments that will be ignored."""
 
   def __init__(self, type_comments):
+    super().__init__()
     self._type_comments = type_comments
     # Lines will be removed from this set during visiting. Any lines that remain
     # at the end are type comments that will be ignored.
     self._ignored_type_lines = set(type_comments)
 
   def visit_code(self, code):
-    """Interface for pyc.visit."""
     for op in code.code_iter:
       # Make sure we have attached the type comment to an opcode.
       if isinstance(op, blocks.STORE_OPCODES):