Merge branch 'enh-tuple-getitem' of https://github.com/jvkersch/pyston

src/runtime/tuple.cpp

@@ -33,26 +33,71 @@ extern "C" Box* createTuple(int64_t nelts, Box** elts) {
     return new BoxedTuple(std::move(velts));
 }
 
-Box* tupleGetitem(BoxedTuple* self, Box* slice) {
-    assert(self->cls == tuple_cls);
+Box* _tupleSlice(BoxedTuple* self, i64 start, i64 stop, i64 step) {
 
     i64 size = self->elts.size();
+    assert(step != 0);
+    if (step > 0) {
+        assert(0 <= start);
+        assert(stop <= size);
+    } else {
+        assert(start < size);
+        assert(-1 <= stop);
+    }
 
-    if (slice->cls == int_cls) {
-        i64 n = static_cast<BoxedInt*>(slice)->n;
+    // 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 (n < 0)
-            n = size - n;
-        if (n < 0 || n >= size) {
-            fprintf(stderr, "IndexError: tuple index out of range\n");
-            raiseExcHelper(IndexError, "");
-        }
+    if (slicelength < 0)
+        slicelength = 0;
 
-        Box* rtn = self->elts[n];
-        return rtn;
-    } else {
-        RELEASE_ASSERT(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];
+
+    return new BoxedTuple(std::move(velts));
+}
+
+Box* tupleGetitemInt(BoxedTuple* self, BoxedInt* slice) {
+    i64 n = slice->n;
+    i64 size = self->elts.size();
+
+    if (n < 0)
+        n = size - n;
+    if (n < 0 || n >= size) {
+        fprintf(stderr, "IndexError: tuple index out of range\n");
+        raiseExcHelper(IndexError, "");
     }
+
+    Box* rtn = self->elts[n];
+    return rtn;
+}
+
+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);
+}
+
+Box* tupleGetitem(BoxedTuple* self, Box* slice) {
+    assert(self->cls == tuple_cls);
+
+    if (slice->cls == int_cls)
+        return tupleGetitemInt(self, static_cast<BoxedInt*>(slice));
+    else if (slice->cls == slice_cls)
+        return tupleGetitemSlice(self, static_cast<BoxedSlice*>(slice));
+    else
+        raiseExcHelper(TypeError, "tuple indices must be integers, not %s", getTypeName(slice)->c_str());
 }
 
 Box* tupleAdd(BoxedTuple* self, Box* rhs) {
@@ -214,7 +259,13 @@ void setupTuple() {
 
     tuple_cls->giveAttr("__name__", boxStrConstant("tuple"));
 
-    tuple_cls->giveAttr("__getitem__", new BoxedFunction(boxRTFunction((void*)tupleGetitem, UNKNOWN, 2)));
+    CLFunction* getitem = createRTFunction(2, 0, 0, 0);
+    addRTFunction(getitem, (void*)tupleGetitemInt, UNKNOWN,
+                  std::vector<ConcreteCompilerType*>{ BOXED_TUPLE, BOXED_INT });
+    addRTFunction(getitem, (void*)tupleGetitemSlice, SLICE, std::vector<ConcreteCompilerType*>{ BOXED_TUPLE, SLICE });
+    addRTFunction(getitem, (void*)tupleGetitem, UNKNOWN, std::vector<ConcreteCompilerType*>{ BOXED_TUPLE, UNKNOWN });
+    tuple_cls->giveAttr("__getitem__", new BoxedFunction(getitem));
+
     tuple_cls->giveAttr("__contains__", new BoxedFunction(boxRTFunction((void*)tupleContains, BOXED_BOOL, 2)));
 
     tuple_cls->giveAttr("__iter__",

test/tests/tuples.py

@@ -76,3 +76,38 @@ print () + ()
 print (1, 2, 3) + ()
 print () + (1, 2, 3)
 print (1, 2) + (2, 3)
+
+# __getitem__
+t = (1, "2")
+print t[0]
+print t[1]
+
+t = (1, 2, 'a', 'b', 'c')
+print t[::-1]
+print t[:-1]
+print t[0:2]
+print t[-5:]
+print t[-5:3]
+print t[-5:10]
+print t[:-5]
+print t[:3]
+print t[:10]
+print t[1:3:-1]
+print t[3:1:-1]
+print t[1:3:1]
+print t[1:3:2]
+print t[1:5:3]
+print t[5:1:-1]
+print t[5:1:-2]
+print t[5:1:-5]
+print t[5:1]
+
+try:
+    t[None]
+except TypeError as e:
+    print e
+
+try:
+    t[(1, 2)]
+except TypeError as e:
+    print e