All domain Registry Traversal per Backend

include/drjit-core/jit.h

@@ -494,9 +494,15 @@ extern JIT_EXPORT void jit_registry_remove(const void *ptr);
 extern JIT_EXPORT uint32_t jit_registry_id(const void *ptr);
 
 /// Return the largest instance ID for the given domain
+/// If the \c domain is a nullptr, it returns the number of active entries in
+/// all domains for the given backend
 extern JIT_EXPORT uint32_t jit_registry_id_bound(JitBackend backend,
                                                  const char *domain);
 
+/// Fills the \c dest pointer array with all pointers registered in the registry
+/// \c dest has to point to an array with \c jit_registry_id_bound(backend, nullptr) entries
+extern JIT_EXPORT void jit_registry_get_pointers(JitBackend backend, void **dest);
+
 /// Return the pointer value associated with a given instance ID
 extern JIT_EXPORT void *jit_registry_ptr(JitBackend backend,
                                          const char *domain, uint32_t id);

src/api.cpp

@@ -983,6 +983,11 @@ uint32_t jit_registry_id_bound(JitBackend backend, const char *domain) {
     return jitc_registry_id_bound(backend, domain);
 }
 
+void jit_registry_get_pointers(JitBackend backend, void **dest) {
+    lock_guard guard(state.lock);
+    return jitc_registry_get_pointers(backend, dest);
+}
+
 void *jit_registry_ptr(JitBackend backend, const char *domain, uint32_t id) {
     lock_guard guard(state.lock);
     return jitc_registry_ptr(backend, domain, id);

src/registry.cpp

@@ -39,6 +39,7 @@ struct Ptr {
 // Per-domain information: a forward map (index-> pointer) and a list of unused entries
 struct Domain {
     const char *name;
+    JitBackend backend;
     uint32_t id_bound;
     std::vector<Ptr> fwd_map;
     std::priority_queue<uint32_t, std::vector<uint32_t>, std::greater<uint32_t>>
@@ -76,6 +77,7 @@ uint32_t jitc_registry_put(JitBackend backend, const char *domain_name, void *pt
         r.domains.emplace_back();
         Domain &domain = r.domains.back();
         domain.name = domain_name;
+        domain.backend = backend;
         domain.id_bound = 0;
     }
 
@@ -149,13 +151,38 @@ uint32_t jitc_registry_id(const void *ptr) {
 
 uint32_t jitc_registry_id_bound(JitBackend backend, const char *domain) {
     Registry &r = registry;
+    if (!domain) {
+        uint32_t n = 0;
+        for (Domain &domain : r.domains) {
+            if (domain.backend == backend)
+                for (auto ptr : domain.fwd_map)
+                    if (ptr.active)
+                        n++;
+        }
+        return n;
+    }
     auto it = r.domain_ids.find(DomainKey{ backend, domain });
     if (it == r.domain_ids.end())
         return 0;
     else
         return r.domains[it->second].id_bound;
 }
 
+void jitc_registry_get_pointers(JitBackend backend, void **dest) {
+    const Registry &r = registry;
+
+    uint32_t n = 0;
+    for (const Domain &domain : r.domains) {
+        if (domain.backend == backend)
+            for (auto ptr : domain.fwd_map) {
+                if (ptr.active) {
+                    dest[n] = ptr.ptr;
+                    n++;
+                }
+            }
+    }
+}
+
 void *jitc_registry_ptr(JitBackend backend, const char *domain_name, uint32_t id) {
     if (id == 0)
         return nullptr;

src/registry.h

@@ -22,8 +22,14 @@ extern void jitc_registry_remove(const void *ptr);
 extern uint32_t jitc_registry_id(const void *ptr);
 
 /// Return the largest instance ID for the given domain
+/// If the \c domain is a nullptr, it returns the number of active entries in
+/// all domains for the given backend
 extern uint32_t jitc_registry_id_bound(JitBackend backend, const char *domain);
 
+/// Fills the \c dest pointer array with all pointers registered in the registry
+/// \c dest has to point to an array with \c jit_registry_id_bound(backend, nullptr) entries
+void extern jitc_registry_get_pointers(JitBackend backend, void **dest);
+
 /// Return the pointer value associated with a given instance ID
 extern void *jitc_registry_ptr(JitBackend backend, const char *domain, uint32_t id);