modelverse/kernel/modelverse_jit/intrinsics.py

import jit
import tree_ir

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' : '-'
}

# 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):
    tmp = tree_ir.StoreLocalInstruction(None, a)
    return tree_ir.create_block(
        tmp,
        tree_ir.CreateEdgeInstruction(tmp.create_load(), b),
        tmp.create_load())

def __dict_add(a, b, c):
    a_tmp = tree_ir.StoreLocalInstruction(None, a)
    b_tmp = tree_ir.StoreLocalInstruction(None, b)
    return tree_ir.create_block(
        a_tmp,
        b_tmp,
        tree_ir.CreateEdgeInstruction(
            tree_ir.CreateEdgeInstruction(a_tmp.create_load(), c),
            b_tmp.create_load()),
        a_tmp.create_load())

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_root
    'read_root' :
        lambda:
        tree_ir.LoadIndexInstruction(
            tree_ir.LoadLocalInstruction(jit.KWARGS_PARAMETER_NAME),
            tree_ir.LiteralInstruction('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,

    # Set operations
    'set_add' : __set_add
}

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 MISC_INTRINSICS.items():
        target_jit.register_intrinsic(key, value)