Move loading of shader programs into shaders module

kitty/shaders.py

@@ -2,11 +2,34 @@
 # License: GPLv3 Copyright: 2023, Kovid Goyal <kovid at kovidgoyal.net>
 
 import re
-from functools import lru_cache
-from typing import Callable, Iterator, Optional
+from functools import lru_cache, partial
+from typing import Any, Callable, Iterator, Optional
 
 from .constants import read_kitty_resource
-from .fast_data_types import GLSL_VERSION, compile_program
+from .fast_data_types import (
+    BGIMAGE_PROGRAM,
+    BLIT_PROGRAM,
+    CELL_BG_PROGRAM,
+    CELL_FG_PROGRAM,
+    CELL_PROGRAM,
+    CELL_SPECIAL_PROGRAM,
+    DECORATION,
+    DECORATION_MASK,
+    DIM,
+    GLSL_VERSION,
+    GRAPHICS_ALPHA_MASK_PROGRAM,
+    GRAPHICS_PREMULT_PROGRAM,
+    GRAPHICS_PROGRAM,
+    MARK,
+    MARK_MASK,
+    NUM_UNDERLINE_STYLES,
+    REVERSE,
+    STRIKETHROUGH,
+    TINT_PROGRAM,
+    compile_program,
+    get_options,
+    init_cell_program,
+)
 
 
 def identity(x: str) -> str:
@@ -55,3 +78,95 @@ class Program:
 @lru_cache(maxsize=64)
 def program_for(name: str) -> Program:
     return Program(name)
+
+
+def multi_replace(src: str, **replacements: Any) -> str:
+    r = {k: str(v) for k, v in replacements.items()}
+
+    def sub(m: 're.Match[str]') -> str:
+        return r.get(m.group(1), m.group(1))
+
+    return re.sub(r'\{([A-Z_]+)\}', sub, src)
+
+
+class LoadShaderPrograms:
+
+    text_fg_override_threshold: float = 0
+    text_old_gamma: bool = False
+    semi_transparent: bool = False
+
+    @property
+    def needs_recompile(self) -> bool:
+        opts = get_options()
+        return opts.text_fg_override_threshold != self.text_fg_override_threshold or (opts.text_composition_strategy == 'legacy') != self.text_old_gamma
+
+    def recompile_if_needed(self) -> None:
+        if self.needs_recompile:
+            self(self.semi_transparent, allow_recompile=True)
+
+    def __call__(self, semi_transparent: bool = False, allow_recompile: bool = False) -> None:
+        self.semi_transparent = semi_transparent
+        opts = get_options()
+        self.text_old_gamma = opts.text_composition_strategy == 'legacy'
+        self.text_fg_override_threshold = max(0, min(opts.text_fg_override_threshold, 100)) * 0.01
+        program_for('blit').compile(BLIT_PROGRAM, allow_recompile)
+        cell = program_for('cell')
+
+        def resolve_cell_vertex_defines(which: str, v: str) -> str:
+            v = multi_replace(
+                v,
+                WHICH_PROGRAM=which,
+                REVERSE_SHIFT=REVERSE,
+                STRIKE_SHIFT=STRIKETHROUGH,
+                DIM_SHIFT=DIM,
+                DECORATION_SHIFT=DECORATION,
+                MARK_SHIFT=MARK,
+                MARK_MASK=MARK_MASK,
+                DECORATION_MASK=DECORATION_MASK,
+                STRIKE_SPRITE_INDEX=NUM_UNDERLINE_STYLES + 1,
+            )
+            if semi_transparent:
+                v = v.replace('#define NOT_TRANSPARENT', '#define TRANSPARENT')
+            return v
+
+        def resolve_cell_fragment_defines(which: str, f: str) -> str:
+            f = f.replace('{WHICH_PROGRAM}', which)
+            if self.text_fg_override_threshold != 0.:
+                f = f.replace('#define NO_FG_OVERRIDE', f'#define FG_OVERRIDE {self.text_fg_override_threshold}')
+            if self.text_old_gamma:
+                f = f.replace('#define TEXT_NEW_GAMMA', '#define TEXT_OLD_GAMMA')
+            if semi_transparent:
+                f = f.replace('#define NOT_TRANSPARENT', '#define TRANSPARENT')
+            return f
+
+        for which, p in {
+            'SIMPLE': CELL_PROGRAM,
+            'BACKGROUND': CELL_BG_PROGRAM,
+            'SPECIAL': CELL_SPECIAL_PROGRAM,
+            'FOREGROUND': CELL_FG_PROGRAM,
+        }.items():
+            cell.apply_to_sources(
+                vertex=partial(resolve_cell_vertex_defines, which),
+                frag=partial(resolve_cell_fragment_defines, which),
+            )
+            cell.compile(p, allow_recompile)
+
+        graphics = program_for('graphics')
+
+        def resolve_graphics_fragment_defines(which: str, f: str) -> str:
+            return f.replace('ALPHA_TYPE', which)
+
+        for which, p in {
+            'SIMPLE': GRAPHICS_PROGRAM,
+            'PREMULT': GRAPHICS_PREMULT_PROGRAM,
+            'ALPHA_MASK': GRAPHICS_ALPHA_MASK_PROGRAM,
+        }.items():
+            graphics.apply_to_sources(frag=partial(resolve_cell_fragment_defines, which))
+            graphics.compile(p, allow_recompile)
+
+        program_for('bgimage').compile(BGIMAGE_PROGRAM, allow_recompile)
+        program_for('tint').compile(TINT_PROGRAM)
+        init_cell_program()
+
+
+load_shader_programs = LoadShaderPrograms()