python · corona10 · Nov 6, 2024 · Nov 5, 2024 · Nov 6, 2024 · markshannon
diff --git a/Lib/test/test_dis.py b/Lib/test/test_dis.py
@@ -1335,6 +1335,27 @@ def test_call_specialize(self):
        got = self.get_disassembly(co, adaptive=True)
        self.do_disassembly_compare(got, call_quicken)

+    @cpython_only
+    @requires_specialization_ft
+    def test_contains_specialize(self):
+        contains_op_quicken = """\
+  0           RESUME_CHECK             0
+
+  1           LOAD_NAME                0 (a)
+              LOAD_NAME                1 (b)
+              %s
+              RETURN_VALUE
+"""
+        co_dict = compile('a in b', "<dict>", "eval")
+        self.code_quicken(lambda: exec(co_dict, {}, {'a': 1, 'b': {1: 5}}))
+        got = self.get_disassembly(co_dict, adaptive=True)
+        self.do_disassembly_compare(got, contains_op_quicken % "CONTAINS_OP_DICT         0 (in)")
+
+        co_set = compile('a in b', "<set>", "eval")
+        self.code_quicken(lambda: exec(co_set, {}, {'a': 1.0, 'b': {1, 2, 3}}))
+        got = self.get_disassembly(co_set, adaptive=True)
+        self.do_disassembly_compare(got, contains_op_quicken % "CONTAINS_OP_SET          0 (in)")
+
    @cpython_only
    @requires_specialization
    def test_loop_quicken(self):

@@ -2508,7 +2508,7 @@ dummy_func(
        }

        specializing op(_SPECIALIZE_CONTAINS_OP, (counter/1, left, right -- left, right)) {
-            #if ENABLE_SPECIALIZATION
+            #if ENABLE_SPECIALIZATION_FT
            if (ADAPTIVE_COUNTER_TRIGGERS(counter)) {
                next_instr = this_instr;
                _Py_Specialize_ContainsOp(right, next_instr);

diff --git a/Python/generated_cases.c.h b/Python/generated_cases.c.h
diff --git a/Python/specialize.c b/Python/specialize.c
@@ -2747,25 +2747,27 @@ _Py_Specialize_ContainsOp(_PyStackRef value_st, _Py_CODEUNIT *instr)
 {
    PyObject *value = PyStackRef_AsPyObjectBorrow(value_st);

-    assert(ENABLE_SPECIALIZATION);
+    assert(ENABLE_SPECIALIZATION_FT);
    assert(_PyOpcode_Caches[CONTAINS_OP] == INLINE_CACHE_ENTRIES_COMPARE_OP);
+    uint8_t specialized_op;
    _PyContainsOpCache *cache = (_PyContainsOpCache  *)(instr + 1);
    if (PyDict_CheckExact(value)) {
-        instr->op.code = CONTAINS_OP_DICT;
+        specialized_op = CONTAINS_OP_DICT;
        goto success;
    }
    if (PySet_CheckExact(value) || PyFrozenSet_CheckExact(value)) {
-        instr->op.code = CONTAINS_OP_SET;
+        specialized_op = CONTAINS_OP_SET;
        goto success;
    }

    SPECIALIZATION_FAIL(CONTAINS_OP, containsop_fail_kind(value));
    STAT_INC(CONTAINS_OP, failure);
-    instr->op.code = CONTAINS_OP;
+    SET_OPCODE_OR_RETURN(instr, CONTAINS_OP);
    cache->counter = adaptive_counter_backoff(cache->counter);
    return;
 success:
    STAT_INC(CONTAINS_OP, success);
+    SET_OPCODE_OR_RETURN(instr, specialized_op);
    cache->counter = adaptive_counter_cooldown();
 }