From ac3db9b636f7c33aecb955a6eeaece9f76ad2b63 Mon Sep 17 00:00:00 2001
From: Grant Sanderson <grant@3blue1brown.com>
Date: Thu, 2 Feb 2023 21:13:18 -0800
Subject: [PATCH] Add set_program_uniform function

---
 manimlib/shader_wrapper.py | 10 ++--------
 manimlib/utils/shaders.py  | 31 ++++++++++++++++++++++++++++++-
 2 files changed, 32 insertions(+), 9 deletions(-)

diff --git a/manimlib/shader_wrapper.py b/manimlib/shader_wrapper.py
index f8091271..d47cb92f 100644
--- a/manimlib/shader_wrapper.py
+++ b/manimlib/shader_wrapper.py
@@ -15,6 +15,7 @@ from manimlib.utils.shaders import image_path_to_texture
 from manimlib.utils.shaders import get_texture_id
 from manimlib.utils.shaders import get_fill_canvas
 from manimlib.utils.shaders import release_texture
+from manimlib.utils.shaders import set_program_uniform
 
 from typing import TYPE_CHECKING
 
@@ -224,14 +225,7 @@ class ShaderWrapper(object):
         if self.program is None:
             return
         for name, value in (*self.mobject_uniforms.items(), *camera_uniforms.items()):
-            if name not in self.program:
-                continue
-            if isinstance(value, np.ndarray) and value.ndim > 0:
-                value = tuple(value)
-            if name in camera_uniforms and self.program_uniform_mirror.get(name, None) == value:
-                continue
-            self.program[name].value = value
-            self.program_uniform_mirror[name] = value
+            set_program_uniform(self.program, name, value)
 
     def get_vertex_buffer_object(self, refresh: bool = True):
         if refresh:
diff --git a/manimlib/utils/shaders.py b/manimlib/utils/shaders.py
index 2a24fd76..3e4f67cf 100644
--- a/manimlib/utils/shaders.py
+++ b/manimlib/utils/shaders.py
@@ -9,7 +9,6 @@ import numpy as np
 
 from manimlib.config import parse_cli
 from manimlib.config import get_configuration
-from manimlib.utils.customization import get_customization
 from manimlib.utils.directories import get_shader_dir
 from manimlib.utils.file_ops import find_file
 
@@ -17,11 +16,13 @@ from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
     from typing import Sequence, Optional, Tuple
+    from manimlib.typing import UniformDict
     from moderngl.vertex_array import VertexArray
     from moderngl.framebuffer import Framebuffer
 
 
 ID_TO_TEXTURE: dict[int, moderngl.Texture] = dict()
+PROGRAM_UNIFORM_MIRRORS: dict[int, dict[str, float | tuple]] = dict()
 
 
 @lru_cache()
@@ -63,6 +64,34 @@ def get_shader_program(
     )
 
 
+def set_program_uniform(
+    program: moderngl.Program,
+    name: str,
+    value: float | tuple | np.ndarray
+) -> bool:
+    """
+    Sets a program uniform, and also keeps track of a dictionary
+    of previously set uniforms for that program so that it
+    doesn't needlessly reset it, requiring an exchange with gpu
+    memory, if it sees the same value again.
+    
+    Returns True if changed the program, False if it left it as is.
+    """
+    if name not in program:
+        return False
+    pid = id(program)
+    if pid not in PROGRAM_UNIFORM_MIRRORS:
+        PROGRAM_UNIFORM_MIRRORS[pid] = dict()
+    uniform_mirror = PROGRAM_UNIFORM_MIRRORS[pid]
+    if isinstance(value, np.ndarray) and value.ndim > 0:
+        value = tuple(value)
+    if uniform_mirror.get(name, None) == value:
+        return False
+    uniform_mirror[name] = value
+    program[name].value = value
+    return True
+
+
 @lru_cache()
 def get_shader_code_from_file(filename: str) -> str | None:
     if not filename: