Add type specific quicksort for strings and unicode in sort and argsort methods.

1 files changed, 211 insertions, 33 deletions

diff --git a/numpy/core/src/_sortmodule.c.src b/numpy/core/src/_sortmodule.c.src
index 573ee0f9c..547c560fd 100644
--- a/numpy/core/src/_sortmodule.c.src
+++ b/numpy/core/src/_sortmodule.c.src
@@ -41,42 +41,10 @@
 #define STRING_LT(pa, pb, len) (compare_string(pa, pb, len) < 0)
 #define STRING_LE(pa, pb, len) (compare_string(pa, pb, len) <= 0)
 #define STRING_EQ(pa, pb, len) (compare_string(pa, pb, len) == 0)
-#define STRING_SWAP(pa, pb, n) {int i; for(i = 0; i < n; ++i) SWAP(pa[i], pb[i]);}
-#define STRING_COPY(pa, pb, n) {int i; for(i = 0; i < n; ++i) pa[i] = pb[i];}
 #define UNICODE_LT(pa, pb, len) (compare_ucs4(pa, pb, len) < 0)
 #define UNICODE_LE(pa, pb, len) (compare_ucs4(pa, pb, len) <= 0)
 #define UNICODE_EQ(pa, pb, len) (compare_ucs4(pa, pb, len) == 0)
 
-#include <stdio.h>
-
-static int inline
-compare_string(char *s1, char *s2, size_t len)
-{
-    const unsigned char *c1 = (unsigned char *)s1;
-    const unsigned char *c2 = (unsigned char *)s2;
-    size_t i;
-
-    for(i = 0; i < len; ++i) {
-        if (c1[i] != c2[i]) {
-            return (c1[i] > c2[i]) ? 1 : -1;
-        }
-    }
-    return 0;
-}
-
-
-static int inline
-compare_ucs4(npy_ucs4 *s1, npy_ucs4 *s2, size_t len)
-{
-    size_t i;
-
-    for(i = 0; i < len; ++i) {
-        if (s1[i] != s2[i]) {
-            return (s1[i] > s2[i]) ? 1 : -1;
-        }
-    }
-    return 0;
-}
 
 /**begin repeat
    #TYPE=BOOL,BYTE,UBYTE,SHORT,USHORT,INT,UINT,LONG,ULONG,LONGLONG,ULONGLONG,FLOAT,DOUBLE,LONGDOUBLE,CFLOAT,CDOUBLE,CLONGDOUBLE#
@@ -407,12 +375,212 @@ static int
 }
 /**end repeat**/
 
+/*
+ * Subroutines that will hopefully be inlined when the code
+ * for strings and unicode is compiled with proper flags.
+ */
+
+static int
+compare_string(char *s1, char *s2, size_t len)
+{
+    const unsigned char *c1 = (unsigned char *)s1;
+    const unsigned char *c2 = (unsigned char *)s2;
+    size_t i;
+
+    for(i = 0; i < len; ++i) {
+        if (c1[i] != c2[i]) {
+            return (c1[i] > c2[i]) ? 1 : -1;
+        }
+    }
+    return 0;
+}
+
+static void
+copy_string(char *s1, char *s2, size_t len)
+{
+    while(len--) {
+        *s1++ = *s2++;
+    }
+}
+
+static void
+swap_string(char *s1, char *s2, size_t len)
+{
+    while(len--) {
+        const char t = *s1;
+        *s1++ = *s2;
+        *s2++ = t;
+    }
+}
+
+
+static int
+compare_ucs4(npy_ucs4 *s1, npy_ucs4 *s2, size_t len)
+{
+    size_t i;
+
+    for(i = 0; i < len; ++i) {
+        if (s1[i] != s2[i]) {
+            return (s1[i] > s2[i]) ? 1 : -1;
+        }
+    }
+    return 0;
+}
+
+
+static void
+copy_ucs4(npy_ucs4 *s1, npy_ucs4 *s2, size_t len)
+{
+    while(len--) {
+        *s1++ = *s2++;
+    }
+}
+
+
+static void
+swap_ucs4(npy_ucs4 *s1, npy_ucs4 *s2, size_t len)
+{
+    while(len--) {
+        const npy_ucs4 t = *s1;
+        *s1++ = *s2;
+        *s2++ = t;
+    }
+}
+
 /**begin repeat
    #TYPE=STRING, UNICODE#
    #type=char, PyArray_UCS4#
    #lessthan=STRING_LT, UNICODE_LT#
    #lessequal=STRING_LE, UNICODE_LE#
+   #swap=swap_string, swap_ucs4#
+   #copy=copy_string, copy_ucs4#
 **/
+
+static int
+@TYPE@_quicksort(@type@ *start, intp num, PyArrayObject *arr)
+{
+    const size_t len = arr->descr->elsize/sizeof(@type@);
+    @type@ *vp = malloc(arr->descr->elsize);
+    @type@ *pl = start;
+    @type@ *pr = start + (num - 1)*len;
+    @type@ *stack[PYA_QS_STACK], **sptr = stack, *pm, *pi, *pj, *pk;
+
+    for(;;) {
+        while ((pr - pl) > 5*len) {
+            /* quicksort partition */
+            pm = pl + (((pr - pl)/len) >> 1)*len;
+            if (@lessthan@(pm, pl, len)) @swap@(pm, pl, len);
+            if (@lessthan@(pr, pm, len)) @swap@(pr, pm, len);
+            if (@lessthan@(pm, pl, len)) @swap@(pm, pl, len);
+            @copy@(vp, pm, len);
+            pi = pl;
+            pj = pr - len;
+            @swap@(pm, pj, len);
+            for(;;) {
+                do pi += len; while (@lessthan@(pi, vp, len));
+                do pj -= len; while (@lessthan@(vp, pj, len));
+                if (pi >= pj)  break;
+                @swap@(pi, pj, len);
+            }
+            pk = pr - len;
+            @swap@(pi, pk, len);
+            /* push largest partition on stack */
+            if (pi - pl < pr - pi) {
+                *sptr++ = pi + len;
+                *sptr++ = pr;
+                pr = pi - len;
+            }
+            else {
+                *sptr++ = pl;
+                *sptr++ = pi - len;
+                pl = pi + len;
+            }
+        }
+
+        /* insertion sort */
+        for(pi = pl + len; pi <= pr; pi += len) {
+            @copy@(vp, pi, len);
+            pj = pi;
+            pk = pi - len;
+            while(pj > pl && @lessthan@(vp, pk, len)) {
+                @copy@(pj, pk, len);
+                pj -= len;
+                pk -= len;
+            }
+            @copy@(pj, vp, len);
+        }
+        if (sptr == stack) break;
+        pr = *(--sptr);
+        pl = *(--sptr);
+    }
+
+    free(vp);
+    return 0;
+}
+
+
+static int
+@TYPE@_aquicksort(@type@ *v, intp* tosort, intp num, PyArrayObject *arr)
+{
+    @type@ *vp;
+    intp *pl = tosort;
+    intp *pr = tosort + num - 1;
+    intp *stack[PYA_QS_STACK];
+    intp **sptr=stack;
+    intp *pm, *pi, *pj, *pk, vi, SWAP_temp;
+    size_t len = arr->descr->elsize/sizeof(@type@);
+
+    for(;;) {
+        while ((pr - pl) > SMALL_QUICKSORT) {
+            /* quicksort partition */
+            pm = pl + ((pr - pl) >> 1);
+            if (@lessthan@(v + (*pm)*len, v + (*pl)*len, len)) SWAP(*pm, *pl);
+            if (@lessthan@(v + (*pr)*len, v + (*pm)*len, len)) SWAP(*pr, *pm);
+            if (@lessthan@(v + (*pm)*len, v + (*pl)*len, len)) SWAP(*pm, *pl);
+            vp = v + (*pm)*len;
+            pi = pl;
+            pj = pr - 1;
+            SWAP(*pm,*pj);
+            for(;;) {
+                do ++pi; while (@lessthan@(v + (*pi)*len, vp, len));
+                do --pj; while (@lessthan@(vp, v + (*pj)*len, len));
+                if (pi >= pj)  break;
+                SWAP(*pi,*pj);
+            }
+            SWAP(*pi,*(pr-1));
+            /* push largest partition on stack */
+            if (pi - pl < pr - pi) {
+                *sptr++ = pi + 1;
+                *sptr++ = pr;
+                pr = pi - 1;
+            }
+            else {
+                *sptr++ = pl;
+                *sptr++ = pi - 1;
+                pl = pi + 1;
+            }
+        }
+
+        /* insertion sort */
+        for(pi = pl + 1; pi <= pr; ++pi) {
+            vi = *pi;
+            vp = v + vi*len;
+            pj = pi;
+            pk = pi -1;
+            while(pj > pl && @lessthan@(vp, v + (*pk)*len, len)) {
+                *pj-- = *pk--;
+            }
+            *pj = vi;
+        }
+        if (sptr == stack) break;
+        pr = *(--sptr);
+        pl = *(--sptr);
+    }
+
+    return 0;
+}
+
+
 static void
 @TYPE@_amergesort0(intp *pl, intp *pr, @type@ *v, intp *pw, int len)
 {
@@ -444,7 +612,7 @@ static void
             vp = v + vi*len;
             pj = pi;
             pk = pi -1;
-            for(; pj > pl && @lessthan@(vp, v + (*pk)*len, len);) {
+            while(pj > pl && @lessthan@(vp, v + (*pk)*len, len)) {
                 *pj-- = *pk--;
             }
             *pj = vi;
@@ -452,6 +620,7 @@ static void
     }
 }
 
+
 static int
 @TYPE@_amergesort(@type@ *v, intp *tosort, intp num, PyArrayObject *arr)
 {
@@ -503,9 +672,18 @@ add_sortfuncs(void)
     descr = PyArray_DescrFromType(PyArray_STRING);
     descr->f->argsort[PyArray_MERGESORT] = \
         (PyArray_ArgSortFunc *)STRING_amergesort;
+    descr->f->argsort[PyArray_QUICKSORT] = \
+        (PyArray_ArgSortFunc *)STRING_aquicksort;
+    descr->f->sort[PyArray_QUICKSORT] = \
+        (PyArray_SortFunc *)STRING_quicksort;
+
     descr = PyArray_DescrFromType(PyArray_UNICODE);
     descr->f->argsort[PyArray_MERGESORT] = \
         (PyArray_ArgSortFunc *)UNICODE_amergesort;
+    descr->f->argsort[PyArray_QUICKSORT] = \
+        (PyArray_ArgSortFunc *)UNICODE_aquicksort;
+    descr->f->sort[PyArray_QUICKSORT] = \
+        (PyArray_SortFunc *)UNICODE_quicksort;
 }
 
 static struct PyMethodDef methods[] = {