From 83fc6ce86e809ac3c86dac544a4848f87bb1aad3 Mon Sep 17 00:00:00 2001 From: Kovid Goyal Date: Tue, 4 Aug 2026 13:04:56 +0530 Subject: [PATCH] allow specialisation of custom shaders at link time --- kitty/shaders/slang.py | 69 +++++++++++++++++++++++++++++++++--------- kitty_tests/slang.py | 28 +++++++++++++++++ 2 files changed, 82 insertions(+), 15 deletions(-) diff --git a/kitty/shaders/slang.py b/kitty/shaders/slang.py index ef0a5f2c1..48d1afed9 100644 --- a/kitty/shaders/slang.py +++ b/kitty/shaders/slang.py @@ -12,7 +12,6 @@ import shutil import subprocess import sys import tempfile -import textwrap import time import types import zlib @@ -1073,17 +1072,41 @@ def is_valid_slot(x: str) -> TypeGuard[Slot]: return x in get_args(Slot) +VALID_VAR_TYPES: frozenset[str] = frozenset({'uint', 'int', 'float', 'double', 'bool'}) +_identifier_re = re.compile(r'^[A-Za-z_][A-Za-z0-9_]*$') + + +def parse_var_directive(parts: list[str]) -> tuple[str, str, str]: + if len(parts) < 4: + raise ValueError('var directive requires: var [=] ') + var_type = parts[1] + if var_type not in VALID_VAR_TYPES: + raise ValueError(f'var type {var_type!r} must be one of: {", ".join(sorted(VALID_VAR_TYPES))}') + var_name = parts[2] + if not _identifier_re.match(var_name): + raise ValueError(f'var name {var_name!r} is not a valid identifier') + if parts[3] == '=': + if len(parts) < 5: + raise ValueError('var directive missing value after =') + value = parts[4] + else: + value = parts[3] + return var_type, var_name, value + + class Group(TypedDict): viewport_pos: tuple[float, float] viewport_size: tuple[float, float] output_texture: NamedTexture shaders: tuple[str, ...] + vars: dict[str, tuple[str, str]] class Pipeline(TypedDict): slot: Slot textures: tuple[NamedTexture, ...] groups: tuple[Group, ...] + vars: dict[str, tuple[str, str]] def parse_pipeline_definition(lines: Iterable[str]) -> Pipeline: @@ -1091,6 +1114,7 @@ def parse_pipeline_definition(lines: Iterable[str]) -> Pipeline: textures: tuple[NamedTexture, ...] = () groups: list[Group] = [] current_group: Group | None = None + pipeline_vars: dict[str, tuple[str, str]] = {} def unit_float(x: str) -> float: return max(0, min(float(x), 1)) @@ -1116,7 +1140,10 @@ def parse_pipeline_definition(lines: Iterable[str]) -> Pipeline: textures = tuple(map(NamedTexture, parts[1:])) case 'startgroup': commit_group() - current_group = {'viewport_pos': (0, 0), 'viewport_size': (1, 1), 'output_texture': NamedTexture.default, 'shaders': ()} + current_group = {'viewport_pos': (0, 0), 'viewport_size': (1, 1), 'output_texture': NamedTexture.default, 'shaders': (), 'vars': {}} + case 'var': + var_type, var_name, value = parse_var_directive(parts) + pipeline_vars[var_name] = (var_type, value) case _: raise ValueError(f'Unknown key {parts[0]}') else: @@ -1131,6 +1158,9 @@ def parse_pipeline_definition(lines: Iterable[str]) -> Pipeline: current_group['output_texture'] = parts[1] case 'shaders': current_group['shaders'] += tuple(parts[1:]) + case 'var': + var_type, var_name, value = parse_var_directive(parts) + current_group['vars'][var_name] = (var_type, value) case 'endgroup': commit_group() case _: @@ -1144,7 +1174,7 @@ def parse_pipeline_definition(lines: Iterable[str]) -> Pipeline: if groups[-1]['viewport_pos'] != (0, 0) or groups[-1]['viewport_size'] != (1, 1): raise ValueError('The final group must not specify a viewport') - return {'slot': slot, 'textures': textures, 'groups': tuple(groups)} + return {'slot': slot, 'textures': textures, 'groups': tuple(groups), 'vars': pipeline_vars} @lru_cache(maxsize=32) @@ -1186,11 +1216,15 @@ def build_custom_shader_pipeline_ir(pipeline: Pipeline, cache_dir: str) -> tuple module_names = {} shaders_content_key = b'' - flat_shader_list = [] - for group in pipeline['groups']: + for n, (t, v) in pipeline['vars'].items(): + shaders_content_key += f':pvar:{t}:{n}:{v}'.encode() + flat_shader_list: list[tuple[int, str]] = [] + for g_idx, group in enumerate(pipeline['groups']): shaders_content_key += b'::' + for n, (t, v) in group['vars'].items(): + shaders_content_key += f':gvar:{t}:{n}:{v}'.encode() for name in group['shaders']: - flat_shader_list.append(name) + flat_shader_list.append((g_idx, name)) path, import_dir, src, content_key = custom_shader(name) if import_dir and import_dir not in import_dirs: import_dirs.append(import_dir) @@ -1220,17 +1254,22 @@ def build_custom_shader_pipeline_ir(pipeline: Pipeline, cache_dir: str) -> tuple wrappers = {} entry_points = [] - for i, name in enumerate(flat_shader_list): + for i, (g_idx, name) in enumerate(flat_shader_list): + group = pipeline['groups'][g_idx] + merged_vars = dict(pipeline['vars']) + merged_vars.update(group['vars']) + specialization_decls = ''.join(f'\nstatic const {t} {n} = {v};' for n, (t, v) in merged_vars.items()) + if specialization_decls: + specialization_decls += '\n' module_name = module_names[name] entry_point = f'fragment_main{i}' - wrapper_src = textwrap.dedent(f""" - #language slang 2026 - implementing {slot_module_name}; - import kitty_custom_shader_types; - import {module_name}; - - public float4 {entry_point}(float4 inp, KittyTextures t, KittyCustomShaderData d, float4 viewport) {{ return fragment_main(inp, t, d, viewport); }} - """) + wrapper_src = f"""#language slang 2026 +implementing {slot_module_name}; +import kitty_custom_shader_types; +import {module_name}; +{specialization_decls} +public float4 {entry_point}(float4 inp, KittyTextures t, KittyCustomShaderData d, float4 viewport) {{ return fragment_main(inp, t, d, viewport); }} +""" wrappers[f'wrapper{i}.slang'] = wrapper_src entry_points.append(entry_point) diff --git a/kitty_tests/slang.py b/kitty_tests/slang.py index dc6ce80f3..557643c57 100644 --- a/kitty_tests/slang.py +++ b/kitty_tests/slang.py @@ -16,6 +16,7 @@ from kitty.shaders.slang import ( clear_caches, parse_pipeline_definition, parse_slang_text, + parse_var_directive, topological_layers, topological_sort, ) @@ -243,6 +244,33 @@ void vsMain() {} # Empty graph self.assertEqual(topological_layers({}), []) + def test_parse_var_directive(self): + self.assertEqual(parse_var_directive(['var', 'uint', 'algo', '=', '1']), ('uint', 'algo', '1')) + self.assertEqual(parse_var_directive(['var', 'float', 'intensity', '=', '0.5']), ('float', 'intensity', '0.5')) + self.assertEqual(parse_var_directive(['var', 'bool', 'flag', 'true']), ('bool', 'flag', 'true')) + self.assertRaises(ValueError, parse_var_directive, ['var', 'badtype', 'x', '=', '1']) + self.assertRaises(ValueError, parse_var_directive, ['var', 'uint', '123bad', '=', '1']) + self.assertRaises(ValueError, parse_var_directive, ['var', 'uint']) + + def test_parse_pipeline_definition_vars(self): + p = parse_pipeline_definition(''' + var uint algo = 1 + var float intensity = 0.5 + startgroup + shaders sample + endgroup + startgroup + var uint algo = 2 + shaders sample + endgroup + '''.splitlines()) + self.assertEqual(p['vars'], {'algo': ('uint', '1'), 'intensity': ('float', '0.5')}) + self.assertEqual(p['groups'][0]['vars'], {}) + self.assertEqual(p['groups'][1]['vars'], {'algo': ('uint', '2')}) + # Groups still carry shaders correctly + self.assertEqual(p['groups'][0]['shaders'], ('sample',)) + self.assertEqual(p['groups'][1]['shaders'], ('sample',)) + def test_build_custom_shader_pipeline_glsl(self): if not shutil.which(slangc()[0]): self.skipTest(f'slangc ({slangc()[0]}) not found in PATH')