Address code review

corona10 · Nov 10, 2024 · 23e190e · 23e190e
1 parent 1f4a350
commit 23e190e
Show file tree

Hide file tree

Showing 3 changed files with 52 additions and 18 deletions.
diff --git a/Include/internal/pycore_typeobject.h b/Include/internal/pycore_typeobject.h
@@ -269,6 +269,11 @@ extern unsigned int _PyType_GetVersionForCurrentState(PyTypeObject *tp);
 PyAPI_FUNC(void) _PyType_SetVersion(PyTypeObject *tp, unsigned int version);
 PyTypeObject *_PyType_LookupByVersion(unsigned int version);
 
+// Returns 0 on success or caller-specific error on failure.
+typedef int (*_py_validate_type)(PyTypeObject *);
+// Returns 0 on success, -1 if no type version could be assigned, or the error returned by validate
+extern int _PyType_Validate(PyTypeObject *ty, _py_validate_type validate, unsigned int *tp_version);
+
 #ifdef __cplusplus
 }
 #endif

diff --git a/Objects/typeobject.c b/Objects/typeobject.c
@@ -5645,6 +5645,24 @@ _PyType_SetFlags(PyTypeObject *self, unsigned long mask, unsigned long flags)
     END_TYPE_LOCK();
 }
 
+int
+_PyType_Validate(PyTypeObject *ty, _py_validate_type validate, unsigned int *tp_version)
+{
+    int err;
+    BEGIN_TYPE_LOCK();
+    err = validate(ty);
+    if (!err) {
+        if(assign_version_tag(_PyInterpreterState_GET(), ty)) {
+            *tp_version = ty->tp_version_tag;
+        }
+        else {
+            err = -1;
+        }
+    }
+    END_TYPE_LOCK();
+    return err;
+}
+
 static void
 set_flags_recursive(PyTypeObject *self, unsigned long mask, unsigned long flags)
 {

diff --git a/Python/specialize.c b/Python/specialize.c
@@ -2690,6 +2690,25 @@ to_bool_fail_kind(PyObject *value)
     return SPEC_FAIL_OTHER;
 }
 
+static int
+check_type_always_true(PyTypeObject *ty)
+{
+    PyNumberMethods *nb = ty->tp_as_number;
+    if (nb && nb->nb_bool) {
+        return SPEC_FAIL_TO_BOOL_NUMBER;
+    }
+    PyMappingMethods *mp = ty->tp_as_mapping;
+    if (mp && mp->mp_length) {
+        return SPEC_FAIL_TO_BOOL_MAPPING;
+    }
+    PySequenceMethods *sq = ty->tp_as_sequence;
+    if (sq && sq->sq_length) {
+      return SPEC_FAIL_TO_BOOL_SEQUENCE;
+    }
+    return 0;
+}
+
+
 void
 _Py_Specialize_ToBool(_PyStackRef value_o, _Py_CODEUNIT *instr)
 {
@@ -2718,33 +2737,25 @@ _Py_Specialize_ToBool(_PyStackRef value_o, _Py_CODEUNIT *instr)
         return;
     }
     if (PyType_HasFeature(Py_TYPE(value), Py_TPFLAGS_HEAPTYPE)) {
-        PyNumberMethods *nb = Py_TYPE(value)->tp_as_number;
-        if (nb && nb->nb_bool) {
-            unspecialize(instr, SPEC_FAIL_TO_BOOL_NUMBER);
-            return;
-        }
-        PyMappingMethods *mp = Py_TYPE(value)->tp_as_mapping;
-        if (mp && mp->mp_length) {
-            unspecialize(instr, SPEC_FAIL_TO_BOOL_MAPPING);
-            return;
-        }
-        PySequenceMethods *sq = Py_TYPE(value)->tp_as_sequence;
-        if (sq && sq->sq_length) {
-            unspecialize(instr, SPEC_FAIL_TO_BOOL_SEQUENCE);
+        unsigned int version = 0;
+        int err = _PyType_Validate(Py_TYPE(value), check_type_always_true, &version);
+        if (err < 0) {
+            unspecialize(instr, SPEC_FAIL_OUT_OF_VERSIONS);
             return;
         }
-        if (!PyUnstable_Type_AssignVersionTag(Py_TYPE(value))) {
-            unspecialize(instr, SPEC_FAIL_OUT_OF_VERSIONS);
+        else if (err > 0) {
+            unspecialize(instr, err);
             return;
         }
-        uint32_t version = Py_TYPE(value)->tp_version_tag;
+
+        assert(err == 0);
         if (version == 0) {
             unspecialize(instr, SPEC_FAIL_OUT_OF_VERSIONS);
             return;
         }
-        specialize(instr, TO_BOOL_ALWAYS_TRUE);
-        write_u32(cache->version, version);
         assert(version);
+        write_u32(cache->version, version);
+        specialize(instr, TO_BOOL_ALWAYS_TRUE);
         return;
     }
     unspecialize(instr, to_bool_fail_kind(value));