Skip to main content

miden_core/deferred/
precompile_registry.rs

1//! Registry that routes deferred tags to their owning precompile.
2
3use alloc::{collections::BTreeMap, sync::Arc, vec::Vec};
4
5use super::precompile::{Precompile, precompile_id};
6use crate::{
7    Felt,
8    deferred::{DeferredContext, Node, NodeType, PrecompileError, Tag},
9};
10
11/// Installed set of precompiles for deferred-node validation and evaluation.
12///
13/// Routing is entirely id-based. The empty registry is valid but rejects every precompile-owned
14/// tag, which is useful for programs that do not use precompile-backed deferred nodes.
15#[derive(Clone, Default)]
16pub struct PrecompileRegistry {
17    precompiles: BTreeMap<Felt, Arc<dyn Precompile>>,
18}
19
20impl core::fmt::Debug for PrecompileRegistry {
21    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
22        f.debug_struct("PrecompileRegistry")
23            .field(
24                "precompiles",
25                &self.precompiles.iter().map(|(id, p)| (id, p.name())).collect::<Vec<_>>(),
26            )
27            .finish()
28    }
29}
30
31impl PrecompileRegistry {
32    /// Creates an empty precompile registry.
33    pub const fn new() -> Self {
34        Self { precompiles: BTreeMap::new() }
35    }
36
37    /// Returns whether this registry contains no installed precompiles.
38    pub fn is_empty(&self) -> bool {
39        self.precompiles.is_empty()
40    }
41
42    /// Adds a precompile to the registry and returns `self` for chaining.
43    ///
44    /// Panics on setup errors: id drift, a framework-reserved id, or a duplicate id.
45    pub fn with_precompile<P: Precompile + 'static>(mut self, precompile: P) -> Self {
46        self.insert_precompile(Arc::new(precompile));
47        self
48    }
49
50    /// Merges another registry into this one.
51    ///
52    /// Panics on duplicate ids, preserving [`Self::with_precompile`]'s setup-failure behavior.
53    pub fn merge(&mut self, registry: Self) -> &mut Self {
54        for precompile in registry.precompiles.into_values() {
55            self.insert_precompile(precompile);
56        }
57        self
58    }
59
60    fn insert_precompile(&mut self, precompile: Arc<dyn Precompile>) {
61        let id = precompile.id();
62        validate_precompile_id(precompile.name(), id, precompile_id(precompile.name()));
63        let name = precompile.name();
64        if let Some(prev) = self.precompiles.get(&id) {
65            panic!("duplicate precompile id in registry (`{}` and `{name}`)", prev.name());
66        }
67        self.precompiles.insert(id, precompile);
68    }
69
70    /// Returns all precompile initialization nodes in deterministic registry id order.
71    ///
72    /// [`DeferredState`](super::DeferredState) loads the full returned set before evaluating each
73    /// init node, so init nodes may depend on TRUE or on any node in the complete init set. Within
74    /// one precompile, nodes retain the order returned by
75    /// [`Precompile::init`](super::Precompile::init).
76    pub(crate) fn init_nodes(&self) -> Vec<Node> {
77        let mut nodes = Vec::new();
78        for precompile in self.precompiles.values() {
79            nodes.extend(precompile.init());
80        }
81        nodes
82    }
83
84    /// Decodes a precompile-owned tag by routing its local arguments to the owning precompile.
85    ///
86    /// Unknown ids are registry failures; recognized ids whose arguments are invalid are
87    /// attributed to the owning precompile. Framework tags are handled by the internal
88    /// framework-aware decoder and rejected here. [`NodeType::True`] is reserved for the framework
89    /// TRUE sentinel, so a precompile that returns it is rejected as an invalid node.
90    pub fn decode_precompile_tag(&self, tag: Tag) -> Result<NodeType, PrecompileError> {
91        if tag.is_framework_reserved() {
92            return Err(PrecompileError::InvalidNode);
93        }
94        let precompile = self.precompiles.get(&tag.id()).ok_or(PrecompileError::InvalidNode)?;
95        let invalid =
96            || PrecompileError::with_precompile(precompile.name(), PrecompileError::InvalidNode);
97        match precompile.decode(tag.args()).ok_or_else(invalid)? {
98            NodeType::True => Err(invalid()),
99            node_type => Ok(node_type),
100        }
101    }
102
103    /// Decodes either a framework-owned tag or a precompile-owned tag.
104    pub(crate) fn decode_node_type(&self, tag: Tag) -> Result<NodeType, PrecompileError> {
105        if tag == Tag::TRUE {
106            Ok(NodeType::True)
107        } else if tag == Tag::AND {
108            Ok(NodeType::Join)
109        } else if tag == Tag::CHUNKS {
110            Ok(NodeType::Data)
111        } else {
112            self.decode_precompile_tag(tag)
113        }
114    }
115
116    /// Validates a node's tag and payload shape under this registry.
117    pub(crate) fn validate_node(&self, node: &Node) -> Result<NodeType, PrecompileError> {
118        let node_type = self.decode_node_type(node.tag())?;
119        node_type.validate_node(node).map_err(|_| PrecompileError::InvalidNode)?;
120        Ok(node_type)
121    }
122
123    /// Evaluates a node through the precompile selected by its tag id.
124    ///
125    /// Failures are wrapped with the owning precompile's name so callers can distinguish routing
126    /// from precompile-local validation.
127    pub(crate) fn evaluate(
128        &self,
129        node: &Node,
130        context: &mut DeferredContext<'_>,
131    ) -> Result<Node, PrecompileError> {
132        let tag = node.tag();
133        if tag.is_framework_reserved() {
134            return Err(PrecompileError::InvalidNode);
135        }
136        let precompile = self.precompiles.get(&tag.id()).ok_or(PrecompileError::InvalidNode)?;
137        precompile
138            .evaluate(tag.args(), node.payload(), context)
139            .map_err(|source| PrecompileError::with_precompile(precompile.name(), source))
140    }
141}
142
143fn validate_precompile_id(name: &'static str, id: Felt, derived: Felt) {
144    assert!(
145        id == derived,
146        "precompile `{name}` declares an id inconsistent with its name derivation"
147    );
148    assert!(
149        !Tag::is_framework_reserved_id(id),
150        "precompile `{name}` derives a framework-reserved id"
151    );
152}
153
154#[cfg(test)]
155mod tests {
156
157    use super::*;
158    use crate::{
159        ONE, ZERO,
160        deferred::{DeferredState, Payload},
161    };
162
163    /// Minimal honest precompile fixture for registry-routing tests.
164    ///
165    /// Names control ids, so duplicate names exercise duplicate-id handling. Non-zero arguments
166    /// are rejected by the fixture, not by the framework.
167    #[derive(Debug, Clone, Copy)]
168    struct Fixture {
169        name: &'static str,
170    }
171
172    impl Fixture {
173        fn new(name: &'static str) -> Self {
174            Self { name }
175        }
176        fn tag(&self) -> Tag {
177            Tag::precompile(self.id(), [ZERO; 3]).expect("fixture id is precompile-owned")
178        }
179    }
180
181    impl Precompile for Fixture {
182        fn name(&self) -> &'static str {
183            self.name
184        }
185        fn id(&self) -> Felt {
186            precompile_id(self.name())
187        }
188        fn decode(&self, args: [Felt; 3]) -> Option<NodeType> {
189            if args != [ZERO; 3] {
190                return None;
191            }
192            Some(NodeType::Data)
193        }
194        fn evaluate(
195            &self,
196            args: [Felt; 3],
197            payload: &Payload,
198            _context: &mut DeferredContext<'_>,
199        ) -> Result<Node, PrecompileError> {
200            let chunk = payload.as_value()?;
201            Ok(Node::value(
202                Tag::precompile(self.id(), args).expect("fixture id is precompile-owned"),
203                *chunk,
204            )?)
205        }
206    }
207
208    #[derive(Debug, Clone, Copy)]
209    struct MaliciousTrue;
210
211    impl Precompile for MaliciousTrue {
212        fn name(&self) -> &'static str {
213            "malicious-true"
214        }
215        fn id(&self) -> Felt {
216            precompile_id(self.name())
217        }
218        fn decode(&self, _args: [Felt; 3]) -> Option<NodeType> {
219            Some(NodeType::True)
220        }
221        fn evaluate(
222            &self,
223            _args: [Felt; 3],
224            _payload: &Payload,
225            _context: &mut DeferredContext<'_>,
226        ) -> Result<Node, PrecompileError> {
227            unreachable!("registry must reject precompile-owned NodeType::True")
228        }
229    }
230
231    #[test]
232    fn dispatches_by_id_across_inserted_and_merged_registries() {
233        let a = Fixture::new("fixture-a");
234        let b = Fixture::new("fixture-b");
235        let tag_a = a.tag();
236        let tag_b = b.tag();
237        let mut registry = PrecompileRegistry::default().with_precompile(a);
238        registry.merge(PrecompileRegistry::default().with_precompile(b));
239
240        assert_eq!(registry.decode_precompile_tag(tag_a).unwrap(), NodeType::Data);
241        assert_eq!(registry.decode_precompile_tag(tag_b).unwrap(), NodeType::Data);
242    }
243
244    #[test]
245    fn registry_decodes_exact_framework_chunks_as_data_only() {
246        let registry = PrecompileRegistry::new();
247        assert_eq!(registry.decode_node_type(Tag::CHUNKS).unwrap(), NodeType::Data);
248        assert!(matches!(
249            registry.decode_precompile_tag(Tag::CHUNKS),
250            Err(PrecompileError::InvalidNode)
251        ));
252
253        let malformed = Tag::from_word([Tag::CHUNKS.id(), ONE, ZERO, ZERO]);
254        assert!(matches!(
255            registry.decode_node_type(malformed),
256            Err(PrecompileError::InvalidNode)
257        ));
258    }
259
260    #[test]
261    fn registry_rejects_precompile_owned_true_shape() {
262        let registry = PrecompileRegistry::default().with_precompile(MaliciousTrue);
263        let tag =
264            Tag::precompile(MaliciousTrue.id(), [ZERO; 3]).expect("test id is precompile-owned");
265        assert!(matches!(
266            registry.decode_precompile_tag(tag),
267            Err(PrecompileError::Precompile { .. })
268        ));
269        assert!(matches!(
270            registry.decode_precompile_tag(tag).unwrap_err().root(),
271            PrecompileError::InvalidNode
272        ));
273    }
274
275    #[test]
276    fn unknown_id_rejected() {
277        let registry = PrecompileRegistry::default().with_precompile(Fixture::new("known"));
278        let bogus = Tag::precompile(Felt::new_unchecked(9999), [ZERO; 3])
279            .expect("bogus id is not framework-reserved");
280        // Unknown id is rejected by the registry itself (not a precompile), so it is *not*
281        // name-wrapped.
282        assert!(matches!(
283            registry.decode_precompile_tag(bogus),
284            Err(PrecompileError::InvalidNode)
285        ));
286    }
287
288    #[test]
289    fn fixture_rejects_nonzero_immediate() {
290        let f = Fixture::new("f");
291        let tag = Tag::precompile(f.id(), [ZERO, ZERO, Felt::new_unchecked(1)])
292            .expect("fixture id is precompile-owned");
293        let registry = PrecompileRegistry::default().with_precompile(f);
294        // The fixture chose to reject the immediate, so the registry name-wraps the cause.
295        assert!(matches!(
296            registry.decode_precompile_tag(tag).unwrap_err().root(),
297            PrecompileError::InvalidNode
298        ));
299    }
300
301    #[test]
302    #[should_panic(expected = "framework-reserved id")]
303    fn true_id_is_reserved_for_framework() {
304        validate_precompile_id("reserved-true", Tag::TRUE.id(), Tag::TRUE.id());
305    }
306
307    #[test]
308    #[should_panic(expected = "framework-reserved id")]
309    fn and_id_is_reserved_for_framework() {
310        validate_precompile_id("reserved-and", Tag::AND.id(), Tag::AND.id());
311    }
312
313    #[test]
314    #[should_panic(expected = "framework-reserved id")]
315    fn chunks_id_is_reserved_for_framework() {
316        validate_precompile_id("reserved-chunks", Tag::CHUNKS.id(), Tag::CHUNKS.id());
317    }
318
319    #[test]
320    #[should_panic(expected = "duplicate precompile id in registry")]
321    fn duplicate_id_panics() {
322        let _ = PrecompileRegistry::default()
323            .with_precompile(Fixture::new("dup"))
324            .with_precompile(Fixture::new("dup"));
325    }
326
327    #[test]
328    fn evaluate_dispatches_to_owning_precompile() {
329        let f = Fixture::new("r");
330        let tag = f.tag();
331        let registry = Arc::new(PrecompileRegistry::default().with_precompile(f));
332        let node = Node::value(tag, [ZERO; 8]).unwrap();
333        let mut state = DeferredState::new(Arc::clone(&registry), usize::MAX).unwrap();
334        // Use the framework's evaluation path so we exercise dispatch end-to-end.
335        let digest = state.register(node.clone()).unwrap();
336        let (canonical_digest, canonical_node) = state.require_canonical_node(digest).unwrap();
337        assert_eq!(canonical_digest, node.digest());
338        assert_eq!(canonical_node, &node);
339    }
340}