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
43 fn get_function_ap_change(
45 &self,
46 function_id: &FunctionId,
47 ) -> Result<SierraApChange, SpecializationError> {
48 self.try_get_function_ap_change(function_id)
49 .ok_or_else(|| SpecializationError::MissingFunction(function_id.clone()))
50 }
51
52 fn get_wrapped_concrete_type(
54 &self,
55 id: GenericTypeId,
56 wrapped: ConcreteTypeId,
57 ) -> Result<ConcreteTypeId, SpecializationError> {
58 self.get_concrete_type(id, &[GenericArg::Type(wrapped)])
59 }
60
61 fn as_type_specialization_context(&self) -> &dyn TypeSpecializationContext;
63}
64
65pub trait SpecializationContext: SignatureSpecializationContext {
67 fn upcast(&self) -> &dyn SignatureSpecializationContext;
70
71 fn try_get_function(&self, function_id: &FunctionId) -> Option<Function>;
73
74 fn get_function(&self, function_id: &FunctionId) -> Result<Function, SpecializationError> {
76 self.try_get_function(function_id)
77 .ok_or_else(|| SpecializationError::MissingFunction(function_id.clone()))
78 }
79}
80
81pub trait GenericLibfunc: Sized {
83 type Concrete: ConcreteLibfunc;
84
85 fn supported_ids() -> Vec<GenericLibfuncId>;
88
89 fn by_id(id: &GenericLibfuncId) -> Option<Self>;
91
92 fn specialize_signature(
94 &self,
95 context: &dyn SignatureSpecializationContext,
96 args: &[GenericArg],
97 ) -> Result<LibfuncSignature, SpecializationError>;
98
99 fn specialize(
101 &self,
102 context: &dyn SpecializationContext,
103 args: &[GenericArg],
104 ) -> Result<Self::Concrete, SpecializationError>;
105}
106
107pub trait GenericLibfuncEx: GenericLibfunc {
109 fn specialize_signature_by_id(
110 context: &dyn SignatureSpecializationContext,
111 libfunc_id: &GenericLibfuncId,
112 args: &[GenericArg],
113 ) -> Result<LibfuncSignature, ExtensionError>;
114
115 fn specialize_by_id(
116 context: &dyn SpecializationContext,
117 libfunc_id: &GenericLibfuncId,
118 args: &[GenericArg],
119 ) -> Result<Self::Concrete, ExtensionError>;
120}
121impl<TGenericLibfunc: GenericLibfunc> GenericLibfuncEx for TGenericLibfunc {
122 fn specialize_signature_by_id(
123 context: &dyn SignatureSpecializationContext,
124 libfunc_id: &GenericLibfuncId,
125 generic_args: &[GenericArg],
126 ) -> Result<LibfuncSignature, ExtensionError> {
127 if let Some(generic_libfunc) = Self::by_id(libfunc_id) {
128 generic_libfunc.specialize_signature(context, generic_args)
129 } else {
130 Err(SpecializationError::UnsupportedId(libfunc_id.0.clone()))
131 }
132 .map_err(move |error| ExtensionError::LibfuncSpecialization {
133 libfunc_id: libfunc_id.clone(),
134 generic_args: generic_args.iter().cloned().collect_vec(),
135 error,
136 })
137 }
138
139 fn specialize_by_id(
140 context: &dyn SpecializationContext,
141 libfunc_id: &GenericLibfuncId,
142 generic_args: &[GenericArg],
143 ) -> Result<TGenericLibfunc::Concrete, ExtensionError> {
144 if let Some(generic_libfunc) = Self::by_id(libfunc_id) {
145 generic_libfunc.specialize(context, generic_args)
146 } else {
147 Err(SpecializationError::UnsupportedId(libfunc_id.0.clone()))
148 }
149 .map_err(move |error| ExtensionError::LibfuncSpecialization {
150 libfunc_id: libfunc_id.clone(),
151 generic_args: generic_args.iter().cloned().collect_vec(),
152 error,
153 })
154 }
155}
156
157pub trait NamedLibfunc: Default {
159 type Concrete: ConcreteLibfunc;
160 const STR_ID: &'static str;
161
162 fn specialize_signature(
164 &self,
165 context: &dyn SignatureSpecializationContext,
166 args: &[GenericArg],
167 ) -> Result<LibfuncSignature, SpecializationError>;
168
169 fn specialize(
171 &self,
172 context: &dyn SpecializationContext,
173 args: &[GenericArg],
174 ) -> Result<Self::Concrete, SpecializationError>;
175}
176impl<TNamedLibfunc: NamedLibfunc> GenericLibfunc for TNamedLibfunc {
177 type Concrete = <Self as NamedLibfunc>::Concrete;
178
179 fn supported_ids() -> Vec<GenericLibfuncId> {
180 vec![GenericLibfuncId::from(Self::STR_ID)]
181 }
182
183 fn by_id(id: &GenericLibfuncId) -> Option<Self> {
184 if Self::STR_ID == id.0 { Some(Self::default()) } else { None }
185 }
186
187 fn specialize_signature(
188 &self,
189 context: &dyn SignatureSpecializationContext,
190 args: &[GenericArg],
191 ) -> Result<LibfuncSignature, SpecializationError> {
192 self.specialize_signature(context, args)
193 }
194
195 fn specialize(
196 &self,
197 context: &dyn SpecializationContext,
198 args: &[GenericArg],
199 ) -> Result<Self::Concrete, SpecializationError> {
200 self.specialize(context, args)
201 }
202}
203
204pub trait SignatureOnlyGenericLibfunc: Default {
206 const STR_ID: &'static str;
207
208 fn specialize_signature(
209 &self,
210 context: &dyn SignatureSpecializationContext,
211 args: &[GenericArg],
212 ) -> Result<LibfuncSignature, SpecializationError>;
213}
214
215impl<T: SignatureOnlyGenericLibfunc> NamedLibfunc for T {
216 type Concrete = SignatureOnlyConcreteLibfunc;
217 const STR_ID: &'static str = <Self as SignatureOnlyGenericLibfunc>::STR_ID;
218
219 fn specialize_signature(
220 &self,
221 context: &dyn SignatureSpecializationContext,
222 args: &[GenericArg],
223 ) -> Result<LibfuncSignature, SpecializationError> {
224 self.specialize_signature(context, args)
225 }
226
227 fn specialize(
228 &self,
229 context: &dyn SpecializationContext,
230 args: &[GenericArg],
231 ) -> Result<Self::Concrete, SpecializationError> {
232 Ok(SignatureOnlyConcreteLibfunc {
233 signature: self.specialize_signature(context.upcast(), args)?,
234 })
235 }
236}
237
238pub trait SignatureAndTypeGenericLibfunc: Default {
241 const STR_ID: &'static str;
242
243 fn specialize_signature(
244 &self,
245 context: &dyn SignatureSpecializationContext,
246 ty: ConcreteTypeId,
247 ) -> Result<LibfuncSignature, SpecializationError>;
248}
249
250#[derive(Default)]
252pub struct WrapSignatureAndTypeGenericLibfunc<T: SignatureAndTypeGenericLibfunc>(T);
253
254impl<T: SignatureAndTypeGenericLibfunc> NamedLibfunc for WrapSignatureAndTypeGenericLibfunc<T> {
255 type Concrete = SignatureAndTypeConcreteLibfunc;
256 const STR_ID: &'static str = <T as SignatureAndTypeGenericLibfunc>::STR_ID;
257
258 fn specialize_signature(
259 &self,
260 context: &dyn SignatureSpecializationContext,
261 args: &[GenericArg],
262 ) -> Result<LibfuncSignature, SpecializationError> {
263 self.0.specialize_signature(context, args_as_single_type(args)?)
264 }
265
266 fn specialize(
267 &self,
268 context: &dyn SpecializationContext,
269 args: &[GenericArg],
270 ) -> Result<Self::Concrete, SpecializationError> {
271 let ty = args_as_single_type(args)?;
272 Ok(SignatureAndTypeConcreteLibfunc {
273 ty: ty.clone(),
274 signature: self.0.specialize_signature(context.upcast(), ty)?,
275 })
276 }
277}
278
279pub trait NoGenericArgsGenericLibfunc: Default {
281 const STR_ID: &'static str;
282
283 fn specialize_signature(
284 &self,
285 context: &dyn SignatureSpecializationContext,
286 ) -> Result<LibfuncSignature, SpecializationError>;
287}
288impl<T: NoGenericArgsGenericLibfunc> SignatureOnlyGenericLibfunc for T {
289 const STR_ID: &'static str = <Self as NoGenericArgsGenericLibfunc>::STR_ID;
290
291 fn specialize_signature(
292 &self,
293 context: &dyn SignatureSpecializationContext,
294 args: &[GenericArg],
295 ) -> Result<LibfuncSignature, SpecializationError> {
296 if args.is_empty() {
297 self.specialize_signature(context)
298 } else {
299 Err(SpecializationError::WrongNumberOfGenericArgs)
300 }
301 }
302}
303
304#[derive(Clone)]
306pub struct ParamSignature {
307 pub ty: ConcreteTypeId,
309 pub allow_deferred: bool,
312 pub allow_add_const: bool,
314 pub allow_const: bool,
316}
317impl ParamSignature {
318 pub fn new(ty: ConcreteTypeId) -> Self {
320 Self { ty, allow_add_const: false, allow_deferred: false, allow_const: false }
321 }
322
323 pub fn with_allow_deferred(mut self) -> Self {
325 self.allow_deferred = true;
326 self
327 }
328
329 pub fn with_allow_add_const(mut self) -> Self {
331 self.allow_add_const = true;
332 self
333 }
334
335 pub fn with_allow_const(mut self) -> Self {
337 self.allow_const = true;
338 self
339 }
340
341 pub fn with_allow_all(mut self) -> Self {
343 self.allow_add_const = true;
344 self.allow_deferred = true;
345 self.allow_const = true;
346 self
347 }
348}
349impl From<ConcreteTypeId> for ParamSignature {
350 fn from(ty: ConcreteTypeId) -> Self {
351 Self::new(ty)
352 }
353}
354
355#[derive(Debug, Clone)]
360pub enum OutputVarReferenceInfo {
361 SameAsParam { param_idx: usize },
363 PartialParam { param_idx: usize },
369 NewTempVar {
372 idx: usize,
375 },
376 NewLocalVar,
378 Deferred(DeferredOutputKind),
381 SimpleDerefs,
383 ZeroSized,
385}
386
387#[derive(Clone, Debug, Eq, PartialEq)]
389pub enum DeferredOutputKind {
390 Const,
392 AddConst { param_idx: usize },
394 Generic,
397}
398
399#[derive(Debug, Clone)]
401pub struct OutputVarInfo {
402 pub ty: ConcreteTypeId,
403 pub ref_info: OutputVarReferenceInfo,
404}
405impl OutputVarInfo {
406 pub fn new_builtin(builtin: ConcreteTypeId, param_idx: usize) -> Self {
408 Self {
409 ty: builtin,
410 ref_info: OutputVarReferenceInfo::Deferred(DeferredOutputKind::AddConst { param_idx }),
411 }
412 }
413}
414
415#[derive(Debug)]
420pub struct BranchSignature {
421 pub vars: Vec<OutputVarInfo>,
423 pub ap_change: SierraApChange,
425}
426
427#[derive(Clone, Debug, Eq, PartialEq)]
429pub enum SierraApChange {
430 Unknown,
432 Known {
434 new_vars_only: bool,
438 },
439 BranchAlign,
442}
443pub trait ConcreteLibfunc {
445 fn param_signatures(&self) -> &[ParamSignature];
448 fn branch_signatures(&self) -> &[BranchSignature];
450 fn fallthrough(&self) -> Option<usize>;
452
453 fn output_types(&self) -> Vec<Vec<ConcreteTypeId>> {
455 self.branch_signatures()
456 .iter()
457 .map(|branch_info| {
458 branch_info.vars.iter().map(|var_info| var_info.ty.clone()).collect()
459 })
460 .collect()
461 }
462}
463
464pub struct LibfuncSignature {
466 pub param_signatures: Vec<ParamSignature>,
469 pub branch_signatures: Vec<BranchSignature>,
472 pub fallthrough: Option<usize>,
474}
475impl LibfuncSignature {
476 pub fn new_non_branch(
478 input_types: Vec<ConcreteTypeId>,
479 output_info: Vec<OutputVarInfo>,
480 ap_change: SierraApChange,
481 ) -> Self {
482 Self::new_non_branch_ex(
483 input_types.into_iter().map(ParamSignature::new).collect(),
484 output_info,
485 ap_change,
486 )
487 }
488
489 pub fn new_non_branch_ex(
492 param_signatures: Vec<ParamSignature>,
493 output_info: Vec<OutputVarInfo>,
494 ap_change: SierraApChange,
495 ) -> LibfuncSignature {
496 Self {
497 param_signatures,
498 branch_signatures: vec![BranchSignature { vars: output_info, ap_change }],
499 fallthrough: Some(0),
500 }
501 }
502}
503
504pub trait SignatureBasedConcreteLibfunc {
507 fn signature(&self) -> &LibfuncSignature;
508}
509
510impl<TSignatureBasedConcreteLibfunc: SignatureBasedConcreteLibfunc> ConcreteLibfunc
511 for TSignatureBasedConcreteLibfunc
512{
513 fn param_signatures(&self) -> &[ParamSignature] {
514 &self.signature().param_signatures
515 }
516 fn branch_signatures(&self) -> &[BranchSignature] {
517 &self.signature().branch_signatures
518 }
519 fn fallthrough(&self) -> Option<usize> {
520 self.signature().fallthrough
521 }
522}
523
524pub struct SignatureAndTypeConcreteLibfunc {
526 pub ty: ConcreteTypeId,
527 pub signature: LibfuncSignature,
528}
529impl SignatureBasedConcreteLibfunc for SignatureAndTypeConcreteLibfunc {
530 fn signature(&self) -> &LibfuncSignature {
531 &self.signature
532 }
533}
534
535pub struct SignatureOnlyConcreteLibfunc {
538 pub signature: LibfuncSignature,
539}
540impl SignatureBasedConcreteLibfunc for SignatureOnlyConcreteLibfunc {
541 fn signature(&self) -> &LibfuncSignature {
542 &self.signature
543 }
544}
545
546#[macro_export]
559macro_rules! define_concrete_libfunc_hierarchy {
560 (pub enum $name:ident $(<
561 $generic_arg:ident : $generic_arg_first_req:ident $(+ $generic_arg_other_reqs:ident)*
562 >)? {
563 $($variant_name:ident ($variant:ty),)*
564 }) => {
565 #[allow(clippy::enum_variant_names)]
566 pub enum $name $(< $generic_arg : $generic_arg_first_req $(+ $generic_arg_other_reqs)* >)? {
567 $($variant_name ($variant),)*
568 }
569 impl $(< $generic_arg : $generic_arg_first_req $(+ $generic_arg_other_reqs)* >)?
570 $crate::extensions::ConcreteLibfunc for $name $(< $generic_arg >)? {
571 $crate::extensions::lib_func::concrete_method_impl! {
572 fn param_signatures(&self) -> &[$crate::extensions::lib_func::ParamSignature] {
573 $($variant_name => $variant,)*
574 }
575 }
576 $crate::extensions::lib_func::concrete_method_impl!{
577 fn branch_signatures(&self) -> &[$crate::extensions::lib_func::BranchSignature] {
578 $($variant_name => $variant,)*
579 }
580 }
581 $crate::extensions::lib_func::concrete_method_impl!{
582 fn fallthrough(&self) -> Option<usize> {
583 $($variant_name => $variant,)*
584 }
585 }
586 }
587 }
588}
589
590macro_rules! concrete_method_impl {
593 (fn $method_name:ident(&self $(,$var_name:ident : $var:ty)*) -> $ret_type:ty {
594 $($variant_name:ident => $variant:ty,)*
595 }) => {
596 fn $method_name(&self $(,$var_name:ident : $var:ty)*) -> $ret_type {
597 match self {
598 $(Self::$variant_name(value) => value.$method_name()),*
599 }
600 }
601 }
602}
603pub(crate) use concrete_method_impl;
604
605#[macro_export]
618macro_rules! define_libfunc_hierarchy {
619 (pub enum $name:ident $(<
620 $generic_arg:ident : $generic_arg_first_req:ident $(+ $generic_arg_other_reqs:ident)*
621 >)? {
622 $($variant_name:ident ($variant:ty),)*
623 },
624 $concrete_name:ident) => {
625 #[allow(clippy::enum_variant_names)]
626 pub enum $name $(< $generic_arg : $generic_arg_first_req $(+ $generic_arg_other_reqs)* >)? {
627 $($variant_name ($variant)),*
628 }
629
630 impl $(< $generic_arg : $generic_arg_first_req $(+ $generic_arg_other_reqs)* >)?
631 $crate::extensions::GenericLibfunc for $name $(< $generic_arg >)? {
632 type Concrete = $concrete_name $(< $generic_arg >)?;
633 fn supported_ids() -> Vec<$crate::ids::GenericLibfuncId> {
634 itertools::chain!(
635 $(
636 <$variant as $crate::extensions::GenericLibfunc>::supported_ids()
637 ),*
638 ).collect()
639 }
640 fn by_id(id: &$crate::ids::GenericLibfuncId) -> Option<Self> {
641 $(
642 if let Some(res) = <$variant>::by_id(id){
643 return Some(Self::$variant_name(res));
644 }
645 )*
646 None
647 }
648 fn specialize_signature(
649 &self,
650 context: &dyn $crate::extensions::lib_func::SignatureSpecializationContext,
651 args: &[$crate::program::GenericArg],
652 ) -> Result<
653 $crate::extensions::lib_func::LibfuncSignature,
654 $crate::extensions::SpecializationError
655 >{
656 match self {
657 $(
658 Self::$variant_name(value) => {
659 <$variant as $crate::extensions::GenericLibfunc>::specialize_signature(
660 value, context, args,
661 )
662 }
663 ),*
664 }
665 }
666 fn specialize(
667 &self,
668 context: &dyn $crate::extensions::lib_func::SpecializationContext,
669 args: &[$crate::program::GenericArg],
670 ) -> Result<Self::Concrete, $crate::extensions::SpecializationError>{
671 match self {
672 $(
673 Self::$variant_name(value) => {
674 Ok(Self::Concrete::$variant_name(
675 <$variant as $crate::extensions::GenericLibfunc>::specialize(
676 value, context, args,
677 )?
678 .into(),
679 ))
680 }
681 ),*
682 }
683 }
684 }
685
686 $crate::define_concrete_libfunc_hierarchy! {
687 pub enum $concrete_name $(<
688 $generic_arg : $generic_arg_first_req $(+ $generic_arg_other_reqs)*
689 >)? {
690 $($variant_name (<$variant as $crate::extensions::GenericLibfunc> ::Concrete),)*
691 }
692 }
693 }
694}