Parser for deko3d shader files

dksh.py

#!/usr/bin/env python3

# See https://github.com/devkitPro/deko3d/blob/master/source/dksh.h

import os, sys
import ctypes, enum

class DkProgramType(enum.Enum):
    Vertex   = 0
    Fragment = 1
    Geometry = 2
    TessCtrl = 3
    TessEval = 4
    Compute  = 5

class DkshHeader(ctypes.Structure):
    _fields_ = [
        ("magic",        ctypes.c_uint32),
        ("header_sz",    ctypes.c_uint32),
        ("control_sz",   ctypes.c_uint32),
        ("code_sz",      ctypes.c_uint32),
        ("programs_off", ctypes.c_uint32),
        ("num_programs", ctypes.c_uint32),
    ]

    MAGIC = 0x48534B44

assert ctypes.sizeof(DkshHeader) == 24

class DkshProgramHeader(ctypes.Structure):
    class ProgramInfo(ctypes.Union):
        class VertInfo(ctypes.Structure):
            _fields_ = [
                ("alt_entrypoint", ctypes.c_uint32),
                ("alt_num_gprs",   ctypes.c_uint32),
            ]

        class FragInfo(ctypes.Structure):
            _fields_ = [
                ("has_table_3d1",        ctypes.c_bool),
                ("early_fragment_tests", ctypes.c_bool),
                ("post_depth_coverage",  ctypes.c_bool),
                ("persample_invocation", ctypes.c_bool),
                ("table_3d1",            ctypes.c_uint32 * 4),
                ("param_d8",             ctypes.c_uint32),
                ("param_65b",            ctypes.c_uint16),
                ("param_489",            ctypes.c_uint16),
            ]

        class GeomInfo(ctypes.Structure):
            _fields_ = [
                ("flag_47c",      ctypes.c_bool),
                ("has_table_490", ctypes.c_bool),
                ("_padding",      ctypes.c_bool * 2),
                ("table",         ctypes.c_uint32 * 8),
            ]

        class TessEvalInfo(ctypes.Structure):
            _fields_ = [
                ("param_c8", ctypes.c_uint32),
            ]

        class CompInfo(ctypes.Structure):
            _fields_ = [
                ("block_dims",       ctypes.c_uint32 * 3),
                ("shared_mem_sz",    ctypes.c_uint32),
                ("local_pos_mem_sz", ctypes.c_uint32),
                ("local_neg_mem_sz", ctypes.c_uint32),
                ("crs_sz",           ctypes.c_uint32),
                ("num_barriers",     ctypes.c_uint32),
            ]

        _fields_ = [
            ("vert",      VertInfo),
            ("frag",      FragInfo),
            ("geom",      GeomInfo),
            ("tess_eval", TessEvalInfo),
            ("comp",      CompInfo),
        ]

    _fields_ = [
        ("type",                ctypes.c_uint32),
        ("entrypoint",          ctypes.c_uint32),
        ("num_gprs",            ctypes.c_uint32),
        ("constbuf1_off",       ctypes.c_uint32),
        ("constbuf1_sz",        ctypes.c_uint32),
        ("per_warp_scratch_sz", ctypes.c_uint32),
        ("prog_info",           ProgramInfo),
        ("reserved",            ctypes.c_uint32),
    ]

assert ctypes.sizeof(DkshProgramHeader) == 64

with open("color_fsh.dksh", "rb") as fp:
    hdr = DkshHeader()
    fp.readinto(hdr)

    assert hdr.magic == DkshHeader.MAGIC

    proghdr = DkshProgramHeader()
    fp.readinto(proghdr)

print("Program type:", DkProgramType(proghdr.type).name)
print("Entrypoint:", hex(proghdr.entrypoint))
print("Num registers:", proghdr.num_gprs)