LlvmGenerator.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
    OpenGEODE - A tiny SDL Editor for TASTE

    This module generates LLVM IR code from SDL process models, allowing
    generation of a binary application without an intermediate language.
    LLVM also allows for various code verification, analysis, and optimization.

    The design is based on the Ada code generator. Check it for details.

    Copyright (c) 2012-2013 European Space Agency

    Designed and implemented by Maxime Perrotin

    Contact: maxime.perrotin@esa.int
"""

import logging

from singledispatch import singledispatch

from llvm import core as lc
from llvm import ee as le

import ogAST
import Helper

LOG = logging.getLogger(__name__)

__all__ = ['generate']


class Context():
    def __init__(self, process):
        self.name = str(process.processName)
        self.process = process
        self.module = lc.Module.new(self.name)
        self.target_data = le.TargetData.new(self.module.data_layout)
        self.dataview = process.dataview
        self.procedures = process.procedures

        self.scope = Scope(self)
        self.global_scope = self.scope
        self.states = {}
        self.enums = {}
        self.structs = {}
        self.unions = {}
        self.strings = {}
        self.funcs = {}
        self.lltypes = {}
        self.basic_asn1types = {}

        # Initialize built-in types
        self.i1 = lc.Type.int(1)
        self.i8 = lc.Type.int(8)
        self.i32 = lc.Type.int(32)
        self.i64 = lc.Type.int(64)
        self.void = lc.Type.void()
        self.double = lc.Type.double()
        self.i1_ptr = lc.Type.pointer(self.i1)
        self.i8_ptr = lc.Type.pointer(self.i8)
        self.i32_ptr = lc.Type.pointer(self.i32)
        self.i64_ptr = lc.Type.pointer(self.i64)
        self.double_ptr = lc.Type.pointer(self.double)

        # Initialize common constants
        self.zero = lc.Constant.int(self.i32, 0)
        self.one = lc.Constant.int(self.i32, 1)

        # Intialize built-in functions
        ty = lc.Type.function(self.void, [self.i8_ptr], True)
        self.funcs['printf'] = self.module.add_function(ty, 'printf')

        ty = lc.Type.function(self.double, [self.double])
        self.funcs['round'] = self.module.add_function(ty, 'round')

        self.funcs['memcpy'] = lc.Function.intrinsic(
            self.module,
            lc.INTR_MEMCPY,
            [self.i8_ptr, self.i8_ptr, self.i64]
        )

        self.funcs['powi'] = lc.Function.intrinsic(
            self.module,
            lc.INTR_POWI,
            [self.double]
        )

        self.funcs['ceil'] = lc.Function.intrinsic(
            self.module,
            lc.INTR_CEIL,
            [self.double]
        )

        self.funcs['cos'] = lc.Function.intrinsic(
            self.module,
            lc.INTR_COS,
            [self.double]
        )

        self.funcs['fabs'] = lc.Function.intrinsic(
            self.module,
            lc.INTR_FABS,
            [self.double]
        )

        self.funcs['floor'] = lc.Function.intrinsic(
            self.module,
            lc.INTR_FLOOR,
            [self.double]
        )

        self.funcs['sin'] = lc.Function.intrinsic(
            self.module,
            lc.INTR_SIN,
            [self.double]
        )

        self.funcs['sqrt'] = lc.Function.intrinsic(
            self.module,
            lc.INTR_SQRT,
            [self.double]
        )

        self.funcs['trunc'] = lc.Function.intrinsic(
            self.module,
            lc.INTR_TRUNC,
            [self.double]
        )

    def open_scope(self):
        ''' Open a scope '''
        self.scope = Scope(self, self.scope)

    def close_scope(self):
        ''' Close the current scope '''
        self.scope = self.scope.parent

    def basic_asn1type_of(self, asn1ty):
        ''' Return the ASN.1 basic type of a type '''
        if asn1ty.kind != 'ReferenceType':
            return asn1ty

        asn1ty_name = asn1ty.ReferencedTypeName.lower()

        # return the basic type if its cached
        if asn1ty_name in self.basic_asn1types:
            return self.basic_asn1types[asn1ty_name]

        basic_asn1ty = asn1ty
        while basic_asn1ty.kind == 'ReferenceType':
            for typename in self.dataview.viewkeys():
                if typename.lower() == basic_asn1ty.ReferencedTypeName.lower():
                    basic_asn1ty = self.dataview[typename].type
                    break

        # cache the basic type
        self.basic_asn1types[asn1ty_name] = basic_asn1ty

        return basic_asn1ty

    def lltype_of(self, asn1ty):
        ''' Return the LL type of a ASN.1 type '''
        try:
            name = asn1ty.ReferencedTypeName.replace('-', '_')
        except AttributeError:
            name = None

        if name and name in self.lltypes:
            return self.lltypes[name]

        basic_asn1ty = self.basic_asn1type_of(asn1ty)

        if basic_asn1ty.kind == 'IntegerType':
            llty = self.i64
        elif basic_asn1ty.kind == 'Integer32Type':
            llty = self.i32
        elif basic_asn1ty.kind == 'BooleanType':
            llty = self.i1
        elif basic_asn1ty.kind == 'RealType':
            llty = self.double
        elif basic_asn1ty.kind == 'SequenceOfType':
            llty = self._lltype_of_sequenceof(name, basic_asn1ty)
        elif basic_asn1ty.kind == 'SequenceType':
            llty = self._lltype_of_sequence(name, basic_asn1ty)
        elif basic_asn1ty.kind == 'EnumeratedType':
            llty = self.i32
        elif basic_asn1ty.kind == 'ChoiceType':
            llty = self._lltype_of_choice(name, basic_asn1ty)
        elif basic_asn1ty.kind == 'OctetStringType':
            llty = self._lltype_of_octetstring(name, basic_asn1ty)
        elif basic_asn1ty.kind in ('StringType', 'StandardStringType'):
            llty = self.i8_ptr
        else:
            raise CompileError('Unknown basic ASN.1 type "%s"' % basic_asn1ty.kind)

        if name:
            self.lltypes[name] = llty

        return llty

    def _lltype_of_sequenceof(self, name, asn1ty):
        ''' Return the LL type of a SequenceOf ASN.1 type '''
        min_size = int(asn1ty.Min)
        max_size = int(asn1ty.Max)
        is_variable_size = min_size != max_size

        elem_llty = self.lltype_of(asn1ty.type)
        array_llty = lc.Type.array(elem_llty, max_size)

        if is_variable_size:
            struct = self.decl_struct(['nCount', 'arr'], [self.i32, array_llty], name)
        else:
            struct = self.decl_struct(['arr'], [array_llty], name)

        struct_ptr = lc.Type.pointer(struct.llty)
        self.decl_func("asn1Scc%s_Equal" % name, self.i1, [struct_ptr, struct_ptr])

        return struct.llty

    def _lltype_of_sequence(self, name, asn1ty):
        ''' Return the LL type of a Sequence ASN.1 type '''
        field_names = []
        field_lltys = []

        for field_name in Helper.sorted_fields(asn1ty):
            field_names.append(field_name.replace('-', '_'))
            field_lltys.append(self.lltype_of(asn1ty.Children[field_name].type))

        struct = self.decl_struct(field_names, field_lltys, name)

        struct_ptr = lc.Type.pointer(struct.llty)
        self.decl_func("asn1Scc%s_Equal" % name, self.i1, [struct_ptr, struct_ptr])

        return struct.llty

    def _lltype_of_choice(self, name, asn1ty):
        ''' Return the equivalent LL type of a Choice ASN.1 type '''
        field_names = []
        field_lltys = []

        for idx, field_name in enumerate(Helper.sorted_fields(asn1ty)):
            # enum values used in choice determinant/present
            self.enums[field_name.replace('-', '_')] = lc.Constant.int(self.i32, idx)

            field_names.append(field_name.replace('-', '_'))
            field_lltys.append(self.lltype_of(asn1ty.Children[field_name].type))

        union = self.decl_union(field_names, field_lltys, name)

        union_ptr = lc.Type.pointer(union.llty)
        self.decl_func("asn1Scc%s_Equal" % name, self.i1, [union_ptr, union_ptr])

        return union.llty

    def _lltype_of_octetstring(self, name, asn1ty):
        ''' Return the equivalent LL type of a OctetString ASN.1 type '''
        min_size = int(asn1ty.Min)
        max_size = int(asn1ty.Max)
        is_variable_size = min_size != max_size

        array_llty = lc.Type.array(self.i8, max_size)

        if is_variable_size:
            struct = self.decl_struct(['nCount', 'arr'], [self.i32, array_llty], name)
        else:
            struct = self.decl_struct(['arr'], [array_llty], name)

        struct_ptr = lc.Type.pointer(struct.llty)
        self.decl_func("asn1Scc%s_Equal" % name, self.i1, [struct_ptr, struct_ptr])

        return struct.llty

    def string_ptr(self, str):
        ''' Returns a pointer to a global string with the given value '''
        if str in self.strings:
            return self.strings[str].gep([self.zero, self.zero])

        str_val = lc.Constant.stringz(str)
        var_name = '.str%s' % len(self.strings)
        var_ptr = self.module.add_global_variable(str_val.type, var_name)
        var_ptr.initializer = str_val
        self.strings[str] = var_ptr
        return var_ptr.gep([self.zero, self.zero])

    def decl_func(self, name, return_llty, param_lltys, extern=False):
        ''' Declare a function '''
        func_llty = lc.Type.function(return_llty, param_lltys)
        func_name = ("%s_RI_%s" % (self.name, name)) if extern else name
        func = lc.Function.new(self.module, func_llty, func_name)
        self.funcs[name.lower()] = func
        return func

    def decl_struct(self, field_names, field_lltys, name=None):
        ''' Declare a struct '''
        name = name if name else "struct.%s" % len(self.structs)
        name = name.replace('-', '_')
        struct = StructType(name, field_names, field_lltys)
        self.structs[name] = struct
        return struct

    def resolve_struct(self, name):
        ''' Return the struct associated to a name '''
        return self.structs[name.replace('-', '_')]

    def decl_union(self, field_names, field_lltys, name=None):
        name = name if name else "union.%s" % len(self.structs)
        name = name.replace('-', '_')
        union = UnionType(name, field_names, field_lltys, self)
        self.unions[name] = union
        return union

    def resolve_union(self, name):
        ''' Return the union associated to a name '''
        return self.unions[name.replace('-', '_')]


class StructType():
    def __init__(self, name, field_names, field_lltys):
        self.name = name
        self.field_names = field_names
        self.llty = lc.Type.struct(field_lltys, self.name)

    def idx(self, field_name):
        return self.field_names.index(field_name)


class UnionType():
    def __init__(self, name, field_names, field_lltys, ctx):
        self.name = name
        self.field_names = field_names
        self.field_lltys = field_lltys
        # Unions are represented a struct with a field indicating the index of its type
        # and a byte array with the size of the biggest type in the union
        self.size = max([ctx.target_data.size(ty) for ty in field_lltys])
        self.llty = lc.Type.struct([ctx.i32, lc.Type.array(ctx.i8, self.size)], name)

    def kind(self, name):
        idx = self.field_names.index(name)
        return (idx, self.field_lltys[idx])


class Scope:
    def __init__(self, ctx, parent=None):
        self.ctx = ctx
        self.vars = {}
        self.labels = {}
        self.parent = parent

    def define(self, name, var):
        self.vars[name.lower()] = var

    def resolve(self, name):
        var = self.vars.get(name.lower())
        if var:
            return var
        if self.parent:
            return self.parent.resolve(name)
        else:
            raise NameError("name '%s' is not defined" % name)

    def label(self, name):
        name = name.lower()
        label_block = self.labels.get(name)
        if not label_block:
            func = self.ctx.builder.basic_block.function
            label_block = func.append_basic_block('label:%s' % name)
            self.labels[name] = label_block
        return label_block


class CompileError(Exception):
    pass


@singledispatch
def generate(ast, ctx=None):
    ''' Generate the IR for an AST node '''
    raise CompileError('Unsupported AST construct "%s"' % ast.__class__.__name__)


# Processing of the AST
@generate.register(ogAST.Process)
def _process(process, ctx=None):
    ''' Generate the IR for a process '''
    process_name = str(process.processName)
    LOG.info('Generating LLVM IR code for process ' + process_name)

    ctx = Context(process)

    # In case model has nested states, flatten everything
    Helper.flatten(process, '.')

    # Make an maping {input: {state: transition...}} in order to easily
    # generate the lookup tables for the state machine runtime
    mapping = Helper.map_input_state(process)

    # Initialize states
    for name, val in process.mapping.viewitems():
        if not name.endswith('START'):
            cons_val = lc.Constant.int(ctx.i32, len(ctx.states))
            ctx.states[name.lower()] = cons_val
        elif name != 'START':
            cons_val = lc.Constant.int(ctx.i32, val)
            ctx.states[name.lower()] = cons_val

    # Generate state var
    state_cons = ctx.module.add_global_variable(ctx.i32, '.state')
    state_cons.initializer = lc.Constant.int(ctx.i32, -1)
    ctx.scope.define('.state', state_cons)

    # Generare process-level vars
    for name, (asn1ty, expr) in process.variables.viewitems():
        var_llty = ctx.lltype_of(asn1ty)
        global_var = ctx.module.add_global_variable(var_llty, str(name))
        global_var.initializer = lc.Constant.null(var_llty)
        ctx.scope.define(str(name).lower(), global_var)

    # Declare set/reset timer functions
    for timer in process.timers:
        # TODO: Should be uint?
        ctx.decl_func("set_%s" % str(timer), ctx.void, [ctx.i64_ptr], True)
        ctx.decl_func("reset_%s" % str(timer), ctx.void, [], True)

    # Declare output signal functions
    for signal in process.output_signals:
        if 'type' in signal:
            param_lltys = [lc.Type.pointer(ctx.lltype_of(signal['type']))]
        else:
            param_lltys = []
        ctx.decl_func(str(signal['name']), ctx.void, param_lltys, True)

    # Declare external procedures functions
    for proc in [proc for proc in process.procedures if proc.external]:
        param_lltys = [lc.Type.pointer(ctx.lltype_of(p['type'])) for p in proc.fpar]
        ctx.decl_func(str(proc.inputString), ctx.void, param_lltys, True)

    # Generate internal procedures
    for proc in process.content.inner_procedures:
        generate(proc, ctx)

    # Generate process functions
    generate_runtr_func(process, ctx)
    generate_startup_func(process, ctx)

    # Generate input signals
    for signal in process.input_signals:
        generate_input_signal(signal, mapping[signal['name']], ctx)

    # Generate timer signal
    for timer in process.timers:
        generate_input_signal({'name': timer.lower()}, mapping[timer], ctx)

    ctx.module.verify()

    with open(ctx.name + '.ll', 'w') as ll_file:
        ll_file.write(str(ctx.module))


def generate_runtr_func(process, ctx):
    ''' Generate the IR for the run_transition function '''
    func = ctx.decl_func('run_transition', ctx.void, [ctx.i32])

    ctx.open_scope()

    entry_block = func.append_basic_block('runtr:entry')
    cond_block = func.append_basic_block('runtr:cond')
    body_block = func.append_basic_block('runtr:body')
    exit_block = func.append_basic_block('runtr:exit')

    ctx.builder = lc.Builder.new(entry_block)

    # entry
    id_ptr = ctx.builder.alloca(ctx.i32, None, 'id')
    ctx.scope.define('id', id_ptr)
    ctx.builder.store(func.args[0], id_ptr)
    ctx.builder.branch(cond_block)

    # cond
    ctx.builder.position_at_end(cond_block)
    no_tr_cons = lc.Constant.int(ctx.i32, -1)
    id_val = ctx.builder.load(id_ptr)
    cond_val = ctx.builder.icmp(lc.ICMP_NE, id_val, no_tr_cons, 'cond')
    ctx.builder.cbranch(cond_val, body_block, exit_block)

    # body
    ctx.builder.position_at_end(body_block)
    switch = ctx.builder.switch(id_val, exit_block)

    # transitions
    for idx, tr in enumerate(process.transitions):
        tr_block = func.append_basic_block('runtr:tr%d' % idx)
        const = lc.Constant.int(ctx.i32, idx)
        switch.add_case(const, tr_block)
        ctx.builder.position_at_end(tr_block)
        generate(tr, ctx)
        if not ctx.builder.basic_block.terminator:
            ctx.builder.branch(cond_block)

    # exit
    ctx.builder.position_at_end(exit_block)
    ctx.builder.ret_void()

    Helper.inner_labels_to_floating(process)
    for label in process.content.floating_labels:
        generate(label, ctx)

    next_tr_label_block = ctx.scope.label('next_transition')
    ctx.builder.position_at_end(next_tr_label_block)
    ctx.builder.branch(cond_block)

    ctx.close_scope()

    func.verify()
    return func


def generate_startup_func(process, ctx):
    ''' Generate the IR for the startup function '''
    func = ctx.decl_func(ctx.name + '_startup', ctx.void, [])

    ctx.open_scope()

    entry_block = func.append_basic_block('startup:entry')
    ctx.builder = lc.Builder.new(entry_block)

    # Initialize process level variables
    for name, (ty, expr) in process.variables.viewitems():
        if expr:
            global_var = ctx.scope.resolve(str(name))
            sdl_assign(global_var, expression(expr, ctx), ctx)

    sdl_call('run_transition', [lc.Constant.int(ctx.i32, 0)], ctx)
    ctx.builder.ret_void()

    ctx.close_scope()

    func.verify()
    return func


def generate_input_signal(signal, inputs, ctx):
    ''' Generate the IR for an input signal '''
    func_name = ctx.name + "_" + str(signal['name'])
    param_lltys = []
    if 'type' in signal:
        param_lltys.append(lc.Type.pointer(ctx.lltype_of(signal['type'])))

    func = ctx.decl_func(func_name, ctx.void, param_lltys)

    ctx.open_scope()

    entry_block = func.append_basic_block('input:entry')
    exit_block = func.append_basic_block('input:exit')
    ctx.builder = lc.Builder.new(entry_block)

    g_state_val = ctx.builder.load(ctx.global_scope.resolve('.state'))
    switch = ctx.builder.switch(g_state_val, exit_block)

    for state_name, state_id in ctx.states.iteritems():
        if state_name.endswith('start'):
            continue

        state_block = func.append_basic_block('input:state_%s' % str(state_name))
        switch.add_case(state_id, state_block)
        ctx.builder.position_at_end(state_block)

        state_input = inputs.get(state_name)
        if not state_input:
            ctx.builder.ret_void()
            continue

        trans = state_input.transition

        for exit_func_name in exit_list(state_name, ctx.process):
            if trans and all(exit_func_name.startswith(trans_st)
                             for trans_st in trans.possible_states):
                sdl_call(exit_func_name, [], ctx)

        for var_name in state_input.parameters:
            var_ptr = ctx.scope.resolve(str(var_name))
            if is_struct_ptr(var_ptr) or is_array_ptr(var_ptr):
                sdl_assign(var_ptr, func.args[0], ctx)
            else:
                sdl_assign(var_ptr, ctx.builder.load(func.args[0]), ctx)

        if trans:
            id_val = lc.Constant.int(ctx.i32, state_input.transition_id)
            sdl_call('run_transition', [id_val], ctx)

        ctx.builder.ret_void()

    ctx.builder.position_at_end(exit_block)
    ctx.builder.ret_void()

    ctx.close_scope()

    func.verify()


def exit_list(state_name, process):
    ''' Calculate the exit call list of a state '''
    context = process
    exitlist = []
    current = ''
    state_tree = state_name.split('.')

    while state_tree:
        current = current + state_tree.pop(0)
        for comp in context.composite_states:
            if current.lower() == comp.statename.lower():
                if comp.exit_procedure:
                    exitlist.append(current + '.exit')
                context = comp
                current = current + '.'
                break

    return reversed(exitlist)


@generate.register(ogAST.Output)
def _output(output, ctx):
    ''' Generate the IR for an output '''
    for out in output.output:
        name = out['outputName'].lower()
        args = out.get('params', [])

        arg_vals = []
        for arg in args:
            arg_val = expression(arg, ctx)
            # Pass by reference
            if arg_val.type.kind != lc.TYPE_POINTER:
                arg_var = ctx.builder.alloca(arg_val.type, None)
                ctx.builder.store(arg_val, arg_var)
                arg_vals.append(arg_var)
            else:
                arg_vals.append(arg_val)

        sdl_call(str(name).lower(), arg_vals, ctx)


@generate.register(ogAST.ProcedureCall)
def _proc_call(proc_call, ctx):
    ''' Generate the IR for a procedure call '''
    output = proc_call.output[0]

    name = output['outputName'].lower()
    args = output.get('params', [])

    if name == 'write' or name == 'writeln':
        arg_vals = [expression(a, ctx) for a in args]
        arg_asn1tys = [a.exprType for a in args]
        sdl_write(arg_vals, arg_asn1tys, ctx, name == 'writeln')
        return
    elif name == 'reset_timer':
        sdl_reset_timer(args[0].value[0], ctx)
        return
    elif name == 'set_timer':
        timer_expr, timer_id = args
        sdl_set_timer(timer_id.value[0], expression(timer_expr, ctx), ctx)
        return

    proc = None
    for p in ctx.procedures:
        if p.inputString.lower() == name:
            proc = p
            break
    else:
        raise CompileError('Procedure "%s" not found' % name)

    arg_vals = []
    for arg, param in zip(args, proc.fpar):
        if param['direction'] == 'out':
            arg_vals.append(reference(arg, ctx))
        else:
            arg_val = expression(arg, ctx)
            # Pass by reference
            if arg_val.type.kind != lc.TYPE_POINTER:
                arg_var = ctx.builder.alloca(arg_val.type, None)
                ctx.builder.store(arg_val, arg_var)
                arg_vals.append(arg_var)
            else:
                arg_vals.append(arg_val)

    sdl_call(str(name).lower(), arg_vals, ctx)


@generate.register(ogAST.TaskAssign)
def _task_assign(task, ctx):
    ''' Generate the IR for a list of assignments '''
    for expr in task.elems:
        expression(expr, ctx)


@generate.register(ogAST.TaskInformalText)
def _task_informal_text(task, ctx):
    ''' Generate comments for informal text '''
    pass


@generate.register(ogAST.TaskForLoop)
def _task_forloop(task, ctx):
    ''' Generate the IRfor a for loop '''
    for loop in task.elems:
        if loop['range']:
            generate_for_range(loop, ctx)
        else:
            generate_for_iterable(loop, ctx)


def generate_for_range(loop, ctx):
    ''' Generate the IR for a for x in range loop '''
    func = ctx.builder.basic_block.function
    cond_block = func.append_basic_block('for:cond')
    body_block = func.append_basic_block('for:body')
    inc_block = func.append_basic_block('for:inc')
    end_block = func.append_basic_block('for:end')

    ctx.open_scope()

    loop_var = ctx.builder.alloca(ctx.i64, None, str(loop['var']))
    ctx.scope.define(str(loop['var']), loop_var)

    if loop['range']['start']:
        start_val = expression(loop['range']['start'], ctx)
        ctx.builder.store(start_val, loop_var)
    else:
        ctx.builder.store(lc.Constant.int(ctx.i64, 0), loop_var)

    stop_val = expression(loop['range']['stop'], ctx)
    ctx.builder.branch(cond_block)

    ctx.builder.position_at_end(cond_block)
    loop_val = ctx.builder.load(loop_var)
    cond_val = ctx.builder.icmp(lc.ICMP_SLT, loop_val, stop_val)
    ctx.builder.cbranch(cond_val, body_block, end_block)

    ctx.builder.position_at_end(body_block)
    generate(loop['transition'], ctx)
    ctx.builder.branch(inc_block)

    ctx.builder.position_at_end(inc_block)
    step_val = lc.Constant.int(ctx.i64, loop['range']['step'])
    loop_val = ctx.builder.load(loop_var)
    temp_val = ctx.builder.add(loop_val, step_val)
    ctx.builder.store(temp_val, loop_var)
    ctx.builder.branch(cond_block)

    ctx.builder.position_at_end(end_block)

    ctx.close_scope()


def generate_for_iterable(loop, ctx):
    ''' Generate the IR for a for x in iterable loop '''
    seqof_asn1ty = ctx.basic_asn1type_of(loop['list'].exprType)