mirage/tools/compiler.py

#!/usr/bin/env python3
"""
着色器编译模块

负责调用glslc编译器将GLSL源代码编译为SPIR-V二进制。
"""

from __future__ import annotations

import logging
import shutil
import subprocess
import tempfile
from pathlib import Path
from typing import List

from .constants import GLSLC_STD, GLSLC_TARGET_ENV
from .spirv_parser import extract_spirv_reflection
from .types import CompilationResult, ShaderMetadata, ToolError

LOG = logging.getLogger("glsl2spirv")


class ShaderCompiler:
    """GLSL到SPIR-V编译器"""

    def __init__(self, verbose: bool = False, dry_run: bool = False, debug_spirv: bool = False):
        self.verbose = verbose
        self.dry_run = dry_run
        self.debug_spirv = debug_spirv
        self.glslc_path = self._find_glslc()

    def _find_glslc(self) -> str:
        """查找glslc编译器"""
        glslc = shutil.which("glslc")
        if glslc:
            return glslc

        glslc = shutil.which("glslc.exe")
        if glslc:
            return glslc

        raise ToolError("Unable to locate glslc compiler. Please install Vulkan SDK.")

    def _build_glslc_command(
        self,
        shader_path: Path,
        output_path: Path,
        metadata: ShaderMetadata,
        shader_type: str,
        debug_mode: bool = False,
    ) -> List[str]:
        """构建glslc命令行
        
        Args:
            shader_path: 着色器源文件路径
            output_path: 输出SPIR-V文件路径
            metadata: 着色器元数据
            shader_type: 着色器类型（vert, frag, comp等）
            debug_mode: 是否启用debug模式（用于提取完整反射数据）
        """
        stage_map = {
            'vert': 'vertex',
            'frag': 'fragment',
            'comp': 'compute',
            'geom': 'geometry',
            'tesc': 'tesscontrol',
            'tese': 'tesseval',
        }
        
        cmd = [self.glslc_path]
        
        if shader_type in stage_map:
            cmd.append(f"-fshader-stage={stage_map[shader_type]}")
        
        cmd.extend([
            str(shader_path),
            "-o", str(output_path),
            "--target-env=" + GLSLC_TARGET_ENV,
            "-std=" + GLSLC_STD,
        ])

        # debug模式：启用调试信息和禁用优化以保留完整的反射数据
        if debug_mode:
            cmd.extend(["-g", "-O0"])

        for inc_path in metadata.include_paths:
            cmd.append(f"-I{inc_path}")

        for macro_name, macro_value in metadata.defines.items():
            if macro_value:
                cmd.append(f"-D{macro_name}={macro_value}")
            else:
                cmd.append(f"-D{macro_name}")

        return cmd

    def _run_glslc(self, cmd: List[str], shader_path: Path, shader_type: str) -> None:
        """执行glslc命令"""
        try:
            result = subprocess.run(cmd, capture_output=True, text=True, check=False)
            if result.returncode != 0:
                raise ToolError(
                    f"glslc compilation failed for {shader_path.name} ({shader_type}):\n{result.stderr}"
                )
        except FileNotFoundError:
            raise ToolError(f"glslc executable not found at {self.glslc_path}")

    def compile_shader(
        self,
        shader_path: Path,
        metadata: ShaderMetadata,
        shader_type: str,
    ) -> CompilationResult:
        """编译单个着色器并提取反射信息
        
        编译流程：
        1. 先用debug模式编译（-g -O0）以获取完整的反射数据
        2. 如果启用了debug_spirv，则最终SPIR-V使用debug版本
           否则用正常模式编译生成最终的优化SPIR-V
        
        这样可以确保绑定代码生成使用完整的变量名信息，
        同时在非debug构建时最终的着色器仍然是优化过的版本。
        """
        with tempfile.TemporaryDirectory() as tmpdir:
            debug_output_path = Path(tmpdir) / f"{shader_path.stem}_debug.spv"
            release_output_path = Path(tmpdir) / f"{shader_path.stem}.spv"

            if not self.dry_run:
                # 步骤1：用debug模式编译以提取完整反射数据
                debug_cmd = self._build_glslc_command(
                    shader_path, debug_output_path, metadata, shader_type, debug_mode=True
                )

                if self.verbose:
                    LOG.debug("Running glslc (debug): %s", " ".join(debug_cmd))

                self._run_glslc(debug_cmd, shader_path, shader_type)
                
                debug_spirv_data = debug_output_path.read_bytes()
                
                # 从debug版本的SPIR-V提取完整反射信息
                reflection_result = extract_spirv_reflection(debug_spirv_data, shader_type)
                entry_point = reflection_result['entry_point']
                bindings = reflection_result['bindings']
                reflection = reflection_result.get('reflection')

                # 步骤2：根据debug_spirv标志决定最终SPIR-V版本
                if self.debug_spirv:
                    # 使用debug版本作为最终输出
                    spirv_data = debug_spirv_data
                    if self.verbose:
                        LOG.debug("Using debug SPIR-V as final output for %s", shader_path.name)
                else:
                    # 用正常模式编译生成最终的优化SPIR-V
                    release_cmd = self._build_glslc_command(
                        shader_path, release_output_path, metadata, shader_type, debug_mode=False
                    )

                    if self.verbose:
                        LOG.debug("Running glslc (release): %s", " ".join(release_cmd))

                    self._run_glslc(release_cmd, shader_path, shader_type)
                    
                    spirv_data = release_output_path.read_bytes()
            else:
                spirv_data = b""
                entry_point = "main"
                bindings = []
                reflection = None

        return CompilationResult(
            source_file=shader_path,
            shader_type=shader_type,
            spirv_data=spirv_data,
            entry_point=entry_point,
            bindings=bindings,
            reflection=reflection,
        )