457 lines
16 KiB
Python
457 lines
16 KiB
Python
import arm.utils
|
|
|
|
if arm.is_reload(__name__):
|
|
arm.utils = arm.reload_module(arm.utils)
|
|
else:
|
|
arm.enable_reload(__name__)
|
|
|
|
# Type aliases for type hints to make it easier to see which kind of
|
|
# shader data type is stored in a string
|
|
floatstr = str
|
|
vec2str = str
|
|
vec3str = str
|
|
vec4str = str
|
|
|
|
|
|
class ShaderData:
|
|
|
|
def __init__(self, material):
|
|
self.material = material
|
|
self.contexts = []
|
|
self.global_elems = [] # bone, weight, ipos, irot, iscl
|
|
self.sd = {}
|
|
self.data = {'shader_datas': [self.sd]}
|
|
self.matname = arm.utils.safesrc(arm.utils.asset_name(material))
|
|
self.sd['name'] = self.matname + '_data'
|
|
self.sd['contexts'] = []
|
|
|
|
def add_context(self, props) -> 'ShaderContext':
|
|
con = ShaderContext(self.material, self.sd, props)
|
|
if con not in self.sd['contexts']:
|
|
for elem in self.global_elems:
|
|
con.add_elem(elem['name'], elem['data'])
|
|
self.sd['contexts'].append(con.get())
|
|
return con
|
|
|
|
def get(self):
|
|
return self.data
|
|
|
|
class ShaderContext:
|
|
|
|
def __init__(self, material, shader_data, props):
|
|
self.vert = None
|
|
self.frag = None
|
|
self.geom = None
|
|
self.tesc = None
|
|
self.tese = None
|
|
self.material = material
|
|
self.matname = arm.utils.safesrc(arm.utils.asset_name(material))
|
|
self.shader_data = shader_data
|
|
self.data = {}
|
|
self.data['name'] = props['name']
|
|
self.data['depth_write'] = props['depth_write']
|
|
self.data['compare_mode'] = props['compare_mode']
|
|
self.data['cull_mode'] = props['cull_mode']
|
|
if 'vertex_elements' in props:
|
|
self.data['vertex_elements'] = props['vertex_elements']
|
|
else:
|
|
self.data['vertex_elements'] = [{'name': 'pos', 'data': 'short4norm'}, {'name': 'nor', 'data': 'short2norm'}] # (p.xyz, n.z), (n.xy)
|
|
if 'blend_source' in props:
|
|
self.data['blend_source'] = props['blend_source']
|
|
if 'blend_destination' in props:
|
|
self.data['blend_destination'] = props['blend_destination']
|
|
if 'blend_operation' in props:
|
|
self.data['blend_operation'] = props['blend_operation']
|
|
if 'alpha_blend_source' in props:
|
|
self.data['alpha_blend_source'] = props['alpha_blend_source']
|
|
if 'alpha_blend_destination' in props:
|
|
self.data['alpha_blend_destination'] = props['alpha_blend_destination']
|
|
if 'alpha_blend_operation' in props:
|
|
self.data['alpha_blend_operation'] = props['alpha_blend_operation']
|
|
if 'color_writes_red' in props:
|
|
self.data['color_writes_red'] = props['color_writes_red']
|
|
if 'color_writes_green' in props:
|
|
self.data['color_writes_green'] = props['color_writes_green']
|
|
if 'color_writes_blue' in props:
|
|
self.data['color_writes_blue'] = props['color_writes_blue']
|
|
if 'color_writes_alpha' in props:
|
|
self.data['color_writes_alpha'] = props['color_writes_alpha']
|
|
if 'color_attachments' in props:
|
|
self.data['color_attachments'] = props['color_attachments']
|
|
|
|
self.data['texture_units'] = []
|
|
self.tunits = self.data['texture_units']
|
|
self.data['constants'] = []
|
|
self.constants = self.data['constants']
|
|
|
|
def add_elem(self, name, data):
|
|
elem = { 'name': name, 'data': data }
|
|
if elem not in self.data['vertex_elements']:
|
|
self.data['vertex_elements'].append(elem)
|
|
self.sort_vs()
|
|
|
|
def sort_vs(self):
|
|
vs = []
|
|
ar = ['pos', 'nor', 'tex', 'tex1', 'morph', 'col', 'tang', 'bone', 'weight', 'ipos', 'irot', 'iscl']
|
|
for ename in ar:
|
|
elem = self.get_elem(ename)
|
|
if elem != None:
|
|
vs.append(elem)
|
|
self.data['vertex_elements'] = vs
|
|
|
|
def is_elem(self, name):
|
|
for elem in self.data['vertex_elements']:
|
|
if elem['name'] == name:
|
|
return True
|
|
return False
|
|
|
|
def get_elem(self, name):
|
|
for elem in self.data['vertex_elements']:
|
|
if elem['name'] == name:
|
|
return elem
|
|
return None
|
|
|
|
def get(self):
|
|
return self.data
|
|
|
|
def add_constant(self, ctype, name, link=None, default_value=None, is_arm_mat_param=None):
|
|
for c in self.constants:
|
|
if c['name'] == name:
|
|
return
|
|
|
|
c = { 'name': name, 'type': ctype}
|
|
if link is not None:
|
|
c['link'] = link
|
|
if default_value is not None:
|
|
if ctype == 'float':
|
|
c['float'] = default_value
|
|
if ctype == 'vec3':
|
|
c['vec3'] = default_value
|
|
if is_arm_mat_param is not None:
|
|
c['is_arm_parameter'] = 'true'
|
|
self.constants.append(c)
|
|
|
|
def add_texture_unit(self, name, link=None, is_image=None,
|
|
addr_u=None, addr_v=None,
|
|
filter_min=None, filter_mag=None, mipmap_filter=None,
|
|
default_value=None, is_arm_mat_param=None):
|
|
for c in self.tunits:
|
|
if c['name'] == name:
|
|
return
|
|
|
|
c = {'name': name}
|
|
if link is not None:
|
|
c['link'] = link
|
|
if is_image is not None:
|
|
c['is_image'] = is_image
|
|
if addr_u is not None:
|
|
c['addressing_u'] = addr_u
|
|
if addr_v is not None:
|
|
c['addressing_v'] = addr_v
|
|
if filter_min is not None:
|
|
c['filter_min'] = filter_min
|
|
if filter_mag is not None:
|
|
c['filter_mag'] = filter_mag
|
|
if mipmap_filter is not None:
|
|
c['mipmap_filter'] = mipmap_filter
|
|
if default_value is not None:
|
|
c['default_image_file'] = default_value
|
|
if is_arm_mat_param is not None:
|
|
c['is_arm_parameter'] = 'true'
|
|
|
|
self.tunits.append(c)
|
|
|
|
def make_vert(self, custom_name: str = None):
|
|
if custom_name is None:
|
|
self.data['vertex_shader'] = self.matname + '_' + self.data['name'] + '.vert'
|
|
else:
|
|
self.data['vertex_shader'] = custom_name + '.vert'
|
|
self.vert = Shader(self, 'vert')
|
|
return self.vert
|
|
|
|
def make_frag(self, custom_name: str = None):
|
|
if custom_name is None:
|
|
self.data['fragment_shader'] = self.matname + '_' + self.data['name'] + '.frag'
|
|
else:
|
|
self.data['fragment_shader'] = custom_name + '.frag'
|
|
self.frag = Shader(self, 'frag')
|
|
return self.frag
|
|
|
|
def make_geom(self, custom_name: str = None):
|
|
if custom_name is None:
|
|
self.data['geometry_shader'] = self.matname + '_' + self.data['name'] + '.geom'
|
|
else:
|
|
self.data['geometry_shader'] = custom_name + '.geom'
|
|
self.geom = Shader(self, 'geom')
|
|
return self.geom
|
|
|
|
def make_tesc(self, custom_name: str = None):
|
|
if custom_name is None:
|
|
self.data['tesscontrol_shader'] = self.matname + '_' + self.data['name'] + '.tesc'
|
|
else:
|
|
self.data['tesscontrol_shader'] = custom_name + '.tesc'
|
|
self.tesc = Shader(self, 'tesc')
|
|
return self.tesc
|
|
|
|
def make_tese(self, custom_name: str = None):
|
|
if custom_name is None:
|
|
self.data['tesseval_shader'] = self.matname + '_' + self.data['name'] + '.tese'
|
|
else:
|
|
self.data['tesseval_shader'] = custom_name + '.tese'
|
|
self.tese = Shader(self, 'tese')
|
|
return self.tese
|
|
|
|
class Shader:
|
|
|
|
def __init__(self, context, shader_type):
|
|
self.context = context
|
|
self.shader_type = shader_type
|
|
self.includes = []
|
|
self.ins = []
|
|
self.outs = []
|
|
self.uniforms_top = []
|
|
self.uniforms = []
|
|
self.constants = []
|
|
self.functions = {}
|
|
self.main = ''
|
|
self.main_init = ''
|
|
self.main_normal = ''
|
|
self.main_textures = ''
|
|
self.main_attribs = ''
|
|
self.header = ''
|
|
self.write_pre = False
|
|
self.write_normal = 0
|
|
self.write_textures = 0
|
|
self.tab = 1
|
|
self.vstruct_as_vsin = True
|
|
self.lock = False
|
|
self.geom_passthrough = False
|
|
self.is_linked = False # Use already generated shader
|
|
self.noprocessing = False
|
|
|
|
def has_include(self, s):
|
|
return s in self.includes
|
|
|
|
def add_include(self, s):
|
|
if not self.has_include(s):
|
|
self.includes.append(s)
|
|
|
|
def add_include_front(self, s):
|
|
if not self.has_include(s):
|
|
pos = 0
|
|
# make sure compiled.inc is always on top
|
|
if len(self.includes) > 0 and self.includes[0] == 'compiled.inc':
|
|
pos = 1
|
|
self.includes.insert(pos, s)
|
|
|
|
def add_in(self, s):
|
|
if s not in self.ins:
|
|
self.ins.append(s)
|
|
|
|
def add_out(self, s):
|
|
if s not in self.outs:
|
|
self.outs.append(s)
|
|
|
|
def add_uniform(self, s, link=None, included=False, top=False,
|
|
tex_addr_u=None, tex_addr_v=None,
|
|
tex_filter_min=None, tex_filter_mag=None,
|
|
tex_mipmap_filter=None, default_value=None, is_arm_mat_param=None):
|
|
ar = s.split(' ')
|
|
# layout(RGBA8) image3D voxels
|
|
utype = ar[-2]
|
|
uname = ar[-1]
|
|
if utype.startswith('sampler') or utype.startswith('image') or utype.startswith('uimage'):
|
|
is_image = True if (utype.startswith('image') or utype.startswith('uimage')) else None
|
|
if uname[-1] == ']': # Array of samplers - sampler2D mySamplers[2]
|
|
# Add individual units - mySamplers[0], mySamplers[1]
|
|
for i in range(int(uname[-2])):
|
|
uname_array = uname[:-2] + str(i) + ']'
|
|
self.context.add_texture_unit(
|
|
uname_array, link, is_image,
|
|
tex_addr_u, tex_addr_v,
|
|
tex_filter_min, tex_filter_mag, tex_mipmap_filter)
|
|
else:
|
|
self.context.add_texture_unit(
|
|
uname, link, is_image,
|
|
tex_addr_u, tex_addr_v,
|
|
tex_filter_min, tex_filter_mag, tex_mipmap_filter,
|
|
default_value=default_value, is_arm_mat_param=is_arm_mat_param)
|
|
else:
|
|
# Prefer vec4[] for d3d to avoid padding
|
|
if ar[0] == 'float' and '[' in ar[1]:
|
|
ar[0] = 'floats'
|
|
ar[1] = ar[1].split('[', 1)[0]
|
|
elif ar[0] == 'vec4' and '[' in ar[1]:
|
|
ar[0] = 'floats'
|
|
ar[1] = ar[1].split('[', 1)[0]
|
|
self.context.add_constant(ar[0], ar[1], link=link, default_value=default_value, is_arm_mat_param=is_arm_mat_param)
|
|
if top:
|
|
if not included and s not in self.uniforms_top:
|
|
self.uniforms_top.append(s)
|
|
elif not included and s not in self.uniforms:
|
|
self.uniforms.append(s)
|
|
|
|
def add_const(self, type_str: str, name: str, value_str: str, array_size: int = 0):
|
|
"""
|
|
Add a global constant to the shader.
|
|
|
|
Parameters
|
|
----------
|
|
type_str: str
|
|
The name of the type, like 'float' or 'vec3'. If the
|
|
constant is an array, there is no need to add `[]` to the
|
|
type
|
|
name: str
|
|
The name of the variable
|
|
value_str: str
|
|
The value of the constant as a string
|
|
array_size: int
|
|
If not 0 (default value), create an array with the given size
|
|
"""
|
|
if array_size == 0:
|
|
self.constants.append(f'{type_str} {name} = {value_str}')
|
|
elif array_size > 0:
|
|
self.constants.append(f'{type_str} {name}[{array_size}] = {type_str}[]({value_str})')
|
|
|
|
def add_function(self, s):
|
|
fname = s.split('(', 1)[0]
|
|
if fname in self.functions:
|
|
return
|
|
self.functions[fname] = s
|
|
|
|
def contains(self, s):
|
|
return s in self.main or \
|
|
s in self.main_init or \
|
|
s in self.main_normal or \
|
|
s in self.ins or \
|
|
s in self.main_textures or \
|
|
s in self.main_attribs
|
|
|
|
def replace(self, old, new):
|
|
self.main = self.main.replace(old, new)
|
|
self.main_init = self.main_init.replace(old, new)
|
|
self.main_normal = self.main_normal.replace(old, new)
|
|
self.main_textures = self.main_textures.replace(old, new)
|
|
self.main_attribs = self.main_attribs.replace(old, new)
|
|
self.uniforms = [u.replace(old, new) for u in self.uniforms]
|
|
|
|
def write_init(self, s, unique=True):
|
|
"""Prepend to the main function. If `unique` is true (default), look for other occurences first."""
|
|
if unique and self.contains(s):
|
|
return
|
|
|
|
self.main_init = '\t' + s + '\n' + self.main_init
|
|
|
|
def write(self, s):
|
|
if self.lock:
|
|
return
|
|
if self.write_textures > 0:
|
|
self.main_textures += '\t' * 1 + s + '\n'
|
|
elif self.write_normal > 0:
|
|
self.main_normal += '\t' * 1 + s + '\n'
|
|
elif self.write_pre:
|
|
self.main_init += '\t' * 1 + s + '\n'
|
|
else:
|
|
self.main += '\t' * self.tab + s + '\n'
|
|
|
|
def write_header(self, s):
|
|
self.header += s + '\n'
|
|
|
|
def write_attrib(self, s):
|
|
self.main_attribs += '\t' + s + '\n'
|
|
|
|
def is_equal(self, sh):
|
|
self.vstruct_to_vsin()
|
|
return self.ins == sh.ins and \
|
|
self.main == sh.main and \
|
|
self.main_normal == sh.main_normal and \
|
|
self.main_init == sh.main_init and \
|
|
self.main_textures == sh.main_textures and \
|
|
self.main_attribs == sh.main_attribs
|
|
|
|
def data_size(self, data):
|
|
if data == 'float1':
|
|
return '1'
|
|
elif data == 'float2':
|
|
return '2'
|
|
elif data == 'float3':
|
|
return '3'
|
|
elif data == 'float4':
|
|
return '4'
|
|
elif data == 'short2norm':
|
|
return '2'
|
|
elif data == 'short4norm':
|
|
return '4'
|
|
|
|
def vstruct_to_vsin(self):
|
|
if self.shader_type != 'vert' or self.ins != [] or not self.vstruct_as_vsin: # Vertex structure as vertex shader input
|
|
return
|
|
vs = self.context.data['vertex_elements']
|
|
for e in vs:
|
|
self.add_in('vec' + self.data_size(e['data']) + ' ' + e['name'])
|
|
|
|
def get(self):
|
|
if self.noprocessing:
|
|
return self.main
|
|
|
|
s = '#version 450\n'
|
|
|
|
s += self.header
|
|
|
|
in_ext = ''
|
|
out_ext = ''
|
|
|
|
if self.shader_type == 'vert':
|
|
self.vstruct_to_vsin()
|
|
|
|
elif self.shader_type == 'tesc':
|
|
in_ext = '[]'
|
|
out_ext = '[]'
|
|
s += 'layout(vertices = 3) out;\n'
|
|
# Gen outs
|
|
for sin in self.ins:
|
|
ar = sin.rsplit(' ', 1) # vec3 wnormal
|
|
tc_s = 'tc_' + ar[1]
|
|
self.add_out(ar[0] + ' ' + tc_s)
|
|
# Pass data
|
|
self.write('{0}[gl_InvocationID] = {1}[gl_InvocationID];'.format(tc_s, ar[1]))
|
|
|
|
elif self.shader_type == 'tese':
|
|
in_ext = '[]'
|
|
s += 'layout(triangles, equal_spacing, ccw) in;\n'
|
|
|
|
elif self.shader_type == 'geom':
|
|
in_ext = '[]'
|
|
s += 'layout(triangles) in;\n'
|
|
if not self.geom_passthrough:
|
|
s += 'layout(triangle_strip) out;\n'
|
|
s += 'layout(max_vertices=3) out;\n'
|
|
|
|
for a in self.uniforms_top:
|
|
s += 'uniform ' + a + ';\n'
|
|
for a in self.includes:
|
|
s += '#include "' + a + '"\n'
|
|
if self.geom_passthrough:
|
|
s += 'layout(passthrough) in gl_PerVertex { vec4 gl_Position; } gl_in[];\n'
|
|
for a in self.ins:
|
|
if self.geom_passthrough:
|
|
s += 'layout(passthrough) '
|
|
s += 'in {0}{1};\n'.format(a, in_ext)
|
|
for a in self.outs:
|
|
if not self.geom_passthrough:
|
|
s += 'out {0}{1};\n'.format(a, out_ext)
|
|
for a in self.uniforms:
|
|
s += 'uniform ' + a + ';\n'
|
|
for c in self.constants:
|
|
s += 'const ' + c + ';\n'
|
|
for f in self.functions:
|
|
s += self.functions[f] + '\n'
|
|
s += 'void main() {\n'
|
|
s += self.main_attribs
|
|
s += self.main_textures
|
|
s += self.main_normal
|
|
s += self.main_init
|
|
s += self.main
|
|
s += '}\n'
|
|
return s
|