nemo_flow/
registry.rs

1// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Priority-sorted named registry.
5//!
6//! [`SortedRegistry`] is the backbone data structure for all guardrail and intercept
7//! registries in the NeMo Flow runtime. It stores entries by unique name and provides
8//! iteration in ascending priority order, with eager re-sorting on every mutation.
9
10use std::collections::HashMap;
11
12/// A named registry that maintains a sorted order by priority.
13///
14/// Items are stored by unique string name and sorted by an integer priority
15/// extracted via a caller-provided function. The sort is performed eagerly:
16/// every [`register`](SortedRegistry::register) or
17/// [`deregister`](SortedRegistry::deregister) call re-sorts immediately, so
18/// [`sorted_values`](SortedRegistry::sorted_values) is a read-only lookup.
19///
20/// # Priority ordering
21///
22/// Entries are sorted in **ascending** priority order (lower numbers run first).
23/// This means a guardrail with priority `1` executes before one with priority `10`.
24///
25/// # Uniqueness
26///
27/// Names must be unique within a registry. Attempting to [`register`](SortedRegistry::register)
28/// a duplicate name returns an error. Use [`deregister`](SortedRegistry::deregister) first
29/// to remove an existing entry before re-registering.
30pub struct SortedRegistry<T> {
31    entries: HashMap<String, T>,
32    sorted_keys: Vec<String>,
33    priority_fn: fn(&T) -> i32,
34}
35
36impl<T> SortedRegistry<T> {
37    /// Create a new empty registry with the given priority extraction function.
38    ///
39    /// The runtime calls `priority_fn` on each stored entry to determine its
40    /// sort key. Lower values are ordered first.
41    ///
42    /// # Parameters
43    /// - `priority_fn`: Function used to extract the integer priority from a
44    ///   stored entry.
45    ///
46    /// # Returns
47    /// A new empty [`SortedRegistry`] with no entries.
48    pub fn new(priority_fn: fn(&T) -> i32) -> Self {
49        Self {
50            entries: HashMap::new(),
51            sorted_keys: Vec::new(),
52            priority_fn,
53        }
54    }
55
56    /// Re-sorts the cached key order by priority. Called eagerly on every mutation.
57    fn resort(&mut self) {
58        let pf = self.priority_fn;
59        let entries = &self.entries;
60        let mut keys: Vec<String> = entries.keys().cloned().collect();
61        keys.sort_by_key(|k| pf(entries.get(k).unwrap()));
62        self.sorted_keys = keys;
63    }
64
65    /// Register a new entry under a unique name.
66    ///
67    /// # Parameters
68    /// - `name`: Unique name used to address the entry later.
69    /// - `entry`: Value to store in the registry.
70    ///
71    /// # Returns
72    /// `Ok(())` when the entry was inserted.
73    ///
74    /// # Errors
75    /// Returns `Err(String)` when `name` is already present in the registry.
76    ///
77    /// # Notes
78    /// Successful registration eagerly re-sorts the cached priority order.
79    pub fn register(&mut self, name: String, entry: T) -> Result<(), String> {
80        if self.entries.contains_key(&name) {
81            return Err(format!("{name} already exists"));
82        }
83        self.entries.insert(name, entry);
84        self.resort();
85        Ok(())
86    }
87
88    /// Deregister an entry by name.
89    ///
90    /// # Parameters
91    /// - `name`: Name of the entry to remove.
92    ///
93    /// # Returns
94    /// `true` when an entry was removed and `false` when `name` was not
95    /// present.
96    ///
97    /// # Notes
98    /// Successful removal eagerly re-sorts the cached priority order.
99    pub fn deregister(&mut self, name: &str) -> bool {
100        if self.entries.remove(name).is_some() {
101            self.resort();
102            true
103        } else {
104            false
105        }
106    }
107
108    /// Return entries sorted by priority (ascending).
109    ///
110    /// This is a read-only operation — the sort order is maintained eagerly
111    /// on every [`register`](SortedRegistry::register) / [`deregister`](SortedRegistry::deregister) call.
112    ///
113    /// # Returns
114    /// A newly allocated [`Vec`] of shared references ordered from lowest
115    /// priority to highest priority.
116    pub fn sorted_values(&self) -> Vec<&T> {
117        self.sorted_keys
118            .iter()
119            .filter_map(|k| self.entries.get(k))
120            .collect()
121    }
122
123    /// Return a shared reference to an entry by name.
124    ///
125    /// # Parameters
126    /// - `name`: Name of the entry to resolve.
127    ///
128    /// # Returns
129    /// `Some(&T)` when an entry exists under `name`, otherwise `None`.
130    pub fn get(&self, name: &str) -> Option<&T> {
131        self.entries.get(name)
132    }
133
134    /// Remove and return an entry by name.
135    ///
136    /// # Parameters
137    /// - `name`: Name of the entry to remove.
138    ///
139    /// # Returns
140    /// `Some(T)` when an entry was removed, otherwise `None`.
141    ///
142    /// # Notes
143    /// Successful removal eagerly re-sorts the cached priority order.
144    pub fn remove(&mut self, name: &str) -> Option<T> {
145        let removed = self.entries.remove(name);
146        if removed.is_some() {
147            self.resort();
148        }
149        removed
150    }
151
152    /// Report whether an entry with the given name exists.
153    ///
154    /// # Parameters
155    /// - `name`: Name to test for membership.
156    ///
157    /// # Returns
158    /// `true` when the registry contains `name`, otherwise `false`.
159    pub fn contains(&self, name: &str) -> bool {
160        self.entries.contains_key(name)
161    }
162}
163
164#[cfg(test)]
165#[path = "../tests/unit/registry_tests.rs"]
166mod tests;
nemo_flow/registry.rs

nemo_flow/
registry.rs
