pypy / pypy / module / micronumpy / arrayimpl / sort.py

from pypy.rpython.lltypesystem import rffi, lltype
from pypy.rlib.listsort import make_timsort_class
from pypy.rlib.objectmodel import specialize
from pypy.rlib.rawstorage import raw_storage_getitem, raw_storage_setitem
from pypy.module.micronumpy.base import W_NDimArray
from pypy.module.micronumpy import interp_dtype

INT_SIZE = rffi.sizeof(lltype.Signed)

@specialize.memo()
def make_sort_classes(space, TP):
    class ArgArrayRepresentation(object):
        def __init__(self, itemsize, size, values, indexes):
            self.itemsize = itemsize
            self.size = size
            self.values = values
            self.indexes = indexes

        def getitem(self, item):
            return (raw_storage_getitem(TP, self.values, item * self.itemsize),
                    raw_storage_getitem(lltype.Signed, self.indexes,
                                        item * INT_SIZE))

        def setitem(self, idx, item):
            raw_storage_setitem(self.values, idx * self.itemsize,
                                rffi.cast(TP, item[0]))
            raw_storage_setitem(self.indexes, idx * INT_SIZE, item[1])

    def arg_getitem(lst, item):
        return lst.getitem(item)

    def arg_setitem(lst, item, value):
        lst.setitem(item, value)

    def arg_length(lst):
        return lst.size

    def arg_getitem_slice(lst, start, stop):
        xxx

    def arg_lt(a, b):
        return a[0] < b[0]

    ArgSort = make_timsort_class(arg_getitem, arg_setitem, arg_length,
                                 arg_getitem_slice, arg_lt)

    return ArgArrayRepresentation, ArgSort

def sort_array(arr, space):
    itemsize = arr.dtype.itemtype.get_element_size()
    # create array of indexes
    dtype = interp_dtype.get_dtype_cache(space).w_longdtype
    indexes = W_NDimArray.from_shape([arr.get_size()], dtype)
    storage = indexes.implementation.get_storage()
    for i in range(arr.get_size()):
        raw_storage_setitem(storage, i * INT_SIZE, i)
    Repr, Sort = make_sort_classes(space, arr.dtype.itemtype.T)
    r = Repr(itemsize, arr.get_size(), arr.get_storage(),
             indexes.implementation.get_storage())
    Sort(r).sort()
    return indexes
Tip: Filter by directory path e.g. /media app.js to search for public/media/app.js.
Tip: Use camelCasing e.g. ProjME to search for ProjectModifiedEvent.java.
Tip: Filter by extension type e.g. /repo .js to search for all .js files in the /repo directory.
Tip: Separate your search with spaces e.g. /ssh pom.xml to search for src/ssh/pom.xml.
Tip: Use ↑ and ↓ arrow keys to navigate and return to view the file.
Tip: You can also navigate files with Ctrl+j (next) and Ctrl+k (previous) and view the file with Ctrl+o.
Tip: You can also navigate files with Alt+j (next) and Alt+k (previous) and view the file with Alt+o.