fix tuple indexing bug (#21521)

Elias Ellison · facebook-github-bot · commit 881adb5bcdfc · 2019-06-07T11:17:59.000-07:00
Summary: lower tuples pass didn't check bounds for tuple index Pull Request resolved: #21521 Differential Revision: D15716813 Pulled By: eellison fbshipit-source-id: 8eead98c2c63118e7d24a8c8bf6184b02afb7dcd
diff --git a/test/test_jit.py b/test/test_jit.py
@@ -10536,6 +10536,30 @@ def test_indexing_out_of_bounds_neg():
         self.checkScriptRaisesRegex(test_indexing_out_of_bounds_pos, (), Exception,
                                     "out of range")
 
+        def negative_index():
+            tup = (1, 2, 3, 4)
+            return tup[-1]
+
+        self.checkScript(negative_index, [])
+
+        def really_negative_index():
+            tup = (1, 2, 3, 4)
+            return tup[-100]
+
+        self.checkScriptRaisesRegex(really_negative_index, [], Exception, "index out of range")
+
+        def negative_slice():
+            tup = (1, 2, 3, 4)
+            return tup[-3:4]
+
+        self.checkScript(negative_slice, [])
+
+        def really_slice_out_of_bounds():
+            tup = (1, 2, 3, 4)
+            return tup[-300:4000]
+
+        self.checkScript(really_slice_out_of_bounds, [])
+
     def test_namedtuple_attr(self):
         def f(x):
             return x.max(dim=1).indices + torch.max(x, dim=1).indices
diff --git a/torch/csrc/jit/passes/lower_tuples.cpp b/torch/csrc/jit/passes/lower_tuples.cpp
@@ -48,7 +48,16 @@ void removeTupleNodes(Node* n, bool must_remove_tuples) {
       }
       return;
     }
-    n->output()->replaceAllUsesWith(construct->inputs().at(*maybe_int));
+    auto int_idx = *maybe_int;
+    auto len = construct->output()->type()->containedTypes().size();
+    if (int_idx < 0) {
+      int_idx += len;
+    }
+    // currently, we allow non-constant tuple index if the tuple is of one type.
+    // so we need to check bounds here
+    if (int_idx >= 0 && static_cast<size_t>(int_idx) < len) {
+      n->output()->replaceAllUsesWith(construct->inputs().at(int_idx));
+    }
   } else if (n->kind() == prim::TupleSlice) {
     std::vector<Value*> values;
     int64_t beg = n->i(attr::beg);
