oh_generate_vectors.py

# Usage: py -3 ./oh_generate_vectors.py [header filename] [header guard name]
#     eg py -3 ./oh_generate_vectors.py oh_generate_vectors.py.h OH_GENERATE_VECTORS_PY_H

from collections import namedtuple
from itertools import chain, product
from contextlib import contextmanager
import re
import sys

Prim = namedtuple('prim', 
    [        'name', 'min',       'max',       'is_float', 'is_bits', 'is_negated', 'is_bool', 'caps'])
# NOTE: min and max are unused (not sure what I had in mind for them)
primitives = {prim.name : prim for prim in [
        Prim('u8t',   '0',         'UINT8_MAX', False,      True,      False       , False,     'U8t'),
        Prim('i32t',  'INT32_MIN', 'INT32_MAX', False,      True,      True        , False,     'I32t'),
        Prim('f32t',  '-INFINITY', 'INFINITY',  True,       False,     True        , False,     'F32t'),
        Prim('f64t',  '-INFINITY', 'INFINITY',  True,       False,     True        , False,     'F64t'),
        Prim('b32t',  None,        None,        False,      False,     False       , True,      'B32t'),
    ]}

def get_prims():
    return [v for (k,v) in primitives.items()]

def convert_func_name(pFrom, pTo):
    return f'Convert{pFrom.caps}_{pTo.caps}' if (pFrom != pTo) else ''

subscript_schemes = [
        [],
        [ ('x',            ), ('r',            ), ],
        [ ('x','y',        ), ('u','v',        ), ],
        [ ('x','y','z',    ), ('r','g','b',    ), ('h','s','v',           ), ('cie_L','cie_a','cie_b',           ) ],
        [ ('x','y','z','w',), ('r','g','b','a',), ('h','s','v','__pun1_a',), ('cie_L','cie_a','cie_b','__pun2_a',) ],
    ]
int_prim = primitives['i32t']
bool_prim = primitives['b32t']

'''
  160:             vector.prim = prim
  161:             vector.n = n
  162:             vector.subSchemes = getSubscriptSchemes(n)
  163:             vector.subs = vector.subSchemes[0]
  164              vecList.append(vector)
  165:             vector.nameByLang = nameVectorByLanguage(prim, n)
  166:             vector.declareVecsInLangs = prim.declareVecsInLangs
  167:             vector.isFloatType = prim.isFloatType
  168:             vector.castUpTo = None
'''
def nljoin(l):
    return '\n'.join(l)

def cmjoin(l):
    return ', '.join(l)

def white_pfx(strings):
    strings = [s for s in strings if s.strip() != '']
    if not strings: return ''
    pat = r'\s*'
    prefix = re.match(pat, strings[0]).group(0)
    for s in strings:
        for i in range(len(prefix)):
            if prefix[i] != s[i]:
                prefix = prefix[:i]
                break
    return prefix

def reindent_txt(txt, newIndent=''):
    lines = deindent_txt_to_lines(txt)
    builder = []
    for l in lines:
        builder.append(newIndent + l)
    return nljoin(builder)

def deindent_txt_to_lines(txt):
    lines = txt.split('\n')
    oldIndent = white_pfx(lines)
    lines_out = []
    for l in lines:
        if oldIndent:
            l = l.replace(oldIndent, '', 1)
        lines_out.append(l)
    return lines_out

class CFmtr(object):
    def __init__(self):
        self._indc = 0
        self._builder = []
        return

    @property
    def _indstr(self):
        return ' ' * 4 * self._indc

    def line(self, s=''):
        self._builder.append(self._indstr + s)

    def triple(self, txt):
        lines = deindent_txt_to_lines(txt)
        self.lines(lines)


    def lines(self, l):
        for line in l:
            self.line(line)

    @contextmanager
    def document(self, s):
        self._builder.append(self._indstr + '// ' + s)
        yield

    @contextmanager
    def indent(self, txt=None, i=1):
        if txt:
            self.line(txt)
        self._indc += i
        yield
        self._indc -= i

    @contextmanager
    def scope(self, txt='', knr=True, sc=False):
        if knr:
            self.line(txt + ' {')
        else:
            self.line(txt)
            self.line('{')
        with self.indent(None):
            yield
        self.line('}' + (';' if sc else ''))

    @property
    def text(self):
        return nljoin(self._builder)

    @property
    def oneliner(self):
        return ' '.join(' '.join(l.split()) for l in self._builder)

class Vector(object):
    def __init__(self, T, n):
        self.T = T
        self.n = n

    def __eq__(self, other):
        return (self.T == other.T) and (self.n == other.n)

    @property
    def sub_schemes(self):
        return subscript_schemes[self.n]

    @property
    def subs(self):
        return self.sub_schemes[0]

    def __str__(self):
        return f"{self.T.name}{self.n}"

    @property
    def caps(self):
        return f"{self.T.caps}{self.n}"

    @property
    def ctor_name(self):
        return 'Mk'+self.caps
    

    def struct_for_subs(self, subs, fmtr=None):
        fmtr = fmtr or CFmtr()
        elements = ' '.join(f'{self.T.name} {ss};' for ss in subs)
        fmtr.line(f'struct {{ {elements} }};')
        return fmtr.text

    #to get this
    def union_def(self, fmtr=None):
        fmtr = fmtr or CFmtr()
        with fmtr.scope(f'union {self}', sc=True):
            for subs in self.sub_schemes:
                self.struct_for_subs(subs, fmtr)
            fmtr.line(f'{self.T.name} E[{self.n}];')
        return fmtr.text

    def ctor_n_prims(self, prim, fmtr=None):
        fmtr = fmtr or CFmtr()
        param_str = ', '.join(f'{prim.name} {ss}' for ss in self.subs)
        with fmtr.scope(f'inline {self} {self.ctor_name}({param_str})'):
            fmtr.line(f'{self} V;')
            fmtr.lines([f'V.{ss} = {convert_func_name(prim, self.T)}({ss});'
                for ss in self.subs])
            fmtr.line('return V;')
        return fmtr.text

    def ctor_single_prim(self, prim, fmtr=None):
        fmtr = fmtr or CFmtr()
        if False: # Multi-line
            with fmtr.scope(f'inline {self} {self.ctor_name}({prim.name} A)'):
                fmtr.line(f'{self} V;')
                fmtr.lines([f'V.{ss} = {convert_func_name(prim, self.T)}(A);'
                    for ss in self.subs])
                fmtr.line('return V;')
            return fmtr.text
        else:
            with fmtr.scope(f'inline {self} {self.ctor_name}({prim.name} A)'):
                terms = ', '.join(f'{convert_func_name(prim, self.T)}(A)' for ss in self.subs)
                fmtr.line(f'return {self.ctor_name}({terms});')
            return fmtr.oneliner

    def ctor_other_vec(self, vec, fmtr=None):
        fmtr = fmtr or CFmtr()
        if self.n != vec.n:
            raise AssertionError("Can't construct from vector of different length")
        if False: # Multi-line
            with fmtr.scope(f'inline {self} {self.ctor_name}({vec} A)'):
                fmtr.line(f'{self} V;')
                fmtr.lines([f'V.{ss} = {convert_func_name(vec.T, self.T)}(A.{ss});'
                    for ss in self.subs])
                fmtr.line('return V;')
        else:
            with fmtr.scope(f'inline {self} {self.ctor_name}({vec} A)'):
                terms = ', '.join(f'A.{ss}' for ss in self.subs)
                fmtr.line(f'return {self.ctor_name}({terms});')
            return fmtr.oneliner

    def ctor_direction(self, fmtr=None):
        fmtr = fmtr or CFmtr()
        with fmtr.document(f'Construct from card_dir'):
            with fmtr.scope(f'inline {self} {self.ctor_name}(card_dir D)'):
                fmtr.line(f'{self} V;')
                fmtr.line(f'V.E[D/2] = {convert_func_name(int_prim, self.T)}(D%2 ? 1 : -1);')
                fmtr.line('return V;')
        return fmtr.text

    def unary_el_op(self, name, op, fmtr=None):
        fmtr = fmtr or CFmtr()
        with fmtr.scope(f'inline {self} {name}({self} A)'):
            with fmtr.indent(f'return {self.ctor_name}('):
                fmtr.lines([op(f'A.{ss}') + (',' if ss != self.subs[-1] else ');')
                    for ss in self.subs])
        return fmtr.oneliner

    def bin_el_op(self, name, op, rVec=None, fmtr=None):
        rVec = rVec or self
        fmtr = fmtr or CFmtr()
        with fmtr.scope(f'inline {rVec} {name}({self} A, {self} B)'):
            with fmtr.indent(f'return {rVec.ctor_name}('):
                fmtr.lines([op(f'A.{ss}', f'B.{ss}') + (',' if ss != self.subs[-1] else ');') 
                    for ss in self.subs])
        return fmtr.oneliner

    def scal_op_r(self, name, op, rVec=None, fmtr=None):
        rVec = rVec or self
        fmtr = fmtr or CFmtr()
        with fmtr.scope(f'inline {rVec} {name}({self} V, {self.T.name} A)'):
            with fmtr.indent(f'return {rVec.ctor_name}('):
                fmtr.lines([op(f'V.{ss}', 'A') + (',' if ss != self.subs[-1] else ');')
                    for ss in self.subs])
        return fmtr.oneliner

    def scal_op_l(self, name, op, rVec=None, fmtr=None):
        rVec = rVec or self
        fmtr = fmtr or CFmtr()
        with fmtr.scope(f'inline {rVec} {name}({self.T.name} A, {self} V)'):
            with fmtr.indent(f'return {rVec.ctor_name}('):
                fmtr.lines([op('A', f'V.{ss}') + (',' if ss != self.subs[-1] else ');')
                    for ss in self.subs])
        return fmtr.oneliner

    def bin_and_scalars(self, name, op, rVec=None):
        rVec = rVec or self
        return '\n'.join([
            self.bin_el_op(name, op, rVec),
            self.scal_op_r(name, op, rVec),
            self.scal_op_l(name, op, rVec),
        ])

    def compound(self, s):
        return '\n'.join([
            f'inline {self}& operator{s}=({self}& Me, {self} A) {{ return (Me = Me {s} A); }}',
            f'inline {self}& operator{s}=({self}& Me, {self.T.name} A) {{ return (Me = Me {s} A); }}'
            ])

    def horiz_op(self, name, op, rtype=None, fmtr=None):
        rtype = rtype or str(self.T.name)
        fmtr = fmtr or CFmtr()
        with fmtr.scope(f'inline {rtype} {name}({self} V)'):
            fmtr.line('return ' + op([f'V.{ss}' for ss in self.subs]) + ';')
        return fmtr.oneliner

    # NOTE there's so many of these it's better to just have it format a block of one-liners *shrug*
    def swizzles(self, fmtr=None):
        fmtr = fmtr or CFmtr()
        for nTo in range(2,5):
            vTo = Vector(self.T, nTo)
            for ssCombo in product(list(self.subs) + ['_'], repeat=vTo.n):
                if (ssCombo.count('_') != len(ssCombo)):
                    params = ', '.join((f'({vTo.T.name})0' if (ss == '_') else f'V.{ss}') for ss in ssCombo)
                    name = 'Swizz' + ''.join(ssCombo).upper()
                    fmtr.line(f'inline {vTo} {name}({self} V) {{ return {vTo.ctor_name}({params}); }}')
        return fmtr.text


def print_vector_h(filename, header_guard):
    file_chunks = []
    vector_types = [Vector(p, n) for p in get_prims() for n in range(2,5)]
    nonbool_vector_types = [v for v in vector_types if not v.T.is_bool]
    bool_vector_types =    [v for v in vector_types if v.T.is_bool]
    negated_vector_types = [v for v in vector_types if v.T.is_negated]
    float_vector_types =   [v for v in vector_types if v.T.is_float]
    int_vector_types =     [v for v in vector_types if v.T == int_prim]
    nonbool_prims = [p for p in get_prims() if not p.is_bool]
    
    for v in vector_types:
        file_chunks.append(v.union_def() + '\n')
    
    for v in nonbool_vector_types:
        for p in nonbool_prims:
            file_chunks.append(v.ctor_n_prims(p) + '\n')
    
    for v in bool_vector_types:
        file_chunks.append(v.ctor_n_prims(bool_prim) + '\n')
    
    for v in negated_vector_types:
        file_chunks.append(v.ctor_direction() + '\n')
    
    for v in nonbool_vector_types:
        for p in nonbool_prims:
            file_chunks.append(v.ctor_single_prim(p))
            file_chunks.append(v.ctor_other_vec(Vector(p,v.n)))
            file_chunks.append('')
    
    for v in bool_vector_types:
        file_chunks.append(v.ctor_single_prim(bool_prim))
        file_chunks.append('')
    
    for v in vector_types:
        file_chunks.append(v.bin_and_scalars('operator!=', lambda a,b: f'{a} != {b}', Vector(bool_prim, v.n)))
        file_chunks.append(v.bin_and_scalars('operator==', lambda a,b: f'{a} == {b}', Vector(bool_prim, v.n)))
        file_chunks.append('')
    
    for v in bool_vector_types:
        file_chunks.append(v.unary_el_op('operator!', lambda a:   f'!{a}' ))
        file_chunks.append(v.bin_and_scalars('operator&&', lambda a,b: f'{a} && {b}'))
        file_chunks.append(v.bin_and_scalars('operator||', lambda a,b: f'{a} || {b}'))
        file_chunks.append(v.horiz_op('HzAnd', lambda l: '&&'.join(l)))
        file_chunks.append(v.horiz_op('HzOr',  lambda l: '||'.join(l)))
        file_chunks.append('')
    
    for v in nonbool_vector_types:
        file_chunks.append(v.unary_el_op('operator-', lambda a:   f'-{a}' ))
        file_chunks.append(v.bin_and_scalars('operator+',  lambda a,b: f'{a} + {b}'))
        file_chunks.append(v.bin_and_scalars('operator-',  lambda a,b: f'{a} - {b}'))
        file_chunks.append(v.bin_and_scalars('operator*',  lambda a,b: f'{a} * {b}'))
        file_chunks.append(v.bin_and_scalars('operator/',  lambda a,b: f'Div({a},{b})'))
        file_chunks.append(v.bin_and_scalars('operator<',  lambda a,b: f'{a} < {b}',  Vector(bool_prim, v.n)))
        file_chunks.append(v.bin_and_scalars('operator<=', lambda a,b: f'{a} <= {b}', Vector(bool_prim, v.n)))
        file_chunks.append(v.bin_and_scalars('operator>',  lambda a,b: f'{a} > {b}',  Vector(bool_prim, v.n)))
        file_chunks.append(v.bin_and_scalars('operator>=', lambda a,b: f'{a} >= {b}', Vector(bool_prim, v.n)))
        file_chunks.append(v.bin_and_scalars('ElMin',      lambda a,b: f'Min({a},{b})'))
        file_chunks.append(v.bin_and_scalars('ElMax',      lambda a,b: f'Max({a},{b})'))
        file_chunks.append('')
    
    for v in nonbool_vector_types:
        file_chunks.append(v.compound('+'))
        file_chunks.append(v.compound('-'))
        file_chunks.append(v.compound('*'))
        file_chunks.append(v.compound('/'))
        file_chunks.append(v.horiz_op('HzAdd', lambda l: '+'.join(l)))
        file_chunks.append(v.horiz_op('HzMlt', lambda l: '*'.join(l)))
        file_chunks.append(v.horiz_op('HzMin', lambda l: f'VARF_{v.n}(Min, {cmjoin(l)});'))
        file_chunks.append(v.horiz_op('HzMax', lambda l: f'VARF_{v.n}(Max, {cmjoin(l)});'))
        file_chunks.append('')
    
    for v in int_vector_types:
        file_chunks.append(v.bin_and_scalars('ElCeilDiv',  lambda a,b: f'CeilDiv({a},{b})'))
        file_chunks.append(v.bin_and_scalars('ElPow2CDiv', lambda a,b: f'Pow2CDiv({a},{b})'))
        file_chunks.append(v.bin_and_scalars('ElPow2Div',  lambda a,b: f'Pow2Div({a},{b})'))
        file_chunks.append(v.bin_and_scalars('ElPow2Mlt',  lambda a,b: f'Pow2Mlt({a},{b})'))
        file_chunks.append(v.bin_and_scalars('ElPow2Rem',  lambda a,b: f'Pow2Rem({a},{b})'))
        file_chunks.append(v.bin_and_scalars('ElRem',      lambda a,b: f'Rem({a},{b})'))
        file_chunks.append('')
    
    for v in float_vector_types:
        file_chunks.append(v.unary_el_op('Ceil',      lambda a:   f'Ceil({a})' ))
        file_chunks.append(f'inline {v.T.name} Len({v} V) {{ return Sqrt(HzAdd(V*V)); }}')
        file_chunks.append(f'inline {v} Norm({v} V) {{ return V/Len(V); }}')
        file_chunks.append('')
    
    for v in nonbool_vector_types:
        file_chunks.append(f'inline {bool_prim.name} AllLt({v} A, {v} B)  {{ return HzAnd(A < B); }}')
        file_chunks.append(f'inline {bool_prim.name} AllLte({v} A, {v} B) {{ return HzAnd(A <= B); }}')
        file_chunks.append(f'inline {bool_prim.name} AllGt({v} A, {v} B)  {{ return HzAnd(A > B); }}')
        file_chunks.append(f'inline {bool_prim.name} AllGte({v} A, {v} B) {{ return HzAnd(A >= B); }}')
        file_chunks.append(f'inline {v.T.name} DotP({v} A, {v} B) {{ return HzAdd(A * B); }}')
        file_chunks.append(f'inline {v.T.name} SqLen({v} V) {{ return HzAdd(V*V); }}')
        file_chunks.append('')
    
    for v in nonbool_vector_types:
        if v.n == 3:
            file_chunks.append(f'inline {v} CrossP({v} A, {v} B)  {{ return {v.ctor_name}' \
                + '( (A.y * B.z) - (A.z * B.y), -(A.x * B.z) + (A.z * B.x), (A.x * B.y) - (A.y * B.x) )' \
                + '; }')
    file_chunks.append('')
    for v in vector_types:
        file_chunks.append(v.swizzles())
        file_chunks.append('\n')

    with open(filename, 'w') as file:
        file.write(f'#ifndef {header_guard}\n')
        file.write(f'#define {header_guard}\n')
        file.write('#pragma warning(disable: 4201)\n')
        code = '\n'.join(file_chunks)
        file.write(code + '\n')
        file.write(f'#endif // {header_guard}')

if __name__ == "__main__":
   print_vector_h(sys.argv[1], sys.argv[2])