Port ManimGL's Mobject family memoization

manim/mobject/mobject.py

@@ -20,11 +20,10 @@
 
 import numpy as np
 
+from manim import config, logger
+from manim.constants import *
 from manim.mobject.opengl.opengl_compatibility import ConvertToOpenGL
-
-from .. import config, logger
-from ..constants import *
-from ..utils.color import (
+from manim.utils.color import (
     BLACK,
     WHITE,
     YELLOW_C,
@@ -33,14 +32,15 @@
     color_gradient,
     interpolate_color,
 )
-from ..utils.exceptions import MultiAnimationOverrideException
-from ..utils.iterables import list_update, remove_list_redundancies
-from ..utils.paths import straight_path
-from ..utils.space_ops import angle_between_vectors, normalize, rotation_matrix
+from manim.utils.exceptions import MultiAnimationOverrideException
+from manim.utils.iterables import list_update, remove_list_redundancies, tuplify
+from manim.utils.paths import straight_path
+from manim.utils.space_ops import angle_between_vectors, normalize, rotation_matrix
 
 if TYPE_CHECKING:
     from typing_extensions import Self, TypeAlias
 
+    from manim.animation.animation import Animation
     from manim.typing import (
         FunctionOverride,
         Image,
@@ -53,8 +53,6 @@
         Vector3D,
     )
 
-    from ..animation.animation import Animation
-
     TimeBasedUpdater: TypeAlias = Callable[["Mobject", float], object]
     NonTimeBasedUpdater: TypeAlias = Callable[["Mobject"], object]
     Updater: TypeAlias = NonTimeBasedUpdater | TimeBasedUpdater
@@ -70,7 +68,16 @@ class Mobject:
     Attributes
     ----------
     submobjects : List[:class:`Mobject`]
-        The contained objects.
+        The contained objects, or "children" of this Mobject.
+    parents : List[:class:`Mobject`]
+        The Mobjects which contain this as part of their submobjects. It is
+        important to have a backreference to the parents in case this Mobject's
+        family changes in order to notify them of that change, because this
+        Mobject is also a part of their families.
+    family : List[:class:`Mobject`] | None
+        An optional, memoized list containing this Mobject, its submobjects,
+        those submobjects' submobjects, and so on. If the family must be
+        recalculated for any reason, this attribute is set to None.
     points : :class:`numpy.ndarray`
         The points of the objects.
 
@@ -106,7 +113,9 @@ def __init__(
         self.target = target
         self.z_index = z_index
         self.point_hash = None
-        self.submobjects = []
+        self._submobjects: list[Mobject] = []
+        self.parents: list[Mobject] = []
+        self.family: list[Mobject] | None = [self]
         self.updaters: list[Updater] = []
         self.updating_suspended = False
         self.color = ManimColor.parse(color)
@@ -115,6 +124,48 @@ def __init__(
         self.generate_points()
         self.init_colors()
 
+    def _assert_valid_submobjects(self, submobjects: list[Mobject]) -> None:
+        """Check that all submobjects are actually values of type Mobject, and
+        that none of them is ``self`` (a :class:`Mobject` cannot contain
+        itself).
+
+        This is an auxiliary function called when adding Mobjects to the
+        :attr:`submobjects` list.
+
+        Parameters
+        ----------
+        submobjects
+            The list containing values which should be Mobjects.
+
+        Raises
+        ------
+        TypeError
+            If any of the values in `submobjects` is not a :class:`Mobject`.
+        ValueError
+            If there was an attempt to add a :class:`Mobject` as its own
+            submobject.
+        """
+        self._assert_valid_submobjects_internal(submobjects, Mobject)
+
+    def _assert_valid_submobjects_internal(
+        self, submobjects: list[Mobject], mob_class: type
+    ):
+        for submob in submobjects:
+            if not isinstance(submob, mob_class):
+                raise TypeError(f"All submobjects must be of type {mob_class.__name__}")
+            if submob is self:
+                raise ValueError("Mobject cannot contain self")
+
+    @property
+    def submobjects(self) -> list[Mobject]:
+        return self._submobjects
+
+    @submobjects.setter
+    def submobjects(self, new_submobjects: list[Mobject]) -> None:
+        self._assert_valid_submobjects(new_submobjects)
+        self._submobjects = new_submobjects
+        self.note_changed_family()
+
     @classmethod
     def animation_override_for(
         cls,
@@ -337,6 +388,12 @@ def __deepcopy__(self, clone_from_id) -> Self:
         result = cls.__new__(cls)
         clone_from_id[id(self)] = result
         for k, v in self.__dict__.items():
+            # This must be set manually because result has no attributes,
+            # and specifically INSIDE the loop to preserve attribute order,
+            # or test_hash_consistency() will fail!
+            if k == "parents":
+                result.parents = []
+                continue
             setattr(result, k, copy.deepcopy(v, clone_from_id))
         result.original_id = str(id(self))
         return result
@@ -432,12 +489,7 @@ def add(self, *mobjects: Mobject) -> Self:
             [child]
 
         """
-        for m in mobjects:
-            if not isinstance(m, Mobject):
-                raise TypeError("All submobjects must be of type Mobject")
-            if m is self:
-                raise ValueError("Mobject cannot contain self")
-
+        self._assert_valid_submobjects(mobjects)
         unique_mobjects = remove_list_redundancies(mobjects)
         if len(mobjects) != len(unique_mobjects):
             logger.warning(
@@ -446,6 +498,10 @@ def add(self, *mobjects: Mobject) -> Self:
             )
 
         self.submobjects = list_update(self.submobjects, unique_mobjects)
+        for mobject in unique_mobjects:
+            if self not in mobject.parents:
+                mobject.parents.append(self)
+        self.note_changed_family()
         return self
 
     def insert(self, index: int, mobject: Mobject) -> None:
@@ -463,11 +519,13 @@ def insert(self, index: int, mobject: Mobject) -> None:
         mobject
             The mobject to be inserted.
         """
-        if not isinstance(mobject, Mobject):
-            raise TypeError("All submobjects must be of type Mobject")
-        if mobject is self:
-            raise ValueError("Mobject cannot contain self")
+        self._assert_valid_submobjects([mobject])
+        # TODO: should verify that subsequent submobjects are not repeated
         self.submobjects.insert(index, mobject)
+        if self not in mobject.parents:
+            mobject.parents.append(self)
+        self.note_changed_family()
+        # TODO: should return Self instead of None?
 
     def __add__(self, mobject: Mobject):
         raise NotImplementedError
@@ -519,16 +577,14 @@ def add_to_back(self, *mobjects: Mobject) -> Self:
         :meth:`add`
 
         """
-        if self in mobjects:
-            raise ValueError("A mobject shouldn't contain itself")
-
-        for mobject in mobjects:
-            if not isinstance(mobject, Mobject):
-                raise TypeError("All submobjects must be of type Mobject")
-
+        self._assert_valid_submobjects(mobjects)
         self.remove(*mobjects)
         # dict.fromkeys() removes duplicates while maintaining order
         self.submobjects = list(dict.fromkeys(mobjects)) + self.submobjects
+        for mobject in mobjects:
+            if self not in mobject.parents:
+                mobject.parents.append(self)
+        self.note_changed_family()
         return self
 
     def remove(self, *mobjects: Mobject) -> Self:
@@ -556,6 +612,9 @@ def remove(self, *mobjects: Mobject) -> Self:
         for mobject in mobjects:
             if mobject in self.submobjects:
                 self.submobjects.remove(mobject)
+            if self in mobject.parents:
+                mobject.parents.remove(self)
+        self.note_changed_family()
         return self
 
     def __sub__(self, other):
@@ -2261,14 +2320,90 @@ def split(self) -> list[Self]:
         result = [self] if len(self.points) > 0 else []
         return result + self.submobjects
 
+    def note_changed_family(self, only_changed_order=False) -> Self:
+        """Indicates that this Mobject's family should be recalculated, by
+        setting it to None to void the previous computation. If this Mobject
+        has parents, it is also part of their respective families, so they must
+        be notified as well.
+
+        This method must be called after any change which involves modifying
+        some Mobject's submobjects, such as a call to Mobject.add.
+
+        Parameters
+        ----------
+        only_changed_order
+            If True, indicate that the family still contains the same Mobjects,
+            only in a different order. This prevents recalculating bounding
+            boxes and updater statuses, because they remain the same. If False,
+            indicate that some Mobjects were added or removed to the family,
+            and trigger the aforementioned recalculations. Default is False.
+
+        Returns
+        -------
+        :class:`Mobject`
+            The Mobject itself.
+        """
+        self.family = None
+        # TODO: Implement when bounding boxes and updater statuses are implemented
+        # if not only_changed_order:
+        #     self.refresh_has_updater_status()
+        #     self.refresh_bounding_box()
+        for parent in self.parents:
+            parent.note_changed_family(only_changed_order)
+        return self
+
     def get_family(self, recurse: bool = True) -> list[Self]:
-        sub_families = [x.get_family() for x in self.submobjects]
-        all_mobjects = [self] + list(it.chain(*sub_families))
-        return remove_list_redundancies(all_mobjects)
+        """Obtain the family of this Mobject, consisting of itself, its
+        submobjects, the submobjects' submobjects, and so on. If the
+        family was calculated previously and memoized into the :attr:`family`
+        attribute as a list, return that list. Otherwise, if the attribute is
+        None, calculate and memoize it now.
+
+        Parameters
+        ----------
+        recurse
+            If True, explore this Mobject's submobjects and so on, to
+            compute the full family. Otherwise, stop at this Mobject and
+            return a list containing only this one. Default is True.
+
+        Returns
+        -------
+        list[:class:`Mobject`]
+            The family of this Mobject.
+        """
+        if not recurse:
+            return [self]
+        if self.family is None:
+            # Reconstruct and save
+            sub_families = (sm.get_family() for sm in self.submobjects)
+            family = [self, *it.chain(*sub_families)]
+            self.family = remove_list_redundancies(family)
+        return self.family
 
     def family_members_with_points(self) -> list[Self]:
         return [m for m in self.get_family() if m.get_num_points() > 0]
 
+    def get_ancestors(self, extended: bool = False) -> list[Mobject]:
+        """
+        Returns parents, grandparents, etc.
+        Order of result should be from higher members of the hierarchy down.
+
+        If extended is set to true, it includes the ancestors of all family members,
+        e.g. any other parents of a submobject.
+        """
+        ancestors = []
+        to_process = self.get_family(recurse=extended).copy()
+        excluded = set(to_process)
+        while to_process:
+            for p in to_process.pop().parents:
+                if p not in excluded:
+                    ancestors.append(p)
+                    to_process.append(p)
+        # Ensure mobjects highest in the hierarchy show up first
+        ancestors.reverse()
+        # Remove list redundancies while preserving order
+        return list(dict.fromkeys(ancestors))
+
     def arrange(
         self,
         direction: Vector3D = RIGHT,
@@ -2552,6 +2687,7 @@ def submob_func(m: Mobject):
                 return point_to_num_func(m.get_center())
 
         self.submobjects.sort(key=submob_func)
+        self.note_changed_family(only_changed_order=True)
         return self
 
     def shuffle(self, recursive: bool = False) -> None:
@@ -2560,6 +2696,8 @@ def shuffle(self, recursive: bool = False) -> None:
             for submob in self.submobjects:
                 submob.shuffle(recursive=True)
         random.shuffle(self.submobjects)
+        self.note_changed_family(only_changed_order=True)
+        # TODO: should return Self instead of None?
 
     def invert(self, recursive: bool = False) -> None:
         """Inverts the list of :attr:`submobjects`.
@@ -2586,6 +2724,8 @@ def construct(self):
             for submob in self.submobjects:
                 submob.invert(recursive=True)
         self.submobjects.reverse()
+        self.note_changed_family(only_changed_order=True)
+        # TODO: should return Self instead of None?
 
     # Just here to keep from breaking old scenes.
     def arrange_submobjects(self, *args, **kwargs) -> Self:
@@ -2715,6 +2855,8 @@ def add_n_more_submobjects(self, n: int) -> Self | None:
         if curr == 0:
             # If empty, simply add n point mobjects
             self.submobjects = [self.get_point_mobject() for k in range(n)]
+            self.note_changed_family()
+            # TODO: shouldn't this return Self instead?
             return None
 
         target = curr + n
@@ -2728,6 +2870,7 @@ def add_n_more_submobjects(self, n: int) -> Self | None:
             for _ in range(1, sf):
                 new_submobs.append(submob.copy().fade(1))
         self.submobjects = new_submobs
+        self.note_changed_family()
         return self
 
     def repeat_submobject(self, submob: Mobject) -> Self:

manim/mobject/types/vectorized_mobject.py

@@ -168,7 +168,8 @@ def __init__(
         self.n_points_per_cubic_curve: int = n_points_per_cubic_curve
         self.cap_style: CapStyleType = cap_style
         super().__init__(**kwargs)
-        self.submobjects: list[VMobject]
+        self._submobjects: list[VMobject]
+        self.family: list[VMobject] | None
 
         # TODO: Find where color overwrites are happening and remove the color doubling
         # if "color" in kwargs:
@@ -179,6 +180,9 @@ def __init__(
         if stroke_color is not None:
             self.stroke_color = ManimColor.parse(stroke_color)
 
+    def _assert_valid_submobjects(self, submobjects: list[VMobject]):
+        self._assert_valid_submobjects_internal(submobjects, VMobject)
+
     # OpenGL compatibility
     @property
     def n_points_per_curve(self) -> int: