cairo_lang_sierra/
program_registry.rs

1#[expect(clippy::disallowed_types)]
2use std::collections::HashMap;
3use std::collections::hash_map::Entry;
4
5use itertools::{chain, izip};
6use thiserror::Error;
7
8use crate::extensions::lib_func::{
9    SierraApChange, SignatureSpecializationContext, SpecializationContext,
10};
11use crate::extensions::type_specialization_context::TypeSpecializationContext;
12use crate::extensions::types::TypeInfo;
13use crate::extensions::{
14    ConcreteLibfunc, ConcreteType, ExtensionError, GenericLibfunc, GenericLibfuncEx, GenericType,
15    GenericTypeEx,
16};
17use crate::ids::{ConcreteLibfuncId, ConcreteTypeId, FunctionId, GenericTypeId};
18use crate::program::{
19    BranchTarget, DeclaredTypeInfo, Function, FunctionSignature, GenericArg, Program, Statement,
20    StatementIdx, TypeDeclaration,
21};
22
23#[cfg(test)]
24#[path = "program_registry_test.rs"]
25mod test;
26
27/// Errors encountered in the program registry.
28#[derive(Error, Debug, Eq, PartialEq)]
29pub enum ProgramRegistryError {
30    #[error("used the same function id twice `{0}`.")]
31    FunctionIdAlreadyExists(FunctionId),
32    #[error("Could not find the requested function `{0}`.")]
33    MissingFunction(FunctionId),
34    #[error("Error during type specialization of `{concrete_id}`: {error}")]
35    TypeSpecialization { concrete_id: ConcreteTypeId, error: ExtensionError },
36    #[error("Used concrete type id `{0}` twice")]
37    TypeConcreteIdAlreadyExists(ConcreteTypeId),
38    #[error("Declared concrete type `{0}` twice")]
39    TypeAlreadyDeclared(Box<TypeDeclaration>),
40    #[error("Could not find requested type `{0}`.")]
41    MissingType(ConcreteTypeId),
42    #[error("Error during libfunc specialization of {concrete_id}: {error}")]
43    LibfuncSpecialization { concrete_id: ConcreteLibfuncId, error: ExtensionError },
44    #[error("Used concrete libfunc id `{0}` twice.")]
45    LibfuncConcreteIdAlreadyExists(ConcreteLibfuncId),
46    #[error("Could not find requested libfunc `{0}`.")]
47    MissingLibfunc(ConcreteLibfuncId),
48    #[error("Type info declaration mismatch for `{0}`.")]
49    TypeInfoDeclarationMismatch(ConcreteTypeId),
50    #[error("Function `{func_id}`'s parameter type `{ty}` is not storable.")]
51    FunctionWithUnstorableType { func_id: FunctionId, ty: ConcreteTypeId },
52    #[error("Function `{0}` points to non existing entry point statement.")]
53    FunctionNonExistingEntryPoint(FunctionId),
54    #[error("#{0}: Libfunc invocation input count mismatch")]
55    LibfuncInvocationInputCountMismatch(StatementIdx),
56    #[error("#{0}: Libfunc invocation branch count mismatch")]
57    LibfuncInvocationBranchCountMismatch(StatementIdx),
58    #[error("#{0}: Libfunc invocation branch #{1} result count mismatch")]
59    LibfuncInvocationBranchResultCountMismatch(StatementIdx, usize),
60    #[error("#{0}: Libfunc invocation branch #{1} target mismatch")]
61    LibfuncInvocationBranchTargetMismatch(StatementIdx, usize),
62    #[error("#{src}: Branch jump backwards to {dst}")]
63    BranchBackwards { src: StatementIdx, dst: StatementIdx },
64    #[error("#{src}: Branch jump to a non-branch align statement #{dst}")]
65    BranchNotToBranchAlign { src: StatementIdx, dst: StatementIdx },
66    #[error("#{src1}, #{src2}: Jump to the same statement #{dst}")]
67    MultipleJumpsToSameStatement { src1: StatementIdx, src2: StatementIdx, dst: StatementIdx },
68    #[error("#{0}: Jump out of range")]
69    JumpOutOfRange(StatementIdx),
70    #[error("Type size computation failed for `{ty}`: missing size information for `{dep}`")]
71    TypeSizeDependencyMissing { ty: ConcreteTypeId, dep: ConcreteTypeId },
72    #[error("Type size computation failed for `{0}`: size overflow.")]
73    TypeSizeOverflow(ConcreteTypeId),
74}
75
76#[expect(clippy::disallowed_types)]
77type TypeMap<TType> = HashMap<ConcreteTypeId, TType>;
78#[expect(clippy::disallowed_types)]
79type LibfuncMap<TLibfunc> = HashMap<ConcreteLibfuncId, TLibfunc>;
80#[expect(clippy::disallowed_types)]
81type FunctionMap = HashMap<FunctionId, Function>;
82/// Mapping from the arguments for generating a concrete type (the generic-id and the arguments) to
83/// the concrete-id that points to it.
84#[expect(clippy::disallowed_types)]
85type ConcreteTypeIdMap<'a> = HashMap<(GenericTypeId, &'a [GenericArg]), ConcreteTypeId>;
86
87/// Registry for the data of the compiler, for all program specific data.
88pub struct ProgramRegistry<TType: GenericType, TLibfunc: GenericLibfunc> {
89    /// Mapping ids to the corresponding user function declaration from the program.
90    functions: FunctionMap,
91    /// Mapping ids to the concrete types represented by them.
92    concrete_types: TypeMap<TType::Concrete>,
93    /// Mapping ids to the concrete libfuncs represented by them.
94    concrete_libfuncs: LibfuncMap<TLibfunc::Concrete>,
95}
96#[expect(clippy::disallowed_types)]
97impl<TType: GenericType, TLibfunc: GenericLibfunc> ProgramRegistry<TType, TLibfunc> {
98    /// Create a registry for the program.
99    pub fn new(
100        program: &Program,
101    ) -> Result<ProgramRegistry<TType, TLibfunc>, Box<ProgramRegistryError>> {
102        let functions = get_functions(program)?;
103        let (concrete_types, concrete_type_ids) = get_concrete_types_maps::<TType>(program)?;
104        let concrete_libfuncs = get_concrete_libfuncs::<TType, TLibfunc>(
105            program,
106            &SpecializationContextForRegistry {
107                functions: &functions,
108                concrete_type_ids: &concrete_type_ids,
109                concrete_types: &concrete_types,
110            },
111        )?;
112        let registry = ProgramRegistry { functions, concrete_types, concrete_libfuncs };
113        registry.validate(program)?;
114        Ok(registry)
115    }
116
117    /// Gets a function from the input program.
118    pub fn get_function<'a>(
119        &'a self,
120        id: &FunctionId,
121    ) -> Result<&'a Function, Box<ProgramRegistryError>> {
122        self.functions
123            .get(id)
124            .ok_or_else(|| Box::new(ProgramRegistryError::MissingFunction(id.clone())))
125    }
126    /// Gets a type from the input program.
127    pub fn get_type<'a>(
128        &'a self,
129        id: &ConcreteTypeId,
130    ) -> Result<&'a TType::Concrete, Box<ProgramRegistryError>> {
131        self.concrete_types
132            .get(id)
133            .ok_or_else(|| Box::new(ProgramRegistryError::MissingType(id.clone())))
134    }
135    /// Gets a libfunc from the input program.
136    pub fn get_libfunc<'a>(
137        &'a self,
138        id: &ConcreteLibfuncId,
139    ) -> Result<&'a TLibfunc::Concrete, Box<ProgramRegistryError>> {
140        self.concrete_libfuncs
141            .get(id)
142            .ok_or_else(|| Box::new(ProgramRegistryError::MissingLibfunc(id.clone())))
143    }
144
145    /// Checks the validity of the [ProgramRegistry] and runs validations on the program.
146    ///
147    /// Later compilation stages may perform more validations as well as repeat these validations.
148    fn validate(&self, program: &Program) -> Result<(), Box<ProgramRegistryError>> {
149        // Check that all the parameter and return types are storable.
150        for func in self.functions.values() {
151            for ty in chain!(func.signature.param_types.iter(), func.signature.ret_types.iter()) {
152                if !self.get_type(ty)?.info().storable {
153                    return Err(Box::new(ProgramRegistryError::FunctionWithUnstorableType {
154                        func_id: func.id.clone(),
155                        ty: ty.clone(),
156                    }));
157                }
158            }
159            if func.entry_point.0 >= program.statements.len() {
160                return Err(Box::new(ProgramRegistryError::FunctionNonExistingEntryPoint(
161                    func.id.clone(),
162                )));
163            }
164        }
165        // A branch map, mapping from a destination statement to the statement that jumps to it.
166        // A branch is considered a branch only if it has more than one target.
167        // Assuming branches into branch alignments only, this should be a bijection.
168        let mut branches: HashMap<StatementIdx, StatementIdx> =
169            HashMap::<StatementIdx, StatementIdx>::default();
170        for (i, statement) in program.statements.iter().enumerate() {
171            self.validate_statement(program, StatementIdx(i), statement, &mut branches)?;
172        }
173        Ok(())
174    }
175
176    /// Checks the validity of a statement.
177    fn validate_statement(
178        &self,
179        program: &Program,
180        index: StatementIdx,
181        statement: &Statement,
182        branches: &mut HashMap<StatementIdx, StatementIdx>,
183    ) -> Result<(), Box<ProgramRegistryError>> {
184        let Statement::Invocation(invocation) = statement else {
185            return Ok(());
186        };
187        let libfunc = self.get_libfunc(&invocation.libfunc_id)?;
188        if invocation.args.len() != libfunc.param_signatures().len() {
189            return Err(Box::new(ProgramRegistryError::LibfuncInvocationInputCountMismatch(index)));
190        }
191        let libfunc_branches = libfunc.branch_signatures();
192        if invocation.branches.len() != libfunc_branches.len() {
193            return Err(Box::new(ProgramRegistryError::LibfuncInvocationBranchCountMismatch(
194                index,
195            )));
196        }
197        let libfunc_fallthrough = libfunc.fallthrough();
198        for (branch_index, (invocation_branch, libfunc_branch)) in
199            izip!(&invocation.branches, libfunc_branches).enumerate()
200        {
201            if invocation_branch.results.len() != libfunc_branch.vars.len() {
202                return Err(Box::new(
203                    ProgramRegistryError::LibfuncInvocationBranchResultCountMismatch(
204                        index,
205                        branch_index,
206                    ),
207                ));
208            }
209            if matches!(libfunc_fallthrough, Some(target) if target == branch_index)
210                != (invocation_branch.target == BranchTarget::Fallthrough)
211            {
212                return Err(Box::new(ProgramRegistryError::LibfuncInvocationBranchTargetMismatch(
213                    index,
214                    branch_index,
215                )));
216            }
217            if !matches!(libfunc_branch.ap_change, SierraApChange::BranchAlign)
218                && let Some(prev) = branches.get(&index)
219            {
220                return Err(Box::new(ProgramRegistryError::BranchNotToBranchAlign {
221                    src: *prev,
222                    dst: index,
223                }));
224            }
225            let next = index.next(invocation_branch.target);
226            if next.0 >= program.statements.len() {
227                return Err(Box::new(ProgramRegistryError::JumpOutOfRange(index)));
228            }
229            if libfunc_branches.len() > 1 {
230                if next.0 < index.0 {
231                    return Err(Box::new(ProgramRegistryError::BranchBackwards {
232                        src: index,
233                        dst: next,
234                    }));
235                }
236                match branches.entry(next) {
237                    Entry::Occupied(e) => {
238                        return Err(Box::new(ProgramRegistryError::MultipleJumpsToSameStatement {
239                            src1: *e.get(),
240                            src2: index,
241                            dst: next,
242                        }));
243                    }
244                    Entry::Vacant(e) => {
245                        e.insert(index);
246                    }
247                }
248            }
249        }
250        Ok(())
251    }
252}
253
254/// Creates the functions map.
255fn get_functions(program: &Program) -> Result<FunctionMap, Box<ProgramRegistryError>> {
256    let mut functions = FunctionMap::with_capacity(program.funcs.len());
257    for func in &program.funcs {
258        match functions.entry(func.id.clone()) {
259            Entry::Occupied(_) => {
260                Err(ProgramRegistryError::FunctionIdAlreadyExists(func.id.clone()))
261            }
262            Entry::Vacant(entry) => Ok(entry.insert(func.clone())),
263        }?;
264    }
265    Ok(functions)
266}
267
268struct TypeSpecializationContextForRegistry<'a, TType: GenericType> {
269    pub concrete_types: &'a TypeMap<TType::Concrete>,
270    pub declared_type_info: &'a TypeMap<TypeInfo>,
271}
272impl<TType: GenericType> TypeSpecializationContext
273    for TypeSpecializationContextForRegistry<'_, TType>
274{
275    fn try_get_type_info(&self, id: &ConcreteTypeId) -> Option<&TypeInfo> {
276        self.declared_type_info.get(id).or_else(|| self.concrete_types.get(id).map(|ty| ty.info()))
277    }
278}
279
280/// Creates the type-id to concrete type map, and the reverse map from generic-id and arguments to
281/// concrete-id.
282#[expect(clippy::disallowed_types)]
283fn get_concrete_types_maps<TType: GenericType>(
284    program: &Program,
285) -> Result<(TypeMap<TType::Concrete>, ConcreteTypeIdMap<'_>), Box<ProgramRegistryError>> {
286    let mut concrete_types = HashMap::with_capacity(program.type_declarations.len());
287    let mut concrete_type_ids =
288        HashMap::<(GenericTypeId, &[GenericArg]), ConcreteTypeId>::with_capacity(
289            program.type_declarations.len(),
290        );
291    let declared_type_info = program
292        .type_declarations
293        .iter()
294        .filter_map(|declaration| {
295            let TypeDeclaration { id, long_id, declared_type_info } = declaration;
296            let DeclaredTypeInfo { storable, droppable, duplicatable, zero_sized } =
297                declared_type_info.as_ref().cloned()?;
298            Some((
299                id.clone(),
300                TypeInfo {
301                    long_id: long_id.clone(),
302                    storable,
303                    droppable,
304                    duplicatable,
305                    zero_sized,
306                },
307            ))
308        })
309        .collect();
310    for declaration in &program.type_declarations {
311        let concrete_type = TType::specialize_by_id(
312            &TypeSpecializationContextForRegistry::<TType> {
313                concrete_types: &concrete_types,
314                declared_type_info: &declared_type_info,
315            },
316            &declaration.long_id.generic_id,
317            &declaration.long_id.generic_args,
318        )
319        .map_err(|error| {
320            Box::new(ProgramRegistryError::TypeSpecialization {
321                concrete_id: declaration.id.clone(),
322                error,
323            })
324        })?;
325        // Check that the info is consistent with declaration.
326        if let Some(declared_info) = declared_type_info.get(&declaration.id)
327            && concrete_type.info() != declared_info
328        {
329            return Err(Box::new(ProgramRegistryError::TypeInfoDeclarationMismatch(
330                declaration.id.clone(),
331            )));
332        }
333
334        match concrete_types.entry(declaration.id.clone()) {
335            Entry::Occupied(_) => Err(Box::new(ProgramRegistryError::TypeConcreteIdAlreadyExists(
336                declaration.id.clone(),
337            ))),
338            Entry::Vacant(entry) => Ok(entry.insert(concrete_type)),
339        }?;
340        match concrete_type_ids
341            .entry((declaration.long_id.generic_id.clone(), &declaration.long_id.generic_args[..]))
342        {
343            Entry::Occupied(_) => Err(Box::new(ProgramRegistryError::TypeAlreadyDeclared(
344                Box::new(declaration.clone()),
345            ))),
346            Entry::Vacant(entry) => Ok(entry.insert(declaration.id.clone())),
347        }?;
348    }
349    Ok((concrete_types, concrete_type_ids))
350}
351
352/// Context required for specialization process.
353pub struct SpecializationContextForRegistry<'a, TType: GenericType> {
354    pub functions: &'a FunctionMap,
355    pub concrete_type_ids: &'a ConcreteTypeIdMap<'a>,
356    pub concrete_types: &'a TypeMap<TType::Concrete>,
357}
358impl<TType: GenericType> TypeSpecializationContext for SpecializationContextForRegistry<'_, TType> {
359    fn try_get_type_info(&self, id: &ConcreteTypeId) -> Option<&TypeInfo> {
360        self.concrete_types.get(id).map(|ty| ty.info())
361    }
362}
363impl<TType: GenericType> SignatureSpecializationContext
364    for SpecializationContextForRegistry<'_, TType>
365{
366    fn try_get_concrete_type(
367        &self,
368        id: GenericTypeId,
369        generic_args: &[GenericArg],
370    ) -> Option<ConcreteTypeId> {
371        self.concrete_type_ids.get(&(id, generic_args)).cloned()
372    }
373
374    fn try_get_function_signature(&self, function_id: &FunctionId) -> Option<FunctionSignature> {
375        self.try_get_function(function_id).map(|f| f.signature)
376    }
377}
378impl<TType: GenericType> SpecializationContext for SpecializationContextForRegistry<'_, TType> {
379    fn try_get_function(&self, function_id: &FunctionId) -> Option<Function> {
380        self.functions.get(function_id).cloned()
381    }
382}
383
384/// Creates the libfuncs map.
385#[expect(clippy::disallowed_types)]
386fn get_concrete_libfuncs<TType: GenericType, TLibfunc: GenericLibfunc>(
387    program: &Program,
388    context: &SpecializationContextForRegistry<'_, TType>,
389) -> Result<LibfuncMap<TLibfunc::Concrete>, Box<ProgramRegistryError>> {
390    let mut concrete_libfuncs = HashMap::with_capacity(program.libfunc_declarations.len());
391    for declaration in &program.libfunc_declarations {
392        let concrete_libfunc = TLibfunc::specialize_by_id(
393            context,
394            &declaration.long_id.generic_id,
395            &declaration.long_id.generic_args,
396        )
397        .map_err(|error| ProgramRegistryError::LibfuncSpecialization {
398            concrete_id: declaration.id.clone(),
399            error,
400        })?;
401        match concrete_libfuncs.entry(declaration.id.clone()) {
402            Entry::Occupied(_) => {
403                Err(ProgramRegistryError::LibfuncConcreteIdAlreadyExists(declaration.id.clone()))
404            }
405            Entry::Vacant(entry) => Ok(entry.insert(concrete_libfunc)),
406        }?;
407    }
408    Ok(concrete_libfuncs)
409}
cairo_lang_sierra/program_registry.rs

cairo_lang_sierra/
program_registry.rs
