Aluode/PerceptionLabPortable

1"""
2Implement slices and various slice computations.
3"""
4
5from itertools import zip_longest
6
7from llvmlite import ir
8from numba.core import cgutils, types, typing, utils
9from numba.core.imputils import (impl_ret_borrowed, impl_ret_new_ref,
10                                 impl_ret_untracked, iternext_impl,
11                                 lower_builtin, lower_cast, lower_constant,
12                                 lower_getattr)
13
14
15def fix_index(builder, idx, size):
16    """
17    Fix negative index by adding *size* to it.  Positive
18    indices are left untouched.
19    """
20    is_negative = builder.icmp_signed('<', idx, ir.Constant(size.type, 0))
21    wrapped_index = builder.add(idx, size)
22    return builder.select(is_negative, wrapped_index, idx)
23
24
25def fix_slice(builder, slice, size):
26    """
27    Fix *slice* start and stop to be valid (inclusive and exclusive, resp)
28    indexing bounds for a sequence of the given *size*.
29    """
30    # See PySlice_GetIndicesEx()
31    zero = ir.Constant(size.type, 0)
32    minus_one = ir.Constant(size.type, -1)
33
34    def fix_bound(bound_name, lower_repl, upper_repl):
35        bound = getattr(slice, bound_name)
36        bound = fix_index(builder, bound, size)
37        # Store value
38        setattr(slice, bound_name, bound)
39        # Still negative? => clamp to lower_repl
40        underflow = builder.icmp_signed('<', bound, zero)
41        with builder.if_then(underflow, likely=False):
42            setattr(slice, bound_name, lower_repl)
43        # Greater than size? => clamp to upper_repl
44        overflow = builder.icmp_signed('>=', bound, size)
45        with builder.if_then(overflow, likely=False):
46            setattr(slice, bound_name, upper_repl)
47
48    with builder.if_else(cgutils.is_neg_int(builder, slice.step)) as (if_neg_step, if_pos_step):
49        with if_pos_step:
50            # < 0 => 0; >= size => size
51            fix_bound('start', zero, size)
52            fix_bound('stop', zero, size)
53        with if_neg_step:
54            # < 0 => -1; >= size => size - 1
55            lower = minus_one
56            upper = builder.add(size, minus_one)
57            fix_bound('start', lower, upper)
58            fix_bound('stop', lower, upper)
59
60
61def get_slice_length(builder, slicestruct):
62    """
63    Given a slice, compute the number of indices it spans, i.e. the
64    number of iterations that for_range_slice() will execute.
65
66    Pseudo-code:
67        assert step != 0
68        if step > 0:
69            if stop <= start:
70                return 0
71            else:
72                return (stop - start - 1) // step + 1
73        else:
74            if stop >= start:
75                return 0
76            else:
77                return (stop - start + 1) // step + 1
78
79    (see PySlice_GetIndicesEx() in CPython)
80    """
81    start = slicestruct.start
82    stop = slicestruct.stop
83    step = slicestruct.step
84    one = ir.Constant(start.type, 1)
85    zero = ir.Constant(start.type, 0)
86
87    is_step_negative = cgutils.is_neg_int(builder, step)
88    delta = builder.sub(stop, start)
89
90    # Nominal case
91    pos_dividend = builder.sub(delta, one)
92    neg_dividend = builder.add(delta, one)
93    dividend  = builder.select(is_step_negative, neg_dividend, pos_dividend)
94    nominal_length = builder.add(one, builder.sdiv(dividend, step))
95
96    # Catch zero length
97    is_zero_length = builder.select(is_step_negative,
98                                    builder.icmp_signed('>=', delta, zero),
99                                    builder.icmp_signed('<=', delta, zero))
100
101    # Clamp to 0 if is_zero_length
102    return builder.select(is_zero_length, zero, nominal_length)
103
104
105def get_slice_bounds(builder, slicestruct):
106    """
107    Return the [lower, upper) indexing bounds of a slice.
108    """
109    start = slicestruct.start
110    stop = slicestruct.stop
111    zero = start.type(0)
112    one = start.type(1)
113    # This is a bit pessimal, e.g. it will return [1, 5) instead
114    # of [1, 4) for `1:5:2`
115    is_step_negative = builder.icmp_signed('<', slicestruct.step, zero)
116    lower = builder.select(is_step_negative,
117                           builder.add(stop, one), start)
118    upper = builder.select(is_step_negative,
119                           builder.add(start, one), stop)
120    return lower, upper
121
122
123def fix_stride(builder, slice, stride):
124    """
125    Fix the given stride for the slice's step.
126    """
127    return builder.mul(slice.step, stride)
128
129def guard_invalid_slice(context, builder, typ, slicestruct):
130    """
131    Guard against *slicestruct* having a zero step (and raise ValueError).
132    """
133    if typ.has_step:
134        cgutils.guard_null(context, builder, slicestruct.step,
135                           (ValueError, "slice step cannot be zero"))
136
137
138def get_defaults(context):
139    """
140    Get the default values for a slice's members:
141    (start for positive step, start for negative step,
142     stop for positive step, stop for negative step, step)
143    """
144    maxint = (1 << (context.address_size - 1)) - 1
145    return (0, maxint, maxint, - maxint - 1, 1)
146
147
148#---------------------------------------------------------------------------
149# The slice structure
150
151@lower_builtin(slice, types.VarArg(types.Any))
152def slice_constructor_impl(context, builder, sig, args):
153    (
154        default_start_pos,
155        default_start_neg,
156        default_stop_pos,
157        default_stop_neg,
158        default_step,
159    ) = [context.get_constant(types.intp, x) for x in get_defaults(context)]
160
161    slice_args = [None] * 3
162
163    # Fetch non-None arguments
164    if len(args) == 1 and sig.args[0] is not types.none:
165        slice_args[1] = args[0]
166    else:
167        for i, (ty, val) in enumerate(zip(sig.args, args)):
168            if ty is not types.none:
169                slice_args[i] = val
170
171    # Fill omitted arguments
172    def get_arg_value(i, default):
173        val = slice_args[i]
174        if val is None:
175            return default
176        else:
177            return val
178
179    step = get_arg_value(2, default_step)
180    is_step_negative = builder.icmp_signed('<', step,
181                                           context.get_constant(types.intp, 0))
182    default_stop = builder.select(is_step_negative,
183                                  default_stop_neg, default_stop_pos)
184    default_start = builder.select(is_step_negative,
185                                   default_start_neg, default_start_pos)
186    stop = get_arg_value(1, default_stop)
187    start = get_arg_value(0, default_start)
188
189    ty = sig.return_type
190    sli = context.make_helper(builder, sig.return_type)
191    sli.start = start
192    sli.stop = stop
193    sli.step = step
194
195    res = sli._getvalue()
196    return impl_ret_untracked(context, builder, sig.return_type, res)
197
198
199@lower_getattr(types.SliceType, "start")
200def slice_start_impl(context, builder, typ, value):
201    sli = context.make_helper(builder, typ, value)
202    return sli.start
203
204@lower_getattr(types.SliceType, "stop")
205def slice_stop_impl(context, builder, typ, value):
206    sli = context.make_helper(builder, typ, value)
207    return sli.stop
208
209@lower_getattr(types.SliceType, "step")
210def slice_step_impl(context, builder, typ, value):
211    if typ.has_step:
212        sli = context.make_helper(builder, typ, value)
213        return sli.step
214    else:
215        return context.get_constant(types.intp, 1)
216
217
218@lower_builtin("slice.indices", types.SliceType, types.Integer)
219def slice_indices(context, builder, sig, args):
220    length = args[1]
221    sli = context.make_helper(builder, sig.args[0], args[0])
222
223    with builder.if_then(cgutils.is_neg_int(builder, length), likely=False):
224        context.call_conv.return_user_exc(
225            builder, ValueError,
226            ("length should not be negative",)
227        )
228    with builder.if_then(cgutils.is_scalar_zero(builder, sli.step), likely=False):
229        context.call_conv.return_user_exc(
230            builder, ValueError,
231            ("slice step cannot be zero",)
232        )
233
234    fix_slice(builder, sli, length)
235
236    return context.make_tuple(
237        builder,
238        sig.return_type,
239        (sli.start, sli.stop, sli.step)
240    )
241
242
243def make_slice_from_constant(context, builder, ty, pyval):
244    sli = context.make_helper(builder, ty)
245    lty = context.get_value_type(types.intp)
246
247    (
248        default_start_pos,
249        default_start_neg,
250        default_stop_pos,
251        default_stop_neg,
252        default_step,
253    ) = [context.get_constant(types.intp, x) for x in get_defaults(context)]
254
255    step = pyval.step
256    if step is None:
257        step_is_neg = False
258        step = default_step
259    else:
260        step_is_neg = step < 0
261        step = lty(step)
262
263    start = pyval.start
264    if start is None:
265        if step_is_neg:
266            start = default_start_neg
267        else:
268            start = default_start_pos
269    else:
270        start = lty(start)
271
272    stop = pyval.stop
273    if stop is None:
274        if step_is_neg:
275            stop = default_stop_neg
276        else:
277            stop = default_stop_pos
278    else:
279        stop = lty(stop)
280
281    sli.start = start
282    sli.stop = stop
283    sli.step = step
284
285    return sli._getvalue()
286
287
288@lower_constant(types.SliceType)
289def constant_slice(context, builder, ty, pyval):
290    if isinstance(ty, types.Literal):
291        typ = ty.literal_type
292    else:
293        typ = ty
294
295    return make_slice_from_constant(context, builder, typ, pyval)
296
297
298@lower_cast(types.misc.SliceLiteral, types.SliceType)
299def cast_from_literal(context, builder, fromty, toty, val):
300    return make_slice_from_constant(
301        context, builder, toty, fromty.literal_value,
302    )
303