Merge pull request #218 from undingen/fix_slice

Fix slicing issues

src/runtime/list.cpp

@@ -105,7 +105,7 @@ extern "C" Box* listLen(BoxedList* self) {
     return boxInt(self->size);
 }
 
-Box* _listSlice(BoxedList* self, i64 start, i64 stop, i64 step) {
+Box* _listSlice(BoxedList* self, i64 start, i64 stop, i64 step, i64 length) {
     // printf("%ld %ld %ld\n", start, stop, step);
     assert(step != 0);
     if (step > 0) {
@@ -117,17 +117,11 @@ Box* _listSlice(BoxedList* self, i64 start, i64 stop, i64 step) {
     }
 
     BoxedList* rtn = new BoxedList();
-
-    if ((step == 1) && ((stop - start) > 0)) {
-        listAppendArrayInternal(rtn, &self->elts->elts[start], stop - start);
-    } else {
-        int cur = start;
-        while ((step > 0 && cur < stop) || (step < 0 && cur > stop)) {
-            listAppendInternal(rtn, self->elts->elts[cur]);
-            cur += step;
-        }
+    if (length > 0) {
+        rtn->ensure(length);
+        copySlice(&rtn->elts->elts[0], &self->elts->elts[0], start, step, length);
+        rtn->size += length;
     }
-
     return rtn;
 }
 
@@ -165,9 +159,9 @@ extern "C" Box* listGetitemSlice(BoxedList* self, BoxedSlice* slice) {
 
     assert(self->cls == list_cls);
     assert(slice->cls == slice_cls);
-    i64 start, stop, step;
-    parseSlice(slice, self->size, &start, &stop, &step);
-    return _listSlice(self, start, stop, step);
+    i64 start, stop, step, length;
+    parseSlice(slice, self->size, &start, &stop, &step, &length);
+    return _listSlice(self, start, stop, step, length);
 }
 
 extern "C" Box* listGetitem(BoxedList* self, Box* slice) {

src/runtime/str.cpp

@@ -423,7 +423,7 @@ extern "C" Box* basestringNew(BoxedClass* cls, Box* args, Box* kwargs) {
     raiseExcHelper(TypeError, "The basestring type cannot be instantiated");
 }
 
-Box* _strSlice(BoxedString* self, i64 start, i64 stop, i64 step) {
+Box* _strSlice(BoxedString* self, i64 start, i64 stop, i64 step, i64 length) {
     assert(self->cls == str_cls);
 
     const std::string& s = self->s;
@@ -437,14 +437,12 @@ Box* _strSlice(BoxedString* self, i64 start, i64 stop, i64 step) {
         assert(-1 <= stop);
     }
 
-    std::vector<char> chars;
-    int cur = start;
-    while ((step > 0 && cur < stop) || (step < 0 && cur > stop)) {
-        chars.push_back(s[cur]);
-        cur += step;
+    std::string chars;
+    if (length > 0) {
+        chars.resize(length);
+        copySlice(&chars[0], &s[0], start, step, length);
     }
-    // TODO too much copying
-    return boxString(std::string(chars.begin(), chars.end()));
+    return boxString(std::move(chars));
 }
 
 Box* strIsAlpha(BoxedString* self) {
@@ -870,9 +868,9 @@ extern "C" Box* strGetitem(BoxedString* self, Box* slice) {
     } else if (slice->cls == slice_cls) {
         BoxedSlice* sslice = static_cast<BoxedSlice*>(slice);
 
-        i64 start, stop, step;
-        parseSlice(sslice, self->s.size(), &start, &stop, &step);
-        return _strSlice(self, start, stop, step);
+        i64 start, stop, step, length;
+        parseSlice(sslice, self->s.size(), &start, &stop, &step, &length);
+        return _strSlice(self, start, stop, step, length);
     } else {
         raiseExcHelper(TypeError, "string indices must be integers, not %s", getTypeName(slice)->c_str());
     }

src/runtime/tuple.cpp

@@ -32,7 +32,7 @@ extern "C" Box* createTuple(int64_t nelts, Box** elts) {
     return new BoxedTuple(std::move(velts));
 }
 
-Box* _tupleSlice(BoxedTuple* self, i64 start, i64 stop, i64 step) {
+Box* _tupleSlice(BoxedTuple* self, i64 start, i64 stop, i64 step, i64 length) {
 
     i64 size = self->elts.size();
     assert(step != 0);
@@ -44,23 +44,10 @@ Box* _tupleSlice(BoxedTuple* self, i64 start, i64 stop, i64 step) {
         assert(-1 <= stop);
     }
 
-    // This is adapted from CPython's PySlice_GetIndicesEx.
-    i64 slicelength;
-    if (step < 0)
-        slicelength = (stop - start + 1) / (step) + 1;
-    else
-        slicelength = (stop - start - 1) / (step) + 1;
-
-    if (slicelength < 0)
-        slicelength = 0;
-
     // FIXME: No need to initialize with 0.
-    BoxedTuple::GCVector velts(slicelength, 0);
-
-    i64 curr, i;
-    for (curr = start, i = 0; i < slicelength; curr += step, i++)
-        velts[i] = self->elts[curr];
-
+    BoxedTuple::GCVector velts(length, 0);
+    if (length > 0)
+        copySlice(&velts[0], &self->elts[0], start, step, length);
     return new BoxedTuple(std::move(velts));
 }
 
@@ -94,9 +81,9 @@ Box* tupleGetitemSlice(BoxedTuple* self, BoxedSlice* slice) {
     assert(self->cls == tuple_cls);
     assert(slice->cls == slice_cls);
 
-    i64 start, stop, step;
-    parseSlice(slice, self->elts.size(), &start, &stop, &step);
-    return _tupleSlice(self, start, stop, step);
+    i64 start, stop, step, length;
+    parseSlice(slice, self->elts.size(), &start, &stop, &step, &length);
+    return _tupleSlice(self, start, stop, step, length);
 }
 
 extern "C" PyObject* PyTuple_GetSlice(PyObject* p, Py_ssize_t low, Py_ssize_t high) {

src/runtime/util.cpp

@@ -20,7 +20,7 @@
 
 namespace pyston {
 
-void parseSlice(BoxedSlice* slice, int size, i64* out_start, i64* out_stop, i64* out_step) {
+void parseSlice(BoxedSlice* slice, int size, i64* out_start, i64* out_stop, i64* out_step, i64* out_length) {
     BoxedSlice* sslice = static_cast<BoxedSlice*>(slice);
 
     Box* start = sslice->start;
@@ -40,6 +40,9 @@ void parseSlice(BoxedSlice* slice, int size, i64* out_start, i64* out_stop, i64*
 
     if (step->cls == int_cls) {
         istep = static_cast<BoxedInt*>(step)->n;
+        if (istep == 0) {
+            raiseExcHelper(ValueError, "slice step cannot be zero");
+        }
     }
 
     if (start->cls == int_cls) {
@@ -63,11 +66,6 @@ void parseSlice(BoxedSlice* slice, int size, i64* out_start, i64* out_stop, i64*
             istop = -1;
     }
 
-    if (istep == 0) {
-        fprintf(stderr, "ValueError: slice step cannot be zero\n");
-        raiseExcHelper(ValueError, "");
-    }
-
     if (istep > 0) {
         if (istart < 0)
             istart = 0;
@@ -83,5 +81,18 @@ void parseSlice(BoxedSlice* slice, int size, i64* out_start, i64* out_stop, i64*
     *out_start = istart;
     *out_stop = istop;
     *out_step = istep;
+
+    if (out_length) {
+        // This is adapted from CPython's PySlice_GetIndicesEx.
+        if ((istep < 0 && istop >= istart) || (istep > 0 && istart >= istop))
+            *out_length = 0;
+        else if (istep < 0)
+            *out_length = (istop - istart + 1) / istep + 1;
+        else
+            *out_length = (istop - istart - 1) / istep + 1;
+
+        if (*out_length < 0)
+            *out_length = 0;
+    }
 }
 }

src/runtime/util.h

@@ -21,6 +21,15 @@ namespace pyston {
 
 class BoxedSlice;
 
-void parseSlice(BoxedSlice* slice, int size, i64* out_start, i64* out_stop, i64* out_end);
+void parseSlice(BoxedSlice* slice, int size, i64* out_start, i64* out_stop, i64* out_end, i64* out_length = nullptr);
+
+template <typename T> void copySlice(T* __restrict__ dst, const T* __restrict__ src, i64 start, i64 step, i64 length) {
+    if (step == 1) {
+        memcpy(dst, &src[start], length * sizeof(T));
+    } else {
+        for (i64 curr = start, i = 0; i < length; curr += step, ++i)
+            dst[i] = src[curr];
+    }
+}
 }
 #endif

test/tests/tuples.py

@@ -143,6 +143,8 @@ print t[5:1:-1]
 print t[5:1:-2]
 print t[5:1:-5]
 print t[5:1]
+print t[-1:-1:-5]
+print t[-5:-4:-5]
 
 try:
     t[None]