test_sort.py

import copy
import itertools
import math
import random
import sys
from typing import KeysView

import numpy as np

from numba.core.compiler import compile_isolated, Flags
from numba import jit, njit
from numba.core import types, utils, errors
import unittest
from numba import testing
from numba.tests.support import TestCase, MemoryLeakMixin, tag

from numba.misc.quicksort import make_py_quicksort, make_jit_quicksort
from numba.misc.mergesort import make_jit_mergesort
from numba.misc.timsort import make_py_timsort, make_jit_timsort, MergeRun


def make_temp_list(keys, n):
    return [keys[0]] * n

def make_temp_array(keys, n):
    return np.empty(n, keys.dtype)


py_list_timsort = make_py_timsort(make_temp_list)

py_array_timsort = make_py_timsort(make_temp_array)

jit_list_timsort = make_jit_timsort(make_temp_list)

jit_array_timsort = make_jit_timsort(make_temp_array)

py_quicksort = make_py_quicksort()

jit_quicksort = make_jit_quicksort()


def sort_usecase(val):
    val.sort()

def argsort_usecase(val):
    return val.argsort()

def argsort_kind_usecase(val, is_stable=False):
    if is_stable:
        return val.argsort(kind='mergesort')
    else:
        return val.argsort(kind='quicksort')

def sorted_usecase(val):
    return sorted(val)

def sorted_reverse_usecase(val, b):
    return sorted(val, reverse=b)

def np_sort_usecase(val):
    return np.sort(val)

def np_argsort_usecase(val):
    return np.argsort(val)

def np_argsort_kind_usecase(val, is_stable=False):
    if is_stable:
        return np.argsort(val, kind='mergesort')
    else:
        return np.argsort(val, kind='quicksort')

def list_sort_usecase(n):
    np.random.seed(42)
    l = []
    for i in range(n):
        l.append(np.random.random())
    ll = l[:]
    ll.sort()
    return l, ll

def list_sort_reverse_usecase(n, b):
    np.random.seed(42)
    l = []
    for i in range(n):
        l.append(np.random.random())
    ll = l[:]
    ll.sort(reverse=b)
    return l, ll


class BaseSortingTest(object):

    def random_list(self, n, offset=10):
        random.seed(42)
        l = list(range(offset, offset + n))
        random.shuffle(l)
        return l

    def sorted_list(self, n, offset=10):
        return list(range(offset, offset + n))

    def revsorted_list(self, n, offset=10):
        return list(range(offset, offset + n))[::-1]

    def initially_sorted_list(self, n, m=None, offset=10):
        if m is None:
            m = n // 2
        l = self.sorted_list(m, offset)
        l += self.random_list(n - m, offset=l[-1] + offset)
        return l

    def duprandom_list(self, n, factor=None, offset=10):
        random.seed(42)
        if factor is None:
            factor = int(math.sqrt(n))
        l = (list(range(offset, offset + (n // factor) + 1)) * (factor + 1))[:n]
        assert len(l) == n
        random.shuffle(l)
        return l

    def dupsorted_list(self, n, factor=None, offset=10):
        if factor is None:
            factor = int(math.sqrt(n))
        l = (list(range(offset, offset + (n // factor) + 1)) * (factor + 1))[:n]
        assert len(l) == n, (len(l), n)
        l.sort()
        return l

    def assertSorted(self, orig, result):
        self.assertEqual(len(result), len(orig))
        # sorted() returns a list, so make sure we compare to another list
        self.assertEqual(list(result), sorted(orig))

    def assertSortedValues(self, orig, orig_values, result, result_values):
        self.assertEqual(len(result), len(orig))
        self.assertEqual(list(result), sorted(orig))
        zip_sorted = sorted(zip(orig, orig_values), key=lambda x: x[0])
        zip_result = list(zip(result, result_values))
        self.assertEqual(zip_sorted, zip_result)
        # Check stability
        for i in range(len(zip_result) - 1):
            (k1, v1), (k2, v2) = zip_result[i], zip_result[i + 1]
            if k1 == k2:
                # Assuming values are unique, which is enforced by the tests
                self.assertLess(orig_values.index(v1), orig_values.index(v2))

    def fibo(self):
        a = 1
        b = 1
        while True:
            yield a
            a, b = b, a + b

    def make_sample_sorted_lists(self, n):
        lists = []
        for offset in (20, 120):
            lists.append(self.sorted_list(n, offset))
            lists.append(self.dupsorted_list(n, offset))
        return lists

    def make_sample_lists(self, n):
        lists = []
        for offset in (20, 120):
            lists.append(self.sorted_list(n, offset))
            lists.append(self.dupsorted_list(n, offset))
            lists.append(self.revsorted_list(n, offset))
            lists.append(self.duprandom_list(n, offset))
        return lists


class BaseTimsortTest(BaseSortingTest):

    def merge_init(self, keys):
        f = self.timsort.merge_init
        return f(keys)

    def test_binarysort(self):
        n = 20
        def check(l, n, start=0):
            res = self.array_factory(l)
            f(res, res, 0, n, start)
            self.assertSorted(l, res)

        f = self.timsort.binarysort
        l = self.sorted_list(n)
        check(l, n)
        check(l, n, n//2)
        l = self.revsorted_list(n)
        check(l, n)
        l = self.initially_sorted_list(n, n//2)
        check(l, n)
        check(l, n, n//2)
        l = self.revsorted_list(n)
        check(l, n)
        l = self.random_list(n)
        check(l, n)
        l = self.duprandom_list(n)
        check(l, n)

    def test_binarysort_with_values(self):
        n = 20
        v = list(range(100, 100+n))

        def check(l, n, start=0):
            res = self.array_factory(l)
            res_v = self.array_factory(v)
            f(res, res_v, 0, n, start)
            self.assertSortedValues(l, v, res, res_v)

        f = self.timsort.binarysort
        l = self.sorted_list(n)
        check(l, n)
        check(l, n, n//2)
        l = self.revsorted_list(n)
        check(l, n)
        l = self.initially_sorted_list(n, n//2)
        check(l, n)
        check(l, n, n//2)
        l = self.revsorted_list(n)
        check(l, n)
        l = self.random_list(n)
        check(l, n)
        l = self.duprandom_list(n)
        check(l, n)

    def test_count_run(self):
        n = 16
        f = self.timsort.count_run

        def check(l, lo, hi):
            n, desc = f(self.array_factory(l), lo, hi)
            # Fully check invariants
            if desc:
                for k in range(lo, lo + n - 1):
                    a, b = l[k], l[k + 1]
                    self.assertGreater(a, b)
                if lo + n < hi:
                    self.assertLessEqual(l[lo + n - 1], l[lo + n])
            else:
                for k in range(lo, lo + n - 1):
                    a, b = l[k], l[k + 1]
                    self.assertLessEqual(a, b)
                if lo + n < hi:
                    self.assertGreater(l[lo + n - 1], l[lo + n], l)


        l = self.sorted_list(n, offset=100)
        check(l, 0, n)
        check(l, 1, n - 1)
        check(l, 1, 2)
        l = self.revsorted_list(n, offset=100)
        check(l, 0, n)
        check(l, 1, n - 1)
        check(l, 1, 2)
        l = self.random_list(n, offset=100)
        for i in range(len(l) - 1):
            check(l, i, n)
        l = self.duprandom_list(n, offset=100)
        for i in range(len(l) - 1):
            check(l, i, n)

    def test_gallop_left(self):
        n = 20
        f = self.timsort.gallop_left

        def check(l, key, start, stop, hint):
            k = f(key, l, start, stop, hint)
            # Fully check invariants
            self.assertGreaterEqual(k, start)
            self.assertLessEqual(k, stop)
            if k > start:
                self.assertLess(l[k - 1], key)
            if k < stop:
                self.assertGreaterEqual(l[k], key)

        def check_all_hints(l, key, start, stop):
            for hint in range(start, stop):
                check(l, key, start, stop, hint)

        def check_sorted_list(l):
            l = self.array_factory(l)
            for key in (l[5], l[15], l[0], -1000, l[-1], 1000):
                check_all_hints(l, key, 0, n)
                check_all_hints(l, key, 1, n - 1)
                check_all_hints(l, key, 8, n - 8)

        l = self.sorted_list(n, offset=100)
        check_sorted_list(l)
        l = self.dupsorted_list(n, offset=100)
        check_sorted_list(l)

    def test_gallop_right(self):
        n = 20
        f = self.timsort.gallop_right

        def check(l, key, start, stop, hint):
            k = f(key, l, start, stop, hint)
            # Fully check invariants
            self.assertGreaterEqual(k, start)
            self.assertLessEqual(k, stop)
            if k > start:
                self.assertLessEqual(l[k - 1], key)
            if k < stop:
                self.assertGreater(l[k], key)

        def check_all_hints(l, key, start, stop):
            for hint in range(start, stop):
                check(l, key, start, stop, hint)

        def check_sorted_list(l):
            l = self.array_factory(l)
            for key in (l[5], l[15], l[0], -1000, l[-1], 1000):
                check_all_hints(l, key, 0, n)
                check_all_hints(l, key, 1, n - 1)
                check_all_hints(l, key, 8, n - 8)

        l = self.sorted_list(n, offset=100)
        check_sorted_list(l)
        l = self.dupsorted_list(n, offset=100)
        check_sorted_list(l)

    def test_merge_compute_minrun(self):
        f = self.timsort.merge_compute_minrun

        for i in range(0, 64):
            self.assertEqual(f(i), i)
        for i in range(6, 63):
            if 2**i > sys.maxsize:
                break
            self.assertEqual(f(2**i), 32)
        for i in self.fibo():
            if i < 64:
                continue
            if i >= sys.maxsize:
                break
            k = f(i)
            self.assertGreaterEqual(k, 32)
            self.assertLessEqual(k, 64)
            if i > 500:
                # i/k is close to, but strictly less than, an exact power of 2
                quot = i // k
                p = 2 ** utils.bit_length(quot)
                self.assertLess(quot, p)
                self.assertGreaterEqual(quot, 0.9 * p)

    def check_merge_lo_hi(self, func, a, b):
        na = len(a)
        nb = len(b)

        # Add sentinels at start and end, to check they weren't moved
        orig_keys = [42] + a + b + [-42]
        keys = self.array_factory(orig_keys)
        ms = self.merge_init(keys)
        ssa = 1
        ssb = ssa + na

        #new_ms = func(ms, keys, [], ssa, na, ssb, nb)
        new_ms = func(ms, keys, keys, ssa, na, ssb, nb)
        self.assertEqual(keys[0], orig_keys[0])
        self.assertEqual(keys[-1], orig_keys[-1])
        self.assertSorted(orig_keys[1:-1], keys[1:-1])
        # Check the MergeState result
        self.assertGreaterEqual(len(new_ms.keys), len(ms.keys))
        self.assertGreaterEqual(len(new_ms.values), len(ms.values))
        self.assertIs(new_ms.pending, ms.pending)
        self.assertGreaterEqual(new_ms.min_gallop, 1)

    def test_merge_lo_hi(self):
        f_lo = self.timsort.merge_lo
        f_hi = self.timsort.merge_hi

        # The larger sizes exercise galloping
        for (na, nb) in [(12, 16), (40, 40), (100, 110), (1000, 1100)]:
            for a, b in itertools.product(self.make_sample_sorted_lists(na),
                                          self.make_sample_sorted_lists(nb)):
                self.check_merge_lo_hi(f_lo, a, b)
                self.check_merge_lo_hi(f_hi, b, a)

    def check_merge_at(self, a, b):
        f = self.timsort.merge_at
        # Prepare the array to be sorted
        na = len(a)
        nb = len(b)
        # Add sentinels at start and end, to check they weren't moved
        orig_keys = [42] + a + b + [-42]
        ssa = 1
        ssb = ssa + na

        stack_sentinel = MergeRun(-42, -42)

        def run_merge_at(ms, keys, i):
            new_ms = f(ms, keys, keys, i)
            self.assertEqual(keys[0], orig_keys[0])
            self.assertEqual(keys[-1], orig_keys[-1])
            self.assertSorted(orig_keys[1:-1], keys[1:-1])
            # Check stack state
            self.assertIs(new_ms.pending, ms.pending)
            self.assertEqual(ms.pending[i], (ssa, na + nb))
            self.assertEqual(ms.pending[0], stack_sentinel)
            return new_ms

        # First check with i == len(stack) - 2
        keys = self.array_factory(orig_keys)
        ms = self.merge_init(keys)
        # Push sentinel on stack, to check it wasn't touched
        ms = self.timsort.merge_append(ms, stack_sentinel)
        i = ms.n
        ms = self.timsort.merge_append(ms, MergeRun(ssa, na))
        ms = self.timsort.merge_append(ms, MergeRun(ssb, nb))
        ms = run_merge_at(ms, keys, i)
        self.assertEqual(ms.n, i + 1)

        # Now check with i == len(stack) - 3
        keys = self.array_factory(orig_keys)
        ms = self.merge_init(keys)
        # Push sentinel on stack, to check it wasn't touched
        ms = self.timsort.merge_append(ms, stack_sentinel)
        i = ms.n
        ms = self.timsort.merge_append(ms, MergeRun(ssa, na))
        ms = self.timsort.merge_append(ms, MergeRun(ssb, nb))
        # A last run (trivial here)
        last_run = MergeRun(ssb + nb, 1)
        ms = self.timsort.merge_append(ms, last_run)
        ms = run_merge_at(ms, keys, i)
        self.assertEqual(ms.n, i + 2)
        self.assertEqual(ms.pending[ms.n - 1], last_run)

    def test_merge_at(self):
        # The larger sizes exercise galloping
        for (na, nb) in [(12, 16), (40, 40), (100, 110), (500, 510)]:
            for a, b in itertools.product(self.make_sample_sorted_lists(na),
                                          self.make_sample_sorted_lists(nb)):
                self.check_merge_at(a, b)
                self.check_merge_at(b, a)

    def test_merge_force_collapse(self):
        f = self.timsort.merge_force_collapse

        # Test with runs of ascending sizes, then descending sizes
        sizes_list = [(8, 10, 15, 20)]
        sizes_list.append(sizes_list[0][::-1])

        for sizes in sizes_list:
            for chunks in itertools.product(*(self.make_sample_sorted_lists(n)
                                              for n in sizes)):
                # Create runs of the given sizes
                orig_keys = sum(chunks, [])
                keys = self.array_factory(orig_keys)
                ms = self.merge_init(keys)
                pos = 0
                for c in chunks:
                    ms = self.timsort.merge_append(ms, MergeRun(pos, len(c)))
                    pos += len(c)
                # Sanity check
                self.assertEqual(sum(ms.pending[ms.n - 1]), len(keys))
                # Now merge the runs
                ms = f(ms, keys, keys)
                # Remaining run is the whole list
                self.assertEqual(ms.n, 1)
                self.assertEqual(ms.pending[0], MergeRun(0, len(keys)))
                # The list is now sorted
                self.assertSorted(orig_keys, keys)

    def test_run_timsort(self):
        f = self.timsort.run_timsort

        for size_factor in (1, 10):
            # Make lists to be sorted from three chunks of different kinds.
            sizes = (15, 30, 20)

            all_lists = [self.make_sample_lists(n * size_factor) for n in sizes]
            for chunks in itertools.product(*all_lists):
                orig_keys = sum(chunks, [])
                keys = self.array_factory(orig_keys)
                f(keys)
                # The list is now sorted
                self.assertSorted(orig_keys, keys)

    def test_run_timsort_with_values(self):
        # Run timsort, but also with a values array
        f = self.timsort.run_timsort_with_values

        for size_factor in (1, 5):
            chunk_size = 80 * size_factor
            a = self.dupsorted_list(chunk_size)
            b = self.duprandom_list(chunk_size)
            c = self.revsorted_list(chunk_size)
            orig_keys = a + b + c
            orig_values = list(range(1000, 1000 + len(orig_keys)))

            keys = self.array_factory(orig_keys)
            values = self.array_factory(orig_values)
            f(keys, values)
            # This checks sort stability
            self.assertSortedValues(orig_keys, orig_values, keys, values)


class TestTimsortPurePython(BaseTimsortTest, TestCase):

    timsort = py_list_timsort

    # Much faster than a Numpy array in pure Python
    array_factory = list


class TestTimsortArraysPurePython(BaseTimsortTest, TestCase):

    timsort = py_array_timsort

    def array_factory(self, lst):
        return np.array(lst, dtype=np.int32)


class JITTimsortMixin(object):

    timsort = jit_array_timsort

    test_merge_at = None
    test_merge_force_collapse = None

    def wrap_with_mergestate(self, timsort, func, _cache={}):
        """
        Wrap *func* into another compiled function inserting a runtime-created
        mergestate as the first function argument.
        """
        key = timsort, func
        if key in _cache:
            return _cache[key]

        merge_init = timsort.merge_init

        @timsort.compile
        def wrapper(keys, values, *args):
            ms = merge_init(keys)
            res = func(ms, keys, values, *args)
            return res

        _cache[key] = wrapper
        return wrapper


class TestTimsortArrays(JITTimsortMixin, BaseTimsortTest, TestCase):

    def array_factory(self, lst):
        return np.array(lst, dtype=np.int32)

    def check_merge_lo_hi(self, func, a, b):
        na = len(a)
        nb = len(b)

        func = self.wrap_with_mergestate(self.timsort, func)

        # Add sentinels at start and end, to check they weren't moved
        orig_keys = [42] + a + b + [-42]
        keys = self.array_factory(orig_keys)
        ssa = 1
        ssb = ssa + na

        new_ms = func(keys, keys, ssa, na, ssb, nb)
        self.assertEqual(keys[0], orig_keys[0])
        self.assertEqual(keys[-1], orig_keys[-1])
        self.assertSorted(orig_keys[1:-1], keys[1:-1])



class BaseQuicksortTest(BaseSortingTest):

    # NOTE these tests assume a non-argsort quicksort.

    def test_insertion_sort(self):
        n = 20
        def check(l, n):
            res = self.array_factory([9999] + l + [-9999])
            f(res, res, 1, n)
            self.assertEqual(res[0], 9999)
            self.assertEqual(res[-1], -9999)
            self.assertSorted(l, res[1:-1])

        f = self.quicksort.insertion_sort
        l = self.sorted_list(n)
        check(l, n)
        l = self.revsorted_list(n)
        check(l, n)
        l = self.initially_sorted_list(n, n//2)
        check(l, n)
        l = self.revsorted_list(n)
        check(l, n)
        l = self.random_list(n)
        check(l, n)
        l = self.duprandom_list(n)
        check(l, n)

    def test_partition(self):
        n = 20
        def check(l, n):
            res = self.array_factory([9999] + l + [-9999])
            index = f(res, res, 1, n)
            self.assertEqual(res[0], 9999)
            self.assertEqual(res[-1], -9999)
            pivot = res[index]
            for i in range(1, index):
                self.assertLessEqual(res[i], pivot)
            for i in range(index + 1, n):
                self.assertGreaterEqual(res[i], pivot)

        f = self.quicksort.partition
        l = self.sorted_list(n)
        check(l, n)
        l = self.revsorted_list(n)
        check(l, n)
        l = self.initially_sorted_list(n, n//2)
        check(l, n)
        l = self.revsorted_list(n)
        check(l, n)
        l = self.random_list(n)
        check(l, n)
        l = self.duprandom_list(n)
        check(l, n)

    def test_partition3(self):
        # Test the unused partition3() function
        n = 20
        def check(l, n):
            res = self.array_factory([9999] + l + [-9999])
            lt, gt = f(res, 1, n)
            self.assertEqual(res[0], 9999)
            self.assertEqual(res[-1], -9999)
            pivot = res[lt]
            for i in range(1, lt):
                self.assertLessEqual(res[i], pivot)
            for i in range(lt, gt + 1):
                self.assertEqual(res[i], pivot)
            for i in range(gt + 1, n):
                self.assertGreater(res[i], pivot)

        f = self.quicksort.partition3
        l = self.sorted_list(n)
        check(l, n)
        l = self.revsorted_list(n)
        check(l, n)
        l = self.initially_sorted_list(n, n//2)
        check(l, n)
        l = self.revsorted_list(n)
        check(l, n)
        l = self.random_list(n)
        check(l, n)
        l = self.duprandom_list(n)
        check(l, n)

    def test_run_quicksort(self):
        f = self.quicksort.run_quicksort

        for size_factor in (1, 5):
            # Make lists to be sorted from two chunks of different kinds.
            sizes = (15, 20)

            all_lists = [self.make_sample_lists(n * size_factor) for n in sizes]
            for chunks in itertools.product(*all_lists):
                orig_keys = sum(chunks, [])
                keys = self.array_factory(orig_keys)
                f(keys)
                # The list is now sorted
                self.assertSorted(orig_keys, keys)

    def test_run_quicksort_lt(self):
        def lt(a, b):
            return a > b

        f = self.make_quicksort(lt=lt).run_quicksort

        for size_factor in (1, 5):
            # Make lists to be sorted from two chunks of different kinds.
            sizes = (15, 20)

            all_lists = [self.make_sample_lists(n * size_factor) for n in sizes]
            for chunks in itertools.product(*all_lists):
                orig_keys = sum(chunks, [])
                keys = self.array_factory(orig_keys)
                f(keys)
                # The list is now rev-sorted
                self.assertSorted(orig_keys, keys[::-1])

        # An imperfect comparison function, as LT(a, b) does not imply not LT(b, a).
        # The sort should handle it gracefully.
        def lt_floats(a, b):
            return math.isnan(b) or a < b

        f = self.make_quicksort(lt=lt_floats).run_quicksort

        np.random.seed(42)
        for size in (5, 20, 50, 500):
            orig = np.random.random(size=size) * 100
            orig[np.random.random(size=size) < 0.1] = float('nan')
            orig_keys = list(orig)
            keys = self.array_factory(orig_keys)
            f(keys)
            non_nans = orig[~np.isnan(orig)]
            # Non-NaNs are sorted at the front
            self.assertSorted(non_nans, keys[:len(non_nans)])


class TestQuicksortPurePython(BaseQuicksortTest, TestCase):

    quicksort = py_quicksort
    make_quicksort = staticmethod(make_py_quicksort)

    # Much faster than a Numpy array in pure Python
    array_factory = list


class TestQuicksortArrays(BaseQuicksortTest, TestCase):

    quicksort = jit_quicksort
    make_quicksort = staticmethod(make_jit_quicksort)

    def array_factory(self, lst):
        return np.array(lst, dtype=np.float64)

class TestQuicksortMultidimensionalArrays(BaseSortingTest, TestCase):

    quicksort = make_jit_quicksort(is_np_array=True)
    make_quicksort = staticmethod(make_jit_quicksort)

    def assertSorted(self, orig, result):
        self.assertEqual(orig.shape, result.shape)
        self.assertPreciseEqual(orig, result)

    def array_factory(self, lst, shape=None):
        array = np.array(lst, dtype=np.float64)
        if shape is None:
            return array.reshape(-1, array.shape[0])
        else:
            return array.reshape(shape)

    def get_shapes(self, n):
        shapes = []
        if n == 1:
            return shapes

        for i in range(2, int(math.sqrt(n)) + 1):
            if n % i == 0:
                shapes.append((n // i, i))
                shapes.append((i, n // i))
                _shapes = self.get_shapes(n // i)
                for _shape in _shapes:
                    shapes.append((i,) + _shape)
                    shapes.append(_shape + (i,))

        return shapes

    def test_run_quicksort(self):
        f = self.quicksort.run_quicksort

        for size_factor in (1, 5):
            # Make lists to be sorted from two chunks of different kinds.
            sizes = (15, 20)

            all_lists = [self.make_sample_lists(n * size_factor) for n in sizes]
            for chunks in itertools.product(*all_lists):
                orig_keys = sum(chunks, [])
                shape_list = self.get_shapes(len(orig_keys))
                shape_list.append(None)
                for shape in shape_list:
                    keys = self.array_factory(orig_keys, shape=shape)
                    keys_copy = self.array_factory(orig_keys, shape=shape)
                    f(keys)
                    keys_copy.sort()
                    # The list is now sorted
                    self.assertSorted(keys_copy, keys)

    def test_run_quicksort_lt(self):
        def lt(a, b):
            return a > b

        f = self.make_quicksort(lt=lt, is_np_array=True).run_quicksort

        for size_factor in (1, 5):
            # Make lists to be sorted from two chunks of different kinds.
            sizes = (15, 20)

            all_lists = [self.make_sample_lists(n * size_factor) for n in sizes]
            for chunks in itertools.product(*all_lists):
                orig_keys = sum(chunks, [])
                shape_list = self.get_shapes(len(orig_keys))
                shape_list.append(None)
                for shape in shape_list:
                    keys = self.array_factory(orig_keys, shape=shape)
                    keys_copy = -self.array_factory(orig_keys, shape=shape)
                    f(keys)
                    # The list is now rev-sorted
                    keys_copy.sort()
                    keys_copy = -keys_copy
                    self.assertSorted(keys_copy, keys)

        # An imperfect comparison function, as LT(a, b) does not imply not LT(b, a).
        # The sort should handle it gracefully.
        def lt_floats(a, b):
            return math.isnan(b) or a < b

        f = self.make_quicksort(lt=lt_floats, is_np_array=True).run_quicksort

        np.random.seed(42)
        for size in (5, 20, 50, 500):
            orig = np.random.random(size=size) * 100
            orig[np.random.random(size=size) < 0.1] = float('nan')
            orig_keys = list(orig)
            shape_list = self.get_shapes(len(orig_keys))
            shape_list.append(None)
            for shape in shape_list:
                keys = self.array_factory(orig_keys, shape=shape)
                keys_copy = self.array_factory(orig_keys, shape=shape)
                f(keys)
                keys_copy.sort()
                # Non-NaNs are sorted at the front
                self.assertSorted(keys_copy, keys)

class TestNumpySort(TestCase):

    def setUp(self):
        np.random.seed(42)

    def int_arrays(self):
        for size in (5, 20, 50, 500):
            yield np.random.randint(99, size=size)

    def float_arrays(self):
        for size in (5, 20, 50, 500):
            yield np.random.random(size=size) * 100
        # Now with NaNs.  Numpy sorts them at the end.
        for size in (5, 20, 50, 500):
            orig = np.random.random(size=size) * 100
            orig[np.random.random(size=size) < 0.1] = float('nan')
            yield orig
        # 90% of values are NaNs.
        for size in (50, 500):
            orig = np.random.random(size=size) * 100
            orig[np.random.random(size=size) < 0.9] = float('nan')
            yield orig

    def has_duplicates(self, arr):
        """
        Whether the array has duplicates.  Takes NaNs into account.
        """
        if np.count_nonzero(np.isnan(arr)) > 1:
            return True
        if np.unique(arr).size < arr.size:
            return True
        return False

    def check_sort_inplace(self, pyfunc, cfunc, val):
        expected = copy.copy(val)
        got = copy.copy(val)
        pyfunc(expected)
        cfunc(got)
        self.assertPreciseEqual(got, expected)

    def check_sort_copy(self, pyfunc, cfunc, val):
        orig = copy.copy(val)
        expected = pyfunc(val)
        got = cfunc(val)
        self.assertPreciseEqual(got, expected)
        # The original wasn't mutated
        self.assertPreciseEqual(val, orig)

    def check_argsort(self, pyfunc, cfunc, val, kwargs={}):
        orig = copy.copy(val)
        expected = pyfunc(val, **kwargs)
        got = cfunc(val, **kwargs)
        self.assertPreciseEqual(orig[got], np.sort(orig),
                                msg="the array wasn't argsorted")
        # Numba and Numpy results may differ if there are duplicates
        # in the array
        if not self.has_duplicates(orig):
            self.assertPreciseEqual(got, expected)
        # The original wasn't mutated
        self.assertPreciseEqual(val, orig)

    def test_array_sort_int(self):
        pyfunc = sort_usecase
        cfunc = jit(nopython=True)(pyfunc)

        for orig in self.int_arrays():
            self.check_sort_inplace(pyfunc, cfunc, orig)

    def test_array_sort_float(self):
        pyfunc = sort_usecase
        cfunc = jit(nopython=True)(pyfunc)

        for orig in self.float_arrays():
            self.check_sort_inplace(pyfunc, cfunc, orig)

    def test_np_sort_int(self):
        pyfunc = np_sort_usecase
        cfunc = jit(nopython=True)(pyfunc)

        for orig in self.int_arrays():
            self.check_sort_copy(pyfunc, cfunc, orig)

    def test_np_sort_float(self):
        pyfunc = np_sort_usecase
        cfunc = jit(nopython=True)(pyfunc)

        for size in (5, 20, 50, 500):
            orig = np.random.random(size=size) * 100
            orig[np.random.random(size=size) < 0.1] = float('nan')
            self.check_sort_copy(pyfunc, cfunc, orig)

    def test_argsort_int(self):
        def check(pyfunc):
            cfunc = jit(nopython=True)(pyfunc)
            for orig in self.int_arrays():
                self.check_argsort(pyfunc, cfunc, orig)

        check(argsort_usecase)
        check(np_argsort_usecase)

    def test_argsort_kind_int(self):
        def check(pyfunc, is_stable):
            cfunc = jit(nopython=True)(pyfunc)
            for orig in self.int_arrays():
                self.check_argsort(pyfunc, cfunc, orig,
                                   dict(is_stable=is_stable))

        check(argsort_kind_usecase, is_stable=True)
        check(np_argsort_kind_usecase, is_stable=True)
        check(argsort_kind_usecase, is_stable=False)
        check(np_argsort_kind_usecase, is_stable=False)

    def test_argsort_float(self):
        def check(pyfunc):
            cfunc = jit(nopython=True)(pyfunc)
            for orig in self.float_arrays():
                self.check_argsort(pyfunc, cfunc, orig)

        check(argsort_usecase)
        check(np_argsort_usecase)

    def test_argsort_float(self):
        def check(pyfunc, is_stable):
            cfunc = jit(nopython=True)(pyfunc)
            for orig in self.float_arrays():
                self.check_argsort(pyfunc, cfunc, orig,
                                   dict(is_stable=is_stable))

        check(argsort_kind_usecase, is_stable=True)
        check(np_argsort_kind_usecase, is_stable=True)
        check(argsort_kind_usecase, is_stable=False)
        check(np_argsort_kind_usecase, is_stable=False)

    def test_bad_array(self):
        cfunc = jit(nopython=True)(np_sort_usecase)
        msg = '.*Argument "a" must be array-like.*'
        with self.assertRaisesRegex(errors.TypingError, msg) as raises:
            cfunc(None)


class TestPythonSort(TestCase):

    def test_list_sort(self):
        pyfunc = list_sort_usecase
        cfunc = jit(nopython=True)(pyfunc)

        for size in (20, 50, 500):
            orig, ret = cfunc(size)
            self.assertEqual(sorted(orig), ret)
            self.assertNotEqual(orig, ret)   # sanity check

    def test_list_sort_reverse(self):
        pyfunc = list_sort_reverse_usecase
        cfunc = jit(nopython=True)(pyfunc)

        for size in (20, 50, 500):
            for b in (False, True):
                orig, ret = cfunc(size, b)
                self.assertEqual(sorted(orig, reverse=b), ret)
                self.assertNotEqual(orig, ret)   # sanity check

    def test_sorted(self):
        pyfunc = sorted_usecase
        cfunc = jit(nopython=True)(pyfunc)

        for size in (20, 50, 500):
            orig = np.random.random(size=size) * 100
            expected = sorted(orig)
            got = cfunc(orig)
            self.assertPreciseEqual(got, expected)
            self.assertNotEqual(list(orig), got)   # sanity check

    def test_sorted_reverse(self):
        pyfunc = sorted_reverse_usecase
        cfunc = jit(nopython=True)(pyfunc)
        size = 20

        orig = np.random.random(size=size) * 100
        for b in (False, True):
            expected = sorted(orig, reverse=b)
            got = cfunc(orig, b)
            self.assertPreciseEqual(got, expected)
            self.assertNotEqual(list(orig), got)   # sanity check


class TestMergeSort(TestCase):
    def setUp(self):
        np.random.seed(321)

    def check_argsort_stable(self, sorter, low, high, count):
        # make data with high possibility of duplicated key
        data = np.random.randint(low, high, count)
        expect = np.argsort(data, kind='mergesort')
        got = sorter(data)
        np.testing.assert_equal(expect, got)

    def test_argsort_stable(self):
        arglist = [
            (-2, 2, 5),
            (-5, 5, 10),
            (0, 10, 101),
            (0, 100, 1003),
        ]
        imp = make_jit_mergesort(is_argsort=True)
        toplevel = imp.run_mergesort
        sorter = njit(lambda arr: toplevel(arr))
        for args in arglist:
            self.check_argsort_stable(sorter, *args)


nop_compiler = lambda x:x


class TestSortSlashSortedWithKey(MemoryLeakMixin, TestCase):

    def test_01(self):

        a = [3, 1, 4, 1, 5, 9]

        @njit
        def external_key(z):
            return 1. / z

        @njit
        def foo(x, key=None):
            new_x = x[:]
            new_x.sort(key=key)
            return sorted(x[:], key=key), new_x

        self.assertPreciseEqual(foo(a[:]), foo.py_func(a[:]))
        self.assertPreciseEqual(foo(a[:], external_key),
                                foo.py_func(a[:], external_key))

    def test_02(self):

        a = [3, 1, 4, 1, 5, 9]

        @njit
        def foo(x):
            def closure_key(z):
                return 1. / z
            new_x = x[:]
            new_x.sort(key=closure_key)
            return sorted(x[:], key=closure_key), new_x

        self.assertPreciseEqual(foo(a[:]), foo.py_func(a[:]))

    def test_03(self):

        a = [3, 1, 4, 1, 5, 9]

        def gen(compiler):

            @compiler
            def bar(x, func):
                new_x = x[:]
                new_x.sort(key=func)
                return sorted(x[:], key=func), new_x

            @compiler
            def foo(x):
                def closure_escapee_key(z):
                    return 1. / z
                return bar(x, closure_escapee_key)

            return foo

        self.assertPreciseEqual(gen(njit)(a[:]), gen(nop_compiler)(a[:]))

    def test_04(self):

        a = ['a','b','B','b','C','A']

        @njit
        def external_key(z):
            return z.upper()

        @njit
        def foo(x, key=None):
            new_x = x[:]
            new_x.sort(key=key)
            return sorted(x[:], key=key), new_x

        self.assertPreciseEqual(foo(a[:]), foo.py_func(a[:]))
        self.assertPreciseEqual(foo(a[:], external_key),
                                foo.py_func(a[:], external_key))

    def test_05(self):

        a = ['a','b','B','b','C','A']

        @njit
        def external_key(z):
            return z.upper()

        @njit
        def foo(x, key=None, reverse=False):
            new_x = x[:]
            new_x.sort(key=key, reverse=reverse)
            return (sorted(x[:], key=key, reverse=reverse), new_x)

        for key, rev in itertools.product((None, external_key),
                                          (True, False, 1, -12, 0)):
            self.assertPreciseEqual(foo(a[:], key, rev),
                                    foo.py_func(a[:], key, rev))

    def test_optional_on_key(self):
        a = [3, 1, 4, 1, 5, 9]

        @njit
        def foo(x, predicate):
            if predicate:
                def closure_key(z):
                    return 1. / z
            else:
                closure_key = None

            new_x = x[:]
            new_x.sort(key=closure_key)

            return (sorted(x[:], key=closure_key), new_x)

        with self.assertRaises(errors.TypingError) as raises:
            TF = True
            foo(a[:], TF)

        msg = "Key must concretely be None or a Numba JIT compiled function"
        self.assertIn(msg, str(raises.exception))

    def test_exceptions_sorted(self):

        @njit
        def foo_sorted(x, key=None, reverse=False):
            return sorted(x[:], key=key, reverse=reverse)

        @njit
        def foo_sort(x, key=None, reverse=False):
            new_x = x[:]
            new_x.sort(key=key, reverse=reverse)
            return new_x

        @njit
        def external_key(z):
            return 1. / z

        a = [3, 1, 4, 1, 5, 9]

        for impl in (foo_sort, foo_sorted):

            # check illegal key
            with self.assertRaises(errors.TypingError) as raises:
                impl(a, key="illegal")

            expect = "Key must be None or a Numba JIT compiled function"
            self.assertIn(expect, str(raises.exception))

            # check illegal reverse
            with self.assertRaises(errors.TypingError) as raises:
                impl(a, key=external_key, reverse="go backwards")

            expect = "an integer is required for 'reverse'"
            self.assertIn(expect, str(raises.exception))


class TestArrayArgsort(MemoryLeakMixin, TestCase):
    """Tests specific to array.argsort"""

    def test_exceptions(self):

        @njit
        def nonliteral_kind(kind):
            np.arange(5).argsort(kind=kind)

        # check non-literal kind
        with self.assertRaises(errors.TypingError) as raises:
            # valid spelling but not literal
            nonliteral_kind('quicksort')

        expect = '"kind" must be a string literal'
        self.assertIn(expect, str(raises.exception))

        @njit
        def unsupported_kwarg():
            np.arange(5).argsort(foo='')

        with self.assertRaises(errors.TypingError) as raises:
            unsupported_kwarg()

        expect = "Unsupported keywords: ['foo']"
        self.assertIn(expect, str(raises.exception))


if __name__ == '__main__':
    unittest.main()