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

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

    # Strings
    'string_get' :
        lambda str_node, index_node:
        tree_ir.CreateNodeWithValueInstruction(
            tree_ir.LoadIndexInstruction(
                tree_ir.ReadValueInstruction(str_node),
                tree_ir.ReadValueInstruction(index_node))),
    'string_len' :
        lambda str_node:
        tree_ir.CreateNodeWithValueInstruction(
            tree_ir.CallInstruction(
                tree_ir.LoadGlobalInstruction('len'),
                [tree_ir.ReadValueInstruction(str_node)])),
    'string_join' :
        lambda lhs, rhs:
        tree_ir.CreateNodeWithValueInstruction(
            tree_ir.BinaryInstruction(
                tree_ir.CallInstruction(
                    tree_ir.LoadGlobalInstruction('str'),
                    [tree_ir.ReadValueInstruction(lhs)]),
                '+',
                tree_ir.CallInstruction(
                    tree_ir.LoadGlobalInstruction('str'),
                    [tree_ir.ReadValueInstruction(rhs)]))),

    # State creation
    'create_node' : tree_ir.CreateNodeInstruction,
    'create_edge' : tree_ir.CreateEdgeInstruction,
    'create_value' :
        lambda val:
        tree_ir.CreateNodeWithValueInstruction(
            tree_ir.ReadValueInstruction(val)),

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

    # Dictionary operations
    'dict_read' :
        lambda dict_node, key:
        tree_ir.ReadDictionaryValueInstruction(
            dict_node, tree_ir.ReadValueInstruction(key)),

    'dict_read_edge' :
        lambda dict_node, key:
        tree_ir.ReadDictionaryEdgeInstruction(
            dict_node, tree_ir.ReadValueInstruction(key))
}

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)