1use 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#[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 pub const fn new() -> Self {
34 Self { precompiles: BTreeMap::new() }
35 }
36
37 pub fn is_empty(&self) -> bool {
39 self.precompiles.is_empty()
40 }
41
42 pub fn with_precompile<P: Precompile + 'static>(mut self, precompile: P) -> Self {
46 self.insert_precompile(Arc::new(precompile));
47 self
48 }
49
50 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 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 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 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 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 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 #[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 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 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(®istry), usize::MAX).unwrap();
334 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}