Simplify changes in pybind11/critical_section.h and add test_nullptr_combinations()

Author: rwgk

include/pybind11/critical_section.h

@@ -13,28 +13,33 @@ PYBIND11_NAMESPACE_BEGIN(PYBIND11_NAMESPACE)
 class scoped_critical_section {
 public:
 #ifdef Py_GIL_DISABLED
-    scoped_critical_section(handle obj1, handle obj2) : m_ptr1(obj1.ptr()), m_ptr2(obj2.ptr()) {
-        if (m_ptr1 == nullptr) {
-            std::swap(m_ptr1, m_ptr2);
-        }
-        if (m_ptr2 != nullptr) {
-            PyCriticalSection2_Begin(&section2, m_ptr1, m_ptr2);
-        } else if (m_ptr1 != nullptr) {
-            PyCriticalSection_Begin(&section, m_ptr1);
+    scoped_critical_section(handle obj1, handle obj2) {
+        if (obj1) {
+            if (obj2) {
+                PyCriticalSection2_Begin(&section2, obj1.ptr(), obj2.ptr());
+                rank = 2;
+            } else {
+                PyCriticalSection_Begin(&section, obj1.ptr());
+                rank = 1;
+            }
+        } else if (obj2) {
+            PyCriticalSection_Begin(&section, obj2.ptr());
+            rank = 1;
         }
     }
 
-    explicit scoped_critical_section(handle obj) : m_ptr1(obj.ptr()) {
-        if (m_ptr1 != nullptr) {
-            PyCriticalSection_Begin(&section, m_ptr1);
+    explicit scoped_critical_section(handle obj) {
+        if (obj) {
+            PyCriticalSection_Begin(&section, obj.ptr());
+            rank = 1;
         }
     }
 
     ~scoped_critical_section() {
-        if (m_ptr2 != nullptr) {
-            PyCriticalSection2_End(&section2);
-        } else if (m_ptr1 != nullptr) {
+        if (rank == 1) {
             PyCriticalSection_End(&section);
+        } else if (rank == 2) {
+            PyCriticalSection2_End(&section2);
         }
     }
 #else
@@ -48,8 +53,7 @@ class scoped_critical_section {
 
 private:
 #ifdef Py_GIL_DISABLED
-    PyObject *m_ptr1{nullptr};
-    PyObject *m_ptr2{nullptr};
+    int rank{0};
     union {
         PyCriticalSection section;
         PyCriticalSection2 section2;

tests/test_scoped_critical_section.cpp

@@ -14,6 +14,20 @@
 
 namespace test_scoped_critical_section_ns {
 
+void test_one_nullptr() { py::scoped_critical_section lock{py::handle{}}; }
+
+void test_two_nullptrs() { py::scoped_critical_section lock{py::handle{}, py::handle{}}; }
+
+void test_first_nullptr() {
+    py::dict d;
+    py::scoped_critical_section lock{py::handle{}, d};
+}
+
+void test_second_nullptr() {
+    py::dict d;
+    py::scoped_critical_section lock{d, py::handle{}};
+}
+
 // Referenced test implementation: https://github.com/PyO3/pyo3/blob/v0.25.0/src/sync.rs#L874
 class BoolWrapper {
 public:
@@ -164,12 +178,10 @@ void test_scoped_critical_section2_same_object_no_deadlock(const py::handle &) {
 TEST_SUBMODULE(scoped_critical_section, m) {
     using namespace test_scoped_critical_section_ns;
 
-    m.attr("defined_THREAD_SANITIZER") =
-#if defined(THREAD_SANITIZER)
-        true;
-#else
-        false;
-#endif
+    m.def("test_one_nullptr", test_one_nullptr);
+    m.def("test_two_nullptrs", test_two_nullptrs);
+    m.def("test_first_nullptr", test_first_nullptr);
+    m.def("test_second_nullptr", test_second_nullptr);
 
     auto BoolWrapperClass = py::class_<BoolWrapper>(m, "BoolWrapper")
                                 .def(py::init<bool>())

tests/test_scoped_critical_section.py

@@ -5,6 +5,13 @@
 from pybind11_tests import scoped_critical_section as m
 
 
+def test_nullptr_combinations():
+    m.test_one_nullptr()
+    m.test_two_nullptrs()
+    m.test_first_nullptr()
+    m.test_second_nullptr()
+
+
 @pytest.mark.skipif(not m.has_barrier, reason="no <barrier>")
 def test_scoped_critical_section() -> None:
     for _ in range(64):