rustkernel_core/
registry.rs

1//! Kernel registry with auto-discovery.
2//!
3//! The registry manages all registered kernels and provides lookup functionality.
4//! Kernels can be registered manually or discovered automatically via proc macros.
5
6use crate::domain::Domain;
7use crate::error::{KernelError, Result};
8use crate::kernel::{KernelMetadata, KernelMode};
9use crate::license::{LicenseError, LicenseValidator, SharedLicenseValidator};
10use crate::traits::{BatchKernel, BatchKernelDyn, GpuKernel, RingKernelDyn, TypeErasedBatchKernel};
11use hashbrown::HashMap;
12use std::sync::{Arc, RwLock};
13use tracing::{debug, info, warn};
14
15/// Registry statistics.
16#[derive(Debug, Clone, Default)]
17pub struct RegistryStats {
18    /// Total number of registered kernels.
19    pub total: usize,
20    /// Number of batch kernels.
21    pub batch_kernels: usize,
22    /// Number of ring kernels.
23    pub ring_kernels: usize,
24    /// Kernels by domain.
25    pub by_domain: HashMap<Domain, usize>,
26}
27
28/// Entry for a batch kernel in the registry.
29#[derive(Clone)]
30pub struct BatchKernelEntry {
31    /// Kernel metadata.
32    pub metadata: KernelMetadata,
33    /// Factory function to create the kernel.
34    factory: Arc<dyn Fn() -> Arc<dyn BatchKernelDyn> + Send + Sync>,
35}
36
37impl BatchKernelEntry {
38    /// Create a new batch kernel entry.
39    pub fn new<F>(metadata: KernelMetadata, factory: F) -> Self
40    where
41        F: Fn() -> Arc<dyn BatchKernelDyn> + Send + Sync + 'static,
42    {
43        Self {
44            metadata,
45            factory: Arc::new(factory),
46        }
47    }
48
49    /// Create an instance of the kernel.
50    #[must_use]
51    pub fn create(&self) -> Arc<dyn BatchKernelDyn> {
52        (self.factory)()
53    }
54}
55
56impl std::fmt::Debug for BatchKernelEntry {
57    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
58        f.debug_struct("BatchKernelEntry")
59            .field("metadata", &self.metadata)
60            .finish()
61    }
62}
63
64/// Entry for a ring kernel in the registry.
65#[derive(Clone)]
66pub struct RingKernelEntry {
67    /// Kernel metadata.
68    pub metadata: KernelMetadata,
69    /// Factory function to create the kernel.
70    factory: Arc<dyn Fn() -> Arc<dyn RingKernelDyn> + Send + Sync>,
71}
72
73impl RingKernelEntry {
74    /// Create a new ring kernel entry.
75    pub fn new<F>(metadata: KernelMetadata, factory: F) -> Self
76    where
77        F: Fn() -> Arc<dyn RingKernelDyn> + Send + Sync + 'static,
78    {
79        Self {
80            metadata,
81            factory: Arc::new(factory),
82        }
83    }
84
85    /// Create an instance of the kernel.
86    #[must_use]
87    pub fn create(&self) -> Arc<dyn RingKernelDyn> {
88        (self.factory)()
89    }
90}
91
92impl std::fmt::Debug for RingKernelEntry {
93    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
94        f.debug_struct("RingKernelEntry")
95            .field("metadata", &self.metadata)
96            .finish()
97    }
98}
99
100/// Central registry for all kernels.
101#[derive(Debug)]
102pub struct KernelRegistry {
103    /// Batch kernel entries by ID.
104    batch_kernels: RwLock<HashMap<String, BatchKernelEntry>>,
105    /// Ring kernel entries by ID.
106    ring_kernels: RwLock<HashMap<String, RingKernelEntry>>,
107    /// Metadata-only entries (for kernels without type-erased factories).
108    metadata_only: RwLock<HashMap<String, KernelMetadata>>,
109    /// License validator.
110    license: Option<SharedLicenseValidator>,
111}
112
113impl KernelRegistry {
114    /// Create a new empty registry.
115    #[must_use]
116    pub fn new() -> Self {
117        Self {
118            batch_kernels: RwLock::new(HashMap::new()),
119            ring_kernels: RwLock::new(HashMap::new()),
120            metadata_only: RwLock::new(HashMap::new()),
121            license: None,
122        }
123    }
124
125    /// Create a registry with license validation.
126    #[must_use]
127    pub fn with_license(license: SharedLicenseValidator) -> Self {
128        Self {
129            batch_kernels: RwLock::new(HashMap::new()),
130            ring_kernels: RwLock::new(HashMap::new()),
131            metadata_only: RwLock::new(HashMap::new()),
132            license: Some(license),
133        }
134    }
135
136    /// Set the license validator.
137    pub fn set_license(&mut self, license: SharedLicenseValidator) {
138        self.license = Some(license);
139    }
140
141    /// Register a batch kernel.
142    pub fn register_batch(&self, entry: BatchKernelEntry) -> Result<()> {
143        // Validate license if present
144        if let Some(ref license) = self.license {
145            self.validate_kernel_license(license.as_ref(), &entry.metadata)?;
146        }
147
148        let id = entry.metadata.id.clone();
149        let mut kernels = self.batch_kernels.write().unwrap();
150
151        if kernels.contains_key(&id) {
152            return Err(KernelError::KernelAlreadyRegistered(id));
153        }
154
155        debug!(kernel_id = %id, domain = %entry.metadata.domain, "Registering batch kernel");
156        kernels.insert(id, entry);
157        Ok(())
158    }
159
160    /// Register a ring kernel.
161    pub fn register_ring(&self, entry: RingKernelEntry) -> Result<()> {
162        // Validate license if present
163        if let Some(ref license) = self.license {
164            self.validate_kernel_license(license.as_ref(), &entry.metadata)?;
165        }
166
167        let id = entry.metadata.id.clone();
168        let mut kernels = self.ring_kernels.write().unwrap();
169
170        if kernels.contains_key(&id) {
171            return Err(KernelError::KernelAlreadyRegistered(id));
172        }
173
174        debug!(kernel_id = %id, domain = %entry.metadata.domain, "Registering ring kernel");
175        kernels.insert(id, entry);
176        Ok(())
177    }
178
179    /// Register kernel metadata without a factory.
180    ///
181    /// This is useful for kernels that implement `GpuKernel` but not the type-erased
182    /// `BatchKernelDyn` or `RingKernelDyn` traits. The metadata is stored for discovery
183    /// purposes, but the kernel cannot be instantiated through the registry.
184    pub fn register_metadata(&self, metadata: KernelMetadata) -> Result<()> {
185        // Validate license if present
186        if let Some(ref license) = self.license {
187            self.validate_kernel_license(license.as_ref(), &metadata)?;
188        }
189
190        let id = metadata.id.clone();
191
192        // Check if already registered in any category
193        {
194            let batch = self.batch_kernels.read().unwrap();
195            if batch.contains_key(&id) {
196                return Err(KernelError::KernelAlreadyRegistered(id));
197            }
198        }
199        {
200            let ring = self.ring_kernels.read().unwrap();
201            if ring.contains_key(&id) {
202                return Err(KernelError::KernelAlreadyRegistered(id));
203            }
204        }
205
206        let mut metadata_map = self.metadata_only.write().unwrap();
207        if metadata_map.contains_key(&id) {
208            return Err(KernelError::KernelAlreadyRegistered(id));
209        }
210
211        debug!(kernel_id = %id, domain = %metadata.domain, mode = ?metadata.mode, "Registering kernel metadata");
212        metadata_map.insert(id, metadata);
213        Ok(())
214    }
215
216    /// Validate kernel license.
217    fn validate_kernel_license(
218        &self,
219        license: &dyn LicenseValidator,
220        metadata: &KernelMetadata,
221    ) -> Result<()> {
222        // Validate domain
223        license
224            .validate_domain(metadata.domain)
225            .map_err(KernelError::from)?;
226
227        // Validate GPU-native requirement
228        if metadata.requires_gpu_native && !license.gpu_native_enabled() {
229            return Err(KernelError::from(LicenseError::GpuNativeNotLicensed));
230        }
231
232        Ok(())
233    }
234
235    /// Get a batch kernel by ID.
236    #[must_use]
237    pub fn get_batch(&self, id: &str) -> Option<BatchKernelEntry> {
238        let kernels = self.batch_kernels.read().unwrap();
239        kernels.get(id).cloned()
240    }
241
242    /// Get a ring kernel by ID.
243    #[must_use]
244    pub fn get_ring(&self, id: &str) -> Option<RingKernelEntry> {
245        let kernels = self.ring_kernels.read().unwrap();
246        kernels.get(id).cloned()
247    }
248
249    /// Get any kernel metadata by ID (batch, ring, or metadata-only).
250    #[must_use]
251    pub fn get(&self, id: &str) -> Option<KernelMetadata> {
252        if let Some(entry) = self.get_batch(id) {
253            return Some(entry.metadata);
254        }
255        if let Some(entry) = self.get_ring(id) {
256            return Some(entry.metadata);
257        }
258        let metadata_map = self.metadata_only.read().unwrap();
259        metadata_map.get(id).cloned()
260    }
261
262    /// Get metadata-only entry by ID.
263    #[must_use]
264    pub fn get_metadata_only(&self, id: &str) -> Option<KernelMetadata> {
265        let metadata_map = self.metadata_only.read().unwrap();
266        metadata_map.get(id).cloned()
267    }
268
269    /// Check if a kernel exists.
270    #[must_use]
271    pub fn contains(&self, id: &str) -> bool {
272        let batch = self.batch_kernels.read().unwrap();
273        let ring = self.ring_kernels.read().unwrap();
274        let metadata = self.metadata_only.read().unwrap();
275        batch.contains_key(id) || ring.contains_key(id) || metadata.contains_key(id)
276    }
277
278    /// Get all batch kernel IDs.
279    #[must_use]
280    pub fn batch_kernel_ids(&self) -> Vec<String> {
281        let kernels = self.batch_kernels.read().unwrap();
282        kernels.keys().cloned().collect()
283    }
284
285    /// Get all ring kernel IDs.
286    #[must_use]
287    pub fn ring_kernel_ids(&self) -> Vec<String> {
288        let kernels = self.ring_kernels.read().unwrap();
289        kernels.keys().cloned().collect()
290    }
291
292    /// Get all metadata-only kernel IDs.
293    #[must_use]
294    pub fn metadata_only_ids(&self) -> Vec<String> {
295        let metadata = self.metadata_only.read().unwrap();
296        metadata.keys().cloned().collect()
297    }
298
299    /// Get all kernel IDs.
300    #[must_use]
301    pub fn all_kernel_ids(&self) -> Vec<String> {
302        let mut ids = self.batch_kernel_ids();
303        ids.extend(self.ring_kernel_ids());
304        ids.extend(self.metadata_only_ids());
305        ids
306    }
307
308    /// Get kernels by domain.
309    #[must_use]
310    pub fn by_domain(&self, domain: Domain) -> Vec<KernelMetadata> {
311        let mut result = Vec::new();
312
313        let batch = self.batch_kernels.read().unwrap();
314        for entry in batch.values() {
315            if entry.metadata.domain == domain {
316                result.push(entry.metadata.clone());
317            }
318        }
319
320        let ring = self.ring_kernels.read().unwrap();
321        for entry in ring.values() {
322            if entry.metadata.domain == domain {
323                result.push(entry.metadata.clone());
324            }
325        }
326
327        let metadata = self.metadata_only.read().unwrap();
328        for entry in metadata.values() {
329            if entry.domain == domain {
330                result.push(entry.clone());
331            }
332        }
333
334        result
335    }
336
337    /// Get kernels by mode.
338    #[must_use]
339    pub fn by_mode(&self, mode: KernelMode) -> Vec<KernelMetadata> {
340        let mut result: Vec<KernelMetadata> = match mode {
341            KernelMode::Batch => {
342                let kernels = self.batch_kernels.read().unwrap();
343                kernels.values().map(|e| e.metadata.clone()).collect()
344            }
345            KernelMode::Ring => {
346                let kernels = self.ring_kernels.read().unwrap();
347                kernels.values().map(|e| e.metadata.clone()).collect()
348            }
349        };
350
351        // Include metadata-only entries with matching mode
352        let metadata = self.metadata_only.read().unwrap();
353        for entry in metadata.values() {
354            if entry.mode == mode {
355                result.push(entry.clone());
356            }
357        }
358
359        result
360    }
361
362    /// Get registry statistics.
363    #[must_use]
364    pub fn stats(&self) -> RegistryStats {
365        let batch = self.batch_kernels.read().unwrap();
366        let ring = self.ring_kernels.read().unwrap();
367        let metadata = self.metadata_only.read().unwrap();
368
369        let mut by_domain: HashMap<Domain, usize> = HashMap::new();
370
371        for entry in batch.values() {
372            *by_domain.entry(entry.metadata.domain).or_default() += 1;
373        }
374
375        for entry in ring.values() {
376            *by_domain.entry(entry.metadata.domain).or_default() += 1;
377        }
378
379        // Count metadata-only entries by their mode for batch/ring stats
380        let mut metadata_batch = 0;
381        let mut metadata_ring = 0;
382        for entry in metadata.values() {
383            *by_domain.entry(entry.domain).or_default() += 1;
384            match entry.mode {
385                KernelMode::Batch => metadata_batch += 1,
386                KernelMode::Ring => metadata_ring += 1,
387            }
388        }
389
390        RegistryStats {
391            total: batch.len() + ring.len() + metadata.len(),
392            batch_kernels: batch.len() + metadata_batch,
393            ring_kernels: ring.len() + metadata_ring,
394            by_domain,
395        }
396    }
397
398    /// Register a batch kernel using a typed factory function.
399    ///
400    /// This is the preferred way to register batch kernels. The factory closure
401    /// creates kernel instances on demand, and type erasure is handled automatically
402    /// via [`TypeErasedBatchKernel`].
403    ///
404    /// # Type Inference
405    ///
406    /// Rust infers `I` and `O` from the `BatchKernel<I, O>` implementation on `K`,
407    /// so turbofish syntax is typically not needed:
408    ///
409    /// ```ignore
410    /// registry.register_batch_typed(|| MyKernel::new())?;
411    /// ```
412    ///
413    /// # Errors
414    ///
415    /// Returns an error if the kernel ID is already registered or fails license validation.
416    pub fn register_batch_typed<K, I, O>(
417        &self,
418        factory: impl Fn() -> K + Send + Sync + 'static,
419    ) -> Result<()>
420    where
421        K: BatchKernel<I, O> + 'static,
422        I: serde::de::DeserializeOwned + Send + Sync + 'static,
423        O: serde::Serialize + Send + Sync + 'static,
424    {
425        let sample = factory();
426        let metadata = sample.metadata().clone();
427        drop(sample);
428        let entry = BatchKernelEntry::new(metadata, move || {
429            Arc::new(TypeErasedBatchKernel::new(factory()))
430        });
431        self.register_batch(entry)
432    }
433
434    /// Register a batch kernel's metadata from a factory function.
435    ///
436    /// This is for batch-mode kernels that implement `GpuKernel` but not
437    /// the full `BatchKernel<I, O>` trait. The metadata is stored for
438    /// discovery and health checking.
439    ///
440    /// # Errors
441    ///
442    /// Returns an error if the kernel ID is already registered or fails license validation.
443    pub fn register_batch_metadata_from<K>(&self, factory: impl Fn() -> K) -> Result<()>
444    where
445        K: GpuKernel,
446    {
447        let sample = factory();
448        let metadata = sample.metadata().clone();
449        self.register_metadata(metadata)
450    }
451
452    /// Register a ring kernel's metadata from a factory function.
453    ///
454    /// Ring kernels require the RingKernel runtime for persistent actor execution
455    /// and cannot be invoked directly via REST. This method registers the kernel's
456    /// metadata for discovery and health checking.
457    ///
458    /// For full ring kernel deployment, use the RingKernel runtime directly.
459    ///
460    /// # Errors
461    ///
462    /// Returns an error if the kernel ID is already registered or fails license validation.
463    pub fn register_ring_metadata_from<K>(&self, factory: impl Fn() -> K) -> Result<()>
464    where
465        K: GpuKernel,
466    {
467        let sample = factory();
468        let metadata = sample.metadata().clone();
469        self.register_metadata(metadata)
470    }
471
472    /// Execute a batch kernel by ID with JSON input/output.
473    ///
474    /// Looks up the kernel in the registry, creates an instance, and executes it
475    /// with type-erased JSON serialization.
476    ///
477    /// # Errors
478    ///
479    /// Returns `KernelNotFound` if no batch kernel with this ID exists, or
480    /// propagates any execution error from the kernel.
481    pub async fn execute_batch(&self, kernel_id: &str, input_json: &[u8]) -> Result<Vec<u8>> {
482        let entry = self
483            .get_batch(kernel_id)
484            .ok_or_else(|| KernelError::KernelNotFound(kernel_id.to_string()))?;
485        let kernel = entry.create();
486        kernel.execute_dyn(input_json).await
487    }
488
489    /// Total number of registered kernels.
490    #[must_use]
491    pub fn total_count(&self) -> usize {
492        let batch = self.batch_kernels.read().unwrap();
493        let ring = self.ring_kernels.read().unwrap();
494        let metadata = self.metadata_only.read().unwrap();
495        batch.len() + ring.len() + metadata.len()
496    }
497
498    /// Get all kernel metadata across all categories, sorted by ID.
499    ///
500    /// Returns metadata for batch, ring, and metadata-only kernels.
501    #[must_use]
502    pub fn all_metadata(&self) -> Vec<KernelMetadata> {
503        let mut result = Vec::new();
504
505        let batch = self.batch_kernels.read().unwrap();
506        for entry in batch.values() {
507            result.push(entry.metadata.clone());
508        }
509
510        let ring = self.ring_kernels.read().unwrap();
511        for entry in ring.values() {
512            result.push(entry.metadata.clone());
513        }
514
515        let metadata = self.metadata_only.read().unwrap();
516        for entry in metadata.values() {
517            result.push(entry.clone());
518        }
519
520        result.sort_by(|a, b| a.id.cmp(&b.id));
521        result
522    }
523
524    /// Search kernels by pattern (case-insensitive substring match on ID and description).
525    #[must_use]
526    pub fn search(&self, pattern: &str) -> Vec<KernelMetadata> {
527        let pattern_lower = pattern.to_lowercase();
528        self.all_metadata()
529            .into_iter()
530            .filter(|m| {
531                m.id.to_lowercase().contains(&pattern_lower)
532                    || m.description.to_lowercase().contains(&pattern_lower)
533            })
534            .collect()
535    }
536
537    /// Get all executable batch kernel IDs (kernels with factory functions).
538    ///
539    /// These are the kernels that can be invoked via REST/gRPC through the
540    /// type-erased `BatchKernelDyn` interface.
541    #[must_use]
542    pub fn executable_kernel_ids(&self) -> Vec<String> {
543        self.batch_kernel_ids()
544    }
545
546    /// Check if a kernel is executable via REST/gRPC (has a `BatchKernelDyn` factory).
547    #[must_use]
548    pub fn is_executable(&self, id: &str) -> bool {
549        let batch = self.batch_kernels.read().unwrap();
550        batch.contains_key(id)
551    }
552
553    /// Clear all registered kernels.
554    pub fn clear(&self) {
555        let mut batch = self.batch_kernels.write().unwrap();
556        let mut ring = self.ring_kernels.write().unwrap();
557        let mut metadata = self.metadata_only.write().unwrap();
558        batch.clear();
559        ring.clear();
560        metadata.clear();
561        info!("Cleared kernel registry");
562    }
563
564    /// Unregister a kernel by ID.
565    pub fn unregister(&self, id: &str) -> bool {
566        let mut batch = self.batch_kernels.write().unwrap();
567        if batch.remove(id).is_some() {
568            debug!(kernel_id = %id, "Unregistered batch kernel");
569            return true;
570        }
571
572        let mut ring = self.ring_kernels.write().unwrap();
573        if ring.remove(id).is_some() {
574            debug!(kernel_id = %id, "Unregistered ring kernel");
575            return true;
576        }
577
578        let mut metadata = self.metadata_only.write().unwrap();
579        if metadata.remove(id).is_some() {
580            debug!(kernel_id = %id, "Unregistered metadata-only kernel");
581            return true;
582        }
583
584        warn!(kernel_id = %id, "Attempted to unregister non-existent kernel");
585        false
586    }
587}
588
589impl Default for KernelRegistry {
590    fn default() -> Self {
591        Self::new()
592    }
593}
594
595/// Builder for kernel registry.
596#[derive(Default)]
597pub struct KernelRegistryBuilder {
598    license: Option<SharedLicenseValidator>,
599    batch_entries: Vec<BatchKernelEntry>,
600    ring_entries: Vec<RingKernelEntry>,
601}
602
603impl KernelRegistryBuilder {
604    /// Create a new builder.
605    #[must_use]
606    pub fn new() -> Self {
607        Self::default()
608    }
609
610    /// Set the license validator.
611    #[must_use]
612    pub fn with_license(mut self, license: SharedLicenseValidator) -> Self {
613        self.license = Some(license);
614        self
615    }
616
617    /// Add a batch kernel.
618    #[must_use]
619    pub fn with_batch(mut self, entry: BatchKernelEntry) -> Self {
620        self.batch_entries.push(entry);
621        self
622    }
623
624    /// Add a ring kernel.
625    #[must_use]
626    pub fn with_ring(mut self, entry: RingKernelEntry) -> Self {
627        self.ring_entries.push(entry);
628        self
629    }
630
631    /// Build the registry.
632    ///
633    /// # Errors
634    ///
635    /// Returns an error if any kernel fails license validation.
636    pub fn build(self) -> Result<KernelRegistry> {
637        let registry = match self.license {
638            Some(license) => KernelRegistry::with_license(license),
639            None => KernelRegistry::new(),
640        };
641
642        for entry in self.batch_entries {
643            registry.register_batch(entry)?;
644        }
645
646        for entry in self.ring_entries {
647            registry.register_ring(entry)?;
648        }
649
650        info!(
651            total = registry.total_count(),
652            batch = registry.batch_kernel_ids().len(),
653            ring = registry.ring_kernel_ids().len(),
654            "Built kernel registry"
655        );
656
657        Ok(registry)
658    }
659}
660
661/// Global kernel registry for auto-discovery.
662///
663/// This is used by the `#[gpu_kernel]` proc macro to automatically register kernels.
664static GLOBAL_REGISTRY: std::sync::OnceLock<KernelRegistry> = std::sync::OnceLock::new();
665
666/// Get or initialize the global registry.
667pub fn global_registry() -> &'static KernelRegistry {
668    GLOBAL_REGISTRY.get_or_init(KernelRegistry::new)
669}
670
671/// Initialize the global registry with a license.
672///
673/// Must be called before any kernel registration.
674///
675/// # Panics
676///
677/// Panics if the global registry has already been initialized.
678pub fn init_global_registry(license: SharedLicenseValidator) -> &'static KernelRegistry {
679    let registry = KernelRegistry::with_license(license);
680    GLOBAL_REGISTRY
681        .set(registry)
682        .expect("Global registry already initialized");
683    GLOBAL_REGISTRY.get().unwrap()
684}
685
686#[cfg(test)]
687mod tests {
688    use super::*;
689    use crate::license::DevelopmentLicense;
690
691    fn test_batch_entry() -> BatchKernelEntry {
692        let metadata = KernelMetadata::batch("test-batch", Domain::Core);
693        BatchKernelEntry::new(metadata, || {
694            // Return a mock kernel
695            panic!("Not implemented for tests")
696        })
697    }
698
699    fn test_ring_entry() -> RingKernelEntry {
700        let metadata = KernelMetadata::ring("test-ring", Domain::Core);
701        RingKernelEntry::new(metadata, || {
702            // Return a mock kernel
703            panic!("Not implemented for tests")
704        })
705    }
706
707    #[test]
708    fn test_registry_creation() {
709        let registry = KernelRegistry::new();
710        assert_eq!(registry.total_count(), 0);
711    }
712
713    #[test]
714    fn test_batch_registration() {
715        let registry = KernelRegistry::new();
716        let entry = test_batch_entry();
717
718        registry.register_batch(entry).unwrap();
719        assert_eq!(registry.total_count(), 1);
720        assert!(registry.contains("test-batch"));
721        assert!(registry.get_batch("test-batch").is_some());
722    }
723
724    #[test]
725    fn test_ring_registration() {
726        let registry = KernelRegistry::new();
727        let entry = test_ring_entry();
728
729        registry.register_ring(entry).unwrap();
730        assert_eq!(registry.total_count(), 1);
731        assert!(registry.contains("test-ring"));
732        assert!(registry.get_ring("test-ring").is_some());
733    }
734
735    #[test]
736    fn test_duplicate_registration() {
737        let registry = KernelRegistry::new();
738        let entry1 = test_batch_entry();
739        let entry2 = test_batch_entry();
740
741        registry.register_batch(entry1).unwrap();
742        let result = registry.register_batch(entry2);
743        assert!(result.is_err());
744    }
745
746    #[test]
747    fn test_by_domain() {
748        let registry = KernelRegistry::new();
749
750        let core_entry = test_batch_entry();
751        registry.register_batch(core_entry).unwrap();
752
753        let graph_entry = BatchKernelEntry::new(
754            KernelMetadata::batch("test-graph", Domain::GraphAnalytics),
755            || panic!("Not implemented"),
756        );
757        registry.register_batch(graph_entry).unwrap();
758
759        let core_kernels = registry.by_domain(Domain::Core);
760        assert_eq!(core_kernels.len(), 1);
761
762        let graph_kernels = registry.by_domain(Domain::GraphAnalytics);
763        assert_eq!(graph_kernels.len(), 1);
764    }
765
766    #[test]
767    fn test_stats() {
768        let registry = KernelRegistry::new();
769
770        registry.register_batch(test_batch_entry()).unwrap();
771        registry.register_ring(test_ring_entry()).unwrap();
772
773        let stats = registry.stats();
774        assert_eq!(stats.total, 2);
775        assert_eq!(stats.batch_kernels, 1);
776        assert_eq!(stats.ring_kernels, 1);
777        assert_eq!(stats.by_domain.get(&Domain::Core), Some(&2));
778    }
779
780    #[test]
781    fn test_unregister() {
782        let registry = KernelRegistry::new();
783        registry.register_batch(test_batch_entry()).unwrap();
784
785        assert!(registry.contains("test-batch"));
786        assert!(registry.unregister("test-batch"));
787        assert!(!registry.contains("test-batch"));
788        assert!(!registry.unregister("test-batch"));
789    }
790
791    #[test]
792    fn test_with_license() {
793        let license: SharedLicenseValidator = Arc::new(DevelopmentLicense);
794        let registry = KernelRegistry::with_license(license);
795
796        // Should succeed with dev license
797        registry.register_batch(test_batch_entry()).unwrap();
798        registry.register_ring(test_ring_entry()).unwrap();
799    }
800
801    #[test]
802    fn test_builder() {
803        let registry = KernelRegistryBuilder::new()
804            .with_batch(test_batch_entry())
805            .with_ring(test_ring_entry())
806            .build()
807            .unwrap();
808
809        assert_eq!(registry.total_count(), 2);
810    }
811
812    #[test]
813    fn test_metadata_only_registration() {
814        let registry = KernelRegistry::new();
815        let metadata = KernelMetadata::batch("test-metadata", Domain::GraphAnalytics);
816
817        registry.register_metadata(metadata).unwrap();
818
819        assert_eq!(registry.total_count(), 1);
820        assert!(registry.contains("test-metadata"));
821        assert!(registry.get("test-metadata").is_some());
822        assert!(registry.get_metadata_only("test-metadata").is_some());
823        assert!(registry.get_batch("test-metadata").is_none()); // Not a batch entry
824    }
825
826    #[test]
827    fn test_metadata_only_duplicate() {
828        let registry = KernelRegistry::new();
829        let metadata1 = KernelMetadata::batch("test-dup", Domain::Core);
830        let metadata2 = KernelMetadata::batch("test-dup", Domain::Core);
831
832        registry.register_metadata(metadata1).unwrap();
833        let result = registry.register_metadata(metadata2);
834        assert!(result.is_err());
835    }
836
837    #[test]
838    fn test_metadata_only_conflict_with_batch() {
839        let registry = KernelRegistry::new();
840        registry.register_batch(test_batch_entry()).unwrap();
841
842        // Try to register metadata with same ID
843        let metadata = KernelMetadata::batch("test-batch", Domain::Core);
844        let result = registry.register_metadata(metadata);
845        assert!(result.is_err());
846    }
847
848    #[test]
849    fn test_metadata_only_in_stats() {
850        let registry = KernelRegistry::new();
851
852        let batch_meta = KernelMetadata::batch("meta-batch", Domain::GraphAnalytics);
853        let ring_meta = KernelMetadata::ring("meta-ring", Domain::GraphAnalytics);
854        registry.register_metadata(batch_meta).unwrap();
855        registry.register_metadata(ring_meta).unwrap();
856
857        let stats = registry.stats();
858        assert_eq!(stats.total, 2);
859        assert_eq!(stats.batch_kernels, 1);
860        assert_eq!(stats.ring_kernels, 1);
861        assert_eq!(stats.by_domain.get(&Domain::GraphAnalytics), Some(&2));
862    }
863
864    #[test]
865    fn test_metadata_only_by_domain() {
866        let registry = KernelRegistry::new();
867
868        let graph_meta = KernelMetadata::batch("graph-kernel", Domain::GraphAnalytics);
869        let ml_meta = KernelMetadata::batch("ml-kernel", Domain::StatisticalML);
870        registry.register_metadata(graph_meta).unwrap();
871        registry.register_metadata(ml_meta).unwrap();
872
873        let graph_kernels = registry.by_domain(Domain::GraphAnalytics);
874        assert_eq!(graph_kernels.len(), 1);
875        assert_eq!(graph_kernels[0].id, "graph-kernel");
876
877        let ml_kernels = registry.by_domain(Domain::StatisticalML);
878        assert_eq!(ml_kernels.len(), 1);
879        assert_eq!(ml_kernels[0].id, "ml-kernel");
880    }
881
882    #[test]
883    fn test_unregister_metadata_only() {
884        let registry = KernelRegistry::new();
885        let metadata = KernelMetadata::batch("test-unreg", Domain::Core);
886
887        registry.register_metadata(metadata).unwrap();
888        assert!(registry.contains("test-unreg"));
889
890        assert!(registry.unregister("test-unreg"));
891        assert!(!registry.contains("test-unreg"));
892    }
893}
rustkernel_core/registry.rs

rustkernel_core/
registry.rs
