modelverse/kernel/modelverse_jit/intrinsics.py

import time
import modelverse_jit.jit as jit
import modelverse_jit.tree_ir as tree_ir
import modelverse_jit.cfg_ir as cfg_ir
import modelverse_jit.runtime as jit_runtime

BINARY_INTRINSICS = {
    'value_eq' : '==',
    'value_neq' : '!=',

    'bool_and' : 'and',
    'bool_or' : 'or',

    'integer_addition' : '+',
    'integer_subtraction' : '-',
    'integer_multiplication' : '*',
    'integer_division' : '/',
    'integer_gt' : '>',
    'integer_gte' : '>=',
    'integer_lt' : '<',
    'integer_lte' : '<=',

    'float_addition' : '+',
    'float_subtraction' : '-',
    'float_multiplication' : '*',
    'float_division' : '/',
    'float_gt' : '>',
    'float_gte' : '>=',
    'float_lt' : '<',
    'float_lte' : '<='
}

UNARY_INTRINSICS = {
    'bool_not' : 'not',
    'integer_neg' : '-',
    'float_neg' : '-'
}

CAST_INTRINSICS = {
    'cast_i2f' : float,
    'cast_i2s' : str,
    'cast_i2b' : bool,
    'cast_f2i' : int,
    'cast_f2s' : str,
    'cast_f2b' : bool,
    'cast_s2i' : int,
    'cast_s2f' : float,
    'cast_s2b' : bool,
    'cast_b2i' : int,
    'cast_b2f' : float,
    'cast_b2s' : str
}

def create_get_length(expression):
    """Creates an expression that evaluates the given expression, and then
       computes the length of its result."""
    return tree_ir.CallInstruction(
        tree_ir.LoadGlobalInstruction('len'),
        [expression])

# Don't compain about the variable names, pylint. It's important that we
# get them right.
# pylint: disable=I0011,C0103
def __set_add(a, b):
    store_a, load_a = tree_ir.evaluate_and_load(a)
    return tree_ir.create_block(
        store_a,
        tree_ir.CreateEdgeInstruction(load_a, b),
        load_a)

def __dict_add(a, b, c):
    store_a, load_a = tree_ir.evaluate_and_load(a)
    store_b, load_b = tree_ir.evaluate_and_load(b)
    return tree_ir.create_block(
        store_a,
        store_b,
        tree_ir.CreateEdgeInstruction(
            tree_ir.CreateEdgeInstruction(load_a, c),
            load_b),
        load_a)

def __dict_add_fast(a, b, c):
    # TODO This might be possible to optimize slightly?
    store_a, load_a = tree_ir.evaluate_and_load(a)
    b_val = tree_ir.StoreLocalInstruction(
        None,
        tree_ir.ReadValueInstruction(b))
    store_c, load_c = tree_ir.evaluate_and_load(c)
    return tree_ir.create_block(
        store_a,
        b_val,
        store_c,
        tree_ir.CreateDictionaryEdgeInstruction(
            load_a,
            b_val.create_load(),
            load_c),
        load_a)

def __list_read(a, b):
    # The statements in this function generate the following code:
    #
    # a_tmp = a # To make sure a is evaluated before b.
    # b_value, = yield [("RV", [b])]
    # result, = yield [("RD", [a_tmp, b_value])]
    # if result is None:
    #     raise Exception("List read out of bounds: %s" % b_value)
    # result

    store_a, load_a = tree_ir.evaluate_and_load(a)
    b_val = tree_ir.StoreLocalInstruction(
        None,
        tree_ir.ReadValueInstruction(b))
    result = tree_ir.StoreLocalInstruction(
        None,
        tree_ir.ReadDictionaryValueInstruction(
            load_a.create_load(), b_val.create_load()))

    return tree_ir.create_block(
        store_a,
        b_val,
        result,
        tree_ir.SelectInstruction(
            tree_ir.BinaryInstruction(
                result.create_load(),
                'is',
                tree_ir.LiteralInstruction(None)),
            tree_ir.RaiseInstruction(
                tree_ir.CallInstruction(
                    tree_ir.LoadGlobalInstruction('Exception'),
                    [tree_ir.BinaryInstruction(
                        tree_ir.LiteralInstruction('List read out of bounds: %s'),
                        '%',
                        b_val.create_load())])),
            tree_ir.EmptyInstruction()),
        result.create_load())

def __list_append(a, b):
    # We want to generate code that is more or less equivalent to:
    #
    # a_tmp = a
    # b_tmp = b
    # a_outgoing, = yield [("RO", [a_tmp])]
    # _ = yield [("CD", [a_tmp, len(a_outgoing), b_tmp])]
    # a

    store_a, load_a = tree_ir.evaluate_and_load(a)
    store_b, load_b = tree_ir.evaluate_and_load(b)
    return tree_ir.create_block(
        store_a,
        store_b,
        tree_ir.CreateDictionaryEdgeInstruction(
            load_a,
            create_get_length(
                tree_ir.ReadOutgoingEdgesInstruction(
                    load_a)),
            load_b),
        load_a)

def __log(a):
    # Original definition:
    #
    # def log(a, **remainder):
    #     a_value, = yield [("RV", [a])]
    #     print("== LOG == " + str(a_value))
    #     raise PrimitiveFinished(a)

    store_a, load_a = tree_ir.evaluate_and_load(a)
    return tree_ir.CompoundInstruction(
        tree_ir.create_block(
            store_a,
            tree_ir.PrintInstruction(
                tree_ir.BinaryInstruction(
                    tree_ir.LiteralInstruction("== LOG == "),
                    '+',
                    tree_ir.CallInstruction(
                        tree_ir.LoadGlobalInstruction('str'),
                        [tree_ir.ReadValueInstruction(load_a)])))),
        load_a)

def __read_nr_out(a):
    # Original definition:
    #
    # def read_nr_out(a, **remainder):
    #     outgoing, = yield [("RO", [a])]
    #     result, = yield [("CNV", [len(outgoing)])]
    #     raise PrimitiveFinished(result)

    return tree_ir.CreateNodeWithValueInstruction(
        create_get_length(tree_ir.ReadOutgoingEdgesInstruction(a)))

MISC_INTRINSICS = {
    # Reference equality
    'element_eq' :
        lambda a, b:
        tree_ir.CreateNodeWithValueInstruction(
            tree_ir.BinaryInstruction(a, '==', b)),
    'element_neq' :
        lambda a, b:
        tree_ir.CreateNodeWithValueInstruction(
            tree_ir.BinaryInstruction(a, '!=', b)),

    # Strings
    'string_get' :
        lambda a, b:
        tree_ir.CreateNodeWithValueInstruction(
            tree_ir.LoadIndexInstruction(
                tree_ir.ReadValueInstruction(a),
                tree_ir.ReadValueInstruction(b))),
    'string_len' :
        lambda a:
        tree_ir.CreateNodeWithValueInstruction(
            tree_ir.CallInstruction(
                tree_ir.LoadGlobalInstruction('len'),
                [tree_ir.ReadValueInstruction(a)])),
    'string_join' :
        lambda a, b:
        tree_ir.CreateNodeWithValueInstruction(
            tree_ir.BinaryInstruction(
                tree_ir.CallInstruction(
                    tree_ir.LoadGlobalInstruction('str'),
                    [tree_ir.ReadValueInstruction(a)]),
                '+',
                tree_ir.CallInstruction(
                    tree_ir.LoadGlobalInstruction('str'),
                    [tree_ir.ReadValueInstruction(b)]))),
    'string_startswith' :
        lambda a, b:
        tree_ir.CreateNodeWithValueInstruction(
            tree_ir.CallInstruction(
                tree_ir.LoadMemberInstruction(
                    tree_ir.ReadValueInstruction(a),
                    'startswith'),
                [tree_ir.ReadValueInstruction(b)])),

    # State creation
    'create_node' : tree_ir.CreateNodeInstruction,
    'create_edge' :
        # Lambda is totally necessary here, pylint.
        # You totally dropped the ball on this one.
        # pylint: disable=I0011,W0108
        lambda a, b:
        tree_ir.CreateEdgeInstruction(a, b),
    'create_value' :
        lambda a:
        tree_ir.CreateNodeWithValueInstruction(
            tree_ir.ReadValueInstruction(a)),

    # State reads
    'read_edge_src' :
        lambda a:
        tree_ir.LoadIndexInstruction(
            tree_ir.ReadEdgeInstruction(a),
            tree_ir.LiteralInstruction(0)),
    'read_edge_dst' :
        lambda a:
        tree_ir.LoadIndexInstruction(
            tree_ir.ReadEdgeInstruction(a),
            tree_ir.LiteralInstruction(1)),
    'is_edge' :
        lambda a:
        tree_ir.CreateNodeWithValueInstruction(
            tree_ir.BinaryInstruction(
                tree_ir.LoadIndexInstruction(
                    tree_ir.ReadEdgeInstruction(a),
                    tree_ir.LiteralInstruction(0)),
                'is not',
                tree_ir.LiteralInstruction(None))),

    'read_nr_out' : __read_nr_out,

    # read_root
    'read_root' :
        lambda:
        tree_ir.LoadIndexInstruction(
            tree_ir.LoadLocalInstruction(jit_runtime.KWARGS_PARAMETER_NAME),
            tree_ir.LiteralInstruction('root')),

    # read_taskroot
    'read_taskroot' :
        lambda:
        tree_ir.LoadIndexInstruction(
            tree_ir.LoadLocalInstruction(jit_runtime.KWARGS_PARAMETER_NAME),
            tree_ir.LiteralInstruction('task_root')),

    # Dictionary operations
    'dict_read' :
        lambda a, b:
        tree_ir.ReadDictionaryValueInstruction(
            a, tree_ir.ReadValueInstruction(b)),

    'dict_read_edge' :
        lambda a, b:
        tree_ir.ReadDictionaryEdgeInstruction(
            a, tree_ir.ReadValueInstruction(b)),

    'dict_add' : __dict_add,
    'dict_add_fast' : __dict_add_fast,
    'dict_len' : __read_nr_out,

    # Set operations
    'set_add' : __set_add,

    # List operations
    'list_len' : __read_nr_out,

    'list_read' : __list_read,
    'list_append' : __list_append,

    # log
    'log' : __log
}

def __read_nr_out_cfg(original_def, a):
    # Original definition:
    #
    # def read_nr_out(a, **remainder):
    #     outgoing, = yield [("RO", [a])]
    #     result, = yield [("CNV", [len(outgoing)])]
    #     raise PrimitiveFinished(result)

    original_def.redefine(
        cfg_ir.CreateNode(
            original_def.insert_before(
                cfg_ir.create_pure_simple_call(
                    'len',
                    original_def.insert_before(
                        cfg_ir.create_read_outgoing_edges(a))))))

def __dict_in_cfg(original_def, a, b):
    # Original definition:
    #
    # def dict_in(a, b, **remainder):
    #     b_value, = yield [("RV", [b])]
    #     value, = yield [("RD", [a, b_value])]
    #     is_in = value is not None
    #     result, = yield [("CNV", [is_in])]
    #     raise PrimitiveFinished(result)

    original_def.redefine(
        cfg_ir.CreateNode(
            original_def.insert_before(
                cfg_ir.Binary(
                    original_def.insert_before(
                        cfg_ir.create_read_dict_value(
                            a, original_def.insert_before(cfg_ir.Read(b)))),
                    'is not',
                    original_def.insert_before(cfg_ir.Literal(None))))))

def __dict_in_node_cfg(original_def, a, b):
    # Original definition:
    #
    # def dict_in_node(a, b, **remainder):
    #     value, = yield [("RDN", [a, b])]
    #     result, = yield [("CNV", [value is not None])]
    #     raise PrimitiveFinished(result)

    original_def.redefine(
        cfg_ir.CreateNode(
            original_def.insert_before(
                cfg_ir.Binary(
                    original_def.insert_before(cfg_ir.create_read_dict_node(a, b)),
                    'is not',
                    original_def.insert_before(cfg_ir.Literal(None))))))

def __dict_read_cfg(original_def, a, b):
    # Original definition:
    #
    # def dict_read(a, b, **remainder):
    #     b_value, = yield [("RV", [b])]
    #     result, = yield [("RD", [a, b_value])]
    #     raise PrimitiveFinished(result)

    original_def.redefine(
        cfg_ir.create_read_dict_value(
            a,
            original_def.insert_before(cfg_ir.Read(b))))

MISC_CFG_INTRINSICS = {
    # Reference equality
    'element_eq' :
        lambda original_def, a, b:
        original_def.redefine(
            cfg_ir.CreateNode(
                original_def.insert_before(
                    cfg_ir.Binary(a, '==', b)))),
    'element_neq' :
        lambda original_def, a, b:
        original_def.redefine(
            cfg_ir.CreateNode(
                original_def.insert_before(
                    cfg_ir.Binary(a, '!=', b)))),


    # String operations
    'string_get' :
        lambda original_def, a, b:
        original_def.redefine(
            cfg_ir.CreateNode(
                original_def.insert_before(
                    cfg_ir.create_index(
                        original_def.insert_before(cfg_ir.Read(a)),
                        original_def.insert_before(cfg_ir.Read(b)))))),
    'string_len' :
        lambda original_def, a:
        original_def.redefine(
            cfg_ir.CreateNode(
                original_def.insert_before(
                    cfg_ir.create_pure_simple_call(
                        'len',
                        original_def.insert_before(cfg_ir.Read(a)))))),
    'string_join' :
        lambda original_def, a, b:
        original_def.redefine(
            cfg_ir.CreateNode(
                original_def.insert_before(
                    cfg_ir.Binary(
                        original_def.insert_before(
                            cfg_ir.create_pure_simple_call(
                                'str',
                                original_def.insert_before(cfg_ir.Read(a)))),
                        '+',
                        original_def.insert_before(
                            cfg_ir.create_pure_simple_call(
                                'str',
                                original_def.insert_before(cfg_ir.Read(b)))))))),
    'string_startswith' :
        lambda original_def, a, b:
        original_def.redefine(
            cfg_ir.CreateNode(
                original_def.insert_before(
                    cfg_ir.DirectFunctionCall(
                        'startswith',
                        [('self', original_def.insert_before(cfg_ir.Read(a))),
                         ('substring', original_def.insert_before(cfg_ir.Read(b)))],
                        calling_convention=cfg_ir.SELF_POSITIONAL_CALLING_CONVENTION,
                        has_value=True, has_side_effects=False)))),

    # State creation
    'create_node' :
        lambda original_def:
        original_def.redefine(
            cfg_ir.CreateNode(original_def.insert_before(cfg_ir.Literal(None)))),
    'create_value' :
        lambda original_def, a:
        original_def.redefine(
            cfg_ir.CreateNode(original_def.insert_before(cfg_ir.Read(a)))),

    # State reads
    'read_nr_out' : __read_nr_out_cfg,

    # Dictionary operations
    'dict_len' : __read_nr_out_cfg,
    'dict_read' : __dict_read_cfg,
    'dict_in' : __dict_in_cfg,
    'dict_in_node' : __dict_in_node_cfg,

    # List operations
    'list_len' : __read_nr_out_cfg
}

def register_time_intrinsic(target_jit):
    """Registers the time() intrinsic with the given JIT."""
    import_name = target_jit.import_value(time.time, 'time')
    target_jit.register_intrinsic(
        'time',
        lambda: tree_ir.CreateNodeWithValueInstruction(
            tree_ir.CallInstruction(
                tree_ir.LoadGlobalInstruction(import_name),
                [])))

def register_intrinsics(target_jit):
    """Registers all intrinsics in the module with the given JIT."""
    for (key, value) in BINARY_INTRINSICS.items():
        target_jit.register_binary_intrinsic(key, value)
    for (key, value) in UNARY_INTRINSICS.items():
        target_jit.register_unary_intrinsic(key, value)
    for (key, value) in CAST_INTRINSICS.items():
        target_jit.register_cast_intrinsic(key, value)
    for (key, value) in MISC_INTRINSICS.items():
        target_jit.register_intrinsic(key, value)
    for (key, value) in MISC_CFG_INTRINSICS.items():
        target_jit.register_cfg_intrinsic(key, value)

    register_time_intrinsic(target_jit)