1use itertools::Itertools;
2
3use super::args_as_single_type;
4use super::error::{ExtensionError, SpecializationError};
5use super::type_specialization_context::TypeSpecializationContext;
6use crate::ids::{ConcreteTypeId, FunctionId, GenericLibfuncId, GenericTypeId};
7use crate::program::{Function, FunctionSignature, GenericArg};
8
9pub trait SignatureSpecializationContext: TypeSpecializationContext {
11 fn try_get_concrete_type(
13 &self,
14 id: GenericTypeId,
15 generic_args: &[GenericArg],
16 ) -> Option<ConcreteTypeId>;
17
18 fn get_concrete_type(
20 &self,
21 id: GenericTypeId,
22 generic_args: &[GenericArg],
23 ) -> Result<ConcreteTypeId, SpecializationError> {
24 self.try_get_concrete_type(id.clone(), generic_args)
25 .ok_or_else(|| SpecializationError::TypeWasNotDeclared(id, generic_args.to_vec()))
26 }
27
28 fn try_get_function_signature(&self, function_id: &FunctionId) -> Option<FunctionSignature>;
30
31 fn get_function_signature(
33 &self,
34 function_id: &FunctionId,
35 ) -> Result<FunctionSignature, SpecializationError> {
36 self.try_get_function_signature(function_id)
37 .ok_or_else(|| SpecializationError::MissingFunction(function_id.clone()))
38 }
39
40 fn try_get_function_ap_change(&self, function_id: &FunctionId) -> Option<SierraApChange> {
42 Some(SierraApChange::FunctionCall(function_id.clone()))
43 }
44
45 fn get_function_ap_change(
47 &self,
48 function_id: &FunctionId,
49 ) -> Result<SierraApChange, SpecializationError> {
50 self.try_get_function_ap_change(function_id)
51 .ok_or_else(|| SpecializationError::MissingFunction(function_id.clone()))
52 }
53
54 fn get_wrapped_concrete_type(
56 &self,
57 id: GenericTypeId,
58 wrapped: ConcreteTypeId,
59 ) -> Result<ConcreteTypeId, SpecializationError> {
60 self.get_concrete_type(id, &[GenericArg::Type(wrapped)])
61 }
62}
63
64pub trait SpecializationContext: SignatureSpecializationContext {
66 fn try_get_function(&self, function_id: &FunctionId) -> Option<Function>;
68
69 fn get_function(&self, function_id: &FunctionId) -> Result<Function, SpecializationError> {
71 self.try_get_function(function_id)
72 .ok_or_else(|| SpecializationError::MissingFunction(function_id.clone()))
73 }
74}
75
76pub trait GenericLibfunc: Sized {
78 type Concrete: ConcreteLibfunc;
79
80 fn supported_ids() -> Vec<GenericLibfuncId>;
83
84 fn by_id(id: &GenericLibfuncId) -> Option<Self>;
86
87 fn specialize_signature(
89 &self,
90 context: &dyn SignatureSpecializationContext,
91 args: &[GenericArg],
92 ) -> Result<LibfuncSignature, SpecializationError>;
93
94 fn specialize(
96 &self,
97 context: &dyn SpecializationContext,
98 args: &[GenericArg],
99 ) -> Result<Self::Concrete, SpecializationError>;
100}
101
102pub trait GenericLibfuncEx: GenericLibfunc {
104 fn specialize_signature_by_id(
105 context: &dyn SignatureSpecializationContext,
106 libfunc_id: &GenericLibfuncId,
107 args: &[GenericArg],
108 ) -> Result<LibfuncSignature, ExtensionError>;
109
110 fn specialize_by_id(
111 context: &dyn SpecializationContext,
112 libfunc_id: &GenericLibfuncId,
113 args: &[GenericArg],
114 ) -> Result<Self::Concrete, ExtensionError>;
115}
116impl<TGenericLibfunc: GenericLibfunc> GenericLibfuncEx for TGenericLibfunc {
117 fn specialize_signature_by_id(
118 context: &dyn SignatureSpecializationContext,
119 libfunc_id: &GenericLibfuncId,
120 generic_args: &[GenericArg],
121 ) -> Result<LibfuncSignature, ExtensionError> {
122 if let Some(generic_libfunc) = Self::by_id(libfunc_id) {
123 generic_libfunc.specialize_signature(context, generic_args)
124 } else {
125 Err(SpecializationError::UnsupportedId(libfunc_id.0.clone()))
126 }
127 .map_err(move |error| ExtensionError::LibfuncSpecialization {
128 libfunc_id: libfunc_id.clone(),
129 generic_args: generic_args.iter().cloned().collect_vec(),
130 error,
131 })
132 }
133
134 fn specialize_by_id(
135 context: &dyn SpecializationContext,
136 libfunc_id: &GenericLibfuncId,
137 generic_args: &[GenericArg],
138 ) -> Result<TGenericLibfunc::Concrete, ExtensionError> {
139 if let Some(generic_libfunc) = Self::by_id(libfunc_id) {
140 generic_libfunc.specialize(context, generic_args)
141 } else {
142 Err(SpecializationError::UnsupportedId(libfunc_id.0.clone()))
143 }
144 .map_err(move |error| ExtensionError::LibfuncSpecialization {
145 libfunc_id: libfunc_id.clone(),
146 generic_args: generic_args.iter().cloned().collect_vec(),
147 error,
148 })
149 }
150}
151
152pub trait NamedLibfunc: Default {
154 type Concrete: ConcreteLibfunc;
155 const STR_ID: &'static str;
156
157 fn specialize_signature(
159 &self,
160 context: &dyn SignatureSpecializationContext,
161 args: &[GenericArg],
162 ) -> Result<LibfuncSignature, SpecializationError>;
163
164 fn specialize(
166 &self,
167 context: &dyn SpecializationContext,
168 args: &[GenericArg],
169 ) -> Result<Self::Concrete, SpecializationError>;
170}
171impl<TNamedLibfunc: NamedLibfunc> GenericLibfunc for TNamedLibfunc {
172 type Concrete = <Self as NamedLibfunc>::Concrete;
173
174 fn supported_ids() -> Vec<GenericLibfuncId> {
175 vec![GenericLibfuncId::from(Self::STR_ID)]
176 }
177
178 fn by_id(id: &GenericLibfuncId) -> Option<Self> {
179 if Self::STR_ID == id.0 { Some(Self::default()) } else { None }
180 }
181
182 fn specialize_signature(
183 &self,
184 context: &dyn SignatureSpecializationContext,
185 args: &[GenericArg],
186 ) -> Result<LibfuncSignature, SpecializationError> {
187 self.specialize_signature(context, args)
188 }
189
190 fn specialize(
191 &self,
192 context: &dyn SpecializationContext,
193 args: &[GenericArg],
194 ) -> Result<Self::Concrete, SpecializationError> {
195 self.specialize(context, args)
196 }
197}
198
199pub trait SignatureOnlyGenericLibfunc: Default {
201 const STR_ID: &'static str;
202
203 fn specialize_signature(
204 &self,
205 context: &dyn SignatureSpecializationContext,
206 args: &[GenericArg],
207 ) -> Result<LibfuncSignature, SpecializationError>;
208}
209
210impl<T: SignatureOnlyGenericLibfunc> NamedLibfunc for T {
211 type Concrete = SignatureOnlyConcreteLibfunc;
212 const STR_ID: &'static str = <Self as SignatureOnlyGenericLibfunc>::STR_ID;
213
214 fn specialize_signature(
215 &self,
216 context: &dyn SignatureSpecializationContext,
217 args: &[GenericArg],
218 ) -> Result<LibfuncSignature, SpecializationError> {
219 self.specialize_signature(context, args)
220 }
221
222 fn specialize(
223 &self,
224 context: &dyn SpecializationContext,
225 args: &[GenericArg],
226 ) -> Result<Self::Concrete, SpecializationError> {
227 Ok(SignatureOnlyConcreteLibfunc { signature: self.specialize_signature(context, args)? })
228 }
229}
230
231pub trait SignatureAndTypeGenericLibfunc: Default {
234 const STR_ID: &'static str;
235
236 fn specialize_signature(
237 &self,
238 context: &dyn SignatureSpecializationContext,
239 ty: ConcreteTypeId,
240 ) -> Result<LibfuncSignature, SpecializationError>;
241}
242
243#[derive(Default)]
245pub struct WrapSignatureAndTypeGenericLibfunc<T: SignatureAndTypeGenericLibfunc>(T);
246
247impl<T: SignatureAndTypeGenericLibfunc> NamedLibfunc for WrapSignatureAndTypeGenericLibfunc<T> {
248 type Concrete = SignatureAndTypeConcreteLibfunc;
249 const STR_ID: &'static str = <T as SignatureAndTypeGenericLibfunc>::STR_ID;
250
251 fn specialize_signature(
252 &self,
253 context: &dyn SignatureSpecializationContext,
254 args: &[GenericArg],
255 ) -> Result<LibfuncSignature, SpecializationError> {
256 self.0.specialize_signature(context, args_as_single_type(args)?.clone())
257 }
258
259 fn specialize(
260 &self,
261 context: &dyn SpecializationContext,
262 args: &[GenericArg],
263 ) -> Result<Self::Concrete, SpecializationError> {
264 let ty = args_as_single_type(args)?;
265 Ok(SignatureAndTypeConcreteLibfunc {
266 ty: ty.clone(),
267 signature: self.0.specialize_signature(context, ty.clone())?,
268 })
269 }
270}
271
272pub trait NoGenericArgsGenericLibfunc: Default {
274 const STR_ID: &'static str;
275
276 fn specialize_signature(
277 &self,
278 context: &dyn SignatureSpecializationContext,
279 ) -> Result<LibfuncSignature, SpecializationError>;
280}
281impl<T: NoGenericArgsGenericLibfunc> SignatureOnlyGenericLibfunc for T {
282 const STR_ID: &'static str = <Self as NoGenericArgsGenericLibfunc>::STR_ID;
283
284 fn specialize_signature(
285 &self,
286 context: &dyn SignatureSpecializationContext,
287 args: &[GenericArg],
288 ) -> Result<LibfuncSignature, SpecializationError> {
289 if args.is_empty() {
290 self.specialize_signature(context)
291 } else {
292 Err(SpecializationError::WrongNumberOfGenericArgs)
293 }
294 }
295}
296
297#[derive(Clone, PartialEq, Eq)]
299pub struct ParamSignature {
300 pub ty: ConcreteTypeId,
302 pub allow_deferred: bool,
305 pub allow_add_const: bool,
307 pub allow_const: bool,
309}
310impl ParamSignature {
311 pub fn new(ty: ConcreteTypeId) -> Self {
313 Self { ty, allow_add_const: false, allow_deferred: false, allow_const: false }
314 }
315
316 pub fn with_allow_deferred(mut self) -> Self {
318 self.allow_deferred = true;
319 self
320 }
321
322 pub fn with_allow_add_const(mut self) -> Self {
324 self.allow_add_const = true;
325 self
326 }
327
328 pub fn with_allow_const(mut self) -> Self {
330 self.allow_const = true;
331 self
332 }
333
334 pub fn with_allow_all(mut self) -> Self {
336 self.allow_add_const = true;
337 self.allow_deferred = true;
338 self.allow_const = true;
339 self
340 }
341}
342impl From<ConcreteTypeId> for ParamSignature {
343 fn from(ty: ConcreteTypeId) -> Self {
344 Self::new(ty)
345 }
346}
347
348#[derive(Debug, Clone, PartialEq, Eq)]
353pub enum OutputVarReferenceInfo {
354 SameAsParam { param_idx: usize },
356 PartialParam { param_idx: usize },
362 NewTempVar {
365 idx: usize,
368 },
369 NewLocalVar,
371 Deferred(DeferredOutputKind),
374 SimpleDerefs,
376 ZeroSized,
378}
379
380#[derive(Copy, Clone, Debug, Eq, PartialEq)]
382pub enum DeferredOutputKind {
383 Const,
385 AddConst,
387 Generic,
390}
391
392#[derive(Debug, Clone, Eq, PartialEq)]
394pub struct OutputVarInfo {
395 pub ty: ConcreteTypeId,
396 pub ref_info: OutputVarReferenceInfo,
397}
398impl OutputVarInfo {
399 pub fn new_builtin(builtin: ConcreteTypeId) -> Self {
401 Self {
402 ty: builtin,
403 ref_info: OutputVarReferenceInfo::Deferred(DeferredOutputKind::AddConst),
404 }
405 }
406}
407
408#[derive(Debug, Clone, Eq, PartialEq)]
413pub struct BranchSignature {
414 pub vars: Vec<OutputVarInfo>,
416 pub ap_change: SierraApChange,
418}
419
420#[derive(Clone, Debug, Eq, PartialEq)]
422pub enum SierraApChange {
423 Unknown,
425 Known {
427 new_vars_only: bool,
431 },
432 BranchAlign,
435 FunctionCall(FunctionId),
437}
438pub trait ConcreteLibfunc {
440 fn param_signatures(&self) -> &[ParamSignature];
443 fn branch_signatures(&self) -> &[BranchSignature];
445 fn fallthrough(&self) -> Option<usize>;
447
448 fn output_types(
450 &self,
451 ) -> impl Iterator<Item = impl ExactSizeIterator<Item = &ConcreteTypeId> + DoubleEndedIterator>
452 {
453 self.branch_signatures()
454 .iter()
455 .map(|branch_info| branch_info.vars.iter().map(|var_info| &var_info.ty))
456 }
457}
458
459#[derive(Clone, PartialEq, Eq)]
461pub struct LibfuncSignature {
462 pub param_signatures: Vec<ParamSignature>,
465 pub branch_signatures: Vec<BranchSignature>,
468 pub fallthrough: Option<usize>,
470}
471impl LibfuncSignature {
472 pub fn new_non_branch(
474 input_types: Vec<ConcreteTypeId>,
475 output_info: Vec<OutputVarInfo>,
476 ap_change: SierraApChange,
477 ) -> Self {
478 Self::new_non_branch_ex(
479 input_types.into_iter().map(ParamSignature::new).collect(),
480 output_info,
481 ap_change,
482 )
483 }
484
485 pub fn new_non_branch_ex(
488 param_signatures: Vec<ParamSignature>,
489 output_info: Vec<OutputVarInfo>,
490 ap_change: SierraApChange,
491 ) -> LibfuncSignature {
492 Self {
493 param_signatures,
494 branch_signatures: vec![BranchSignature { vars: output_info, ap_change }],
495 fallthrough: Some(0),
496 }
497 }
498}
499
500pub trait SignatureBasedConcreteLibfunc {
503 fn signature(&self) -> &LibfuncSignature;
504}
505
506impl<TSignatureBasedConcreteLibfunc: SignatureBasedConcreteLibfunc> ConcreteLibfunc
507 for TSignatureBasedConcreteLibfunc
508{
509 fn param_signatures(&self) -> &[ParamSignature] {
510 &self.signature().param_signatures
511 }
512 fn branch_signatures(&self) -> &[BranchSignature] {
513 &self.signature().branch_signatures
514 }
515 fn fallthrough(&self) -> Option<usize> {
516 self.signature().fallthrough
517 }
518}
519
520pub struct SignatureAndTypeConcreteLibfunc {
522 pub ty: ConcreteTypeId,
523 pub signature: LibfuncSignature,
524}
525impl SignatureBasedConcreteLibfunc for SignatureAndTypeConcreteLibfunc {
526 fn signature(&self) -> &LibfuncSignature {
527 &self.signature
528 }
529}
530
531pub struct SignatureOnlyConcreteLibfunc {
534 pub signature: LibfuncSignature,
535}
536impl SignatureBasedConcreteLibfunc for SignatureOnlyConcreteLibfunc {
537 fn signature(&self) -> &LibfuncSignature {
538 &self.signature
539 }
540}
541
542#[macro_export]
555macro_rules! define_concrete_libfunc_hierarchy {
556 (pub enum $name:ident $(<
557 $generic_arg:ident : $generic_arg_first_req:ident $(+ $generic_arg_other_reqs:ident)*
558 >)? {
559 $($variant_name:ident ($variant:ty),)*
560 }) => {
561 #[allow(clippy::enum_variant_names)]
562 pub enum $name $(< $generic_arg : $generic_arg_first_req $(+ $generic_arg_other_reqs)* >)? {
563 $($variant_name ($variant),)*
564 }
565 impl $(< $generic_arg : $generic_arg_first_req $(+ $generic_arg_other_reqs)* >)?
566 $crate::extensions::ConcreteLibfunc for $name $(< $generic_arg >)? {
567 $crate::extensions::lib_func::concrete_method_impl! {
568 fn param_signatures(&self) -> &[$crate::extensions::lib_func::ParamSignature] {
569 $($variant_name => $variant,)*
570 }
571 }
572 $crate::extensions::lib_func::concrete_method_impl!{
573 fn branch_signatures(&self) -> &[$crate::extensions::lib_func::BranchSignature] {
574 $($variant_name => $variant,)*
575 }
576 }
577 $crate::extensions::lib_func::concrete_method_impl!{
578 fn fallthrough(&self) -> Option<usize> {
579 $($variant_name => $variant,)*
580 }
581 }
582 }
583 }
584}
585
586macro_rules! concrete_method_impl {
589 (fn $method_name:ident(&self $(,$var_name:ident : $var:ty)*) -> $ret_type:ty {
590 $($variant_name:ident => $variant:ty,)*
591 }) => {
592 fn $method_name(&self $(,$var_name:ident : $var:ty)*) -> $ret_type {
593 match self {
594 $(Self::$variant_name(value) => value.$method_name()),*
595 }
596 }
597 }
598}
599pub(crate) use concrete_method_impl;
600
601#[macro_export]
614macro_rules! define_libfunc_hierarchy {
615 (pub enum $name:ident $(<
616 $generic_arg:ident : $generic_arg_first_req:ident $(+ $generic_arg_other_reqs:ident)*
617 >)? {
618 $($variant_name:ident ($variant:ty),)*
619 },
620 $concrete_name:ident) => {
621 #[allow(clippy::enum_variant_names)]
622 pub enum $name $(< $generic_arg : $generic_arg_first_req $(+ $generic_arg_other_reqs)* >)? {
623 $($variant_name ($variant)),*
624 }
625
626 impl $(< $generic_arg : $generic_arg_first_req $(+ $generic_arg_other_reqs)* >)?
627 $crate::extensions::GenericLibfunc for $name $(< $generic_arg >)? {
628 type Concrete = $concrete_name $(< $generic_arg >)?;
629 fn supported_ids() -> Vec<$crate::ids::GenericLibfuncId> {
630 itertools::chain!(
631 $(
632 <$variant as $crate::extensions::GenericLibfunc>::supported_ids()
633 ),*
634 ).collect()
635 }
636 fn by_id(id: &$crate::ids::GenericLibfuncId) -> Option<Self> {
637 $(
638 if let Some(res) = <$variant>::by_id(id){
639 return Some(Self::$variant_name(res));
640 }
641 )*
642 None
643 }
644 fn specialize_signature(
645 &self,
646 context: &dyn $crate::extensions::lib_func::SignatureSpecializationContext,
647 args: &[$crate::program::GenericArg],
648 ) -> Result<
649 $crate::extensions::lib_func::LibfuncSignature,
650 $crate::extensions::SpecializationError
651 >{
652 match self {
653 $(
654 Self::$variant_name(value) => {
655 <$variant as $crate::extensions::GenericLibfunc>::specialize_signature(
656 value, context, args,
657 )
658 }
659 ),*
660 }
661 }
662 fn specialize(
663 &self,
664 context: &dyn $crate::extensions::lib_func::SpecializationContext,
665 args: &[$crate::program::GenericArg],
666 ) -> Result<Self::Concrete, $crate::extensions::SpecializationError>{
667 match self {
668 $(
669 Self::$variant_name(value) => {
670 Ok(Self::Concrete::$variant_name(
671 <$variant as $crate::extensions::GenericLibfunc>::specialize(
672 value, context, args,
673 )?
674 .into(),
675 ))
676 }
677 ),*
678 }
679 }
680 }
681
682 $crate::define_concrete_libfunc_hierarchy! {
683 pub enum $concrete_name $(<
684 $generic_arg : $generic_arg_first_req $(+ $generic_arg_other_reqs)*
685 >)? {
686 $($variant_name (<$variant as $crate::extensions::GenericLibfunc> ::Concrete),)*
687 }
688 }
689 }
690}