3b1b-manim/manimlib/utils/shaders.py

import os
import warnings
import re
import moderngl

from manimlib.constants import SHADER_DIR

# Mobjects that should be rendered with
# the same shader will be organized and
# clumped together based on keeping track
# of a dict holding all the relevant information
# to that shader


# TODO, this should all be treated as an object
# This object a shader program instead of the vert,
# geom and frag file names, and it should cache those
# programs in the way currently handled by Camera
# It should replace the Camera.get_shader method with
# its own get_shader_program method, which will take
# in the camera's perspective_uniforms.


SHADER_INFO_KEYS = [
    # Vertex data for the shader (as structured array)
    "vert_data",
    # Index data (if applicable) for the shader
    "index_data",
    # List of variable names corresponding to inputs of vertex shader
    "attributes",
    # Filename of vetex shader
    "vert",
    # Filename of geometry shader, if there is one
    "geom",
    # Filename of fragment shader
    "frag",
    # A dictionary mapping names of uniform variables
    "uniforms",
    # A dictionary mapping names (as they show up in)
    # the shader to filepaths for textures.
    "texture_paths",
    # Whether or not to apply depth test
    "depth_test",
    # E.g. moderngl.TRIANGLE_STRIP
    "render_primative",
]

# Exclude data
SHADER_KEYS_FOR_ID = SHADER_INFO_KEYS[3:]


def get_shader_info(vert_data=None,
                    vert_indices=None,
                    attributes=None,
                    vert_file=None,
                    geom_file=None,
                    frag_file=None,
                    uniforms=None,
                    texture_paths=None,
                    depth_test=False,
                    render_primative=moderngl.TRIANGLE_STRIP,
                    ):
    result = {
        "vert_data": vert_data,
        "vert_indices": vert_indices,
        "attributes": attributes,
        "vert": vert_file,
        "geom": geom_file,
        "frag": frag_file,
        "uniforms": uniforms or dict(),
        "texture_paths": texture_paths or dict(),
        "depth_test": depth_test,
        "render_primative": str(render_primative),
    }
    result["id"] = create_shader_info_id(result)
    result["prog_id"] = create_shader_info_program_id(result)
    return result


def is_valid_shader_info(shader_info):
    vert_data = shader_info["vert_data"]
    return all([
        vert_data is not None,
        shader_info["vert"],
        shader_info["frag"],
    ])


def shader_info_to_id(shader_info):
    return shader_info["id"]


def shader_info_program_id(shader_info):
    return shader_info["prog_id"]


def create_shader_info_id(shader_info):
    # A unique id for a shader
    return "|".join([str(shader_info[key]) for key in SHADER_KEYS_FOR_ID])


def refresh_shader_info_id(shader_info):
    shader_info["id"] = create_shader_info_id(shader_info)


def create_shader_info_program_id(shader_info):
    return "|".join([str(shader_info[key]) for key in ["vert", "geom", "frag"]])


def same_shader_type(info1, info2):
    return info1["id"] == info2["id"]


def shader_info_to_program_code(shader_info):
    return {
        "vertex_shader": get_shader_code_from_file(shader_info["vert"]),
        "geometry_shader": get_shader_code_from_file(shader_info["geom"]),
        "fragment_shader": get_shader_code_from_file(shader_info["frag"]),
    }


def get_shader_code_from_file(filename):
    if not filename:
        return None

    filepath = os.path.join(SHADER_DIR, filename)
    if not os.path.exists(filepath):
        warnings.warn(f"No file at {filepath}")
        return

    with open(filepath, "r") as f:
        result = f.read()

    # To share functionality between shaders, some functions are read in
    # from other files an inserted into the relevant strings before
    # passing to ctx.program for compiling
    # Replace "#INSERT " lines with relevant code
    insertions = re.findall(r"^#INSERT .*\.glsl$", result, flags=re.MULTILINE)
    for line in insertions:
        inserted_code = get_shader_code_from_file(line.replace("#INSERT ", ""))
        result = result.replace(line, inserted_code)
    return result