/pypy/module/micronumpy/loop.py

""" This file is the main run loop as well as evaluation loops for various
operations. This is the place to look for all the computations that iterate
over all the array elements.
"""
import py
from pypy.interpreter.error import oefmt
from rpython.rlib import jit
from rpython.rlib.rstring import StringBuilder
from rpython.rtyper.lltypesystem import lltype, rffi
from pypy.module.micronumpy import support, constants as NPY
from pypy.module.micronumpy.base import W_NDimArray, convert_to_array
from pypy.module.micronumpy.iterators import PureShapeIter, AxisIter, \
    AllButAxisIter, ArrayIter
from pypy.interpreter.argument import Arguments


def call2(space, shape, func, calc_dtype, w_lhs, w_rhs, out):
    if w_lhs.get_size() == 1:
        w_left = w_lhs.get_scalar_value().convert_to(space, calc_dtype)
        left_iter = left_state = None
    else:
        w_left = None
        left_iter, left_state = w_lhs.create_iter(shape)
        left_iter.track_index = False

    if w_rhs.get_size() == 1:
        w_right = w_rhs.get_scalar_value().convert_to(space, calc_dtype)
        right_iter = right_state = None
    else:
        w_right = None
        right_iter, right_state = w_rhs.create_iter(shape)
        right_iter.track_index = False

    out_iter, out_state = out.create_iter(shape)
    shapelen = len(shape)
    res_dtype = out.get_dtype()
    call2_func = try_to_share_iterators_call2(left_iter, right_iter,
            left_state, right_state, out_state)
    params = (space, shapelen, func, calc_dtype, res_dtype, out,
              w_left, w_right, left_iter, right_iter, out_iter,
              left_state, right_state, out_state)
    return call2_func(*params)

def try_to_share_iterators_call2(left_iter, right_iter, left_state, right_state, out_state):
    # these are all possible iterator sharing combinations
    # left == right == out
    # left == right
    # left == out
    # right == out
    right_out_equal = False
    if right_iter:
        # rhs is not a scalar
        if out_state.same(right_state):
            right_out_equal = True
    #
    if not left_iter:
        # lhs is a scalar
        if right_out_equal:
            return call2_advance_out_left
        else:
            # worst case, nothing can be shared and lhs is a scalar
            return call2_advance_out_left_right
    else:
        # lhs is NOT a scalar
        if out_state.same(left_state):
            # (2) out and left are the same -> remove left
            if right_out_equal:
                # the best case
                return call2_advance_out
            else:
                return call2_advance_out_right
        else:
            if right_out_equal:
                # right and out are equal, only advance left and out
                return call2_advance_out_left
            else:
                if right_iter and right_state.same(left_state):
                    # left and right are equal, but still need to advance out
                    return call2_advance_out_left_eq_right
                else:
                    # worst case, nothing can be shared
                    return call2_advance_out_left_right

    assert 0, "logical problem with the selection of the call2 case"

def generate_call2_cases(name, left_state, right_state):
    call2_driver = jit.JitDriver(name='numpy_call2_' + name,
        greens=['shapelen', 'func', 'calc_dtype', 'res_dtype'],
        reds='auto', vectorize=True)
    #
    advance_left_state = left_state == "left_state"
    advance_right_state = right_state == "right_state"
    code = """
    def method(space, shapelen, func, calc_dtype, res_dtype, out,
               w_left, w_right, left_iter, right_iter, out_iter,
               left_state, right_state, out_state):
        while not out_iter.done(out_state):
            call2_driver.jit_merge_point(shapelen=shapelen, func=func,
                    calc_dtype=calc_dtype, res_dtype=res_dtype)
            if left_iter:
                w_left = left_iter.getitem({left_state}).convert_to(space, calc_dtype)
            if right_iter:
                w_right = right_iter.getitem({right_state}).convert_to(space, calc_dtype)
            w_out = func(calc_dtype, w_left, w_right)
            out_iter.setitem(out_state, w_out.convert_to(space, res_dtype))
            out_state = out_iter.next(out_state)
            if advance_left_state and left_iter:
                left_state = left_iter.next(left_state)
            if advance_right_state and right_iter:
                right_state = right_iter.next(right_state)
            #
            # if not set to None, the values will be loop carried
            # (for the var,var case), forcing the vectorization to unpack
            # the vector registers at the end of the loop
            if left_iter:
                w_left = None
            if right_iter:
                w_right = None
        return out
    """
    exec(py.code.Source(code.format(left_state=left_state,right_state=right_state)).compile(), locals())
    method.__name__ = "call2_" + name
    return method

call2_advance_out = generate_call2_cases("inc_out", "out_state", "out_state")
call2_advance_out_left = generate_call2_cases("inc_out_left", "left_state", "out_state")
call2_advance_out_right = generate_call2_cases("inc_out_right", "out_state", "right_state")
call2_advance_out_left_eq_right = generate_call2_cases("inc_out_left_eq_right", "left_state", "left_state")
call2_advance_out_left_right = generate_call2_cases("inc_out_left_right", "left_state", "right_state")

call1_driver = jit.JitDriver(
    name='numpy_call1',
    greens=['shapelen', 'share_iterator', 'func', 'calc_dtype', 'res_dtype'],
    reds='auto', vectorize=True)

def call1(space, shape, func, calc_dtype, w_obj, w_ret):
    obj_iter, obj_state = w_obj.create_iter(shape)
    obj_iter.track_index = False
    out_iter, out_state = w_ret.create_iter(shape)
    shapelen = len(shape)
    res_dtype = w_ret.get_dtype()
    share_iterator = out_state.same(obj_state)
    while not out_iter.done(out_state):
        call1_driver.jit_merge_point(shapelen=shapelen, func=func,
                                     share_iterator=share_iterator,
                                     calc_dtype=calc_dtype, res_dtype=res_dtype)
        if share_iterator:
            # use out state as param to getitem
            elem = obj_iter.getitem(out_state).convert_to(space, calc_dtype)
        else:
            elem = obj_iter.getitem(obj_state).convert_to(space, calc_dtype)
        out_iter.setitem(out_state, func(calc_dtype, elem).convert_to(space, res_dtype))
        if share_iterator:
            # only advance out, they share the same iteration space
            out_state = out_iter.next(out_state)
        else:
            out_state = out_iter.next(out_state)
            obj_state = obj_iter.next(obj_state)
        elem = None
    return w_ret

call_many_to_one_driver = jit.JitDriver(
    name='numpy_call_many_to_one',
    greens=['shapelen', 'nin', 'func', 'in_dtypes', 'res_dtype'],
    reds='auto')

def call_many_to_one(space, shape, func, in_dtypes, res_dtype, in_args, out):
    # out must hav been built. func needs no calc_type, is usually an
    # external ufunc
    nin = len(in_args)
    in_iters = [None] * nin
    in_states = [None] * nin
    for i in range(nin):
        in_i = in_args[i]
        assert isinstance(in_i, W_NDimArray)
        in_iter, in_state = in_i.create_iter(shape)
        in_iters[i] = in_iter
        in_states[i] = in_state
    shapelen = len(shape)
    assert isinstance(out, W_NDimArray)
    out_iter, out_state = out.create_iter(shape)
    vals = [None] * nin
    while not out_iter.done(out_state):
        call_many_to_one_driver.jit_merge_point(shapelen=shapelen, func=func,
                        in_dtypes=in_dtypes, res_dtype=res_dtype, nin=nin)
        for i in range(nin):
            vals[i] = in_dtypes[i].coerce(space, in_iters[i].getitem(in_states[i]))
        w_arglist = space.newlist(vals)
        w_out_val = space.call_args(func, Arguments.frompacked(space, w_arglist))
        out_iter.setitem(out_state, res_dtype.coerce(space, w_out_val))
        for i in range(nin):
            in_states[i] = in_iters[i].next(in_states[i])
        out_state = out_iter.next(out_state)
    return out

call_many_to_many_driver = jit.JitDriver(
    name='numpy_call_many_to_many',
    greens=['shapelen', 'nin', 'nout', 'func', 'in_dtypes', 'out_dtypes'],
    reds='auto')

def call_many_to_many(space, shape, func, in_dtypes, out_dtypes, in_args, out_args):
    # out must have been built. func needs no calc_type, is usually an
    # external ufunc
    nin = len(in_args)
    in_iters = [None] * nin
    in_states = [None] * nin
    nout = len(out_args)
    out_iters = [None] * nout
    out_states = [None] * nout
    for i in range(nin):
        in_i = in_args[i]
        assert isinstance(in_i, W_NDimArray)
        in_iter, in_state = in_i.create_iter(shape)
        in_iters[i] = in_iter
        in_states[i] = in_state
    for i in range(nout):
        out_i = out_args[i]
        assert isinstance(out_i, W_NDimArray)
        out_iter, out_state = out_i.create_iter(shape)
        out_iters[i] = out_iter
        out_states[i] = out_state
    shapelen = len(shape)
    vals = [None] * nin
    test_iter, test_state = in_iters[-1], in_states[-1]
    if nout > 0:
        test_iter, test_state = out_iters[0], out_states[0]
    while not test_iter.done(test_state):
        call_many_to_many_driver.jit_merge_point(shapelen=shapelen, func=func,
                             in_dtypes=in_dtypes, out_dtypes=out_dtypes,
                             nin=nin, nout=nout)
        for i in range(nin):
            vals[i] = in_dtypes[i].coerce(space, in_iters[i].getitem(in_states[i]))
        w_arglist = space.newlist(vals)
        w_outvals = space.call_args(func, Arguments.frompacked(space, w_arglist))
        # w_outvals should be a tuple, but func can return a single value as well 
        if space.isinstance_w(w_outvals, space.w_tuple):
            batch = space.listview(w_outvals)
            for i in range(len(batch)):
                out_iters[i].setitem(out_states[i], out_dtypes[i].coerce(space, batch[i]))
                out_states[i] = out_iters[i].next(out_states[i])
        elif nout > 0:
            out_iters[0].setitem(out_states[0], out_dtypes[0].coerce(space, w_outvals))
            out_states[0] = out_iters[0].next(out_states[0])
        for i in range(nin):
            in_states[i] = in_iters[i].next(in_states[i])
        test_state = test_iter.next(test_state)
    return space.newtuple([convert_to_array(space, o) for o in out_args])

setslice_driver = jit.JitDriver(name='numpy_setslice',
                                greens = ['shapelen', 'dtype'],
                                reds = 'auto', vectorize=True)

def setslice(space, shape, target, source):
    if not shape:
        dtype = target.dtype
        val = source.getitem(source.start)
        if dtype.is_str_or_unicode():
            val = dtype.coerce(space, val)
        else:
            val = val.convert_to(space, dtype)
        target.setitem(target.start, val)
        return target
    return _setslice(space, shape, target, source)

def _setslice(space, shape, target, source):
    # note that unlike everything else, target and source here are
    # array implementations, not arrays
    target_iter, target_state = target.create_iter(shape)
    source_iter, source_state = source.create_iter(shape)
    source_iter.track_index = False
    dtype = target.dtype
    shapelen = len(shape)
    while not target_iter.done(target_state):
        setslice_driver.jit_merge_point(shapelen=shapelen, dtype=dtype)
        val = source_iter.getitem(source_state)
        if dtype.is_str_or_unicode() or dtype.is_record():
            val = dtype.coerce(space, val)
        else:
            val = val.convert_to(space, dtype)
        target_iter.setitem(target_state, val)
        target_state = target_iter.next(target_state)
        source_state = source_iter.next(source_state)
    return target


def split_iter(arr, axis_flags):
    """Prepare 2 iterators for nested iteration over `arr`.

    Arguments:
        arr: instance of BaseConcreteArray
        axis_flags: list of bools, one for each dimension of `arr`.The inner
        iterator operates over the dimensions for which the flag is True
    """
    shape = arr.get_shape()
    strides = arr.get_strides()
    backstrides = arr.get_backstrides()
    shapelen = len(shape)
    assert len(axis_flags) == shapelen
    inner_shape = [-1] * shapelen
    inner_strides = [-1] * shapelen
    inner_backstrides = [-1] * shapelen
    outer_shape = [-1] * shapelen
    outer_strides = [-1] * shapelen
    outer_backstrides = [-1] * shapelen
    for i in range(len(shape)):
        if axis_flags[i]:
            inner_shape[i] = shape[i]
            inner_strides[i] = strides[i]
            inner_backstrides[i] = backstrides[i]
            outer_shape[i] = 1
            outer_strides[i] = 0
            outer_backstrides[i] = 0
        else:
            outer_shape[i] = shape[i]
            outer_strides[i] = strides[i]
            outer_backstrides[i] = backstrides[i]
            inner_shape[i] = 1
            inner_strides[i] = 0
            inner_backstrides[i] = 0
    inner_iter = ArrayIter(arr, support.product(inner_shape),
                           inner_shape, inner_strides, inner_backstrides)
    outer_iter = ArrayIter(arr, support.product(outer_shape),
                           outer_shape, outer_strides, outer_backstrides)
    return inner_iter, outer_iter


reduce_flat_driver = jit.JitDriver(
    name='numpy_reduce_flat',
    greens = ['shapelen', 'func', 'done_func', 'calc_dtype'], reds = 'auto',
    vectorize = True)

def reduce_flat(space, func, w_arr, calc_dtype, done_func, identity):
    obj_iter, obj_state = w_arr.create_iter()
    if identity is None:
        cur_value = obj_iter.getitem(obj_state).convert_to(space, calc_dtype)
        obj_state = obj_iter.next(obj_state)
    else:
        cur_value = identity.convert_to(space, calc_dtype)
    shapelen = len(w_arr.get_shape())
    while not obj_iter.done(obj_state):
        reduce_flat_driver.jit_merge_point(
            shapelen=shapelen, func=func,
            done_func=done_func, calc_dtype=calc_dtype)
        rval = obj_iter.getitem(obj_state).convert_to(space, calc_dtype)
        if done_func is not None and done_func(calc_dtype, rval):
            return rval
        cur_value = func(calc_dtype, cur_value, rval)
        obj_state = obj_iter.next(obj_state)
    return cur_value

reduce_driver = jit.JitDriver(
    name='numpy_reduce',
    greens=['shapelen', 'func', 'dtype'], reds='auto',
    vectorize=True)

def reduce(space, func, w_arr, axis_flags, dtype, out, identity):
    out_iter, out_state = out.create_iter()
    out_iter.track_index = False
    shape = w_arr.get_shape()
    shapelen = len(shape)
    inner_iter, outer_iter = split_iter(w_arr.implementation, axis_flags)
    assert outer_iter.size == out_iter.size

    if identity is not None:
        identity = identity.convert_to(space, dtype)
    outer_state = outer_iter.reset()
    while not outer_iter.done(outer_state):
        inner_state = inner_iter.reset()
        inner_state.offset = outer_state.offset
        if identity is not None:
            w_val = identity
        else:
            w_val = inner_iter.getitem(inner_state).convert_to(space, dtype)
            inner_state = inner_iter.next(inner_state)
        while not inner_iter.done(inner_state):
            reduce_driver.jit_merge_point(
                shapelen=shapelen, func=func, dtype=dtype)
            w_item = inner_iter.getitem(inner_state).convert_to(space, dtype)
            w_val = func(dtype, w_item, w_val)
            inner_state = inner_iter.next(inner_state)
        out_iter.setitem(out_state, w_val)
        out_state = out_iter.next(out_state)
        outer_state = outer_iter.next(outer_state)
    return out

accumulate_flat_driver = jit.JitDriver(
    name='numpy_accumulate_flat',
    greens=['shapelen', 'func', 'dtype', 'out_dtype'],
    reds='auto', vectorize=True)

def accumulate_flat(space, func, w_arr, calc_dtype, w_out, identity):
    arr_iter, arr_state = w_arr.create_iter()
    out_iter, out_state = w_out.create_iter()
    out_iter.track_index = False
    if identity is None:
        cur_value = arr_iter.getitem(arr_state).convert_to(space, calc_dtype)
        out_iter.setitem(out_state, cur_value)
        out_state = out_iter.next(out_state)
        arr_state = arr_iter.next(arr_state)
    else:
        cur_value = identity.convert_to(space, calc_dtype)
    shapelen = len(w_arr.get_shape())
    out_dtype = w_out.get_dtype()
    while not arr_iter.done(arr_state):
        accumulate_flat_driver.jit_merge_point(
            shapelen=shapelen, func=func, dtype=calc_dtype,
            out_dtype=out_dtype)
        w_item = arr_iter.getitem(arr_state).convert_to(space, calc_dtype)
        cur_value = func(calc_dtype, cur_value, w_item)
        out_iter.setitem(out_state, out_dtype.coerce(space, cur_value))
        out_state = out_iter.next(out_state)
        arr_state = arr_iter.next(arr_state)

accumulate_driver = jit.JitDriver(
    name='numpy_accumulate',
    greens=['shapelen', 'func', 'calc_dtype'],
    reds='auto',
    vectorize=True)


def accumulate(space, func, w_arr, axis, calc_dtype, w_out, identity):
    out_iter, out_state = w_out.create_iter()
    arr_shape = w_arr.get_shape()
    temp_shape = arr_shape[:axis] + arr_shape[axis + 1:]
    temp = W_NDimArray.from_shape(space, temp_shape, calc_dtype, w_instance=w_arr)
    temp_iter = AxisIter(temp.implementation, w_arr.get_shape(), axis)
    temp_state = temp_iter.reset()
    arr_iter, arr_state = w_arr.create_iter()
    arr_iter.track_index = False
    if identity is not None:
        identity = identity.convert_to(space, calc_dtype)
    shapelen = len(arr_shape)
    while not out_iter.done(out_state):
        accumulate_driver.jit_merge_point(shapelen=shapelen, func=func,
                                          calc_dtype=calc_dtype)
        w_item = arr_iter.getitem(arr_state).convert_to(space, calc_dtype)
        arr_state = arr_iter.next(arr_state)

        out_indices = out_iter.indices(out_state)
        if out_indices[axis] == 0:
            if identity is not None:
                w_item = func(calc_dtype, identity, w_item)
        else:
            cur_value = temp_iter.getitem(temp_state)
            w_item = func(calc_dtype, cur_value, w_item)

        out_iter.setitem(out_state, w_item)
        out_state = out_iter.next(out_state)
        temp_iter.setitem(temp_state, w_item)
        temp_state = temp_iter.next(temp_state)
    return w_out

def fill(arr, box):
    arr_iter, arr_state = arr.create_iter()
    while not arr_iter.done(arr_state):
        arr_iter.setitem(arr_state, box)
        arr_state = arr_iter.next(arr_state)

def assign(space, arr, seq):
    arr_iter, arr_state = arr.create_iter()
    arr_dtype = arr.get_dtype()
    for item in seq:
        arr_iter.setitem(arr_state, arr_dtype.coerce(space, item))
        arr_state = arr_iter.next(arr_state)

where_driver = jit.JitDriver(name='numpy_where',
                             greens = ['shapelen', 'dtype', 'arr_dtype'],
                             reds = 'auto',
                             vectorize=True)

def where(space, out, shape, arr, x, y, dtype):
    out_iter, out_state = out.create_iter(shape)
    arr_iter, arr_state = arr.create_iter(shape)
    arr_dtype = arr.get_dtype()
    x_iter, x_state = x.create_iter(shape)
    y_iter, y_state = y.create_iter(shape)
    if x.is_scalar():
        if y.is_scalar():
            iter, state = arr_iter, arr_state
        else:
            iter, state = y_iter, y_state
    else:
        iter, state = x_iter, x_state
    out_iter.track_index = x_iter.track_index = False
    arr_iter.track_index = y_iter.track_index = False
    iter.track_index = True
    shapelen = len(shape)
    while not iter.done(state):
        where_driver.jit_merge_point(shapelen=shapelen, dtype=dtype,
                                        arr_dtype=arr_dtype)
        w_cond = arr_iter.getitem(arr_state)
        if arr_dtype.itemtype.bool(w_cond):
            w_val = x_iter.getitem(x_state).convert_to(space, dtype)
        else:
            w_val = y_iter.getitem(y_state).convert_to(space, dtype)
        out_iter.setitem(out_state, w_val)
        out_state = out_iter.next(out_state)
        arr_state = arr_iter.next(arr_state)
        x_state = x_iter.next(x_state)
        y_state = y_iter.next(y_state)
        if x.is_scalar():
            if y.is_scalar():
                state = arr_state
            else:
                state = y_state
        else:
            state = x_state
    return out

def _new_argmin_argmax(op_name):
    arg_driver = jit.JitDriver(name='numpy_' + op_name,
                               greens = ['shapelen', 'dtype'],
                               reds = 'auto')
    arg_flat_driver = jit.JitDriver(name='numpy_flat_' + op_name,
                                    greens = ['shapelen', 'dtype'],
                                    reds = 'auto')

    def argmin_argmax(space, w_arr, w_out, axis):
        from pypy.module.micronumpy.descriptor import get_dtype_cache
        dtype = w_arr.get_dtype()
        shapelen = len(w_arr.get_shape())
        axis_flags = [False] * shapelen
        axis_flags[axis] = True
        inner_iter, outer_iter = split_iter(w_arr.implementation, axis_flags)
        outer_state = outer_iter.reset()
        out_iter, out_state = w_out.create_iter()
        while not outer_iter.done(outer_state):
            inner_state = inner_iter.reset()
            inner_state.offset = outer_state.offset
            cur_best = inner_iter.getitem(inner_state)
            inner_state = inner_iter.next(inner_state)
            result = 0
            idx = 1
            while not inner_iter.done(inner_state):
                arg_driver.jit_merge_point(shapelen=shapelen, dtype=dtype)
                w_val = inner_iter.getitem(inner_state)
                old_best = getattr(dtype.itemtype, op_name)(cur_best, w_val)
                if not old_best:
                    result = idx
                    cur_best = w_val
                inner_state = inner_iter.next(inner_state)
                idx += 1
            result = get_dtype_cache(space).w_longdtype.box(result)
            out_iter.setitem(out_state, result)
            out_state = out_iter.next(out_state)
            outer_state = outer_iter.next(outer_state)
        return w_out

    def argmin_argmax_flat(w_arr):
        result = 0
        idx = 1
        dtype = w_arr.get_dtype()
        iter, state = w_arr.create_iter()
        cur_best = iter.getitem(state)
        state = iter.next(state)
        shapelen = len(w_arr.get_shape())
        while not iter.done(state):
            arg_flat_driver.jit_merge_point(shapelen=shapelen, dtype=dtype)
            w_val = iter.getitem(state)
            old_best = getattr(dtype.itemtype, op_name)(cur_best, w_val)
            if not old_best:
                result = idx
                cur_best = w_val
            state = iter.next(state)
            idx += 1
        return result

    return argmin_argmax, argmin_argmax_flat
argmin, argmin_flat = _new_argmin_argmax('argmin')
argmax, argmax_flat = _new_argmin_argmax('argmax')

dot_driver = jit.JitDriver(name = 'numpy_dot',
                           greens = ['dtype'],
                           reds = 'auto',
                           vectorize=True)

def multidim_dot(space, left, right, result, dtype, right_critical_dim):
    ''' assumes left, right are concrete arrays
    given left.shape == [3, 5, 7],
          right.shape == [2, 7, 4]
    then
     result.shape == [3, 5, 2, 4]
     broadcast shape should be [3, 5, 2, 7, 4]
     result should skip dims 3 which is len(result_shape) - 1
        (note that if right is 1d, result should
                  skip len(result_shape))
     left should skip 2, 4 which is a.ndims-1 + range(right.ndims)
          except where it==(right.ndims-2)
     right should skip 0, 1
    '''
    left_shape = left.get_shape()
    right_shape = right.get_shape()
    left_impl = left.implementation
    right_impl = right.implementation
    assert left_shape[-1] == right_shape[right_critical_dim]
    assert result.get_dtype() == dtype
    outi, outs = result.create_iter()
    outi.track_index = False
    lefti = AllButAxisIter(left_impl, len(left_shape) - 1)
    righti = AllButAxisIter(right_impl, right_critical_dim)
    lefts = lefti.reset()
    rights = righti.reset()
    n = left_impl.shape[-1]
    s1 = left_impl.strides[-1]
    s2 = right_impl.strides[right_critical_dim]
    while not lefti.done(lefts):
        while not righti.done(rights):
            oval = outi.getitem(outs)
            i1 = lefts.offset
            i2 = rights.offset
            i = 0
            while i < n:
                i += 1
                dot_driver.jit_merge_point(dtype=dtype)
                lval = left_impl.getitem(i1).convert_to(space, dtype)
                rval = right_impl.getitem(i2).convert_to(space, dtype)
                oval = dtype.itemtype.add(oval, dtype.itemtype.mul(lval, rval))
                i1 += jit.promote(s1)
                i2 += jit.promote(s2)
            outi.setitem(outs, oval)
            outs = outi.next(outs)
            rights = righti.next(rights)
        rights = righti.reset(rights)
        lefts = lefti.next(lefts)
    return result

count_all_true_driver = jit.JitDriver(name = 'numpy_count',
                                      greens = ['shapelen', 'dtype'],
                                      reds = 'auto',
                                      vectorize=True)

def count_all_true_concrete(impl):
    s = 0
    iter, state = impl.create_iter()
    shapelen = len(impl.shape)
    dtype = impl.dtype
    while not iter.done(state):
        count_all_true_driver.jit_merge_point(shapelen=shapelen, dtype=dtype)
        s += iter.getitem_bool(state)
        state = iter.next(state)
    return s

def count_all_true(arr):
    if arr.is_scalar():
        return arr.get_dtype().itemtype.bool(arr.get_scalar_value())
    else:
        return count_all_true_concrete(arr.implementation)

nonzero_driver = jit.JitDriver(name = 'numpy_nonzero',
                               greens = ['shapelen', 'dims', 'dtype'],
                               reds = 'auto',
                               vectorize=True)

def nonzero(res, arr, box):
    res_iter, res_state = res.create_iter()
    arr_iter, arr_state = arr.create_iter()
    shapelen = len(arr.shape)
    dtype = arr.dtype
    dims = range(shapelen)
    while not arr_iter.done(arr_state):
        nonzero_driver.jit_merge_point(shapelen=shapelen, dims=dims, dtype=dtype)
        if arr_iter.getitem_bool(arr_state):
            arr_indices = arr_iter.indices(arr_state)
            for d in dims:
                res_iter.setitem(res_state, box(arr_indices[d]))
                res_state = res_iter.next(res_state)
        arr_state = arr_iter.next(arr_state)
    return res


getitem_filter_driver = jit.JitDriver(name = 'numpy_getitem_bool',
                                      greens = ['shapelen', 'arr_dtype',
                                                'index_dtype'],
                                      reds = 'auto',
                                      vectorize=True)

def getitem_filter(res, arr, index):
    res_iter, res_state = res.create_iter()
    shapelen = len(arr.get_shape())
    if shapelen > 1 and len(index.get_shape()) < 2:
        index_iter, index_state = index.create_iter(arr.get_shape(), backward_broadcast=True)
    else:
        index_iter, index_state = index.create_iter()
    arr_iter, arr_state = arr.create_iter()
    arr_dtype = arr.get_dtype()
    index_dtype = index.get_dtype()
    # support the deprecated form where arr([True]) will return arr[0, ...]
    # by iterating over res_iter, not index_iter
    while not res_iter.done(res_state):
        getitem_filter_driver.jit_merge_point(shapelen=shapelen,
                                              index_dtype=index_dtype,
                                              arr_dtype=arr_dtype,
                                              )
        if index_iter.getitem_bool(index_state):
            res_iter.setitem(res_state, arr_iter.getitem(arr_state))
            res_state = res_iter.next(res_state)
        index_state = index_iter.next(index_state)
        arr_state = arr_iter.next(arr_state)
    return res

setitem_filter_driver = jit.JitDriver(name = 'numpy_setitem_bool',
                                      greens = ['shapelen', 'arr_dtype',
                                                'index_dtype'],
                                      reds = 'auto',
                                      vectorize=True)

def setitem_filter(space, arr, index, value):
    arr_iter, arr_state = arr.create_iter()
    shapelen = len(arr.get_shape())
    if shapelen > 1 and len(index.get_shape()) < 2:
        index_iter, index_state = index.create_iter(arr.get_shape(), backward_broadcast=True)
    else:
        index_iter, index_state = index.create_iter()
    if value.get_size() == 1:
        value_iter, value_state = value.create_iter(arr.get_shape())
    else:
        value_iter, value_state = value.create_iter()
    index_dtype = index.get_dtype()
    arr_dtype = arr.get_dtype()
    while not index_iter.done(index_state):
        setitem_filter_driver.jit_merge_point(shapelen=shapelen,
                                              index_dtype=index_dtype,
                                              arr_dtype=arr_dtype,
                                             )
        if index_iter.getitem_bool(index_state):
            val = arr_dtype.coerce(space, value_iter.getitem(value_state))
            value_state = value_iter.next(value_state)
            arr_iter.setitem(arr_state, val)
        arr_state = arr_iter.next(arr_state)
        index_state = index_iter.next(index_state)

flatiter_getitem_driver = jit.JitDriver(name = 'numpy_flatiter_getitem',
                                        greens = ['dtype'],
                                        reds = 'auto',
                                        vectorize=True)

def flatiter_getitem(res, base_iter, base_state, step):
    ri, rs = res.create_iter()
    dtype = res.get_dtype()
    while not ri.done(rs):
        flatiter_getitem_driver.jit_merge_point(dtype=dtype)
        ri.setitem(rs, base_iter.getitem(base_state))
        base_state = base_iter.goto(base_state.index + step)
        rs = ri.next(rs)
    return res

flatiter_setitem_driver = jit.JitDriver(name = 'numpy_flatiter_setitem',
                                        greens = ['dtype'],
                                        reds = 'auto',
                                        vectorize=True)

def flatiter_setitem(space, dtype, val, arr_iter, arr_state, step, length):
    val_iter, val_state = val.create_iter()
    while length > 0:
        flatiter_setitem_driver.jit_merge_point(dtype=dtype)
        val = val_iter.getitem(val_state)
        if dtype.is_str_or_unicode():
            val = dtype.coerce(space, val)
        else:
            val = val.convert_to(space, dtype)
        arr_iter.setitem(arr_state, val)
        arr_state = arr_iter.goto(arr_state.index + step)
        val_state = val_iter.next(val_state)
        if val_iter.done(val_state):
            val_state = val_iter.reset(val_state)
        length -= 1

fromstring_driver = jit.JitDriver(name = 'numpy_fromstring',
                                  greens = ['itemsize', 'dtype'],
                                  reds = 'auto')

def fromstring_loop(space, a, dtype, itemsize, s):
    i = 0
    ai, state = a.create_iter()
    while not ai.done(state):
        fromstring_driver.jit_merge_point(dtype=dtype, itemsize=itemsize)
        sub = s[i*itemsize:i*itemsize + itemsize]
        val = dtype.runpack_str(space, sub)
        ai.setitem(state, val)
        state = ai.next(state)
        i += 1

def tostring(space, arr):
    builder = StringBuilder()
    iter, state = arr.create_iter()
    w_res_str = W_NDimArray.from_shape(space, [1], arr.get_dtype())
    itemsize = arr.get_dtype().elsize
    with w_res_str.implementation as storage:
        res_str_casted = rffi.cast(rffi.CArrayPtr(lltype.Char),
                               support.get_storage_as_int(storage))
        while not iter.done(state):
            w_res_str.implementation.setitem(0, iter.getitem(state))
            for i in range(itemsize):
                builder.append(res_str_casted[i])
            state = iter.next(state)
        return builder.build()

getitem_int_driver = jit.JitDriver(name = 'numpy_getitem_int',
                                   greens = ['shapelen', 'indexlen',
                                             'prefixlen', 'dtype'],
                                   reds = 'auto')

def getitem_array_int(space, arr, res, iter_shape, indexes_w, prefix_w):
    shapelen = len(iter_shape)
    prefixlen = len(prefix_w)
    indexlen = len(indexes_w)
    dtype = arr.get_dtype()
    iter = PureShapeIter(iter_shape, indexes_w)
    while not iter.done():
        getitem_int_driver.jit_merge_point(shapelen=shapelen, indexlen=indexlen,
                                           dtype=dtype, prefixlen=prefixlen)
        # prepare the index
        index_w = [None] * indexlen
        for i in range(indexlen):
            if iter.idx_w_i[i] is not None:
                index_w[i] = iter.idx_w_i[i].getitem(iter.idx_w_s[i])
            else:
                index_w[i] = indexes_w[i]
        res.descr_setitem(space, space.newtuple(prefix_w[:prefixlen] +
                                            iter.get_index(space, shapelen)),
                          arr.descr_getitem(space, space.newtuple(index_w)))
        iter.next()
    return res

setitem_int_driver = jit.JitDriver(name = 'numpy_setitem_int',
                                   greens = ['shapelen', 'indexlen',
                                             'prefixlen', 'dtype'],
                                   reds = 'auto')

def setitem_array_int(space, arr, iter_shape, indexes_w, val_arr,
                      prefix_w):
    shapelen = len(iter_shape)
    indexlen = len(indexes_w)
    prefixlen = len(prefix_w)
    dtype = arr.get_dtype()
    iter = PureShapeIter(iter_shape, indexes_w)
    while not iter.done():
        setitem_int_driver.jit_merge_point(shapelen=shapelen, indexlen=indexlen,
                                           dtype=dtype, prefixlen=prefixlen)
        # prepare the index
        index_w = [None] * indexlen
        for i in range(indexlen):
            if iter.idx_w_i[i] is not None:
                index_w[i] = iter.idx_w_i[i].getitem(iter.idx_w_s[i])
            else:
                index_w[i] = indexes_w[i]
        w_idx = space.newtuple(prefix_w[:prefixlen] + iter.get_index(space,
                                                                  shapelen))
        if val_arr.is_scalar():
            w_value = val_arr.get_scalar_value()
        else:
            w_value = val_arr.descr_getitem(space, w_idx)
        arr.descr_setitem(space, space.newtuple(index_w), w_value)
        iter.next()

byteswap_driver = jit.JitDriver(name='numpy_byteswap_driver',
                                greens = ['dtype'],
                                reds = 'auto',
                                vectorize=True)

def byteswap(from_, to):
    dtype = from_.dtype
    from_iter, from_state = from_.create_iter()
    to_iter, to_state = to.create_iter()
    while not from_iter.done(from_state):
        byteswap_driver.jit_merge_point(dtype=dtype)
        val = dtype.itemtype.byteswap(from_iter.getitem(from_state))
        to_iter.setitem(to_state, val)
        to_state = to_iter.next(to_state)
        from_state = from_iter.next(from_state)

choose_driver = jit.JitDriver(name='numpy_choose_driver',
                              greens = ['shapelen', 'mode', 'dtype'],
                              reds = 'auto',
                              vectorize=True)

def choose(space, arr, choices, shape, dtype, out, mode):
    shapelen = len(shape)
    pairs = [a.create_iter(shape) for a in choices]
    iterators = [i[0] for i in pairs]
    states = [i[1] for i in pairs]
    arr_iter, arr_state = arr.create_iter(shape)
    out_iter, out_state = out.create_iter(shape)
    while not arr_iter.done(arr_state):
        choose_driver.jit_merge_point(shapelen=shapelen, dtype=dtype,
                                      mode=mode)
        index = support.index_w(space, arr_iter.getitem(arr_state))
        if index < 0 or index >= len(iterators):
            if mode == NPY.RAISE:
                raise oefmt(space.w_ValueError,
                            "invalid entry in choice array")
            elif mode == NPY.WRAP:
                index = index % (len(iterators))
            else:
                assert mode == NPY.CLIP
                if index < 0:
                    index = 0
                else:
                    index = len(iterators) - 1
        val = iterators[index].getitem(states[index]).convert_to(space, dtype)
        out_iter.setitem(out_state, val)
        for i in range(len(iterators)):
            states[i] = iterators[i].next(states[i])
        out_state = out_iter.next(out_state)
        arr_state = arr_iter.next(arr_state)

clip_driver = jit.JitDriver(name='numpy_clip_driver',
                            greens = ['shapelen', 'dtype'],
                            reds = 'auto',
                            vectorize=True)

def clip(space, arr, shape, min, max, out):
    assert min or max
    arr_iter, arr_state = arr.create_iter(shape)
    if min is not None:
        min_iter, min_state = min.create_iter(shape)
    else:
        min_iter, min_state = None, None
    if max is not None:
        max_iter, max_state = max.create_iter(shape)
    else:
        max_iter, max_state = None, None
    out_iter, out_state = out.create_iter(shape)
    shapelen = len(shape)
    dtype = out.get_dtype()
    while not arr_iter.done(arr_state):
        clip_driver.jit_merge_point(shapelen=shapelen, dtype=dtype)
        w_v = arr_iter.getitem(arr_state).convert_to(space, dtype)
        arr_state = arr_iter.next(arr_state)
        if min_iter is not None:
            w_min = min_iter.getitem(min_state).convert_to(space, dtype)
            if dtype.itemtype.lt(w_v, w_min):
                w_v = w_min
            min_state = min_iter.next(min_state)
        if max_iter is not None:
            w_max = max_iter.getitem(max_state).convert_to(space, dtype)
            if dtype.itemtype.gt(w_v, w_max):
                w_v = w_max
            max_state = max_iter.next(max_state)
        out_iter.setitem(out_state, w_v)
        out_state = out_iter.next(out_state)

round_driver = jit.JitDriver(name='numpy_round_driver',
                             greens = ['shapelen', 'dtype'],
                             reds = 'auto',
                             vectorize=True)

def round(space, arr, dtype, shape, decimals, out):
    arr_iter, arr_state = arr.create_iter(shape)
    out_iter, out_state = out.create_iter(shape)
    shapelen = len(shape)
    while not arr_iter.done(arr_state):
        round_driver.jit_merge_point(shapelen=shapelen, dtype=dtype)
        w_v = arr_iter.getitem(arr_state).convert_to(space, dtype)
        w_v = dtype.itemtype.round(w_v, decimals)
        out_iter.setitem(out_state, w_v)
        arr_state = arr_iter.next(arr_state)
        out_state = out_iter.next(out_state)

diagonal_simple_driver = jit.JitDriver(name='numpy_diagonal_simple_driver',
                                       greens = ['axis1', 'axis2'],
                                       reds = 'auto')

def diagonal_simple(space, arr, out, offset, axis1, axis2, size):
    out_iter, out_state = out.create_iter()
    i = 0
    index = [0] * 2
    while i < size:
        diagonal_simple_driver.jit_merge_point(axis1=axis1, axis2=axis2)
        index[axis1] = i
        index[axis2] = i + offset
        out_iter.setitem(out_state, arr.getitem_index(space, index))
        i += 1
        out_state = out_iter.next(out_state)

def diagonal_array(space, arr, out, offset, axis1, axis2, shape):
    out_iter, out_state = out.create_iter()
    iter = PureShapeIter(shape, [])
    shapelen_minus_1 = len(shape) - 1
    assert shapelen_minus_1 >= 0
    if axis1 < axis2:
        a = axis1
        b = axis2 - 1
    else:
        a = axis2
        b = axis1 - 1
    assert a >= 0
    assert b >= 0
    while not iter.done():
        last_index = iter.indexes[-1]
        if axis1 < axis2:
            indexes = (iter.indexes[:a] + [last_index] +
                       iter.indexes[a:b] + [last_index + offset] +
                       iter.indexes[b:shapelen_minus_1])
        else:
            indexes = (iter.indexes[:a] + [last_index + offset] +
                       iter.indexes[a:b] + [last_index] +
                       iter.indexes[b:shapelen_minus_1])
        out_iter.setitem(out_state, arr.getitem_index(space, indexes))
        iter.next()
        out_state = out_iter.next(out_state)

def _new_binsearch(side, op_name):
    binsearch_driver = jit.JitDriver(name='numpy_binsearch_' + side,
                                     greens=['dtype'],
                                     reds='auto')

    def binsearch(space, arr, key, ret):
        assert len(arr.get_shape()) == 1
        dtype = key.get_dtype()
        op = getattr(dtype.itemtype, op_name)
        key_iter, key_state = key.create_iter()
        ret_iter, ret_state = ret.create_iter()
        ret_iter.track_index = False
        size = arr.get_size()
        min_idx = 0
        max_idx = size
        last_key_val = key_iter.getitem(key_state)
        while not key_iter.done(key_state):
            key_val = key_iter.getitem(key_state)
            if dtype.itemtype.lt(last_key_val, key_val):
                max_idx = size
            else:
                min_idx = 0
                max_idx = max_idx + 1 if max_idx < size else size
            last_key_val = key_val
            while min_idx < max_idx:
                binsearch_driver.jit_merge_point(dtype=dtype)
                mid_idx = min_idx + ((max_idx - min_idx) >> 1)
                mid_val = arr.getitem(space, [mid_idx]).convert_to(space, dtype)
                if op(mid_val, key_val):
                    min_idx = mid_idx + 1
                else:
                    max_idx = mid_idx
            ret_iter.setitem(ret_state, ret.get_dtype().box(min_idx))
            ret_state = ret_iter.next(ret_state)
            key_state = key_iter.next(key_state)
    return binsearch

binsearch_left = _new_binsearch('left', 'lt')
binsearch_right = _new_binsearch('right', 'le')