codekingpro/portable-devtools

1import contextlib
2import itertools
3import operator
4
5import pytest
6
7import numpy as np
8import numpy._core._multiarray_tests as mt
9from numpy.testing import assert_equal, assert_raises
10
11INT64_MAX = np.iinfo(np.int64).max
12INT64_MIN = np.iinfo(np.int64).min
13INT64_MID = 2**32
14
15# int128 is not two's complement, the sign bit is separate
16INT128_MAX = 2**128 - 1
17INT128_MIN = -INT128_MAX
18INT128_MID = 2**64
19
20INT64_VALUES = (
21    [INT64_MIN + j for j in range(20)] +
22    [INT64_MAX - j for j in range(20)] +
23    [INT64_MID + j for j in range(-20, 20)] +
24    [2 * INT64_MID + j for j in range(-20, 20)] +
25    [INT64_MID // 2 + j for j in range(-20, 20)] +
26    list(range(-70, 70))
27)
28
29INT128_VALUES = (
30    [INT128_MIN + j for j in range(20)] +
31    [INT128_MAX - j for j in range(20)] +
32    [INT128_MID + j for j in range(-20, 20)] +
33    [2 * INT128_MID + j for j in range(-20, 20)] +
34    [INT128_MID // 2 + j for j in range(-20, 20)] +
35    list(range(-70, 70)) +
36    [False]  # negative zero
37)
38
39INT64_POS_VALUES = [x for x in INT64_VALUES if x > 0]
40
41
42@contextlib.contextmanager
43def exc_iter(*args):
44    """
45    Iterate over Cartesian product of *args, and if an exception is raised,
46    add information of the current iterate.
47    """
48
49    value = [None]
50
51    def iterate():
52        for v in itertools.product(*args):
53            value[0] = v
54            yield v
55
56    try:
57        yield iterate()
58    except Exception:
59        import traceback
60        msg = f"At: {repr(value[0])!r}\n{traceback.format_exc()}"
61        raise AssertionError(msg)
62
63
64def test_safe_binop():
65    # Test checked arithmetic routines
66
67    ops = [
68        (operator.add, 1),
69        (operator.sub, 2),
70        (operator.mul, 3)
71    ]
72
73    with exc_iter(ops, INT64_VALUES, INT64_VALUES) as it:
74        for xop, a, b in it:
75            pyop, op = xop
76            c = pyop(a, b)
77
78            if not (INT64_MIN <= c <= INT64_MAX):
79                assert_raises(OverflowError, mt.extint_safe_binop, a, b, op)
80            else:
81                d = mt.extint_safe_binop(a, b, op)
82                if c != d:
83                    # assert_equal is slow
84                    assert_equal(d, c)
85
86
87def test_to_128():
88    with exc_iter(INT64_VALUES) as it:
89        for a, in it:
90            b = mt.extint_to_128(a)
91            if a != b:
92                assert_equal(b, a)
93
94
95def test_to_64():
96    with exc_iter(INT128_VALUES) as it:
97        for a, in it:
98            if not (INT64_MIN <= a <= INT64_MAX):
99                assert_raises(OverflowError, mt.extint_to_64, a)
100            else:
101                b = mt.extint_to_64(a)
102                if a != b:
103                    assert_equal(b, a)
104
105
106def test_mul_64_64():
107    with exc_iter(INT64_VALUES, INT64_VALUES) as it:
108        for a, b in it:
109            c = a * b
110            d = mt.extint_mul_64_64(a, b)
111            if c != d:
112                assert_equal(d, c)
113
114
115def test_add_128():
116    with exc_iter(INT128_VALUES, INT128_VALUES) as it:
117        for a, b in it:
118            c = a + b
119            if not (INT128_MIN <= c <= INT128_MAX):
120                assert_raises(OverflowError, mt.extint_add_128, a, b)
121            else:
122                d = mt.extint_add_128(a, b)
123                if c != d:
124                    assert_equal(d, c)
125
126
127def test_sub_128():
128    with exc_iter(INT128_VALUES, INT128_VALUES) as it:
129        for a, b in it:
130            c = a - b
131            if not (INT128_MIN <= c <= INT128_MAX):
132                assert_raises(OverflowError, mt.extint_sub_128, a, b)
133            else:
134                d = mt.extint_sub_128(a, b)
135                if c != d:
136                    assert_equal(d, c)
137
138
139def test_neg_128():
140    with exc_iter(INT128_VALUES) as it:
141        for a, in it:
142            b = -a
143            c = mt.extint_neg_128(a)
144            if b != c:
145                assert_equal(c, b)
146
147
148def test_shl_128():
149    with exc_iter(INT128_VALUES) as it:
150        for a, in it:
151            if a < 0:
152                b = -(((-a) << 1) & (2**128 - 1))
153            else:
154                b = (a << 1) & (2**128 - 1)
155            c = mt.extint_shl_128(a)
156            if b != c:
157                assert_equal(c, b)
158
159
160def test_shr_128():
161    with exc_iter(INT128_VALUES) as it:
162        for a, in it:
163            if a < 0:
164                b = -((-a) >> 1)
165            else:
166                b = a >> 1
167            c = mt.extint_shr_128(a)
168            if b != c:
169                assert_equal(c, b)
170
171
172def test_gt_128():
173    with exc_iter(INT128_VALUES, INT128_VALUES) as it:
174        for a, b in it:
175            c = a > b
176            d = mt.extint_gt_128(a, b)
177            if c != d:
178                assert_equal(d, c)
179
180
181@pytest.mark.slow
182def test_divmod_128_64():
183    with exc_iter(INT128_VALUES, INT64_POS_VALUES) as it:
184        for a, b in it:
185            if a >= 0:
186                c, cr = divmod(a, b)
187            else:
188                c, cr = divmod(-a, b)
189                c = -c
190                cr = -cr
191
192            d, dr = mt.extint_divmod_128_64(a, b)
193
194            if c != d or d != dr or b * d + dr != a:
195                assert_equal(d, c)
196                assert_equal(dr, cr)
197                assert_equal(b * d + dr, a)
198
199
200def test_floordiv_128_64():
201    with exc_iter(INT128_VALUES, INT64_POS_VALUES) as it:
202        for a, b in it:
203            c = a // b
204            d = mt.extint_floordiv_128_64(a, b)
205
206            if c != d:
207                assert_equal(d, c)
208
209
210def test_ceildiv_128_64():
211    with exc_iter(INT128_VALUES, INT64_POS_VALUES) as it:
212        for a, b in it:
213            c = (a + b - 1) // b
214            d = mt.extint_ceildiv_128_64(a, b)
215
216            if c != d:
217                assert_equal(d, c)
218