Skip to content

Commit 3d68529

Browse files
Add constant lowering support for SliceTypes
1 parent f2a2ae1 commit 3d68529

4 files changed

Lines changed: 78 additions & 5 deletions

File tree

‎numba/core/datamodel/models.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1077,6 +1077,7 @@ def __init__(self, dmm, fe_type):
10771077
super(UniTupleIter, self).__init__(dmm, fe_type, members)
10781078

10791079

1080+
@register_default(types.misc.SliceLiteral)
10801081
@register_default(types.SliceType)
10811082
class SliceModel(StructModel):
10821083
def __init__(self, dmm, fe_type):

‎numba/core/types/misc.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -334,6 +334,10 @@ def __init__(self, value):
334334
members = 2 if value.step is None else 3
335335
SliceType.__init__(self, name=name, members=members)
336336

337+
@property
338+
def key(self):
339+
return self.name
340+
337341

338342
Literal.ctor_map[slice] = SliceLiteral
339343

‎numba/cpython/slicing.py‎

Lines changed: 62 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -5,11 +5,11 @@
55
from itertools import zip_longest
66

77
from llvmlite import ir
8-
9-
from numba.core import types, typing, utils, cgutils
10-
from numba.core.imputils import (lower_builtin, lower_getattr,
11-
iternext_impl, impl_ret_borrowed,
12-
impl_ret_new_ref, impl_ret_untracked)
8+
from numba.core import cgutils, types, typing, utils
9+
from 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)
1313

1414

1515
def fix_index(builder, idx, size):
@@ -233,3 +233,60 @@ def slice_indices(context, builder, sig, args):
233233
sig.return_type,
234234
(sli.start, sli.stop, sli.step)
235235
)
236+
237+
238+
def make_slice_from_constant(context, builder, ty, pyval):
239+
sli = context.make_helper(builder, ty)
240+
lty = context.get_value_type(types.intp)
241+
242+
default_start_pos, default_start_neg, default_stop_pos, default_stop_neg, default_step = \
243+
[context.get_constant(types.intp, x) for x in get_defaults(context)]
244+
245+
step = pyval.step
246+
if step is None:
247+
step_is_neg = False
248+
step = default_step
249+
else:
250+
step_is_neg = step < 0
251+
step = lty(step)
252+
253+
start = pyval.start
254+
if start is None:
255+
if step_is_neg:
256+
start = default_start_neg
257+
else:
258+
start = default_start_pos
259+
else:
260+
start = lty(start)
261+
262+
stop = pyval.stop
263+
if stop is None:
264+
if step_is_neg:
265+
stop = default_stop_neg
266+
else:
267+
stop = default_stop_pos
268+
else:
269+
stop = lty(stop)
270+
271+
sli.start = start
272+
sli.stop = stop
273+
sli.step = step
274+
275+
return sli._getvalue()
276+
277+
278+
@lower_constant(types.SliceType)
279+
def constant_slice(context, builder, ty, pyval):
280+
if isinstance(ty, types.Literal):
281+
typ = ty.literal_type
282+
else:
283+
typ = ty
284+
285+
return make_slice_from_constant(context, builder, typ, pyval)
286+
287+
288+
@lower_cast(types.misc.SliceLiteral, types.SliceType)
289+
def cast_from_literal(context, builder, fromty, toty, val):
290+
return make_slice_from_constant(
291+
context, builder, toty, fromty.literal_value,
292+
)

‎numba/tests/test_slices.py‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -185,6 +185,17 @@ def test_slice_indices_examples(self):
185185
str(e.exception)
186186
)
187187

188+
def test_slice_constant_lowering(self):
189+
for ts in [slice(1, 3, None), slice(2, None, None), slice(None, 2, -1),
190+
slice(None, None, None), slice(None, None, -2)]:
191+
test_tuple = (1, 2, 3, 4)
192+
193+
@jit(nopython=True)
194+
def test_fn():
195+
return test_tuple[ts]
196+
197+
assert test_fn() == test_tuple[ts]
198+
188199

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

0 commit comments

Comments
 (0)