LNXSDK/leenkx/blender/lnx/material/make_shader.py

import os
import re
import subprocess
from concurrent.futures import ThreadPoolExecutor
from typing import Dict, List, Tuple

import bpy
from bpy.types import Material
from bpy.types import Object

import lnx.api
import lnx.assets as assets
import lnx.exporter
import lnx.log as log
import lnx.material.cycles as cycles
import lnx.material.make_decal as make_decal
import lnx.material.make_depth as make_depth
import lnx.material.make_mesh as make_mesh
import lnx.material.make_overlay as make_overlay
import lnx.material.make_transluc as make_transluc
import lnx.material.make_refract as make_refract
import lnx.material.make_voxel as make_voxel
import lnx.material.mat_state as mat_state
import lnx.material.mat_utils as mat_utils
from lnx.material.shader import Shader, ShaderContext, ShaderData
import lnx.utils

if lnx.is_reload(__name__):
    lnx.api = lnx.reload_module(lnx.api)
    assets = lnx.reload_module(assets)
    lnx.exporter = lnx.reload_module(lnx.exporter)
    log = lnx.reload_module(log)
    cycles = lnx.reload_module(cycles)
    make_decal = lnx.reload_module(make_decal)
    make_depth = lnx.reload_module(make_depth)
    make_mesh = lnx.reload_module(make_mesh)
    make_overlay = lnx.reload_module(make_overlay)
    make_transluc = lnx.reload_module(make_transluc)
    make_voxel = lnx.reload_module(make_voxel)
    mat_state = lnx.reload_module(mat_state)
    mat_utils = lnx.reload_module(mat_utils)
    lnx.material.shader = lnx.reload_module(lnx.material.shader)
    from lnx.material.shader import Shader, ShaderContext, ShaderData
    lnx.utils = lnx.reload_module(lnx.utils)
else:
    lnx.enable_reload(__name__)

rpass_hook = None


def cache_hit(cached, wrd, full_path, matname):
    mat_state.features = dict(cached['features'])
    mat_state.emission_type = cached['emission_type']
    mat_state.texture_grad = cached['texture_grad']
    mat_state.data = cached['shader_data']

    if cached['world_defs_delta']:
        for define in re.findall(r'_[A-Za-z0-9]+', cached['world_defs_delta']):
            assets.add_world_def(wrd, define)

    for d in cached['khafile_defs_delta']:
        assets.add_khafile_def(d)

    for shader_path in cached['shader_paths']:
        assets.add_shader(shader_path)

    shader_data_name = cached['shader_data_name']
    if wrd.lnx_single_data_file:
        pass
    else:
        shader_data_path = lnx.utils.get_fp_build() + '/compiled/Shaders/' + shader_data_name + '.lnx'
        assets.add_shader_data(shader_data_path)

    return cached['rpasses'], mat_state.data, shader_data_name, cached['bind_constants'], cached['bind_textures']


def build(material: Material, mat_users: Dict[Material, List[Object]], mat_lnxusers) -> Tuple:
    mat_state.mat_users = mat_users
    mat_state.mat_lnxusers = mat_lnxusers
    mat_state.material = material
    mat_state.nodes = material.node_tree.nodes
    mat_state.data = ShaderData(material)
    mat_state.output_node = cycles.node_by_type(mat_state.nodes, 'OUTPUT_MATERIAL')
    if mat_state.output_node is None:
        # Place empty material output to keep compiler happy..
        mat_state.output_node = mat_state.nodes.new('ShaderNodeOutputMaterial')
    
    wrd = bpy.data.worlds['Lnx']
    rpdat = lnx.utils.get_rp()
    rpasses = mat_utils.get_rpasses(material)
    matname = lnx.utils.safesrc(lnx.utils.asset_name(material))
    rel_path = lnx.utils.build_dir() + '/compiled/Shaders/'
    full_path = lnx.utils.get_fp() + '/' + rel_path
    if not os.path.exists(full_path):
        os.makedirs(full_path)

    make_instancing_and_skinning(material, mat_users)

    cache_key = None
    cached = None
    if material.signature:
        global_elems_key = tuple((e['name'], e['data']) for e in mat_state.data.global_elems)
        cache_key = (material.signature, mat_state.uses_instancing, tuple(rpasses), global_elems_key)
        cached = mat_state.material_cache.get(cache_key)

    if cached is not None:
        return cache_hit(cached, wrd, full_path, matname)

    world_defs_before = wrd.world_defs
    khafile_defs_before = set(assets.khafile_defs)

    bind_constants = dict()
    bind_textures = dict()
    all_shader_paths = []

    for rp in rpasses:
        car = []
        bind_constants[rp] = car
        mat_state.bind_constants = car
        tar = []
        bind_textures[rp] = tar
        mat_state.bind_textures = tar

        con = None

        if rpdat.rp_driver != 'Leenkx' and lnx.api.drivers[rpdat.rp_driver]['make_rpass'] is not None:
            con = lnx.api.drivers[rpdat.rp_driver]['make_rpass'](rp)

        if con is not None:
            pass

        elif rp == 'mesh':
            con = make_mesh.make(rp, rpasses)

        elif rp == 'shadowmap':
            con = make_depth.make(rp, rpasses, shadowmap=True)

        elif rp == 'shadowmap_transparent':
            con = make_depth.make(rp, rpasses, shadowmap=True, shadowmap_transparent=True)

        elif rp == 'translucent':
            con = make_transluc.make(rp)

        elif rp == 'refraction':
            con = make_refract.make(rp)

        elif rp == 'overlay':
            con = make_overlay.make(rp)

        elif rp == 'decal':
            con = make_decal.make(rp)

        elif rp == 'depth':
            con = make_depth.make(rp, rpasses)

        elif rp == 'voxel':
            con = make_voxel.make(rp)

        elif rpass_hook is not None:
            con = rpass_hook(rp)

        if con is not None:
            rp_shader_paths = write_shaders(rel_path, con, rp, matname)
            all_shader_paths.extend(rp_shader_paths)

    shader_data_name = matname + '_data'

    if wrd.lnx_single_data_file:
        if 'shader_datas' not in lnx.exporter.current_output:
            lnx.exporter.current_output['shader_datas'] = []
        lnx.exporter.current_output['shader_datas'].append(mat_state.data.get()['shader_datas'][0])
    else:
        lnx.utils.write_lnx(full_path + '/' + matname + '_data.lnx', mat_state.data.get())
        shader_data_path = lnx.utils.get_fp_build() + '/compiled/Shaders/' + shader_data_name + '.lnx'
        assets.add_shader_data(shader_data_path)

    if cache_key is not None:
        world_defs_delta = wrd.world_defs[len(world_defs_before):]
        khafile_defs_delta = [d for d in assets.khafile_defs if d not in khafile_defs_before]
        mat_state.material_cache[cache_key] = {
            'shader_data': mat_state.data,
            'shader_data_name': shader_data_name,
            'bind_constants': bind_constants,
            'bind_textures': bind_textures,
            'features': dict(mat_state.features),
            'emission_type': mat_state.emission_type,
            'texture_grad': mat_state.texture_grad,
            'shader_paths': all_shader_paths,
            'world_defs_delta': world_defs_delta,
            'khafile_defs_delta': khafile_defs_delta,
            'rpasses': rpasses,
        }

    return rpasses, mat_state.data, shader_data_name, bind_constants, bind_textures


def write_shaders(rel_path: str, con: ShaderContext, rpass: str, matname: str) -> List[str]:
    keep_cache = mat_state.material.lnx_cached
    shaders = [con.vert, con.frag, con.geom, con.tesc, con.tese]
    exts = ['vert', 'frag', 'geom', 'tesc', 'tese']
    shader_paths = []

    write_tasks = []
    for shader, ext in zip(shaders, exts):
        if shader is None or shader.is_linked:
            continue

        validation_issues = shader.validate()
        if validation_issues:
            for issue in validation_issues:
                log.warn(f"Shader validation issue in {matname}_{rpass}.{ext}: {issue}. ")

        output_rpass = rpass
        if output_rpass == 'mesh' and mat_state.material.lnx_blending:
            output_rpass = 'blend'

        output_ext = '.glsl'
        output_rel_path = rel_path
        if shader.noprocessing:
            hlsl_dir = lnx.utils.build_dir() + '/compiled/Hlsl/'
            os.makedirs(hlsl_dir, exist_ok=True)
            output_ext = '.hlsl'
            output_rel_path = rel_path.replace('/compiled/Shaders/', '/compiled/Hlsl/')

        shader_file = matname + '_' + output_rpass + '.' + ext + output_ext
        shader_path = lnx.utils.get_fp() + '/' + output_rel_path + '/' + shader_file
        assets.add_shader(shader_path)
        shader_paths.append(shader_path)

        if os.path.isfile(shader_path) and keep_cache:
            continue

        content = shader.get()

        if shader.noprocessing:
            written = lnx.utils.write_if_changed(shader_path, content)
            if written:
                cwd = os.getcwd()
                os.chdir(lnx.utils.get_fp() + '/' + output_rel_path)
                hlslbin_path = lnx.utils.get_sdk_path() + '/lib/leenkx_tools/hlslbin/hlslbin.exe'
                prof = 'vs_5_0' if ext == 'vert' else 'ps_5_0' if ext == 'frag' else 'gs_5_0'
                args = [hlslbin_path.replace('/', '\\').replace('\\\\', '\\'), shader_file, shader_file[:-4] + 'glsl', prof]
                if ext == 'vert':
                    args.append('-i')
                    args.append('pos')
                proc = subprocess.call(args)
                os.chdir(cwd)
        else:
            write_tasks.append((shader_path, content))

    if write_tasks:
        with ThreadPoolExecutor(max_workers=min(5, len(write_tasks))) as pool:
            list(pool.map(lambda t: lnx.utils.write_if_changed(t[0], t[1]), write_tasks))

    return shader_paths


def write_shader(rel_path: str, shader: Shader, ext: str, rpass: str, matname: str, keep_cache=True) -> None:
    if shader is None or shader.is_linked:
        return

    validation_issues = shader.validate()
    if validation_issues:
        for issue in validation_issues:
            log.warn(f"Shader validation issue in {matname}_{rpass}.{ext}: {issue}. ")

    # TODO: blend context
    if rpass == 'mesh' and mat_state.material.lnx_blending:
        rpass = 'blend'

    file_ext = '.glsl'
    if shader.noprocessing:
        # Use hlsl directly
        hlsl_dir = lnx.utils.build_dir() + '/compiled/Hlsl/'
        os.makedirs(hlsl_dir, exist_ok=True)
        file_ext = '.hlsl'
        rel_path = rel_path.replace('/compiled/Shaders/', '/compiled/Hlsl/')

    shader_file = matname + '_' + rpass + '.' + ext + file_ext
    shader_path = lnx.utils.get_fp() + '/' + rel_path + '/' + shader_file
    assets.add_shader(shader_path)
    if not os.path.isfile(shader_path) or not keep_cache:
        written = lnx.utils.write_if_changed(shader_path, shader.get())

        if written and shader.noprocessing:
            cwd = os.getcwd()
            os.chdir(lnx.utils.get_fp() + '/' + rel_path)
            hlslbin_path = lnx.utils.get_sdk_path() + '/lib/leenkx_tools/hlslbin/hlslbin.exe'
            prof = 'vs_5_0' if ext == 'vert' else 'ps_5_0' if ext == 'frag' else 'gs_5_0'
            # noprocessing flag - gets renamed to .d3d11
            args = [hlslbin_path.replace('/', '\\').replace('\\\\', '\\'), shader_file, shader_file[:-4] + 'glsl', prof]
            if ext == 'vert':
                args.append('-i')
                args.append('pos')
            proc = subprocess.call(args)
            os.chdir(cwd)


def make_instancing_and_skinning(mat: Material, mat_users: Dict[Material, List[Object]]) -> None:
    """Build material with instancing or skinning if enabled.
    If the material is a custom material, only validation checks for instancing are performed."""
    global_elems = []
    if mat_users is not None and mat in mat_users:
        # Whether there are both an instanced object and a not instanced object with this material
        instancing_usage = [False, False]
        mat_state.uses_instancing = False

        for bo in mat_users[mat]:
            if mat.lnx_custom_material == '':
                # Morph Targets
                if lnx.utils.export_morph_targets(bo):
                    global_elems.append({'name': 'morph', 'data': 'short2norm'})
                # GPU Skinning
                if lnx.utils.export_bone_data(bo):
                    global_elems.append({'name': 'bone', 'data': 'short4norm'})
                    global_elems.append({'name': 'weight', 'data': 'short4norm'})

            # Instancing
            inst = bo.lnx_instanced
            if inst != 'Off' or (mat.lnx_particle_flag and lnx.utils.get_rp().lnx_particles == 'GPU'):
                instancing_usage[0] = True
                mat_state.uses_instancing = True

                if mat.lnx_custom_material == '':
                    global_elems.append({'name': 'ipos', 'data': 'float3'})
                    if 'Rot' in inst:
                        global_elems.append({'name': 'irot', 'data': 'float3'})
                    if 'Scale' in inst:
                        global_elems.append({'name': 'iscl', 'data': 'float3'})

            elif inst == 'Off':
                # Ignore children of instanced objects, they are instanced even when set to 'Off'
                instancing_usage[1] = bo.parent is None or bo.parent.lnx_instanced == 'Off'

        if instancing_usage[0] and instancing_usage[1]:
            # Display a warning for invalid instancing configurations
            # See https://github.com/leenkx3d/leenkx/issues/2072
            log.warn(f'Material "{mat.name}" has both instanced and not instanced objects, objects might flicker!')

    if mat.lnx_custom_material == '':
        mat_state.data.global_elems = global_elems