Hashing stability for transient objects in manim 0.20.1

manim/utils/hashing.py

@@ -100,20 +100,18 @@ class _Memoizer:
     content-equality detection.
     """
 
-    _already_processed = set()
-
-    # Can be changed to whatever string to help debugging the JSon generation.
+    # Can be changed to whatever string to help debugging the JSON generation.
     ALREADY_PROCESSED_PLACEHOLDER = "AP"
     THRESHOLD_WARNING = 170_000
 
-    @classmethod
-    def reset_already_processed(cls: type[_Memoizer]) -> None:
-        cls._already_processed.clear()
+    def __init__(self) -> None:
+        self._already_processed: set[tuple[str, int]] = set()
+        # Keep tracked objects alive for this serialization operation. CPython can
+        # otherwise reuse a transient object's address before encoding has completed;
+        # this also matters for objects whose hash is derived from their identity.
+        self._keep_alive: dict[int, Any] = {}
 
-    @classmethod
-    def check_already_processed_decorator(
-        cls: type[_Memoizer], is_method: bool = False
-    ) -> Callable:
+    def check_already_processed_decorator(self, is_method: bool = False) -> Callable:
         """Decorator to handle the arguments that goes through the decorated function.
         Returns the value of ALREADY_PROCESSED_PLACEHOLDER if the obj has been processed,
         or lets the decorated function call go ahead.
@@ -128,16 +126,17 @@ def layer(func: Callable[[Any], Any]) -> Callable:
             # NOTE : There is probably a better way to separate both case when func is
             # a method or a function.
             if is_method:
-                return lambda self, obj: cls._handle_already_processed(
+                return lambda obj_self, obj: self._handle_already_processed(
                     obj,
-                    default_function=lambda obj: func(self, obj),
+                    default_function=lambda obj: func(obj_self, obj),
                 )
-            return lambda obj: cls._handle_already_processed(obj, default_function=func)
+            return lambda obj: self._handle_already_processed(
+                obj, default_function=func
+            )
 
         return layer
 
-    @classmethod
-    def check_already_processed(cls: type[_Memoizer], obj: Any) -> Any:
+    def check_already_processed(self, obj: Any) -> Any:
         """Checks if obj has been already processed. Returns itself if it has not been,
         or the value of ALREADY_PROCESSED_PLACEHOLDER if it has.
         Marks the object as processed in the second case.
@@ -153,23 +152,21 @@ def check_already_processed(cls: type[_Memoizer], obj: Any) -> Any:
             Either the object itself or the placeholder.
         """
         # When the object is not memoized, we return the object itself.
-        return cls._handle_already_processed(obj, lambda x: x)
+        return self._handle_already_processed(obj, lambda x: x)
 
-    @classmethod
-    def mark_as_processed(cls: type[_Memoizer], obj: Any) -> None:
+    def mark_as_processed(self, obj: Any) -> None:
         """Marks an object as processed.
 
         Parameters
         ----------
         obj
             The object to mark as processed.
         """
-        cls._handle_already_processed(obj, lambda x: x)
-        return cls._return(obj, id, lambda x: x, memoizing=False)
+        self._handle_already_processed(obj, lambda x: x)
+        self._return(obj, id, lambda x: x, memoizing=False)
 
-    @classmethod
     def _handle_already_processed(
-        cls: type[_Memoizer],
+        self,
         obj: Any,
         default_function: Callable[[Any], Any],
     ) -> str | Any:
@@ -181,41 +178,39 @@ def _handle_already_processed(
                 str,
                 complex,
             ),
-        ) and obj not in [None, cls.ALREADY_PROCESSED_PLACEHOLDER]:
+        ) and obj not in [None, self.ALREADY_PROCESSED_PLACEHOLDER]:
             # It makes no sense (and it'd slower) to memoize objects of these primitive
             # types.  Hence, we simply return the object.
             return obj
         if isinstance(obj, Hashable):
             try:
-                return cls._return(obj, hash, default_function)
+                return self._return(obj, hash, default_function)
             except TypeError:
                 # In case of an error with the hash (eg an object is marked as hashable
                 # but contains a non hashable within it)
                 # Fallback to use the built-in function id instead.
                 pass
-        return cls._return(obj, id, default_function)
+        return self._return(obj, id, default_function)
 
-    @classmethod
     def _return(
-        cls: type[_Memoizer],
+        self,
         obj: Any,
         obj_to_membership_sign: Callable[[Any], int],
         default_func: Callable[[Any], Any],
         memoizing: bool = True,
     ) -> str | Any:
         # hash() values and id() addresses are both raw ints that share no meaning, so entries are tagged with
         # their signature kind: in an untagged set, a hash() value that numerically equals a recorded id() address
-        # (or vice versa) collapses a never-processed object to the placeholder. Whether such a coincidence occurs
-        # varies from process to process (str hashes are randomized), making play hashes unstable across identical
-        # runs and causing unnecessary re-renders of already cached partial movie files.
+        # (or vice versa) collapses a never-processed object to the placeholder.
         sign_kind = "h" if obj_to_membership_sign is hash else "i"
-        obj_membership_sign = (sign_kind, obj_to_membership_sign(obj))
-        if obj_membership_sign in cls._already_processed:
-            return cls.ALREADY_PROCESSED_PLACEHOLDER
+        signature = obj_to_membership_sign(obj)
+        obj_membership_sign = (sign_kind, signature)
+        if obj_membership_sign in self._already_processed:
+            return self.ALREADY_PROCESSED_PLACEHOLDER
         if memoizing:
             if (
                 not config.disable_caching_warning
-                and len(cls._already_processed) == cls.THRESHOLD_WARNING
+                and len(self._already_processed) == self.THRESHOLD_WARNING
             ):
                 logger.warning(
                     "It looks like the scene contains a lot of sub-mobjects. Caching "
@@ -228,18 +223,21 @@ def _return(
                     "to True in your config file.",
                 )
 
-            cls._already_processed.add(obj_membership_sign)
+            self._already_processed.add(obj_membership_sign)
+            self._keep_alive.setdefault(id(obj), obj)
         return default_func(obj)
 
 
 class _CustomEncoder(json.JSONEncoder):
     def __init__(
         self,
         *args: Any,
+        memoizer: _Memoizer,
         include_pixel_array: bool = False,
         **kwargs: Any,
     ) -> None:
         super().__init__(*args, **kwargs)
+        self._memoizer = memoizer
         self.include_pixel_array = include_pixel_array
 
     def default(self, obj: Any) -> Any:
@@ -322,12 +320,12 @@ def _cleaned_iterable(self, iterable):
         """
 
         def _key_to_hash(key: Any) -> int:
-            return zlib.crc32(json.dumps(key, cls=_CustomEncoder).encode())
+            return zlib.crc32(_get_json(key, self._memoizer).encode())
 
         def _iter_check_list(lst: Sequence[Any]) -> list[Any]:
             processed_list = [None] * len(lst)
             for i, el in enumerate(lst):
-                el = _Memoizer.check_already_processed(el)
+                el = self._memoizer.check_already_processed(el)
                 if isinstance(el, (list, tuple)):
                     new_value = _iter_check_list(el)
                 elif isinstance(el, dict):
@@ -340,7 +338,7 @@ def _iter_check_list(lst: Sequence[Any]) -> list[Any]:
         def _iter_check_dict(dct: dict[Any, Any]) -> dict[Any, Any]:
             processed_dict = {}
             for k, v in dct.items():
-                v = _Memoizer.check_already_processed(v)
+                v = self._memoizer.check_already_processed(v)
                 if k in KEYS_TO_FILTER_OUT and not (
                     self.include_pixel_array and k == "pixel_array"
                 ):
@@ -379,12 +377,24 @@ def encode(self, obj: Any) -> str:
         :class:`str`
            The object encoder with the standard json process.
         """
-        _Memoizer.mark_as_processed(obj)
+        self._memoizer.mark_as_processed(obj)
         if isinstance(obj, (dict, list, tuple)):
             return super().encode(self._cleaned_iterable(obj))
         return super().encode(obj)
 
 
+def _get_json(
+    obj: Any, memoizer: _Memoizer, *, include_pixel_array: bool = False
+) -> str:
+    """Serialize ``obj`` using an existing serialization-operation memoizer."""
+    return json.dumps(
+        obj,
+        cls=_CustomEncoder,
+        memoizer=memoizer,
+        include_pixel_array=include_pixel_array,
+    )
+
+
 def get_json(obj: Any, *, include_pixel_array: bool = False) -> str:
     """Recursively serialize `object` to JSON using the :class:`CustomEncoder` class.
 
@@ -400,11 +410,7 @@ def get_json(obj: Any, *, include_pixel_array: bool = False) -> str:
     :class:`str`
         The flattened object
     """
-    return json.dumps(
-        obj,
-        cls=_CustomEncoder,
-        include_pixel_array=include_pixel_array,
-    )
+    return _get_json(obj, _Memoizer(), include_pixel_array=include_pixel_array)
 
 
 def get_hash_from_play_call(
@@ -436,14 +442,16 @@ def get_hash_from_play_call(
     """
     logger.debug("Hashing ...")
     t_start = perf_counter()
-    _Memoizer.mark_as_processed(scene_object)
-    camera_json = get_json(camera_object)
+    memoizer = _Memoizer()
+    memoizer.mark_as_processed(scene_object)
+    camera_json = _get_json(camera_object, memoizer)
     animations_list_json = [
-        get_json(animation, include_pixel_array=True)
+        _get_json(animation, memoizer, include_pixel_array=True)
         for animation in sorted(animations_list, key=str)
     ]
     current_mobjects_list_json = [
-        get_json(mobject, include_pixel_array=True) for mobject in current_mobjects_list
+        _get_json(mobject, memoizer, include_pixel_array=True)
+        for mobject in current_mobjects_list
     ]
     hash_camera, hash_animations, hash_current_mobjects = (
         zlib.crc32(repr(json_val).encode())
@@ -452,7 +460,5 @@ def get_hash_from_play_call(
     hash_complete = f"{hash_camera}_{hash_animations}_{hash_current_mobjects}"
     t_end = perf_counter()
     logger.debug("Hashing done in %(time)s s.", {"time": str(t_end - t_start)[:8]})
-    # End of the hashing for the animation, reset all the memoize.
-    _Memoizer.reset_already_processed()
     logger.debug("Hash generated :  %(h)s", {"h": hash_complete})
     return hash_complete