@@ -720,116 +720,128 @@ Register* simplifyLoadMethod(Env& env, const LoadMethod* load_meth) {
720720 *load_meth->frameState ());
721721}
722722
723- Register* simplifyBinaryOp (Env& env, const BinaryOp* instr) {
723+ Register* simplifySubscript (Env& env, const BinaryOp* instr) {
724724 BinaryOpKind op = instr->op ();
725725 Register* lhs = instr->left ();
726726 Register* rhs = instr->right ();
727+ JIT_CHECK (
728+ op == BinaryOpKind::kSubscript ,
729+ " simplifySubscript trying to optimize {}" ,
730+ op);
727731
728- if (op == BinaryOpKind::kSubscript ) {
729- if (lhs->isA (TDictExact)) {
730- return env.emit <DictSubscr>(lhs, rhs, *instr->frameState ());
732+ if (lhs->isA (TDictExact)) {
733+ return env.emit <DictSubscr>(lhs, rhs, *instr->frameState ());
734+ }
735+ if (!rhs->isA (TLongExact)) {
736+ return nullptr ;
737+ }
738+
739+ Type lhs_type = lhs->type ();
740+ Type rhs_type = rhs->type ();
741+
742+ // Constant tuple subscripted by constant long.
743+ if (lhs_type <= TTupleExact && lhs_type.hasObjectSpec () &&
744+ rhs_type.hasObjectSpec ()) {
745+ int overflow;
746+ Py_ssize_t index =
747+ PyLong_AsLongAndOverflow (rhs_type.objectSpec (), &overflow);
748+ if (!overflow) {
749+ BorrowedRef<> lhs_obj = lhs_type.objectSpec ();
750+ if (index >= 0 && index < PyTuple_GET_SIZE (lhs_obj)) {
751+ BorrowedRef<> item = PyTuple_GET_ITEM (lhs_obj.get (), index);
752+ env.emit <UseType>(lhs, lhs_type);
753+ env.emit <UseType>(rhs, rhs_type);
754+ return env.emit <LoadConst>(
755+ Type::fromObject (env.func .env .addReference (item)));
756+ }
731757 }
732- if (!rhs->isA (TLongExact)) {
733- return nullptr ;
758+ }
759+
760+ // TODO(T255264263). Enable this for FT builds. See P2169673256.
761+ if (!kFreeThreadedBuild && (lhs->isA (TListExact) || lhs->isA (TTupleExact))) {
762+ // TASK(T93509109): Replace TCInt64 with a less platform-specific
763+ // representation of the type, which should be analagous to Py_ssize_t.
764+ env.emit <UseType>(lhs, lhs->isA (TListExact) ? TListExact : TTupleExact);
765+ env.emit <UseType>(rhs, TLongExact);
766+ Register* right_index = env.emit <IndexUnbox>(rhs);
767+ env.emit <IsNegativeAndErrOccurred>(right_index, *instr->frameState ());
768+ Register* adjusted_idx =
769+ env.emit <CheckSequenceBounds>(lhs, right_index, *instr->frameState ());
770+ Py_ssize_t offset = offsetof (PyTupleObject, ob_item);
771+ Register* array = lhs;
772+ // Lists carry a nested array of ob_item whereas tuples are variable-sized
773+ // structs.
774+ if (lhs->isA (TListExact)) {
775+ array = env.emit <LoadField>(
776+ lhs, " ob_item" , offsetof (PyListObject, ob_item), TCPtr);
777+ offset = 0 ;
734778 }
735- Type lhs_type = lhs->type ();
736- Type rhs_type = rhs->type ();
737- if (lhs_type <= TTupleExact && lhs_type.hasObjectSpec () &&
738- rhs_type.hasObjectSpec ()) {
739- int overflow;
740- Py_ssize_t index =
741- PyLong_AsLongAndOverflow (rhs_type.objectSpec (), &overflow);
742- if (!overflow) {
743- PyObject* lhs_obj = lhs_type.objectSpec ();
744- if (index >= 0 && index < PyTuple_GET_SIZE (lhs_obj)) {
745- BorrowedRef<> item = PyTuple_GET_ITEM (lhs_obj, index);
746- env.emit <UseType>(lhs, lhs_type);
747- env.emit <UseType>(rhs, rhs_type);
748- return env.emit <LoadConst>(
749- Type::fromObject (env.func .env .addReference (item)));
750- }
751- // Fallthrough
779+ return env.emit <LoadArrayItem>(array, adjusted_idx, lhs, offset, TObject);
780+ }
781+
782+ // Unicode subscript.
783+ if (lhs_type <= TUnicodeExact && rhs_type <= TLongExact) {
784+ // Constant fold. This isn't safe in multi-threaded compilation because the
785+ // worker doesn't hold the GIL and that's required for creating a new
786+ // string.
787+ if (!getThreadedCompileContext ().compileRunning () &&
788+ lhs_type.hasObjectSpec () && rhs_type.hasObjectSpec ()) {
789+ Py_ssize_t idx = PyLong_AsSsize_t (rhs_type.objectSpec ());
790+ if (idx == -1 && PyErr_Occurred ()) {
791+ PyErr_Clear ();
792+ return nullptr ;
752793 }
753- // Fallthrough
754- }
755- // TODO(T255264263). Enable this again. See P2169673256.
756- #ifndef Py_GIL_DISABLED
757- if (lhs->isA (TListExact) || lhs->isA (TTupleExact)) {
758- // TASK(T93509109): Replace TCInt64 with a less platform-specific
759- // representation of the type, which should be analagous to Py_ssize_t.
760- env.emit <UseType>(lhs, lhs->isA (TListExact) ? TListExact : TTupleExact);
761- env.emit <UseType>(rhs, TLongExact);
762- Register* right_index = env.emit <IndexUnbox>(rhs);
763- env.emit <IsNegativeAndErrOccurred>(right_index, *instr->frameState ());
764- Register* adjusted_idx =
765- env.emit <CheckSequenceBounds>(lhs, right_index, *instr->frameState ());
766- Py_ssize_t offset = offsetof (PyTupleObject, ob_item);
767- Register* array = lhs;
768- // Lists carry a nested array of ob_item whereas tuples are variable-sized
769- // structs.
770- if (lhs->isA (TListExact)) {
771- array = env.emit <LoadField>(
772- lhs, " ob_item" , offsetof (PyListObject, ob_item), TCPtr);
773- offset = 0 ;
794+ Py_ssize_t n = PyUnicode_GetLength (lhs_type.objectSpec ());
795+ if (idx < -n || idx >= n) {
796+ return nullptr ;
797+ }
798+
799+ if (idx < 0 ) {
800+ idx += n;
801+ }
802+
803+ ThreadedCompileSerialize guard;
804+ Py_UCS4 c = PyUnicode_ReadChar (lhs_type.objectSpec (), idx);
805+ PyObject* substr = PyUnicode_FromKindAndData (PyUnicode_4BYTE_KIND, &c, 1 );
806+ if (substr == nullptr ) {
807+ return nullptr ;
774808 }
775- return env.emit <LoadArrayItem>(array, adjusted_idx, lhs, offset, TObject);
809+ PyUnicode_InternInPlace (&substr);
810+ Ref<> result = Ref<>::steal (substr);
811+
812+ // Use exact types since we're relying on the object specializations.
813+ env.emit <UseType>(lhs, lhs_type);
814+ env.emit <UseType>(rhs, rhs_type);
815+ return env.emit <LoadConst>(
816+ Type::fromObject (env.func .env .addReference (std::move (result))));
776817 }
777- #endif
778- if (lhs_type <= TUnicodeExact && rhs_type <= TLongExact) { // Unicode subscr
779- if (lhs_type.hasObjectSpec () && rhs_type.hasObjectSpec ()) {
780- // This isn't safe in the multi-threaded compilation on 3.12 because
781- // we don't hold the GIL which is required for
782- // PyUnicode_InternInPlace.
783- RETURN_MULTITHREADED_COMPILE (nullptr );
784-
785- // Constant propagation
786- Py_ssize_t idx = PyLong_AsSsize_t (rhs_type.objectSpec ());
787- if (idx == -1 && PyErr_Occurred ()) {
788- PyErr_Clear ();
789- return nullptr ;
790- }
791- Py_ssize_t n = PyUnicode_GetLength (lhs_type.objectSpec ());
792818
793- if (idx < -n || idx >= n) {
794- return nullptr ;
795- }
819+ env.emit <UseType>(lhs, TUnicodeExact);
820+ env.emit <UseType>(rhs, TLongExact);
821+ Register* unboxed_idx = env.emit <IndexUnbox>(rhs);
822+ env.emit <IsNegativeAndErrOccurred>(unboxed_idx, *instr->frameState ());
823+ Register* adjusted_idx =
824+ env.emit <CheckSequenceBounds>(lhs, unboxed_idx, *instr->frameState ());
825+ return env.emit <UnicodeSubscr>(lhs, adjusted_idx, *instr->frameState ());
826+ }
796827
797- if (idx < 0 ) {
798- idx += n;
799- }
828+ return nullptr ;
829+ }
800830
801- ThreadedCompileSerialize guard;
802- Py_UCS4 c = PyUnicode_ReadChar (lhs_type.objectSpec (), idx);
803- PyObject* substr =
804- PyUnicode_FromKindAndData (PyUnicode_4BYTE_KIND, &c, 1 );
805- if (substr == nullptr ) {
806- return nullptr ;
807- }
808- PyUnicode_InternInPlace (&substr);
809- Ref<> result = Ref<>::steal (substr);
831+ Register* simplifyBinaryOp (Env& env, const BinaryOp* instr) {
832+ BinaryOpKind op = instr->op ();
833+ Register* lhs = instr->left ();
834+ Register* rhs = instr->right ();
810835
811- // Use exact types since we're relying on the object specializations.
812- env.emit <UseType>(lhs, lhs_type);
813- env.emit <UseType>(rhs, rhs_type);
814- return env.emit <LoadConst>(
815- Type::fromObject (env.func .env .addReference (std::move (result))));
816- } else {
817- env.emit <UseType>(lhs, TUnicodeExact);
818- env.emit <UseType>(rhs, TLongExact);
819- Register* unboxed_idx = env.emit <IndexUnbox>(rhs);
820- env.emit <IsNegativeAndErrOccurred>(unboxed_idx, *instr->frameState ());
821- Register* adjusted_idx = env.emit <CheckSequenceBounds>(
822- lhs, unboxed_idx, *instr->frameState ());
823- return env.emit <UnicodeSubscr>(lhs, adjusted_idx, *instr->frameState ());
824- }
825- }
836+ if (op == BinaryOpKind::kSubscript ) {
837+ return simplifySubscript (env, instr);
826838 }
827839
828840 if (lhs->isA (TLongExact) && rhs->isA (TLongExact)) {
829841 // All binary ops on TLong's return mutable so can be freely simplified with
830842 // no explicit check.
831- if (op == BinaryOpKind::kMatrixMultiply || op == BinaryOpKind:: kSubscript ) {
832- // These will generate an error at runtime.
843+ if (op == BinaryOpKind::kMatrixMultiply ) {
844+ // This will generate an error at runtime.
833845 return nullptr ;
834846 }
835847 env.emit <UseType>(lhs, TLongExact);
0 commit comments