Fix some error handling in Simplify

Summary:
This started as cleanup of some small problems I noticed on
D47854475 after it landed, and snowballed a bit:
- Make sure to check the return value of C-API functions that can raise.
- When emitting UseTypes for types where we care about the specialization, use
  the original type. These types can't come from a guard at the moment anyway,
  but it won't hurt anything.
- `PrimitiveUnbox` should be always followed by `IsNegativeOrErrOccurred`.
- To better match exceptions raised by the interpreter, add `hir::IndexUnbox`,
  which is similar to `PrimitiveUnbox` but raises a different (and
  configurable) exception.

Reviewed By: alexmalyshev

Differential Revision: D50475141

fbshipit-source-id: a9f26d61cde92f54d68b855ce1d4c7a6f7f58cdb

Jit/hir/analysis.cpp

@@ -80,6 +80,7 @@ bool isPassthrough(const Instr& instr) {
     case Opcode::kImportFrom:
     case Opcode::kImportName:
     case Opcode::kInPlaceOp:
+    case Opcode::kIndexUnbox:
     case Opcode::kInitialYield:
     case Opcode::kIntBinaryOp:
     case Opcode::kIntConvert:

Jit/hir/hir.cpp

@@ -142,6 +142,7 @@ bool Instr::isReplayable() const {
     case Opcode::kGuardIs:
     case Opcode::kGuardType:
     case Opcode::kHintType:
+    case Opcode::kIndexUnbox:
     case Opcode::kIntBinaryOp:
     case Opcode::kIntConvert:
     case Opcode::kIsNegativeAndErrOccurred:

Jit/hir/hir.h

@@ -268,6 +268,7 @@ struct FrameState {
   V(ImportName)                        \
   V(InPlaceOp)                         \
   V(Incref)                            \
+  V(IndexUnbox)                        \
   V(InitialYield)                      \
   V(IntBinaryOp)                       \
   V(PrimitiveBoxBool)                  \
@@ -2624,6 +2625,21 @@ class INSTR_CLASS(PrimitiveUnbox, (), HasOutput, Operands<1>) {
   Type type_;
 };
 
+// Similar to PrimitiveUnbox, but uses PyNumber_AsSsize_t() instead of
+// PyLong_AsSize_t() for a different exception and message on overflow.
+class INSTR_CLASS(IndexUnbox, (TLong), HasOutput, Operands<1>) {
+ public:
+  IndexUnbox(Register* dst, Register* value, PyObject* exc = PyExc_IndexError)
+      : InstrT(dst, value), exc_(exc) {}
+
+  PyObject* exception() const {
+    return exc_;
+  }
+
+ private:
+  PyObject* exc_;
+};
+
 class CondBranchBase : public Instr {
  public:
   CondBranchBase(Opcode opcode, BasicBlock* true_bb, BasicBlock* false_bb)

Jit/hir/memory_effects.cpp

@@ -41,6 +41,7 @@ MemoryEffects memoryEffects(const Instr& inst) {
     case Opcode::kGetSecondOutput:
     case Opcode::kGuardType:
     case Opcode::kHintType:
+    case Opcode::kIndexUnbox:
     case Opcode::kIntBinaryOp:
     case Opcode::kIntConvert:
     case Opcode::kIsNegativeAndErrOccurred:

Jit/hir/printer.cpp

@@ -529,6 +529,11 @@ static std::string format_immediates(const Instr& instr) {
       const auto& unbox = static_cast<const PrimitiveUnbox&>(instr);
       return fmt::format("{}", unbox.type());
     }
+    case Opcode::kIndexUnbox: {
+      const auto& unbox = static_cast<const IndexUnbox&>(instr);
+      return fmt::format(
+          "{}", reinterpret_cast<PyTypeObject*>(unbox.exception())->tp_name);
+    }
     case Opcode::kLoadGlobalCached: {
       const auto& load = static_cast<const LoadGlobalCached&>(instr);
       return format_name(load, load.name_idx());

Jit/hir/simplify.cpp

@@ -561,7 +561,8 @@ Register* simplifyBinaryOp(Env& env, const BinaryOp* instr) {
       // representation of the type, which should be analagous to Py_ssize_t.
       env.emit<UseType>(lhs, lhs->isA(TListExact) ? TListExact : TTupleExact);
       env.emit<UseType>(rhs, TLongExact);
-      Register* right_index = env.emit<PrimitiveUnbox>(rhs, TCInt64);
+      Register* right_index = env.emit<IndexUnbox>(rhs);
+      env.emit<IsNegativeAndErrOccurred>(right_index, *instr->frameState());
       Register* adjusted_idx =
           env.emit<CheckSequenceBounds>(lhs, right_index, *instr->frameState());
       ssize_t offset = offsetof(PyTupleObject, ob_item);
@@ -579,6 +580,10 @@ Register* simplifyBinaryOp(Env& env, const BinaryOp* instr) {
       if (lhs_type.hasObjectSpec() && rhs_type.hasObjectSpec()) {
         // Constant propagation
         Py_ssize_t idx = PyLong_AsSsize_t(rhs_type.objectSpec());
+        if (idx == -1 && PyErr_Occurred()) {
+          PyErr_Clear();
+          return nullptr;
+        }
         Py_ssize_t n = PyUnicode_GetLength(lhs_type.objectSpec());
 
         if (idx < -n || idx >= n) {
@@ -593,17 +598,22 @@ Register* simplifyBinaryOp(Env& env, const BinaryOp* instr) {
         Py_UCS4 c = PyUnicode_ReadChar(lhs_type.objectSpec(), idx);
         PyObject* substr =
             PyUnicode_FromKindAndData(PyUnicode_4BYTE_KIND, &c, 1);
+        if (substr == nullptr) {
+          return nullptr;
+        }
         PyUnicode_InternInPlace(&substr);
         Ref<> result = Ref<>::steal(substr);
 
-        env.emit<UseType>(lhs, TUnicodeExact);
-        env.emit<UseType>(rhs, TLongExact);
+        // Use exact types since we're relying on the object specializations.
+        env.emit<UseType>(lhs, lhs_type);
+        env.emit<UseType>(rhs, rhs_type);
         return env.emit<LoadConst>(
             Type::fromObject(env.func.env.addReference(result)));
       } else {
         env.emit<UseType>(lhs, TUnicodeExact);
         env.emit<UseType>(rhs, TLongExact);
-        Register* unboxed_idx = env.emit<PrimitiveUnbox>(rhs, TCInt64);
+        Register* unboxed_idx = env.emit<IndexUnbox>(rhs);
+        env.emit<IsNegativeAndErrOccurred>(unboxed_idx, *instr->frameState());
         Register* adjusted_idx = env.emit<CheckSequenceBounds>(
             lhs, unboxed_idx, *instr->frameState());
         return env.emit<UnicodeSubscr>(lhs, adjusted_idx, *instr->frameState());
@@ -624,7 +634,7 @@ Register* simplifyBinaryOp(Env& env, const BinaryOp* instr) {
   }
   if ((lhs->isA(TUnicodeExact) && rhs->isA(TLongExact)) &&
       (instr->op() == BinaryOpKind::kMultiply)) {
-    Register* unboxed_rhs = env.emit<PrimitiveUnbox>(rhs, TCInt64);
+    Register* unboxed_rhs = env.emit<IndexUnbox>(rhs, PyExc_OverflowError);
     env.emit<IsNegativeAndErrOccurred>(unboxed_rhs, *instr->frameState());
     return env.emit<UnicodeRepeat>(lhs, unboxed_rhs, *instr->frameState());
   }
@@ -717,50 +727,49 @@ Register* simplifyPrimitiveBoxBool(Env& env, const PrimitiveBoxBool* instr) {
   return nullptr;
 }
 
-Register* simplifyPrimitiveUnbox(Env& env, const PrimitiveUnbox* instr) {
-  Register* unboxed_value = instr->GetOperand(0);
-  if (unboxed_value->instr()->IsPrimitiveBox()) {
+Register* simplifyUnbox(Env& env, const Instr* instr) {
+  Register* input_value = instr->GetOperand(0);
+  Type output_type = instr->GetOutput()->type();
+  if (input_value->instr()->IsPrimitiveBox()) {
     // Simplify unbox(box(x)) -> x
-    const PrimitiveBox* box =
-        static_cast<PrimitiveBox*>(unboxed_value->instr());
-    if (box->type() == instr->type()) {
+    const auto box = static_cast<PrimitiveBox*>(input_value->instr());
+    if (box->type() == output_type) {
       // We can't optimize away the potential overflow in unboxing.
       return box->GetOperand(0);
     }
   }
-  Type unbox_output_type = instr->GetOutput()->type();
   // Ensure that we are dealing with either a integer or a double.
-  Type unboxed_value_type = unboxed_value->type();
-  if (!(unboxed_value_type.hasObjectSpec())) {
+  Type input_value_type = input_value->type();
+  if (!(input_value_type.hasObjectSpec())) {
     return nullptr;
   }
-  PyObject* value = unboxed_value_type.objectSpec();
-  if (unbox_output_type <= (TCSigned | TCUnsigned)) {
+  PyObject* value = input_value_type.objectSpec();
+  if (output_type <= (TCSigned | TCUnsigned)) {
     if (!PyLong_Check(value)) {
       return nullptr;
     }
     int overflow = 0;
     long number =
-        PyLong_AsLongAndOverflow(unboxed_value_type.objectSpec(), &overflow);
+        PyLong_AsLongAndOverflow(input_value_type.objectSpec(), &overflow);
     if (overflow != 0) {
       return nullptr;
     }
-    if (unbox_output_type <= TCSigned) {
-      if (!Type::CIntFitsType(number, unbox_output_type)) {
+    if (output_type <= TCSigned) {
+      if (!Type::CIntFitsType(number, output_type)) {
         return nullptr;
       }
-      return env.emit<LoadConst>(Type::fromCInt(number, unbox_output_type));
+      return env.emit<LoadConst>(Type::fromCInt(number, output_type));
     } else {
-      if (!Type::CUIntFitsType(number, unbox_output_type)) {
+      if (!Type::CUIntFitsType(number, output_type)) {
         return nullptr;
       }
-      return env.emit<LoadConst>(Type::fromCUInt(number, unbox_output_type));
+      return env.emit<LoadConst>(Type::fromCUInt(number, output_type));
     }
-  } else if (unbox_output_type <= TCDouble) {
+  } else if (output_type <= TCDouble) {
     if (!PyFloat_Check(value)) {
       return nullptr;
     }
-    double number = PyFloat_AS_DOUBLE(unboxed_value_type.objectSpec());
+    double number = PyFloat_AS_DOUBLE(input_value_type.objectSpec());
     return env.emit<LoadConst>(Type::fromCDouble(number));
   }
   return nullptr;
@@ -1315,9 +1324,9 @@ Register* simplifyInstr(Env& env, const Instr* instr) {
     case Opcode::kPrimitiveBoxBool:
       return simplifyPrimitiveBoxBool(
           env, static_cast<const PrimitiveBoxBool*>(instr));
+    case Opcode::kIndexUnbox:
     case Opcode::kPrimitiveUnbox:
-      return simplifyPrimitiveUnbox(
-          env, static_cast<const PrimitiveUnbox*>(instr));
+      return simplifyUnbox(env, instr);
 
     case Opcode::kIsNegativeAndErrOccurred:
       return simplifyIsNegativeAndErrOccurred(

Jit/hir/ssa.cpp

@@ -639,6 +639,9 @@ Type outputType(
       auto& unbox = static_cast<const PrimitiveUnbox&>(instr);
       return unbox.type();
     }
+    case Opcode::kIndexUnbox: {
+      return TCInt64;
+    }
 
     // Check opcodes return a copy of their input that is statically known to
     // not be null.

Jit/jit_rt.cpp

@@ -1784,6 +1784,7 @@ void JITRT_BatchDecref(PyObject** args, int nargs) {
 }
 
 Py_ssize_t JITRT_CheckSequenceBounds(PyObject* s, Py_ssize_t i) {
+  JIT_DCHECK(!PyErr_Occurred(), "called with error set");
   i = i < 0 ? i + Py_SIZE(s) : i;
   if (i < 0 || i >= Py_SIZE(s)) {
     PyErr_SetString(PyExc_IndexError, "index out of range");

Jit/lir/generator.cpp

@@ -1065,6 +1065,15 @@ LIRGenerator::TranslatedBlock LIRGenerator::TranslateOneBasicBlock(
         }
         break;
       }
+      case Opcode::kIndexUnbox: {
+        auto instr = static_cast<const IndexUnbox*>(&i);
+        bbb.AppendCall(
+            instr->dst(),
+            PyNumber_AsSsize_t,
+            instr->GetOperand(0),
+            instr->exception());
+        break;
+      }
       case Opcode::kPrimitiveUnaryOp: {
         auto instr = static_cast<const PrimitiveUnaryOp*>(&i);
         switch (instr->op()) {

Lib/test/cinderjit_profile_test.py

@@ -6,6 +6,7 @@
 # 'testcinder_jit_profile' make target.
 
 import os
+import sys
 import threading
 import unittest
 from collections import Counter
@@ -104,6 +105,33 @@ def do_pow(a, b):
             with self.assertDeopts({}):
                 self.assertEqual(do_pow(2, 8), 256)
 
+    def assert_index_error(self):
+        return self.assertRaisesRegex(
+            IndexError, "cannot fit 'int' into an index-sized integer"
+        )
+
+    def test_raising_str_subscr(self):
+        s = "abcdefu"
+        with self.assert_index_error():
+            return s[sys.maxsize + 1]
+
+    def test_raising_list_subscr(self):
+        l = [1, 2, 3]
+        with self.assert_index_error():
+            return l[sys.maxsize + 1]
+
+    def test_raising_tuple_subscr(self):
+        t = (1, 2, 3)
+        with self.assert_index_error():
+            return t[sys.maxsize + 1]
+
+    def test_raising_str_repeat(self):
+        s = "123"
+        with self.assertRaisesRegex(
+            OverflowError, "cannot fit 'int' into an index-sized integer"
+        ):
+            return s * (sys.maxsize + 1)
+
 
 class NewThreadTests(ProfileTest):
     def test_create_thread(self):

RuntimeTests/hir_tests/all_passes_static_test.txt

@@ -18,7 +18,7 @@ fun jittestmodule:test {
 
   bb 4 (preds 0, 2) {
     v24:CInt64 = Phi<0, 2> v18 v37
-    v25:OptObject = Phi<0, 2> v16 v47
+    v25:OptObject = Phi<0, 2> v16 v48
     v21:CInt32 = LoadEvalBreaker
     CondBranch<5, 1> v21
   }
@@ -42,34 +42,42 @@ fun jittestmodule:test {
   }
 
   bb 2 (preds 1) {
-    v46:CInt64 = CheckSequenceBounds v15 v24 {
+    v46:CInt64 = IsNegativeAndErrOccurred v24 {
       LiveValues<3> b:v15 s:v24 o:v25
       FrameState {
         NextInstrOffset 24
         Locals<3> v15 v24 v25
         Stack<1> v15
       }
     }
-    v47:Object = LoadArrayItem<Offset[24]> v15 v46 v15
+    v47:CInt64 = CheckSequenceBounds v15 v24 {
+      LiveValues<3> b:v15 s:v24 o:v25
+      FrameState {
+        NextInstrOffset 24
+        Locals<3> v15 v24 v25
+        Stack<1> v15
+      }
+    }
+    v48:Object = LoadArrayItem<Offset[24]> v15 v47 v15
     v36:CInt64[1] = LoadConst<CInt64[1]>
     v37:CInt64 = IntBinaryOp<Add> v24 v36
     v40:OptObject = LoadGlobalCached<0; "print">
     v41:MortalObjectUser[builtin_function_or_method:print:0xdeadbeef] = GuardIs<0xdeadbeef> v40 {
       Descr 'LOAD_GLOBAL: print'
-      LiveValues<6> b:v15 s:v24 o:v25 s:v37 b:v40 b:v47
+      LiveValues<6> b:v15 s:v24 o:v25 s:v37 b:v40 b:v48
       FrameState {
         NextInstrOffset 30
         Locals<3> v15 v24 v25
-        Stack<3> v15 v47 v37
+        Stack<3> v15 v48 v37
       }
     }
-    Incref v47
+    Incref v48
     XDecref v25
-    v43:Object = VectorCall<1> v41 v47 {
-      LiveValues<4> b:v15 s:v37 b:v41 o:v47
+    v43:Object = VectorCall<1> v41 v48 {
+      LiveValues<4> b:v15 s:v37 b:v41 o:v48
       FrameState {
         NextInstrOffset 40
-        Locals<3> v15 v37 v47
+        Locals<3> v15 v37 v48
         Stack<1> v15
       }
     }

RuntimeTests/hir_tests/inliner_test.txt

@@ -134,19 +134,20 @@ fun jittestmodule:test {
     Snapshot
     UseType<ListExact> v8
     UseType<LongExact> v15
-    v33:CInt64[1] = LoadConst<CInt64[1]>
-    v30:CInt64 = CheckSequenceBounds v8 v33 {
+    v34:CInt64[1] = LoadConst<CInt64[1]>
+    v35:CInt64[0] = LoadConst<CInt64[0]>
+    v31:CInt64 = CheckSequenceBounds v8 v34 {
       FrameState {
         NextInstrOffset 6
         Locals<2> v8 v15
       }
     }
-    v31:CPtr = LoadField<ob_item@24, CPtr, borrowed> v8
-    v32:Object = LoadArrayItem v31 v30 v8
+    v32:CPtr = LoadField<ob_item@24, CPtr, borrowed> v8
+    v33:Object = LoadArrayItem v32 v31 v8
     Snapshot
     EndInlinedFunction
     Snapshot
-    Return v32
+    Return v33
   }
 }
 ---

RuntimeTests/hir_tests/profile_data_hir_test.txt

@@ -21,16 +21,22 @@ fun jittestmodule:test {
     }
     UseType<TupleExact> v7
     UseType<LongExact> v8
-    v10:CInt64 = PrimitiveUnbox<CInt64> v8
-    v11:CInt64 = CheckSequenceBounds v7 v10 {
+    v10:CInt64 = IndexUnbox<IndexError> v8
+    v11:CInt64 = IsNegativeAndErrOccurred v10 {
       FrameState {
         NextInstrOffset 6
         Locals<2> v7 v8
       }
     }
-    v12:Object = LoadArrayItem<Offset[24]> v7 v11 v7
+    v12:CInt64 = CheckSequenceBounds v7 v10 {
+      FrameState {
+        NextInstrOffset 6
+        Locals<2> v7 v8
+      }
+    }
+    v13:Object = LoadArrayItem<Offset[24]> v7 v12 v7
     Snapshot
-    Return v12
+    Return v13
   }
 }
 ---