use itertools::Itertools;
use super::args_as_single_type;
use super::error::{ExtensionError, SpecializationError};
use super::type_specialization_context::TypeSpecializationContext;
use crate::ids::{ConcreteTypeId, FunctionId, GenericLibfuncId, GenericTypeId};
use crate::program::{Function, FunctionSignature, GenericArg};
pub trait SignatureSpecializationContext: TypeSpecializationContext {
fn try_get_concrete_type(
&self,
id: GenericTypeId,
generic_args: &[GenericArg],
) -> Option<ConcreteTypeId>;
fn get_concrete_type(
&self,
id: GenericTypeId,
generic_args: &[GenericArg],
) -> Result<ConcreteTypeId, SpecializationError> {
self.try_get_concrete_type(id.clone(), generic_args)
.ok_or_else(|| SpecializationError::TypeWasNotDeclared(id, generic_args.to_vec()))
}
fn try_get_function_signature(&self, function_id: &FunctionId) -> Option<FunctionSignature>;
fn get_function_signature(
&self,
function_id: &FunctionId,
) -> Result<FunctionSignature, SpecializationError> {
self.try_get_function_signature(function_id)
.ok_or_else(|| SpecializationError::MissingFunction(function_id.clone()))
}
fn try_get_function_ap_change(&self, function_id: &FunctionId) -> Option<SierraApChange>;
fn get_function_ap_change(
&self,
function_id: &FunctionId,
) -> Result<SierraApChange, SpecializationError> {
self.try_get_function_ap_change(function_id)
.ok_or_else(|| SpecializationError::MissingFunction(function_id.clone()))
}
fn get_wrapped_concrete_type(
&self,
id: GenericTypeId,
wrapped: ConcreteTypeId,
) -> Result<ConcreteTypeId, SpecializationError> {
self.get_concrete_type(id, &[GenericArg::Type(wrapped)])
}
fn as_type_specialization_context(&self) -> &dyn TypeSpecializationContext;
}
pub trait SpecializationContext: SignatureSpecializationContext {
fn upcast(&self) -> &dyn SignatureSpecializationContext;
fn try_get_function(&self, function_id: &FunctionId) -> Option<Function>;
fn get_function(&self, function_id: &FunctionId) -> Result<Function, SpecializationError> {
self.try_get_function(function_id)
.ok_or_else(|| SpecializationError::MissingFunction(function_id.clone()))
}
}
pub trait GenericLibfunc: Sized {
type Concrete: ConcreteLibfunc;
fn supported_ids() -> Vec<GenericLibfuncId>;
fn by_id(id: &GenericLibfuncId) -> Option<Self>;
fn specialize_signature(
&self,
context: &dyn SignatureSpecializationContext,
args: &[GenericArg],
) -> Result<LibfuncSignature, SpecializationError>;
fn specialize(
&self,
context: &dyn SpecializationContext,
args: &[GenericArg],
) -> Result<Self::Concrete, SpecializationError>;
}
pub trait GenericLibfuncEx: GenericLibfunc {
fn specialize_signature_by_id(
context: &dyn SignatureSpecializationContext,
libfunc_id: &GenericLibfuncId,
args: &[GenericArg],
) -> Result<LibfuncSignature, ExtensionError>;
fn specialize_by_id(
context: &dyn SpecializationContext,
libfunc_id: &GenericLibfuncId,
args: &[GenericArg],
) -> Result<Self::Concrete, ExtensionError>;
}
impl<TGenericLibfunc: GenericLibfunc> GenericLibfuncEx for TGenericLibfunc {
fn specialize_signature_by_id(
context: &dyn SignatureSpecializationContext,
libfunc_id: &GenericLibfuncId,
generic_args: &[GenericArg],
) -> Result<LibfuncSignature, ExtensionError> {
if let Some(generic_libfunc) = Self::by_id(libfunc_id) {
generic_libfunc.specialize_signature(context, generic_args)
} else {
Err(SpecializationError::UnsupportedId(libfunc_id.0.clone()))
}
.map_err(move |error| ExtensionError::LibfuncSpecialization {
libfunc_id: libfunc_id.clone(),
generic_args: generic_args.iter().cloned().collect_vec(),
error,
})
}
fn specialize_by_id(
context: &dyn SpecializationContext,
libfunc_id: &GenericLibfuncId,
generic_args: &[GenericArg],
) -> Result<TGenericLibfunc::Concrete, ExtensionError> {
if let Some(generic_libfunc) = Self::by_id(libfunc_id) {
generic_libfunc.specialize(context, generic_args)
} else {
Err(SpecializationError::UnsupportedId(libfunc_id.0.clone()))
}
.map_err(move |error| ExtensionError::LibfuncSpecialization {
libfunc_id: libfunc_id.clone(),
generic_args: generic_args.iter().cloned().collect_vec(),
error,
})
}
}
pub trait NamedLibfunc: Default {
type Concrete: ConcreteLibfunc;
const STR_ID: &'static str;
fn specialize_signature(
&self,
context: &dyn SignatureSpecializationContext,
args: &[GenericArg],
) -> Result<LibfuncSignature, SpecializationError>;
fn specialize(
&self,
context: &dyn SpecializationContext,
args: &[GenericArg],
) -> Result<Self::Concrete, SpecializationError>;
}
impl<TNamedLibfunc: NamedLibfunc> GenericLibfunc for TNamedLibfunc {
type Concrete = <Self as NamedLibfunc>::Concrete;
fn supported_ids() -> Vec<GenericLibfuncId> {
vec![GenericLibfuncId::from(Self::STR_ID)]
}
fn by_id(id: &GenericLibfuncId) -> Option<Self> {
if Self::STR_ID == id.0 { Some(Self::default()) } else { None }
}
fn specialize_signature(
&self,
context: &dyn SignatureSpecializationContext,
args: &[GenericArg],
) -> Result<LibfuncSignature, SpecializationError> {
self.specialize_signature(context, args)
}
fn specialize(
&self,
context: &dyn SpecializationContext,
args: &[GenericArg],
) -> Result<Self::Concrete, SpecializationError> {
self.specialize(context, args)
}
}
pub trait SignatureOnlyGenericLibfunc: Default {
const STR_ID: &'static str;
fn specialize_signature(
&self,
context: &dyn SignatureSpecializationContext,
args: &[GenericArg],
) -> Result<LibfuncSignature, SpecializationError>;
}
impl<T: SignatureOnlyGenericLibfunc> NamedLibfunc for T {
type Concrete = SignatureOnlyConcreteLibfunc;
const STR_ID: &'static str = <Self as SignatureOnlyGenericLibfunc>::STR_ID;
fn specialize_signature(
&self,
context: &dyn SignatureSpecializationContext,
args: &[GenericArg],
) -> Result<LibfuncSignature, SpecializationError> {
self.specialize_signature(context, args)
}
fn specialize(
&self,
context: &dyn SpecializationContext,
args: &[GenericArg],
) -> Result<Self::Concrete, SpecializationError> {
Ok(SignatureOnlyConcreteLibfunc {
signature: self.specialize_signature(context.upcast(), args)?,
})
}
}
pub trait SignatureAndTypeGenericLibfunc: Default {
const STR_ID: &'static str;
fn specialize_signature(
&self,
context: &dyn SignatureSpecializationContext,
ty: ConcreteTypeId,
) -> Result<LibfuncSignature, SpecializationError>;
}
#[derive(Default)]
pub struct WrapSignatureAndTypeGenericLibfunc<T: SignatureAndTypeGenericLibfunc>(T);
impl<T: SignatureAndTypeGenericLibfunc> NamedLibfunc for WrapSignatureAndTypeGenericLibfunc<T> {
type Concrete = SignatureAndTypeConcreteLibfunc;
const STR_ID: &'static str = <T as SignatureAndTypeGenericLibfunc>::STR_ID;
fn specialize_signature(
&self,
context: &dyn SignatureSpecializationContext,
args: &[GenericArg],
) -> Result<LibfuncSignature, SpecializationError> {
self.0.specialize_signature(context, args_as_single_type(args)?)
}
fn specialize(
&self,
context: &dyn SpecializationContext,
args: &[GenericArg],
) -> Result<Self::Concrete, SpecializationError> {
let ty = args_as_single_type(args)?;
Ok(SignatureAndTypeConcreteLibfunc {
ty: ty.clone(),
signature: self.0.specialize_signature(context.upcast(), ty)?,
})
}
}
pub trait NoGenericArgsGenericLibfunc: Default {
const STR_ID: &'static str;
fn specialize_signature(
&self,
context: &dyn SignatureSpecializationContext,
) -> Result<LibfuncSignature, SpecializationError>;
}
impl<T: NoGenericArgsGenericLibfunc> SignatureOnlyGenericLibfunc for T {
const STR_ID: &'static str = <Self as NoGenericArgsGenericLibfunc>::STR_ID;
fn specialize_signature(
&self,
context: &dyn SignatureSpecializationContext,
args: &[GenericArg],
) -> Result<LibfuncSignature, SpecializationError> {
if args.is_empty() {
self.specialize_signature(context)
} else {
Err(SpecializationError::WrongNumberOfGenericArgs)
}
}
}
#[derive(Clone)]
pub struct ParamSignature {
pub ty: ConcreteTypeId,
pub allow_deferred: bool,
pub allow_add_const: bool,
pub allow_const: bool,
}
impl ParamSignature {
pub fn new(ty: ConcreteTypeId) -> Self {
Self { ty, allow_add_const: false, allow_deferred: false, allow_const: false }
}
pub fn with_allow_deferred(mut self) -> Self {
self.allow_deferred = true;
self
}
pub fn with_allow_add_const(mut self) -> Self {
self.allow_add_const = true;
self
}
pub fn with_allow_const(mut self) -> Self {
self.allow_const = true;
self
}
pub fn with_allow_all(mut self) -> Self {
self.allow_add_const = true;
self.allow_deferred = true;
self.allow_const = true;
self
}
}
impl From<ConcreteTypeId> for ParamSignature {
fn from(ty: ConcreteTypeId) -> Self {
Self::new(ty)
}
}
#[derive(Debug, Clone)]
pub enum OutputVarReferenceInfo {
SameAsParam { param_idx: usize },
PartialParam { param_idx: usize },
NewTempVar {
idx: usize,
},
NewLocalVar,
Deferred(DeferredOutputKind),
SimpleDerefs,
ZeroSized,
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum DeferredOutputKind {
Const,
AddConst { param_idx: usize },
Generic,
}
#[derive(Debug, Clone)]
pub struct OutputVarInfo {
pub ty: ConcreteTypeId,
pub ref_info: OutputVarReferenceInfo,
}
impl OutputVarInfo {
pub fn new_builtin(builtin: ConcreteTypeId, param_idx: usize) -> Self {
Self {
ty: builtin,
ref_info: OutputVarReferenceInfo::Deferred(DeferredOutputKind::AddConst { param_idx }),
}
}
}
#[derive(Debug)]
pub struct BranchSignature {
pub vars: Vec<OutputVarInfo>,
pub ap_change: SierraApChange,
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum SierraApChange {
Unknown,
Known {
new_vars_only: bool,
},
BranchAlign,
}
pub trait ConcreteLibfunc {
fn param_signatures(&self) -> &[ParamSignature];
fn branch_signatures(&self) -> &[BranchSignature];
fn fallthrough(&self) -> Option<usize>;
fn output_types(&self) -> Vec<Vec<ConcreteTypeId>> {
self.branch_signatures()
.iter()
.map(|branch_info| {
branch_info.vars.iter().map(|var_info| var_info.ty.clone()).collect()
})
.collect()
}
}
pub struct LibfuncSignature {
pub param_signatures: Vec<ParamSignature>,
pub branch_signatures: Vec<BranchSignature>,
pub fallthrough: Option<usize>,
}
impl LibfuncSignature {
pub fn new_non_branch(
input_types: Vec<ConcreteTypeId>,
output_info: Vec<OutputVarInfo>,
ap_change: SierraApChange,
) -> Self {
Self::new_non_branch_ex(
input_types.into_iter().map(ParamSignature::new).collect(),
output_info,
ap_change,
)
}
pub fn new_non_branch_ex(
param_signatures: Vec<ParamSignature>,
output_info: Vec<OutputVarInfo>,
ap_change: SierraApChange,
) -> LibfuncSignature {
Self {
param_signatures,
branch_signatures: vec![BranchSignature { vars: output_info, ap_change }],
fallthrough: Some(0),
}
}
}
pub trait SignatureBasedConcreteLibfunc {
fn signature(&self) -> &LibfuncSignature;
}
impl<TSignatureBasedConcreteLibfunc: SignatureBasedConcreteLibfunc> ConcreteLibfunc
for TSignatureBasedConcreteLibfunc
{
fn param_signatures(&self) -> &[ParamSignature] {
&self.signature().param_signatures
}
fn branch_signatures(&self) -> &[BranchSignature] {
&self.signature().branch_signatures
}
fn fallthrough(&self) -> Option<usize> {
self.signature().fallthrough
}
}
pub struct SignatureAndTypeConcreteLibfunc {
pub ty: ConcreteTypeId,
pub signature: LibfuncSignature,
}
impl SignatureBasedConcreteLibfunc for SignatureAndTypeConcreteLibfunc {
fn signature(&self) -> &LibfuncSignature {
&self.signature
}
}
pub struct SignatureOnlyConcreteLibfunc {
pub signature: LibfuncSignature,
}
impl SignatureBasedConcreteLibfunc for SignatureOnlyConcreteLibfunc {
fn signature(&self) -> &LibfuncSignature {
&self.signature
}
}
#[macro_export]
macro_rules! define_concrete_libfunc_hierarchy {
(pub enum $name:ident $(<
$generic_arg:ident : $generic_arg_first_req:ident $(+ $generic_arg_other_reqs:ident)*
>)? {
$($variant_name:ident ($variant:ty),)*
}) => {
#[allow(clippy::enum_variant_names)]
pub enum $name $(< $generic_arg : $generic_arg_first_req $(+ $generic_arg_other_reqs)* >)? {
$($variant_name ($variant),)*
}
impl $(< $generic_arg : $generic_arg_first_req $(+ $generic_arg_other_reqs)* >)?
$crate::extensions::ConcreteLibfunc for $name $(< $generic_arg >)? {
$crate::extensions::lib_func::concrete_method_impl! {
fn param_signatures(&self) -> &[$crate::extensions::lib_func::ParamSignature] {
$($variant_name => $variant,)*
}
}
$crate::extensions::lib_func::concrete_method_impl!{
fn branch_signatures(&self) -> &[$crate::extensions::lib_func::BranchSignature] {
$($variant_name => $variant,)*
}
}
$crate::extensions::lib_func::concrete_method_impl!{
fn fallthrough(&self) -> Option<usize> {
$($variant_name => $variant,)*
}
}
}
}
}
macro_rules! concrete_method_impl {
(fn $method_name:ident(&self $(,$var_name:ident : $var:ty)*) -> $ret_type:ty {
$($variant_name:ident => $variant:ty,)*
}) => {
fn $method_name(&self $(,$var_name:ident : $var:ty)*) -> $ret_type {
match self {
$(Self::$variant_name(value) => value.$method_name()),*
}
}
}
}
pub(crate) use concrete_method_impl;
#[macro_export]
macro_rules! define_libfunc_hierarchy {
(pub enum $name:ident $(<
$generic_arg:ident : $generic_arg_first_req:ident $(+ $generic_arg_other_reqs:ident)*
>)? {
$($variant_name:ident ($variant:ty),)*
},
$concrete_name:ident) => {
#[allow(clippy::enum_variant_names)]
pub enum $name $(< $generic_arg : $generic_arg_first_req $(+ $generic_arg_other_reqs)* >)? {
$($variant_name ($variant)),*
}
impl $(< $generic_arg : $generic_arg_first_req $(+ $generic_arg_other_reqs)* >)?
$crate::extensions::GenericLibfunc for $name $(< $generic_arg >)? {
type Concrete = $concrete_name $(< $generic_arg >)?;
fn supported_ids() -> Vec<$crate::ids::GenericLibfuncId> {
itertools::chain!(
$(
<$variant as $crate::extensions::GenericLibfunc>::supported_ids()
),*
).collect()
}
fn by_id(id: &$crate::ids::GenericLibfuncId) -> Option<Self> {
$(
if let Some(res) = <$variant>::by_id(id){
return Some(Self::$variant_name(res));
}
)*
None
}
fn specialize_signature(
&self,
context: &dyn $crate::extensions::lib_func::SignatureSpecializationContext,
args: &[$crate::program::GenericArg],
) -> Result<
$crate::extensions::lib_func::LibfuncSignature,
$crate::extensions::SpecializationError
>{
match self {
$(
Self::$variant_name(value) => {
<$variant as $crate::extensions::GenericLibfunc>::specialize_signature(
value, context, args,
)
}
),*
}
}
fn specialize(
&self,
context: &dyn $crate::extensions::lib_func::SpecializationContext,
args: &[$crate::program::GenericArg],
) -> Result<Self::Concrete, $crate::extensions::SpecializationError>{
match self {
$(
Self::$variant_name(value) => {
Ok(Self::Concrete::$variant_name(
<$variant as $crate::extensions::GenericLibfunc>::specialize(
value, context, args,
)?
.into(),
))
}
),*
}
}
}
$crate::define_concrete_libfunc_hierarchy! {
pub enum $concrete_name $(<
$generic_arg : $generic_arg_first_req $(+ $generic_arg_other_reqs)*
>)? {
$($variant_name (<$variant as $crate::extensions::GenericLibfunc> ::Concrete),)*
}
}
}
}