"""
Transforms analyzed nodes into vectorized functions.

**The contents of this module are intended only for internal Deephaven use and may change at any time.**
"""

# When copying "ast we want to generate" from the ast.dump method, you just need to add a from import here,
# and then wire the new ast into the correct place of the assembled function.
from _ast import Module, For, Return, Str
from ast import fix_missing_locations, increment_lineno, NodeTransformer

import numba
import numpy
from numba.types import Type

from deephaven.lang.constants import primitive_types, VecMode
from deephaven.lang.tools import param, method_call, field_read, read, write, decorate, new_func, assign, array_access, \
    write_array, new_list, check_if_none, keyword_arg, num_size, int_or_float, itr, values, read_none, arg_or_name
from deephaven.lang.analyzer import Analyzer


# The methods we are overriding have non-standard uppercase in their names.
# noinspection PyPep8Naming
class Transformer(NodeTransformer):
    """Transforms the analyzed node into a vectorized function.
    This operation is ~O(1) if you do not use any assignments (we only visit children if needed)"""

    _key_index = 'i__'
    _key_result_default = 'out'
    _unique = 0

    @classmethod
    def next_func_name(cls):
        """Returns the next name that will be generated by any transformer instance.
        Visible for testing."""
        return '__vec{}'.format(cls._unique)

    def __init__(self, analyzed, return_type, ref_values=None, mode=VecMode.default_mode):
        assert isinstance(analyzed, Analyzer)
        assert analyzed.empty or len(analyzed.all_names) > 0, \
            "Must have at least one column reference in a vectorized function"
        self.analyzed = analyzed
        self.mode = mode
        self.return_type = return_type
        if ref_values is None:
            ref_values = {}
        self.external = {k: v for k, v in itr(ref_values) if k in analyzed.used_refs}

        self.func_name = '__vec{}'.format(Transformer._unique)
        Transformer._unique += 1

        self.key_result = Transformer._key_result_default if analyzed.last_assign is None \
            else analyzed.last_assign.targets[0].id
        self.index = read(Transformer._key_index)
        # we add one line for the function definition,
        # then non-@vectorize uses two lines to create the output array if needed
        self.extra_lines = 1 if VecMode.vectorize else 3

    def visit_Module(self, node):
        """This is the main entry point of the source AST.
        We want to wrap the expression the user sent us
        with a function definition, a loop, and then
        an initialization of all references variables.

        This might be easier to read if we just modified the source,
        but it is far more dependable to work with raw ast nodes,
        and never even think of manually parsing any strings."""

        # we'll assemble an array of statements to form the body of the generated function
        fun_body = []
        for_body = []

        arg_format = '{}' if self.mode is VecMode.vectorize else '{}_'
        # Lets prepare the arguments.  First, all column references are translated from colname into colname_
        args = list(map(lambda x: arg_or_name(arg_format.format(x.id)), values(self.analyzed.all_names)))
        # Next, add any external references that are simply passed through to the function
        args.extend(list(map(lambda x: arg_or_name(x[0]), self.external)))

        if self.mode is not VecMode.vectorize:
            # Add the returned array.  We will give it a default value of None, and initialize it when it is None
            args.append(arg_or_name(self.key_result))

        default_args = [read_none()] if self.mode is VecMode.jit else []

        # Add the desired decorator.
        decorators = [self.choose_decorator()]

        if self.mode is VecMode.jit:
            # only jit-mode has to worry about an empty output array being passed.
            # vectorize and guvectorize both handle output arrays "for free".

            # assemble: dtype=np.return_type
            ret_dtype = keyword_arg('dtype', field_read('np', self.return_type_name()))

            # assemble: np.empty(arg0.shape, dtype=np.return_type)
            create_empty = method_call(
                scope=field_read('np', 'empty'),
                args=[field_read(read(args[0]), 'shape')],
                keywords=[ret_dtype]
            )

            # assemble: out = np.empty(arg0.shape, dtype=np.return_type)
            create_result = assign(self.key_result, create_empty)
            # assemble: if out is None:
            #             out = np.empty(arg0.shape, dtype=np.return_type)
            check_result = check_if_none(self.key_result, [create_result])
            fun_body.append(check_result)

        # Fill in the body of the function.  All types except vectorize use the same code path.
        if self.mode is VecMode.vectorize:
            if self.analyzed.last_assign is None:
                fun_body.extend(node.body[0:-1])
                arr_assign = assign(self.key_result, node.body[-1].value)
                fun_body.append(arr_assign)
            else:
                fun_body.extend(node.body)
        else:
            # jit, guvectorize
            # assemble: for i__ in range(arg0.shape[0]):
            #             a = a__[i__]
            for arg in values(self.analyzed.all_names):
                # Create assignments to the expected variable names w/in the loop
                self.extra_lines += 1
                # assemble: a = a[i__]
                a = assign(arg.id, array_access(arg_format.format(arg.id), self.index))
                for_body.append(a)

            # visit the original body of source code, to do transformation on assignments
            # This O(n) visit is only necessary if the original expression contained an assignment.
            fixed_body = node.body
            if self.analyzed.last_assign is None:
                # The user did not supply an assignment statement; build one for them.
                for_body.extend(fixed_body[0:-1])
                arr_assign = assign(write_array(self.key_result, self.index), fixed_body[-1].value)
                for_body.append(arr_assign)
            else:
                fixed_body = list(map(lambda x: increment_lineno(x, self.extra_lines), fixed_body))
                fixed_body = list(map(lambda x: self.generic_visit(x), fixed_body))
                # append the fixed statements to the new body
                for_body.extend(fixed_body)

            # assemble: r__.shape[0]
            shape = array_access(field_read(self.key_result, 'shape'), 0)
            # assemble: range(r__.shape[0])
            shape_range = method_call('range', [shape])
            # assemble: for i__ in range(r__.shape[0]):
            for_loop = For(target=write(Transformer._key_index), iter=shape_range, body=for_body, orelse=[])
            # print our loop
            fun_body.append(for_loop)

        # print the return statement
        fun_body.append(Return(read(self.key_result)))

        # create the whole function
        func = new_func(self.func_name, args, fun_body, decorators, default_args)

        # build a new module w/ just the import and the function (so we can call it later)
        # if we can forgo the imports, or do them manually to the environment,
        # then we could use `single` mode instead of `exec` mode when compiling,
        # and just return a single function (effectively what the java version did)
        fixed = Module(body=[
            # # import numpy as np
            # Import(names=[alias(name='numpy', asname='np')]),
            # # from numba import jit
            # ImportFrom(module='numba', names=[alias(name='jit', asname=None)], level=0),
            # the function we just built
            func])

        # built-in module to generate source information in the generated results.
        # TODO: verify these are actually correct in unit test
        node = fix_missing_locations(fixed)

        # print dump(node)
        return node

    def choose_decorator(self):
        """
        Choose an appropriate vectorization decorator.
        :return: decorator.
        """
        nopython = True
        for col in self.analyzed.used_cols:
            t = self.analyzed.cols[col]
            ty = t if t in primitive_types else numba.typeof(t)
            from numpy import dtype
            dt = dtype(ty)
            if dt is numpy.object_:
                # TODO: consider warning / reporting this fact somehow.
                nopython = False

        args = {'nopython': '{}'.format(nopython)}

        # Constructs the expected signature; this can save numba time, since it won't have to infer as much;
        # it also allows use to overcome numba translating 32 bit types to 64 types.

        if self.mode is VecMode.vectorize:
            if self.return_type is not None:
                args['ftylist_or_function'] = new_list(self.compute_sigs())
            return decorate('vectorize', args)
        else:
            if self.return_type is not None:
                sigs = self.compute_sigs('[:]')
                if (self.mode is VecMode.jit):
                    # in @jit mode, we have to reduce the amount of types we can accept
                    # in order to be able to generate code that can create an empty output array.
                    if self.analyzed.performed_inspect:
                        # especially when the analyzer actually did work, we can only the inferred output type;
                        # this could likely be simplified, but jit mode isn't even default, so it stays hacky like this
                        sigs = [sigs[-1]]
                    # in order to omit passing the output array, we have to add an additional signature
                    sigs.append(sigs[-1].replace("{}[:])".format(self.return_type_name()), 'Omitted(None))'))

                args['signature_or_function'] = new_list(sigs)

                if self.mode is VecMode.guvectorize or self.mode is VecMode.cuda:
                    if self.mode is VecMode.cuda:
                        args['type'] = Str(s='cuda')
                    # TODO Also add the necessary numba signature
                    return decorate('guvectorize', args)
            return decorate('jit', args)

    def arg_types(self):
        """
        Argument types.
        :return: argument types.
        """
        types = []
        for i in self.analyzed.used_cols:
            type_name = self.analyzed.cols[i]
            types.append(type_name)
        return types

    def generic_visit(self, node):
        node = super(Transformer, self).generic_visit(node)
        return node

    def visit_Name(self, node):
        if self.analyzed.last_assign is not None and node == self.analyzed.last_assign.targets[0]:
            # Need to convert result assignments to an array store (Subscript)
            return write_array(self.key_result, self.index)
        return node

    def return_type_name(self):
        """
        Gets the return type name.
        :return: return type name."""
        return self.return_type.name if isinstance(self.return_type, Type) else numpy.dtype(self.return_type).name

    def compute_sigs(self, suffix='', sig_return=None):
        """
        Compute a list of signatures that are valid for the function we are building.

        for example, int32 + int32 w/out explicit return would yield:
        @vectorize (empty suffix):
        ['int32(int32,int32), int64(int32, int32)']
        @jit (suffix='[:]'):
        ['int32[:](int32[:],int32[:],int32[:])', 'int64[:](int32[:], int32[:],int64[:])']
        @guivectorize (suffix='[:]', sig_return='void'):
        ['void(int32[:],int32[:],int32[:])', 'void(int32[:], int32[:],int64[:])']

        @vectorize will implicitly add optional output array.
        @guivectorize _requires_ explicit output array (void return type)
        @jit we generate code to add optional output array

        :param suffix: A suffix to apply, to explicitly increase vector dimension (turn scalar into vector)
        :param sig_return: The explicit return type to use.  Defaults to the function's expected output type.
        :return: list of valid function signatures.
        """
        arg_list = []
        small_num_sizes = [8, 16, 32]
        has_sized_num = False
        for arg in self.arg_types():
            name = arg.__name__
            arg_list.append(name)
            size = num_size(name)
            if size is not None:
                has_sized_num = True
                small_num_sizes = [x for x in small_num_sizes if x >= size]

        ret_type = self.return_type_name()
        void = sig_return is 'void'
        if sig_return is None:
            sig_return = '{}{}'.format(ret_type, suffix)
        size = num_size(ret_type)
        if size in small_num_sizes:
            small_num_sizes.remove(size)

        if suffix:
            sig = ','.join(map(lambda arg: arg + suffix, arg_list))
            last_arg = ',{}{}'.format(ret_type, suffix)
        else:
            sig = ','.join(arg_list)
            last_arg = ''
        sigs = []
        if not void and has_sized_num:
            type_name = int_or_float(ret_type)
            for i in small_num_sizes:
                my_type = "{}{}".format(type_name, i)
                my_last = ',{}{}'.format(my_type, suffix) if last_arg else ''
                sigs.append("{}{}({}{})".format(my_type, suffix, sig, my_last))

        sigs.append("{}({}{})".format(sig_return, sig, last_arg))
        return sigs

"""
Transforms analyzed nodes into vectorized functions.

**The contents of this module are intended only for internal Deephaven use and may change at any time.**
"""

# When copying "ast we want to generate" from the ast.dump method, you just need to add a from import here,
# and then wire the new ast into the correct place of the assembled function.
from _ast import Module, For, Return, Str
from ast import fix_missing_locations, increment_lineno, NodeTransformer

import numba
import numpy
from numba.types import Type

from deephaven.lang.constants import primitive_types, VecMode
from deephaven.lang.tools import param, method_call, field_read, read, write, decorate, new_func, assign, array_access, \
    write_array, new_list, check_if_none, keyword_arg, num_size, int_or_float, itr, values, read_none, arg_or_name
from deephaven.lang.analyzer import Analyzer


# The methods we are overriding have non-standard uppercase in their names.
# noinspection PyPep8Naming
class Transformer(NodeTransformer):
    """Transforms the analyzed node into a vectorized function.
    This operation is ~O(1) if you do not use any assignments (we only visit children if needed)"""

    _key_index = 'i__'
    _key_result_default = 'out'
    _unique = 0

    @classmethod
    def next_func_name(cls):
        """Returns the next name that will be generated by any transformer instance.
        Visible for testing."""
        return '__vec{}'.format(cls._unique)

    def __init__(self, analyzed, return_type, ref_values=None, mode=VecMode.default_mode):
        assert isinstance(analyzed, Analyzer)
        assert analyzed.empty or len(analyzed.all_names) > 0, \
            "Must have at least one column reference in a vectorized function"
        self.analyzed = analyzed
        self.mode = mode
        self.return_type = return_type
        if ref_values is None:
            ref_values = {}
        self.external = {k: v for k, v in itr(ref_values) if k in analyzed.used_refs}

        self.func_name = '__vec{}'.format(Transformer._unique)
        Transformer._unique += 1

        self.key_result = Transformer._key_result_default if analyzed.last_assign is None \
            else analyzed.last_assign.targets[0].id
        self.index = read(Transformer._key_index)
        # we add one line for the function definition,
        # then non-@vectorize uses two lines to create the output array if needed
        self.extra_lines = 1 if VecMode.vectorize else 3

    def visit_Module(self, node):
        """This is the main entry point of the source AST.
        We want to wrap the expression the user sent us
        with a function definition, a loop, and then
        an initialization of all references variables.

        This might be easier to read if we just modified the source,
        but it is far more dependable to work with raw ast nodes,
        and never even think of manually parsing any strings."""

        # we'll assemble an array of statements to form the body of the generated function
        fun_body = []
        for_body = []

        arg_format = '{}' if self.mode is VecMode.vectorize else '{}_'
        # Lets prepare the arguments.  First, all column references are translated from colname into colname_
        args = list(map(lambda x: arg_or_name(arg_format.format(x.id)), values(self.analyzed.all_names)))
        # Next, add any external references that are simply passed through to the function
        args.extend(list(map(lambda x: arg_or_name(x[0]), self.external)))

        if self.mode is not VecMode.vectorize:
            # Add the returned array.  We will give it a default value of None, and initialize it when it is None
            args.append(arg_or_name(self.key_result))

        default_args = [read_none()] if self.mode is VecMode.jit else []

        # Add the desired decorator.
        decorators = [self.choose_decorator()]

        if self.mode is VecMode.jit:
            # only jit-mode has to worry about an empty output array being passed.
            # vectorize and guvectorize both handle output arrays "for free".

            # assemble: dtype=np.return_type
            ret_dtype = keyword_arg('dtype', field_read('np', self.return_type_name()))

            # assemble: np.empty(arg0.shape, dtype=np.return_type)
            create_empty = method_call(
                scope=field_read('np', 'empty'),
                args=[field_read(read(args[0]), 'shape')],
                keywords=[ret_dtype]
            )

            # assemble: out = np.empty(arg0.shape, dtype=np.return_type)
            create_result = assign(self.key_result, create_empty)
            # assemble: if out is None:
            #             out = np.empty(arg0.shape, dtype=np.return_type)
            check_result = check_if_none(self.key_result, [create_result])
            fun_body.append(check_result)

        # Fill in the body of the function.  All types except vectorize use the same code path.
        if self.mode is VecMode.vectorize:
            if self.analyzed.last_assign is None:
                fun_body.extend(node.body[0:-1])
                arr_assign = assign(self.key_result, node.body[-1].value)
                fun_body.append(arr_assign)
            else:
                fun_body.extend(node.body)
        else:
            # jit, guvectorize
            # assemble: for i__ in range(arg0.shape[0]):
            #             a = a__[i__]
            for arg in values(self.analyzed.all_names):
                # Create assignments to the expected variable names w/in the loop
                self.extra_lines += 1
                # assemble: a = a[i__]
                a = assign(arg.id, array_access(arg_format.format(arg.id), self.index))
                for_body.append(a)

            # visit the original body of source code, to do transformation on assignments
            # This O(n) visit is only necessary if the original expression contained an assignment.
            fixed_body = node.body
            if self.analyzed.last_assign is None:
                # The user did not supply an assignment statement; build one for them.
                for_body.extend(fixed_body[0:-1])
                arr_assign = assign(write_array(self.key_result, self.index), fixed_body[-1].value)
                for_body.append(arr_assign)
            else:
                fixed_body = list(map(lambda x: increment_lineno(x, self.extra_lines), fixed_body))
                fixed_body = list(map(lambda x: self.generic_visit(x), fixed_body))
                # append the fixed statements to the new body
                for_body.extend(fixed_body)

            # assemble: r__.shape[0]
            shape = array_access(field_read(self.key_result, 'shape'), 0)
            # assemble: range(r__.shape[0])
            shape_range = method_call('range', [shape])
            # assemble: for i__ in range(r__.shape[0]):
            for_loop = For(target=write(Transformer._key_index), iter=shape_range, body=for_body, orelse=[])
            # print our loop
            fun_body.append(for_loop)

        # print the return statement
        fun_body.append(Return(read(self.key_result)))

        # create the whole function
        func = new_func(self.func_name, args, fun_body, decorators, default_args)

        # build a new module w/ just the import and the function (so we can call it later)
        # if we can forgo the imports, or do them manually to the environment,
        # then we could use `single` mode instead of `exec` mode when compiling,
        # and just return a single function (effectively what the java version did)
        fixed = Module(body=[
            # # import numpy as np
            # Import(names=[alias(name='numpy', asname='np')]),
            # # from numba import jit
            # ImportFrom(module='numba', names=[alias(name='jit', asname=None)], level=0),
            # the function we just built
            func])

        # built-in module to generate source information in the generated results.
        # TODO: verify these are actually correct in unit test
        node = fix_missing_locations(fixed)

        # print dump(node)
        return node

    def choose_decorator(self):
        """
        Choose an appropriate vectorization decorator.
        :return: decorator.
        """
        nopython = True
        for col in self.analyzed.used_cols:
            t = self.analyzed.cols[col]
            ty = t if t in primitive_types else numba.typeof(t)
            from numpy import dtype
            dt = dtype(ty)
            if dt is numpy.object_:
                # TODO: consider warning / reporting this fact somehow.
                nopython = False

        args = {'nopython': '{}'.format(nopython)}

        # Constructs the expected signature; this can save numba time, since it won't have to infer as much;
        # it also allows use to overcome numba translating 32 bit types to 64 types.

        if self.mode is VecMode.vectorize:
            if self.return_type is not None:
                args['ftylist_or_function'] = new_list(self.compute_sigs())
            return decorate('vectorize', args)
        else:
            if self.return_type is not None:
                sigs = self.compute_sigs('[:]')
                if (self.mode is VecMode.jit):
                    # in @jit mode, we have to reduce the amount of types we can accept
                    # in order to be able to generate code that can create an empty output array.
                    if self.analyzed.performed_inspect:
                        # especially when the analyzer actually did work, we can only the inferred output type;
                        # this could likely be simplified, but jit mode isn't even default, so it stays hacky like this
                        sigs = [sigs[-1]]
                    # in order to omit passing the output array, we have to add an additional signature
                    sigs.append(sigs[-1].replace("{}[:])".format(self.return_type_name()), 'Omitted(None))'))

                args['signature_or_function'] = new_list(sigs)

                if self.mode is VecMode.guvectorize or self.mode is VecMode.cuda:
                    if self.mode is VecMode.cuda:
                        args['type'] = Str(s='cuda')
                    # TODO Also add the necessary numba signature
                    return decorate('guvectorize', args)
            return decorate('jit', args)

    def arg_types(self):
        """
        Argument types.
        :return: argument types.
        """
        types = []
        for i in self.analyzed.used_cols:
            type_name = self.analyzed.cols[i]
            types.append(type_name)
        return types

    def generic_visit(self, node):
        node = super(Transformer, self).generic_visit(node)
        return node

    def visit_Name(self, node):
        if self.analyzed.last_assign is not None and node == self.analyzed.last_assign.targets[0]:
            # Need to convert result assignments to an array store (Subscript)
            return write_array(self.key_result, self.index)
        return node

    def return_type_name(self):
        """
        Gets the return type name.
        :return: return type name."""
        return self.return_type.name if isinstance(self.return_type, Type) else numpy.dtype(self.return_type).name

    def compute_sigs(self, suffix='', sig_return=None):
        """
        Compute a list of signatures that are valid for the function we are building.

        for example, int32 + int32 w/out explicit return would yield:
        @vectorize (empty suffix):
        ['int32(int32,int32), int64(int32, int32)']
        @jit (suffix='[:]'):
        ['int32[:](int32[:],int32[:],int32[:])', 'int64[:](int32[:], int32[:],int64[:])']
        @guivectorize (suffix='[:]', sig_return='void'):
        ['void(int32[:],int32[:],int32[:])', 'void(int32[:], int32[:],int64[:])']

        @vectorize will implicitly add optional output array.
        @guivectorize _requires_ explicit output array (void return type)
        @jit we generate code to add optional output array

        :param suffix: A suffix to apply, to explicitly increase vector dimension (turn scalar into vector)
        :param sig_return: The explicit return type to use.  Defaults to the function's expected output type.
        :return: list of valid function signatures.
        """
        arg_list = []
        small_num_sizes = [8, 16, 32]
        has_sized_num = False
        for arg in self.arg_types():
            name = arg.__name__
            arg_list.append(name)
            size = num_size(name)
            if size is not None:
                has_sized_num = True
                small_num_sizes = [x for x in small_num_sizes if x >= size]

        ret_type = self.return_type_name()
        void = sig_return is 'void'
        if sig_return is None:
            sig_return = '{}{}'.format(ret_type, suffix)
        size = num_size(ret_type)
        if size in small_num_sizes:
            small_num_sizes.remove(size)

        if suffix:
            sig = ','.join(map(lambda arg: arg + suffix, arg_list))
            last_arg = ',{}{}'.format(ret_type, suffix)
        else:
            sig = ','.join(arg_list)
            last_arg = ''
        sigs = []
        if not void and has_sized_num:
            type_name = int_or_float(ret_type)
            for i in small_num_sizes:
                my_type = "{}{}".format(type_name, i)
                my_last = ',{}{}'.format(my_type, suffix) if last_arg else ''
                sigs.append("{}{}({}{})".format(my_type, suffix, sig, my_last))

        sigs.append("{}({}{})".format(sig_return, sig, last_arg))
        return sigs

Transforms analyzed nodes into vectorized functions.  **The contents of this module are intended only for internal Deephaven use and may change at any time.** # When copying "ast we want to generate" from the ast.dump method, you just need to add a from import here, # and then wire the new ast into the correct place of the assembled function. # The methods we are overriding have non-standard uppercase in their names. # noinspection PyPep8Naming Transforms the analyzed node into a vectorized function.     This operation is ~O(1) if you do not use any assignments (we only visit children if needed) Returns the next name that will be generated by any transformer instance.         Visible for testing. # we add one line for the function definition, # then non-@vectorize uses two lines to create the output array if needed This is the main entry point of the source AST.         We want to wrap the expression the user sent us         with a function definition, a loop, and then         an initialization of all references variables.          This might be easier to read if we just modified the source,         but it is far more dependable to work with raw ast nodes,         and never even think of manually parsing any strings. # we'll assemble an array of statements to form the body of the generated function # Lets prepare the arguments.  First, all column references are translated from colname into colname_ # Next, add any external references that are simply passed through to the function # Add the returned array.  We will give it a default value of None, and initialize it when it is None # Add the desired decorator. # only jit-mode has to worry about an empty output array being passed. # vectorize and guvectorize both handle output arrays "for free". # assemble: dtype=np.return_type # assemble: np.empty(arg0.shape, dtype=np.return_type) # assemble: out = np.empty(arg0.shape, dtype=np.return_type) # assemble: if out is None: #             out = np.empty(arg0.shape, dtype=np.return_type) # Fill in the body of the function.  All types except vectorize use the same code path. # jit, guvectorize # assemble: for i__ in range(arg0.shape[0]): #             a = a__[i__] # Create assignments to the expected variable names w/in the loop # assemble: a = a[i__] # visit the original body of source code, to do transformation on assignments # This O(n) visit is only necessary if the original expression contained an assignment. # The user did not supply an assignment statement; build one for them. # append the fixed statements to the new body # assemble: r__.shape[0] # assemble: range(r__.shape[0]) # assemble: for i__ in range(r__.shape[0]): # print our loop # print the return statement # create the whole function # build a new module w/ just the import and the function (so we can call it later) # if we can forgo the imports, or do them manually to the environment, # then we could use `single` mode instead of `exec` mode when compiling, # and just return a single function (effectively what the java version did) # # import numpy as np # Import(names=[alias(name='numpy', asname='np')]), # # from numba import jit # ImportFrom(module='numba', names=[alias(name='jit', asname=None)], level=0), # the function we just built # built-in module to generate source information in the generated results. # TODO: verify these are actually correct in unit test # print dump(node) Choose an appropriate vectorization decorator.         :return: decorator. # TODO: consider warning / reporting this fact somehow. # Constructs the expected signature; this can save numba time, since it won't have to infer as much; # it also allows use to overcome numba translating 32 bit types to 64 types. # in @jit mode, we have to reduce the amount of types we can accept # in order to be able to generate code that can create an empty output array. # especially when the analyzer actually did work, we can only the inferred output type; # this could likely be simplified, but jit mode isn't even default, so it stays hacky like this # in order to omit passing the output array, we have to add an additional signature # TODO Also add the necessary numba signature Argument types.         :return: argument types. # Need to convert result assignments to an array store (Subscript) Gets the return type name.         :return: return type name. Compute a list of signatures that are valid for the function we are building.          for example, int32 + int32 w/out explicit return would yield:         @vectorize (empty suffix):         ['int32(int32,int32), int64(int32, int32)']         @jit (suffix='[:]'):         ['int32[:](int32[:],int32[:],int32[:])', 'int64[:](int32[:], int32[:],int64[:])']         @guivectorize (suffix='[:]', sig_return='void'):         ['void(int32[:],int32[:],int32[:])', 'void(int32[:], int32[:],int64[:])']          @vectorize will implicitly add optional output array.         @guivectorize _requires_ explicit output array (void return type)         @jit we generate code to add optional output array          :param suffix: A suffix to apply, to explicitly increase vector dimension (turn scalar into vector)         :param sig_return: The explicit return type to use.  Defaults to the function's expected output type.         :return: list of valid function signatures.