Added base value test for custom_type_ext traversal

Author: DoeringChristian

tests/custom_type_ext.cpp

@@ -51,12 +51,15 @@ class Object : public drjit::TraversableBase {
 
 template <typename Value>
 class CustomBase : public Object{
+    Value m_base_value;
+
 public:
-    CustomBase() : Object() {}
+    CustomBase(const Value &base_value) : Object(), m_base_value(base_value) {}
 
+    Value &base_value() { return m_base_value; }
     virtual Value &value() = 0;
 
-    DR_TRAVERSE_CB(Object);
+    DR_TRAVERSE_CB(Object, m_base_value);
 };
 
 template <typename Value>
@@ -65,7 +68,7 @@ class PyCustomBase : public CustomBase<Value>{
     using Base = CustomBase<Value>;
     NB_TRAMPOLINE(Base, 1);
 
-    PyCustomBase() : Base() {}
+    PyCustomBase(const Value &base_value) : Base(base_value) {}
 
     Value &value() override { NB_OVERRIDE_PURE(value); }
 
@@ -77,8 +80,7 @@ class CustomA: public CustomBase<Value>{
 public:
     using Base = CustomBase<Value>;
 
-    CustomA() {}
-    CustomA(const Value &v) : m_value(v) {}
+    CustomA(const Value &value, const Value &base_value) : Base(base_value), m_value(value) {}
 
     Value &value() override { return m_value; }
 
@@ -119,13 +121,15 @@ template <JitBackend Backend> void bind(nb::module_ &m) {
         nb::intrusive_ptr<Object>(
             [](Object *o, PyObject *po) noexcept { o->set_self_py(po); }));
 
-    auto base = nb::class_<CustomBase, Object, PyCustomBase>(m, "CustomBase")
-                    .def(nb::init())
-                    .def("value", nb::overload_cast<>(&CustomBase::value));
+    auto base =
+        nb::class_<CustomBase, Object, PyCustomBase>(m, "CustomBase")
+            .def(nb::init<Float>())
+            .def("value", nb::overload_cast<>(&CustomBase::value))
+            .def("base_value", nb::overload_cast<>(&CustomBase::base_value));
 
     drjit::bind_traverse(base);
 
-    auto a = nb::class_<CustomA>(m, "CustomA").def(nb::init<Float>());
+    auto a = nb::class_<CustomA>(m, "CustomA").def(nb::init<Float, Float>());
 
     drjit::bind_traverse(a);
 

tests/test_custom_type_ext.py

@@ -78,10 +78,11 @@ def test04_traverse_opaque(t):
     pkg = get_pkg(t)
     Float = t
 
-    v = dr.arange(Float, 10)
+    value = dr.arange(Float, 10)
+    base_value = dr.arange(Float, 10)
 
-    a = pkg.CustomA(v)
-    assert dr.detail.collect_indices(a) == [v.index]
+    a = pkg.CustomA(value, base_value)
+    assert dr.detail.collect_indices(a) == [base_value.index, value.index]
 
 
 @pytest.test_arrays("float32,-diff,shape=(*),jit")
@@ -119,16 +120,17 @@ def test06_trampoline_traversal(t):
     pkg = get_pkg(t)
     Float = t
 
-    v = dr.opaque(Float, 0, 3)
+    value = dr.opaque(Float, 0, 3)
+    base_value = dr.opaque(Float, 1, 3)
 
     class B(pkg.CustomBase):
-        def __init__(self, v) -> None:
-            super().__init__()
-            self.v = v
+        def __init__(self, value, base_value) -> None:
+            super().__init__(base_value)
+            self._value = value
 
         def value(self):
-            return self.v
+            return self._value
 
-    b = B(v)
+    b = B(value, base_value)
 
-    assert dr.detail.collect_indices(b) == [v.index]
+    assert dr.detail.collect_indices(b) == [base_value.index, value.index]