ncca-ngl 0.3.4__py3-none-any.whl → 0.5.0__py3-none-any.whl

ncca/ngl/webgpu/wip/shader_templates.py

@@ -0,0 +1,563 @@
+"""
+Shader template system for WebGPU pipelines.
+Provides configurable shader templates that can generate single/multi-color variants
+and support user-defined shaders.
+"""
+
+from abc import ABC, abstractmethod
+from enum import Enum
+from typing import Dict, List, Optional, Union
+import dataclasses
+
+
+class ColorMode(Enum):
+    """Color mode for shader generation."""
+
+    SINGLE = "single"
+    MULTI = "multi"
+
+
+class PrimitiveType(Enum):
+    """Primitive type for shader generation."""
+
+    POINT = "point"
+    LINE = "line"
+    TRIANGLE = "triangle"
+
+
+class RenderMode(Enum):
+    """Render mode specialization."""
+
+    STANDARD = "standard"
+    BILLBOARDED = "billboarded"
+    INSTANCED = "instanced"
+
+
+@dataclasses.dataclass
+class UniformField:
+    """Definition of a uniform field in the shader."""
+
+    name: str
+    data_type: str
+    array_size: Optional[int] = None
+    default_value: Optional[Union[float, List[float]]] = None
+
+
+@dataclasses.dataclass
+class VertexAttribute:
+    """Definition of a vertex attribute."""
+
+    name: str
+    data_type: str
+    location: int
+    step_mode: str = "vertex"
+
+
+@dataclasses.dataclass
+class ShaderConfig:
+    """Configuration for shader generation."""
+
+    primitive_type: PrimitiveType
+    color_mode: ColorMode
+    render_mode: RenderMode
+    custom_uniforms: Optional[List[UniformField]] = None
+    custom_attributes: Optional[List[VertexAttribute]] = None
+    custom_vertex_shader: Optional[str] = None
+    custom_fragment_shader: Optional[str] = None
+
+
+class ShaderTemplate(ABC):
+    """Abstract base class for shader templates."""
+
+    @abstractmethod
+    def generate_vertex_shader(self, config: ShaderConfig) -> str:
+        """Generate vertex shader code from configuration."""
+        pass
+
+    @abstractmethod
+    def generate_fragment_shader(self, config: ShaderConfig) -> str:
+        """Generate fragment shader code from configuration."""
+        pass
+
+    def generate_shader_code(self, config: ShaderConfig) -> str:
+        """Generate complete WGSL shader code."""
+        vertex_shader = self.generate_vertex_shader(config)
+        fragment_shader = self.generate_fragment_shader(config)
+
+        # Ensure we have valid shader code
+        if not vertex_shader or not fragment_shader:
+            raise ValueError("Generated vertex or fragment shader is empty")
+
+        return f"{vertex_shader}\n{fragment_shader}"
+
+    def _generate_uniform_struct(self, config: ShaderConfig) -> str:
+        """Generate uniform buffer structure definition."""
+        uniforms = self._get_base_uniforms(config)
+
+        if config.custom_uniforms:
+            uniforms.extend(config.custom_uniforms)
+
+        lines = ["struct Uniforms {"]
+        for uniform in uniforms:
+            if uniform.array_size:
+                line = f"    {uniform.name}: array<{uniform.data_type}, {uniform.array_size}>,"
+            else:
+                line = f"    {uniform.name}: {uniform.data_type},"
+            lines.append(line)
+        lines.append("};")
+
+        return "\n".join(lines)
+
+    def _get_base_uniforms(self, config: ShaderConfig) -> List[UniformField]:
+        """Get base uniform fields for the configuration."""
+        uniforms = [
+            UniformField("MVP", "mat4x4<f32>"),
+        ]
+
+        # Add view matrix for billboarding
+        if config.render_mode == RenderMode.BILLBOARDED:
+            uniforms.append(UniformField("ViewMatrix", "mat4x4<f32>"))
+
+        # Add color/size based on color mode and primitive type
+        if config.color_mode == ColorMode.SINGLE:
+            if config.primitive_type == PrimitiveType.POINT:
+                uniforms.append(UniformField("ColourSize", "vec4<f32>"))
+            else:
+                uniforms.append(UniformField("colour", "vec3<f32>"))
+                uniforms.append(UniformField("padding", "f32"))
+
+        # Add size for points in multi-color mode
+        elif config.primitive_type == PrimitiveType.POINT:
+            uniforms.append(UniformField("size", "f32"))
+            uniforms.extend([UniformField(f"padding{i}", "u32") for i in range(3)])
+        else:
+            # Add padding for alignment
+            uniforms.extend([UniformField(f"padding{i}", "f32") for i in range(4)])
+
+        return uniforms
+
+    def _generate_vertex_input_struct(self, config: ShaderConfig) -> str:
+        """Generate vertex input structure definition."""
+        attributes = []
+
+        # Always include position
+        attributes.append(VertexAttribute("position", "vec3<f32>", 0))
+
+        # Add color for multi-color mode
+        if config.color_mode == ColorMode.MULTI:
+            attributes.append(VertexAttribute("colour", "vec3<f32>", 1))
+
+        # Add custom attributes
+        if config.custom_attributes:
+            attributes.extend(config.custom_attributes)
+
+        lines = ["struct VertexIn {"]
+        for attr in attributes:
+            lines.append(
+                f"    @location({attr.location}) {attr.name}: {attr.data_type},"
+            )
+        lines.append("};")
+
+        return "\n".join(lines)
+
+    def _generate_vertex_output_struct(self, config: ShaderConfig) -> str:
+        """Generate vertex output structure definition."""
+        lines = ["struct VertexOut {"]
+        lines.append("    @builtin(position) position: vec4<f32>,")
+
+        # Add color for multi-color mode
+        if config.color_mode == ColorMode.MULTI:
+            lines.append("    @location(0) fragColour: vec3<f32>,")
+
+        # Add UV for billboarding points
+        if (
+            config.primitive_type == PrimitiveType.POINT
+            and config.render_mode == RenderMode.BILLBOARDED
+        ):
+            location = 1 if config.color_mode == ColorMode.MULTI else 0
+            lines.append(f"    @location({location}) uv: vec2<f32>,")
+
+        # Add custom outputs from custom vertex shader analysis
+        if config.custom_vertex_shader:
+            # Simple parsing for @location annotations in custom vertex shader
+            for line in config.custom_vertex_shader.split("\n"):
+                if "@location(" in line and "output:" in line:
+                    lines.append(f"    {line.strip()},")
+
+        lines.append("};")
+
+        return "\n".join(lines)
+
+
+class PointShaderTemplate(ShaderTemplate):
+    """Shader template for point rendering."""
+
+    def generate_vertex_shader(self, config: ShaderConfig) -> str:
+        """Generate point vertex shader."""
+        uniform_struct = self._generate_uniform_struct(config)
+        vertex_input = self._generate_vertex_input_struct(config)
+        vertex_output = self._generate_vertex_output_struct(config)
+
+        # Custom vertex shader takes precedence
+        if config.custom_vertex_shader:
+            return f"{uniform_struct}\n{vertex_input}\n{vertex_output}\n{config.custom_vertex_shader}"
+
+        # Generate billboarding vertex shader
+        if config.render_mode == RenderMode.BILLBOARDED:
+            return self._generate_billboard_vertex_shader(
+                config, uniform_struct, vertex_input, vertex_output
+            )
+        else:
+            return self._generate_basic_vertex_shader(
+                config, uniform_struct, vertex_input, vertex_output
+            )
+
+    def generate_fragment_shader(self, config: ShaderConfig) -> str:
+        """Generate point fragment shader."""
+        # Custom fragment shader takes precedence
+        if config.custom_fragment_shader:
+            return config.custom_fragment_shader
+
+        # Generate circular point fragment shader for billboarding
+        if config.render_mode == RenderMode.BILLBOARDED:
+            return self._generate_circle_fragment_shader(config)
+        else:
+            return self._generate_basic_fragment_shader(config)
+
+    def _generate_billboard_vertex_shader(
+        self,
+        config: ShaderConfig,
+        uniform_struct: str,
+        vertex_input: str,
+        vertex_output: str,
+    ) -> str:
+        """Generate billboarding vertex shader for points."""
+        return f"""
+{uniform_struct}
+
+{vertex_input}
+
+{vertex_output}
+
+@binding(0) @group(0) var<uniform> uniforms : Uniforms;
+
+@vertex
+fn vertex_main(input: VertexIn, @builtin(vertex_index) vertex_index: u32) -> VertexOut {{
+    var output: VertexOut;
+    let quad_offsets = array<vec2<f32>, 4>(
+        vec2<f32>(-1.0, -1.0), // bottom-left
+        vec2<f32>(1.0, -1.0),  // bottom-right
+        vec2<f32>(-1.0, 1.0),   // top-left
+        vec2<f32>(1.0, 1.0)    // top-right
+    );
+
+    // Extract camera right and up vectors from view matrix
+    let cameraRight = normalize(vec3<f32>(uniforms.ViewMatrix[0][0], uniforms.ViewMatrix[1][0], uniforms.ViewMatrix[2][0]));
+    let cameraUp = normalize(vec3<f32>(uniforms.ViewMatrix[0][1], uniforms.ViewMatrix[1][1], uniforms.ViewMatrix[2][1]));
+
+    // Calculate billboard offset in world space
+    let size = {self._get_size_expression(config)};
+    let offset2D = quad_offsets[vertex_index] * size;
+    let offset3D = cameraRight * offset2D.x + cameraUp * offset2D.y;
+    let worldPos = input.position + offset3D;
+
+    output.position = uniforms.MVP * vec4<f32>(worldPos, 1.0);
+    {self._get_color_passthrough(config, "output")}
+    // convert offset from -1 -> 1 to 0 -> 1 for uv
+    output.uv = quad_offsets[vertex_index] * 0.5 + 0.5;
+
+    return output;
+}}
+"""
+
+    def _generate_basic_vertex_shader(
+        self,
+        config: ShaderConfig,
+        uniform_struct: str,
+        vertex_input: str,
+        vertex_output: str,
+    ) -> str:
+        """Generate basic vertex shader for points."""
+        return f"""
+{uniform_struct}
+
+{vertex_input}
+
+{vertex_output}
+
+@binding(0) @group(0) var<uniform> uniforms : Uniforms;
+
+@vertex
+fn vertex_main(input: VertexIn) -> VertexOut {{
+    var output: VertexOut;
+    output.position = uniforms.MVP * vec4<f32>(input.position, 1.0);
+    {self._get_color_passthrough(config, "output")}
+    return output;
+}}
+"""
+
+    def _generate_circle_fragment_shader(self, config: ShaderConfig) -> str:
+        """Generate fragment shader for circular points."""
+        color_expression = self._get_color_expression(config)
+
+        return f"""
+@fragment
+fn fragment_main(fragData: VertexOut) -> @location(0) vec4<f32>
+{{
+    let center = vec2<f32>(0.5, 0.5); // Center of the quad in UV space
+    let dist = distance(fragData.uv, center); // Distance from center
+    let radius = 0.5; // Circle radius (quad is 1.0 in UV space)
+
+    if (dist > radius)
+    {{
+        discard; // Remove pixels outside the circle
+    }}
+
+    return {color_expression};
+}}
+"""
+
+    def _generate_basic_fragment_shader(self, config: ShaderConfig) -> str:
+        """Generate basic fragment shader for points."""
+        color_expression = self._get_color_expression(config)
+
+        return f"""
+@fragment
+fn fragment_main(fragData: VertexOut) -> @location(0) vec4<f32>
+{{
+    return {color_expression};
+}}
+"""
+
+    def _get_size_expression(self, config: ShaderConfig) -> str:
+        """Get size expression based on color mode."""
+        if config.color_mode == ColorMode.SINGLE:
+            return "uniforms.ColourSize.w"
+        else:
+            return "uniforms.size"
+
+    def _get_color_passthrough(self, config: ShaderConfig, output_var: str) -> str:
+        """Get color passthrough expression for vertex shader."""
+        if config.color_mode == ColorMode.MULTI:
+            return f"{output_var}.fragColour = input.colour;"
+        else:
+            return ""
+
+    def _get_color_expression(self, config: ShaderConfig) -> str:
+        """Get color expression for fragment shader."""
+        # DEBUG: Always return white to test point visibility
+        return "vec4<f32>(1.0, 1.0, 1.0, 1.0)"
+
+        # Original code (commented out for debugging)
+        # if config.color_mode == ColorMode.MULTI:
+        #     return "vec4<f32>(fragData.fragColour, 1.0)"
+        # else:
+        #     return "vec4<f32>(uniforms.colour, 1.0)"
+
+
+class LineShaderTemplate(ShaderTemplate):
+    """Shader template for line rendering."""
+
+    def generate_vertex_shader(self, config: ShaderConfig) -> str:
+        """Generate line vertex shader."""
+        uniform_struct = self._generate_uniform_struct(config)
+        vertex_input = self._generate_vertex_input_struct(config)
+        vertex_output = self._generate_vertex_output_struct(config)
+
+        if config.custom_vertex_shader:
+            return f"{uniform_struct}\n{vertex_input}\n{vertex_output}\n{config.custom_vertex_shader}"
+
+        return f"""
+{uniform_struct}
+
+{vertex_input}
+
+{vertex_output}
+
+@binding(0) @group(0) var<uniform> uniforms : Uniforms;
+
+@vertex
+fn vertex_main(input: VertexIn) -> VertexOut {{
+    var output: VertexOut;
+    output.position = uniforms.MVP * vec4<f32>(input.position, 1.0);
+    {self._get_color_passthrough(config, "output")}
+    return output;
+}}
+"""
+
+    def generate_fragment_shader(self, config: ShaderConfig) -> str:
+        """Generate line fragment shader."""
+        if config.custom_fragment_shader:
+            return config.custom_fragment_shader
+
+        color_expression = self._get_color_expression(config)
+
+        return f"""
+@fragment
+fn fragment_main(fragData: VertexOut) -> @location(0) vec4<f32>
+{{
+    return {color_expression};
+}}
+"""
+
+    def _get_color_passthrough(self, config: ShaderConfig, output_var: str) -> str:
+        """Get color passthrough expression for vertex shader."""
+        if config.color_mode == ColorMode.MULTI:
+            return f"{output_var}.fragColour = input.colour;"
+        else:
+            return ""
+
+    def _get_color_expression(self, config: ShaderConfig) -> str:
+        """Get color expression for fragment shader."""
+        # DEBUG: Always return white to test line visibility
+        return "vec4<f32>(1.0, 1.0, 1.0, 1.0)"
+
+        # Original code (commented out for debugging)
+        # if config.color_mode == ColorMode.MULTI:
+        #     return "vec4<f32>(fragData.fragColour, 1.0)"
+        # else:
+        #     return "vec4<f32>(uniforms.colour, 1.0)"
+
+
+class TriangleShaderTemplate(ShaderTemplate):
+    """Shader template for triangle rendering."""
+
+    def generate_vertex_shader(self, config: ShaderConfig) -> str:
+        """Generate triangle vertex shader."""
+        uniform_struct = self._generate_uniform_struct(config)
+        vertex_input = self._generate_vertex_input_struct(config)
+        vertex_output = self._generate_vertex_output_struct(config)
+
+        if config.custom_vertex_shader:
+            return f"{uniform_struct}\n{vertex_input}\n{vertex_output}\n{config.custom_vertex_shader}"
+
+        return f"""
+{uniform_struct}
+
+{vertex_input}
+
+{vertex_output}
+
+@binding(0) @group(0) var<uniform> uniforms : Uniforms;
+
+@vertex
+fn vertex_main(input: VertexIn) -> VertexOut {{
+    var output: VertexOut;
+    output.position = uniforms.MVP * vec4<f32>(input.position, 1.0);
+    {self._get_color_passthrough(config, "output")}
+    return output;
+}}
+"""
+
+    def generate_fragment_shader(self, config: ShaderConfig) -> str:
+        """Generate triangle fragment shader."""
+        if config.custom_fragment_shader:
+            return config.custom_fragment_shader
+
+        color_expression = self._get_color_expression(config)
+
+        return f"""
+@fragment
+fn fragment_main(fragData: VertexOut) -> @location(0) vec4<f32>
+{{
+    return {color_expression};
+}}
+"""
+
+    def _get_color_passthrough(self, config: ShaderConfig, output_var: str) -> str:
+        """Get color passthrough expression for vertex shader."""
+        if config.color_mode == ColorMode.MULTI:
+            return f"{output_var}.fragColour = input.colour;"
+        else:
+            return ""
+
+    def _get_color_expression(self, config: ShaderConfig) -> str:
+        """Get color expression for fragment shader."""
+        # DEBUG: Always return white to test triangle visibility
+        return "vec4<f32>(1.0, 1.0, 1.0, 1.0)"
+
+        # Original code (commented out for debugging)
+        # if config.color_mode == ColorMode.MULTI:
+        #     return "vec4<f32>(fragData.fragColour, 1.0)"
+        # else:
+        #     return "vec4<f32>(uniforms.colour, 1.0)"
+
+
+class ShaderTemplateRegistry:
+    """Registry for shader templates."""
+
+    def __init__(self):
+        """Initialize registry with default templates."""
+        # Move template instantiation to separate method to avoid forward reference issues
+        self._templates: Dict[PrimitiveType, ShaderTemplate] = {}
+
+    def _initialize_templates(self):
+        """Initialize templates after all classes are defined."""
+        if not self._templates:
+            self._templates = {
+                PrimitiveType.POINT: PointShaderTemplate(),
+                PrimitiveType.LINE: LineShaderTemplate(),
+                PrimitiveType.TRIANGLE: TriangleShaderTemplate(),
+            }
+
+    def register_template(
+        self, primitive_type: PrimitiveType, template: ShaderTemplate
+    ) -> None:
+        """Register a custom shader template."""
+        self._templates[primitive_type] = template
+
+    def get_template(self, primitive_type: PrimitiveType) -> ShaderTemplate:
+        """Get shader template for primitive type."""
+        self._initialize_templates()  # Ensure templates are initialized
+        if primitive_type not in self._templates:
+            raise ValueError(
+                f"No template registered for primitive type: {primitive_type}"
+            )
+        return self._templates[primitive_type]
+
+    def generate_shader(self, config: ShaderConfig) -> str:
+        """Generate shader code from configuration."""
+        template = self.get_template(config.primitive_type)
+        return template.generate_shader_code(config)
+
+
+# Global registry instance
+shader_registry = ShaderTemplateRegistry()
+
+
+def generate_shader_code(config: ShaderConfig) -> str:
+    """Generate shader code from configuration using global registry."""
+    return shader_registry.generate_shader(config)
+
+
+def create_shader_config(
+    primitive_type: PrimitiveType,
+    color_mode: ColorMode,
+    render_mode: RenderMode = RenderMode.STANDARD,
+    custom_vertex_shader: Optional[str] = None,
+    custom_fragment_shader: Optional[str] = None,
+    custom_uniforms: Optional[List[UniformField]] = None,
+    custom_attributes: Optional[List[VertexAttribute]] = None,
+) -> ShaderConfig:
+    """Create a shader configuration."""
+    return ShaderConfig(
+        primitive_type=primitive_type,
+        color_mode=color_mode,
+        render_mode=render_mode,
+        custom_vertex_shader=custom_vertex_shader,
+        custom_fragment_shader=custom_fragment_shader,
+        custom_uniforms=custom_uniforms,
+        custom_attributes=custom_attributes,
+    )
+
+
+# Register default templates
+shader_registry.register_template(PrimitiveType.POINT, PointShaderTemplate())
+shader_registry.register_template(PrimitiveType.LINE, LineShaderTemplate())
+shader_registry.register_template(PrimitiveType.TRIANGLE, TriangleShaderTemplate())
+
+
+# Register default templates
+shader_registry.register_template(PrimitiveType.POINT, PointShaderTemplate())
+shader_registry.register_template(PrimitiveType.LINE, LineShaderTemplate())
+shader_registry.register_template(PrimitiveType.TRIANGLE, TriangleShaderTemplate())